Upd ws, db and import state of changes

app/db.py

@@ -1,54 +1,60 @@
 # app/db.py
 from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorGridFSBucket
-import os
-import logging
 from pymongo.uri_parser import parse_uri
+import os, logging
 
-logger = logging.getLogger("book-query")
+logger          = logging.getLogger("book-query")
+MONGO_URI       = os.getenv("MONGODB_URI")
+MONGO_DB_NAME   = os.getenv("MONGODB_DB", "querysearcher")
+TEXTBOOK_URI    = os.getenv("TEXTBOOK_URI")         
 
-MONGO_URI = os.getenv("MONGODB_URI")
-MONGO_DB_NAME = os.getenv("MONGODB_DB", "querysearcher")
-TEXTBOOK_URI = os.getenv("TEXTBOOK_URI")
-
-# == QUERY ==
-# Return a fresh client for current event loop
+# ────────────────────────────────────────────────────────────────
+# helpers for the main “query-searcher” DB
+# ────────────────────────────────────────────────────────────────
 def get_client():
     return AsyncIOMotorClient(MONGO_URI)
+
 def get_db():
     return get_client()[MONGO_DB_NAME]
+
 def get_gridfs():
     return AsyncIOMotorGridFSBucket(get_db())
 
-# == PDF SAVER ==
+# ────────────────────────────────────────────────────────────────
+# ONE canonical helper for the *textbook* bucket
+# ────────────────────────────────────────────────────────────────
+def _get_textbook_fs() -> AsyncIOMotorGridFSBucket:
+    """
+    Build (and cache) a GridFS bucket that points to the textbook replica.
+    Works whether TEXTBOOK_URI ends with '/<db>' or not.
+    """
+    if not TEXTBOOK_URI:
+        raise RuntimeError("TEXTBOOK_URI not set in environment")
+
+    parsed   = parse_uri(TEXTBOOK_URI)
+    db_name  = parsed.get("database") or "textbooks"      # fallback name
+    client   = AsyncIOMotorClient(TEXTBOOK_URI)
+    return AsyncIOMotorGridFSBucket(client[db_name])
+
+# ── public wrappers ─────────────────────────────────────────────
 async def save_to_textbook_fs(doc_id: str, file_path: str):
     try:
-        uri = TEXTBOOK_URI
-        if not uri:
-            raise ValueError("TEXTBOOK_URI not set")
-        # Parse to correct db
-        parsed = parse_uri(uri)
-        db_name = parsed.get("database")
-        if not db_name:
-            raise ValueError("No default database name defined or provided.")
-        # 
-        textbook_client = AsyncIOMotorClient(uri)
-        textbook_db = textbook_client[db_name]
-        textbook_fs = AsyncIOMotorGridFSBucket(textbook_db)
-        # Read
-        with open(file_path, "rb") as f:
-            await textbook_fs.upload_from_stream(f"{doc_id}.pdf", f)
-        # Log
-        logger.info(f"📦 PDF also stored to textbook bucket at: {TEXTBOOK_URI}")
+        bucket = _get_textbook_fs()
+        with open(file_path, "rb") as fp:
+            await bucket.upload_from_stream(f"{doc_id}.pdf", fp)
+        logger.info(f"📦 textbook PDF stored → {TEXTBOOK_URI}/{bucket.database.name}")
     except Exception as e:
-        logger.warning(f"⚠️ Failed to save to textbook storage: {e}")
+        logger.warning(f"⚠️ textbook GridFS save failed: {e}")
 
-# == PDF FETCHER ==
 async def fetch_textbook_pdf(doc_id: str):
+    bucket = _get_textbook_fs()
+    return await bucket.open_download_stream_by_name(f"{doc_id}.pdf")
+
+async def delete_textbook_pdf(doc_id: str):
+    bucket = _get_textbook_fs()
     try:
-        textbook_client = AsyncIOMotorClient(TEXTBOOK_URI)
-        textbook_db = textbook_client.get_default_database()
-        textbook_fs = AsyncIOMotorGridFSBucket(textbook_db)
-        return await textbook_fs.open_download_stream_by_name(f"{doc_id}.pdf")
-    except Exception as e:
-        logger.warning(f"⚠️ Failed to fetch textbook PDF for {doc_id}: {e}")
-        raise
+        file = await bucket.find({"filename": f"{doc_id}.pdf"}).to_list(1)
+        if file:
+            await bucket.delete(file[0]["_id"])
+    except Exception:
+        pass

app/routers/import_doc.py

@@ -76,7 +76,7 @@ async def import_book(req: ImportRequest):
     doc = {
         "_id": req.candidate_id,
         "title": req.title,
-        "status": "queued",
+        "status": "DOWNLOADING",
         "metadata": result
     }
     db = get_db()
@@ -91,7 +91,7 @@ async def import_book(req: ImportRequest):
         "source": req.source,
         "documentId": req.candidate_id,
         "uri": uri
-}
+    }
 
 
 # Fetch textbook on id

app/routers/ws_progress.py

@@ -1,28 +1,53 @@
-import asyncio
-import logging
-from fastapi import WebSocket
-from app.db import get_db
+# app/routers/ws_progress.py
+import asyncio, logging, contextlib
+from fastapi import WebSocket, WebSocketDisconnect
+from bson import ObjectId
+from app.db import get_db, get_gridfs, delete_textbook_pdf
 
 logger = logging.getLogger("book-query")
 
+async def _delete_everything(doc_id: str):
+    """Remove embeddings, GridFS file & textbook copy if the user vanished."""
+    db           = get_db()
+    gridfs_query = get_gridfs()
+    # remove metadata + embeddings
+    await db.documents.delete_one({"_id": doc_id})
+    await db.embeddings.delete_many({"document_id": doc_id})
+    # remove original PDF from the main bucket
+    with contextlib.suppress(Exception):
+        file = await gridfs_query.find({"filename": f"{doc_id}.pdf"}).to_list(1)
+        if file:
+            await gridfs_query.delete(file[0]["_id"])
+    # remove textbook replica
+    with contextlib.suppress(Exception):
+        await delete_textbook_pdf(doc_id)
+    logger.info(f"🗑️  cleaned up artefacts of {doc_id}")
+
+
 async def forward_progress(websocket: WebSocket, document_id: str):
+    """Handle state change allowing frontend to connect and update seamlessly."""
     db = get_db()
     try:
         while True:
             doc = await db.documents.find_one({"_id": document_id})
             if not doc:
                 await websocket.send_json({"status": "NOT_FOUND"})
-                await asyncio.sleep(2)
-                continue
+            else:
+                status = doc.get("status", "UNKNOWN")
+                await websocket.send_json({"status": status})
+                if status in {"READY", "FAILED"}:
+                    break
+            await asyncio.sleep(2)
 
-            status = doc.get("status", "UNKNOWN")
-            await websocket.send_json({"status": status})
-            if status in {"READY", "FAILED"}:
-                break
+    except WebSocketDisconnect:
+        logger.info(f"💨 client closed WS for {document_id} – deleting artefacts")
+        await _delete_everything(document_id)
 
-            await asyncio.sleep(2)
     except Exception as e:
-        logger.exception(f"📡 WebSocket failed for doc {document_id}: {e}")
-        await websocket.send_json({"status": "ERROR"})
+        logger.exception(f"📡 WS loop crashed for {document_id}: {e}")
+        with contextlib.suppress(Exception):
+            await websocket.send_json({"status": "ERROR"})
+
     finally:
-        await websocket.close()
+        with contextlib.suppress(Exception):
+            await websocket.close()