Update app.py

add status endpoint

app.py

@@ -118,6 +118,7 @@ import logging
 import sys
 from logging.handlers import RotatingFileHandler
 import asyncio
+from datetime import datetime, timezone
 from pydantic import BaseModel
 
 # Import core helpers from the existing codebase
@@ -200,6 +201,36 @@ app.add_middleware(
 SESSIONS = {}
 
 
+def _utc_now_iso() -> str:
+    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
+
+
+def _build_status_payload(session_id: str, sess: dict) -> dict:
+    return {
+        "session_id": session_id,
+        "status": sess.get("status", "queued"),
+        "progress_pct": sess.get("progress_pct", 0),
+        "message": sess.get("message", ""),
+        "doc_count": sess.get("doc_count"),
+        "chunk_count": sess.get("chunk_count"),
+        "started_at": sess.get("started_at"),
+        "updated_at": sess.get("updated_at"),
+        "completed_at": sess.get("completed_at"),
+        "error": sess.get("error"),
+    }
+
+
+def _update_session(session_id: str, **updates):
+    sess = SESSIONS.get(session_id)
+    if not sess:
+        return
+    now = _utc_now_iso()
+    updates["updated_at"] = now
+    if updates.get("status") in {"ready", "failed"} and not updates.get("completed_at"):
+        updates["completed_at"] = now
+    sess.update(updates)
+
+
 async def _background_process(session_id: str, tmpdir: str, paths: list):
     """Background worker to parse uploaded PDFs and build vectorstore/chain.
 
@@ -208,12 +239,32 @@ async def _background_process(session_id: str, tmpdir: str, paths: list):
     try:
         from generate_testset import get_documents_from_pdfs
 
+        _update_session(
+            session_id,
+            status="parsing",
+            progress_pct=10,
+            message="Parsing uploaded PDFs",
+            error=None,
+        )
         documents = await asyncio.to_thread(get_documents_from_pdfs, paths)
         if not documents:
-            SESSIONS[session_id]["status"] = "failed"
-            SESSIONS[session_id]["error"] = "No documents parsed from uploads"
+            _update_session(
+                session_id,
+                status="failed",
+                progress_pct=100,
+                message="Parsing failed: no documents parsed",
+                error="No documents parsed from uploads",
+            )
             return
 
+        _update_session(
+            session_id,
+            status="chunking",
+            progress_pct=35,
+            message="Chunking parsed text",
+            doc_count=len(documents),
+        )
+
         texts, metadatas = get_text_chunks(
             [
                 {
@@ -232,29 +283,56 @@ async def _background_process(session_id: str, tmpdir: str, paths: list):
             ]
         )
 
+        _update_session(
+            session_id,
+            status="indexing",
+            progress_pct=60,
+            message="Building vector store",
+            chunk_count=len(texts),
+        )
+
         vectorstore = await asyncio.to_thread(get_vectorstore, texts, metadatas)
 
+        _update_session(
+            session_id,
+            status="building_chain",
+            progress_pct=85,
+            message="Building conversation chain",
+        )
+
         try:
             conversation = await asyncio.to_thread(get_conversation_chain, vectorstore)
         except Exception as e:
-            SESSIONS[session_id]["status"] = "failed"
-            SESSIONS[session_id]["error"] = f"Conversation creation failed: {e}"
+            _update_session(
+                session_id,
+                status="failed",
+                progress_pct=100,
+                message="Conversation chain creation failed",
+                error=f"Conversation creation failed: {e}",
+            )
             return
 
-        SESSIONS[session_id].update(
-            {
-                "vectorstore": vectorstore,
-                "conversation": conversation,
-                "status": "ready",
-            }
+        _update_session(
+            session_id,
+            vectorstore=vectorstore,
+            conversation=conversation,
+            status="ready",
+            progress_pct=100,
+            message="Session is ready for questions",
+            error=None,
         )
         logger.info("Background processing complete for session %s", session_id)
     except Exception as e:
         logger.exception(
             "Background processing failed for session %s: %s", session_id, e
         )
-        SESSIONS[session_id]["status"] = "failed"
-        SESSIONS[session_id]["error"] = str(e)
+        _update_session(
+            session_id,
+            status="failed",
+            progress_pct=100,
+            message="Background processing failed",
+            error=str(e),
+        )
     finally:
         try:
             shutil.rmtree(tmpdir)
@@ -267,6 +345,14 @@ def health():
     return {"ok": True}
 
 
+@app.get("/status")
+def get_status(session_id: str, api_key: str = Depends(verify_api_key)):
+    sess = SESSIONS.get(session_id)
+    if not sess:
+        raise HTTPException(status_code=404, detail="session_id not found")
+    return _build_status_payload(session_id, sess)
+
+
 @app.post("/process")
 async def process(
     files: List[UploadFile] = File(...), api_key: str = Depends(verify_api_key)
@@ -285,8 +371,26 @@ async def process(
     paths = []
     # create session early and schedule background processing
     session_id = str(uuid.uuid4())
-    SESSIONS[session_id] = {"status": "processing", "tmpdir": tmpdir}
+    now = _utc_now_iso()
+    SESSIONS[session_id] = {
+        "status": "queued",
+        "progress_pct": 0,
+        "message": "Upload received and queued",
+        "started_at": now,
+        "updated_at": now,
+        "completed_at": None,
+        "error": None,
+        "doc_count": None,
+        "chunk_count": None,
+        "tmpdir": tmpdir,
+    }
     try:
+        _update_session(
+            session_id,
+            status="parsing",
+            progress_pct=5,
+            message="Saving uploaded files",
+        )
         for f in files:
             dest = os.path.join(tmpdir, f.filename or f"upload-{uuid.uuid4()}.pdf")
             with open(dest, "wb") as out:
@@ -294,12 +398,29 @@ async def process(
                 out.write(content)
             paths.append(dest)
 
+        _update_session(
+            session_id,
+            status="queued",
+            progress_pct=8,
+            message="Files saved; waiting for background processing",
+            doc_count=len(paths),
+        )
+
         # schedule background processing to build index & chain
         asyncio.create_task(_background_process(session_id, tmpdir, paths))
         logger.info("Scheduled background processing for session %s", session_id)
-        return {"session_id": session_id, "status": "processing"}
+        payload = _build_status_payload(session_id, SESSIONS[session_id])
+        payload["estimated_poll_interval_ms"] = 2500
+        return payload
     except Exception:
         # on unexpected error, clean up tmpdir and re-raise
+        _update_session(
+            session_id,
+            status="failed",
+            progress_pct=100,
+            message="Failed while saving uploads",
+            error="Unexpected error while handling uploaded files",
+        )
         try:
             shutil.rmtree(tmpdir)
         except Exception:
@@ -323,11 +444,14 @@ async def ask(
     if not sess:
         raise HTTPException(status_code=404, detail="session_id not found")
 
-    # If processing not complete, return 202 Accepted with status
+    # If processing not complete, return 202 Accepted with full status payload
     status = sess.get("status")
+    if status == "failed":
+        payload = _build_status_payload(session_id, sess)
+        return JSONResponse(status_code=409, content=payload)
     if status != "ready":
         return JSONResponse(
-            status_code=202, content={"session_id": session_id, "status": status}
+            status_code=202, content=_build_status_payload(session_id, sess)
         )
 
     conv = sess.get("conversation")