Spaces:
Running
Running
File size: 2,958 Bytes
ef4c8c3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
"""
Api.py: FastAPI endpoints for dataset generation, progress polling, and download.
"""
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, StreamingResponse
from .Core import job_manager
from .Progress import progress_tracker
from .Payment import payment_manager
import io
app = FastAPI()
@app.post("/generate-dataset")
async def generate_dataset(request: Request):
user_input = await request.json()
job_id, error = job_manager.start_job(user_input)
if error:
return JSONResponse({"error": error}, status_code=400)
return {"job_id": job_id}
@app.get("/progress/{job_id}")
def get_progress(job_id: str):
progress = progress_tracker.get(job_id)
if not progress:
return JSONResponse({"error": "Job not found"}, status_code=404)
return progress
@app.get("/download/{job_id}")
def download(job_id: str):
job = job_manager.get_job_status(job_id)
if not job or job.get("status") != "complete":
return JSONResponse({"error": "Job not complete"}, status_code=400)
# Payment check
plan = job.get("plan", "free")
tokens = job.get("token_budget", 0)
if payment_manager.requires_payment(plan, tokens):
return JSONResponse({"error": "Payment required", "checkout_url": payment_manager.create_checkout_session(plan, job_id)}, status_code=402)
# In production, use FileResponse to serve the file
return {
"download_url": job["result_path"],
"stats": job.get("stats", {})
}
@app.get("/download-corpus/{job_id}")
def download_corpus(job_id: str):
job = job_manager.get_job_status(job_id)
if not job or job.get("status") != "complete":
return JSONResponse({"error": "Job not complete"}, status_code=400)
if job.get("job_type") != "corpus":
return JSONResponse({"error": "Not a corpus job"}, status_code=400)
plan = job.get("plan", "free")
tokens = job.get("token_budget", 0)
if payment_manager.requires_payment(plan, tokens):
return JSONResponse({"error": "Payment required", "checkout_url": payment_manager.create_checkout_session(plan, job_id)}, status_code=402)
jsonl_lines = job.get("jsonl_lines", [])
stats = job.get("stats", {})
# Stream the JSONL as a file
file_like = io.StringIO("\n".join(jsonl_lines))
headers = {
"Content-Disposition": f"attachment; filename=scientific_corpus_{job_id}.jsonl"
}
return StreamingResponse(file_like, media_type="application/jsonl", headers=headers)
@app.get("/job-stats/{job_id}")
def job_stats(job_id: str):
job = job_manager.get_job_status(job_id)
if not job:
return JSONResponse({"error": "Job not found"}, status_code=404)
return {"stats": job.get("stats", {})}
@app.get("/price/{plan}")
def get_price(plan: str):
price = payment_manager.get_price(plan)
return {"plan": plan, "price": price}
|