update flow poolling

app.py

@@ -2,6 +2,9 @@ import json
 import os
 import shutil
 import uuid
+import time
+import asyncio
+from pathlib import Path
 from contextlib import asynccontextmanager
 from typing import Annotated, Optional
 
@@ -88,8 +91,22 @@ class PredictionResponse(BaseModel):
     )
 
 
+class JobStatus(BaseModel):
+    job_id: str
+    status: str
+    message: str | None = None
+    result: PredictionResponse | None = None
+    created_at: float
+    updated_at: float
+
+
 predictor: G3BatchPredictor
 
+MAX_CONCURRENT = int(os.getenv("MAX_CONCURRENT", "10"))
+jobs: dict[str, dict] = {}
+jobs_lock = asyncio.Lock()
+worker_sem = asyncio.Semaphore(MAX_CONCURRENT)
+
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -114,51 +131,115 @@ app = FastAPI(
 )
 
 
+async def _update_job(job_id: str, **fields) -> dict:
+    async with jobs_lock:
+        job = jobs[job_id]
+        job.update(fields)
+        job["updated_at"] = time.time()
+        return job.copy()
+
+
+async def _get_job(job_id: str) -> dict | None:
+    async with jobs_lock:
+        job = jobs.get(job_id)
+        return None if job is None else job.copy()
+
+
+async def _run_job(job_id: str, job_dir: Path) -> None:
+    await _update_job(job_id, status="running", message=None)
+    async with worker_sem:
+        try:
+            predictor.clear_directories()
+
+            os.makedirs(predictor.input_dir, exist_ok=True)
+            for file_path in job_dir.iterdir():
+                if file_path.is_file():
+                    dest = predictor.input_dir / file_path.name
+                    shutil.copy(file_path, dest)
+
+            response = await predictor.predict(model_name="gemini-2.5-pro")
+            prediction = LocationPredictionResponse(
+                latitude=response.latitude,
+                longitude=response.longitude,
+                location=response.location,
+                evidence=[
+                    EvidenceResponse(analysis=ev.analysis, references=ev.references)
+                    for ev in response.evidence
+                ],
+            )
+            transcript = predictor.get_transcript()
+            images_b64 = load_images_as_base64()
+
+            result = PredictionResponse(
+                prediction=prediction,
+                transcript=transcript,
+                media=images_b64,
+            )
+
+            await _update_job(job_id, status="succeeded", result=result)
+        except Exception as e:
+            await _update_job(job_id, status="failed", message=str(e))
+        finally:
+            shutil.rmtree(job_dir, ignore_errors=True)
+
+
 @app.post(
     "/g3/predict",
-    description="Provide location prediction.",
+    description="Provide location prediction (async, polling via job_id).",
+    response_model=JobStatus,
 )
 async def predict_endpoint(
     files: Annotated[
         list[UploadFile],
         File(description="Input images, videos and metadata json."),
     ],
-) -> PredictionResponse:
-    # Write files to disk
+) -> JobStatus:
+    if not files:
+        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="No files provided")
+
+    job_id = uuid.uuid4().hex
+    job_dir = Path("uploads") / job_id
+    os.makedirs(job_dir, exist_ok=True)
+
     try:
-        predictor.clear_directories()
         for file in files:
             filename = file.filename if file.filename is not None else uuid.uuid4().hex
-            filepath = predictor.input_dir / filename
-            os.makedirs(predictor.input_dir, exist_ok=True)
+            filepath = job_dir / filename
             with open(filepath, "wb") as buffer:
                 shutil.copyfileobj(file.file, buffer)
+
+        now = time.time()
+        async with jobs_lock:
+            jobs[job_id] = {
+                "job_id": job_id,
+                "status": "queued",
+                "message": None,
+                "result": None,
+                "created_at": now,
+                "updated_at": now,
+            }
+
+        asyncio.create_task(_run_job(job_id, job_dir))
+        job = await _get_job(job_id)
+        return job  # type: ignore[return-value]
     except Exception as e:
+        shutil.rmtree(job_dir, ignore_errors=True)
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=f"Failed to save file: {e}",
+            detail=f"Failed to enqueue job: {e}",
         )
 
-    # Get prediction
-    response = await predictor.predict(model_name="gemini-2.5-pro")
-    # response = predictor.get_response(response)
-    prediction = LocationPredictionResponse(
-        latitude=response.latitude,
-        longitude=response.longitude,
-        location=response.location,
-        evidence=[
-            EvidenceResponse(analysis=ev.analysis, references=ev.references)
-            for ev in response.evidence
-        ],
-    )
-    # Get transcript if available
-    transcript = predictor.get_transcript()
-
-    # Get media files if available
-    images_b64 = load_images_as_base64()
 
-    # Clear directories
-    return PredictionResponse(prediction=prediction, transcript=transcript, media=images_b64)
+@app.get(
+    "/g3/predict/{job_id}",
+    description="Get prediction job status/result.",
+    response_model=JobStatus,
+)
+async def get_job_status(job_id: str) -> JobStatus:
+    job = await _get_job(job_id)
+    if job is None:
+        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Job not found")
+    return job  # type: ignore[return-value]
 
 
 @app.get(