Update app.py

app.py

@@ -61,8 +61,6 @@ SESSION_STORES = {}
 STORE_LOCK = threading.RLock()
 CLEANUP_INTERVAL = 1800  # 30 minutes
 STORE_TTL = 1800  # 30 minutes
-# Note: asyncio.to_thread() uses its own internal thread pool, so a dedicated
-# global pool is not strictly necessary unless fine-grained control is needed.
 
 # --- Pydantic Models ---
 class ChatRequest(BaseModel):
@@ -148,6 +146,8 @@ def init_rag_models():
 def load_session_from_mongodb(session_id: str) -> Dict[str, Any]:
     """Load pre-existing session data and embeddings from MongoDB"""
     session_logger = create_session_logger(session_id)
+    session_logger.info(f"Loading session from MongoDB: {session_id}")
+    
     if not DB:
         raise ValueError("Database not connected")
 
@@ -170,11 +170,12 @@ def load_session_from_mongodb(session_id: str) -> Dict[str, Any]:
     groq_api_key = os.getenv("GROQ_API_KEY")
     session_rag = SessionRAG(session_id, groq_api_key)
     
-    # Load the existing chunks and indices (don't recreate embeddings)
+    session_logger.info(f"Loading existing session data with {len(chunks_list)} chunks")
     session_rag.load_existing_session_data(chunks_list)
 
     session_store = {
         "session_rag": session_rag,
+        "indexed": True,  # Add this field for health check
         "metadata": {
             "session_id": session_id,
             "title": session_doc.get("filename", "Document"),
@@ -186,10 +187,10 @@ def load_session_from_mongodb(session_id: str) -> Dict[str, Any]:
     session_logger.info("Session loaded from MongoDB with existing embeddings")
     return session_store
 
-
 def get_chat_history_safely(session_id: str, limit: int = 50) -> List[Dict[str, Any]]:
     """Get chat history with error handling."""
-    if not DB: return []
+    if not DB: 
+        return []
     try:
         chats_cursor = DB.chats.find({"session_id": session_id}).sort("created_at", -1).limit(limit)
         # Reverse to get chronological order [oldest -> newest]
@@ -228,11 +229,21 @@ async def periodic_cleanup():
         logger.info("Running periodic cleanup of expired sessions...")
         cleanup_expired_sessions()
 
+async def save_chat_message_safely(session_id: str, role: str, message: str):
+    """Saves chat messages in a non-blocking way."""
+    if not DB: 
+        return
+    try:
+        await asyncio.to_thread(
+            DB.chats.insert_one,
+            {"session_id": session_id, "role": role, "message": message, "created_at": datetime.utcnow()}
+        )
+    except Exception as e:
+        logger.error(f"Failed to save chat message for session {session_id}: {e}")
 
 # --- Application Lifespan ---
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # --- FIX APPLIED: Variable initialized to None ---
     cleanup_task = None
     APP_STATE["startup_time"] = datetime.utcnow()
     logger.info("Starting Advanced RAG Chat Service...")
@@ -252,7 +263,6 @@ async def lifespan(app: FastAPI):
         MONGO_CLIENT.close()
     logger.info("Shutdown complete.")
 
-
 # --- FastAPI App and Endpoints ---
 app = FastAPI(
     title="Advanced RAG Chat Service",
@@ -263,7 +273,10 @@ app = FastAPI(
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"],
+    allow_origins=["*"], 
+    allow_credentials=True, 
+    allow_methods=["*"], 
+    allow_headers=["*"],
 )
 
 @app.get("/")
@@ -275,7 +288,7 @@ async def health_check():
     uptime = (datetime.utcnow() - APP_STATE["startup_time"]).total_seconds()
     with STORE_LOCK:
         active_sessions = len(SESSION_STORES)
-        indexed_sessions = sum(1 for s in SESSION_STORES.values() if s["indexed"])
+        indexed_sessions = sum(1 for s in SESSION_STORES.values() if s.get("indexed", False))
 
     status = "healthy"
     if not RAG_MODELS_INITIALIZED or not APP_STATE["mongodb_connected"]:
@@ -284,7 +297,7 @@ async def health_check():
     return HealthResponse(
         status=status,
         mongodb_connected=APP_STATE["mongodb_connected"],
-        rag_models_initialized=RAG_MODELS_INITIALIZED, # Correct field name
+        rag_models_initialized=RAG_MODELS_INITIALIZED,
         faiss_available=FAISS_AVAILABLE,
         active_sessions=active_sessions,
         memory_usage={"loaded_sessions": active_sessions, "indexed_sessions": indexed_sessions},
@@ -298,22 +311,35 @@ async def chat_with_document(session_id: str, request: ChatRequest):
     start_time = time.time()
     
     try:
+        session_logger.info(f"Received chat request: {request.message[:100]}...")
+        
         # Check if session is already loaded in memory
         with STORE_LOCK:
             if session_id not in SESSION_STORES:
                 # Lazy load: Load session from MongoDB when first chat request comes
                 session_logger.info("Loading session from MongoDB for first chat request")
-                session_store = await asyncio.to_thread(load_session_from_mongodb, session_id)
-                SESSION_STORES[session_id] = session_store
+                try:
+                    session_store = await asyncio.to_thread(load_session_from_mongodb, session_id)
+                    SESSION_STORES[session_id] = session_store
+                    session_logger.info("Session loaded successfully from MongoDB")
+                except Exception as load_error:
+                    session_logger.error(f"Failed to load session: {load_error}")
+                    raise HTTPException(status_code=404, detail=f"Failed to load session: {str(load_error)}")
             
             # Update last access time
             SESSION_LAST_ACCESS[session_id] = datetime.utcnow()
             session_rag = SESSION_STORES[session_id]["session_rag"]
-            # Process the query
-        session_logger.info(f"Processing query: {request.message[:100]}...")
+        
+        session_logger.info(f"Processing query with SessionRAG...")
+        
+        # Process the query using SessionRAG
         result = await asyncio.to_thread(session_rag.query_documents, request.message, top_k=5)
-        APP_STATE["total_queries"] += 1
         
+        if 'error' in result:
+            session_logger.error(f"Query processing error: {result['error']}")
+            raise HTTPException(status_code=500, detail=result['error'])
+        
+        APP_STATE["total_queries"] += 1
         answer = result.get('answer', 'Unable to generate an answer.')
         
         # Save chat messages asynchronously
@@ -321,7 +347,7 @@ async def chat_with_document(session_id: str, request: ChatRequest):
         asyncio.create_task(save_chat_message_safely(session_id, "assistant", answer))
         
         processing_time = time.time() - start_time
-        session_logger.info(f"Query processed in {processing_time:.2f}s.")
+        session_logger.info(f"Query processed successfully in {processing_time:.2f}s")
         
         return ChatResponse(
             success=True,
@@ -333,16 +359,20 @@ async def chat_with_document(session_id: str, request: ChatRequest):
             query_analysis=result.get('query_analysis'),
             confidence=result.get('confidence')
         )
+        
+    except HTTPException:
+        raise
     except Exception as e:
         session_logger.error(f"Chat processing failed: {e}", exc_info=True)
-        raise HTTPException(status_code=500, detail=f"Chat processing error: {e}")
+        APP_STATE["errors"].append(f"Chat error: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Chat processing error: {str(e)}")
 
 @app.get("/history/{session_id}")
 async def get_session_history(session_id: str):
     """Retrieves chat history for a session."""
     if not DB:
         raise HTTPException(status_code=503, detail="Database not connected")
-    # --- FIX APPLIED: Consistent threading model ---
+    
     history = await asyncio.to_thread(get_chat_history_safely, session_id)
     return {"session_id": session_id, "chat_history": history}
 
@@ -352,19 +382,8 @@ async def cleanup_session(session_id: str):
     cleanup_session_resources(session_id)
     return {"success": True, "message": f"Session {session_id} cleaned up."}
 
-async def save_chat_message_safely(session_id: str, role: str, message: str):
-    """Saves chat messages in a non-blocking way."""
-    if not DB: return
-    try:
-        await asyncio.to_thread(
-            DB.chats.insert_one,
-            {"session_id": session_id, "role": role, "message": message, "created_at": datetime.utcnow()}
-        )
-    except Exception as e:
-        logger.error(f"Failed to save chat message for session {session_id}: {e}")
-
 if __name__ == "__main__":
     import uvicorn
     port = int(os.getenv("PORT", 7861))
     logger.info(f"Starting server on http://0.0.0.0:{port}")
-    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=True)
+    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=True)