Update app.py

app.py

@@ -128,63 +128,66 @@ def init_rag_models():
         return False
 
 def load_session_from_mongodb(session_id: str) -> Dict[str, Any]:
-    """
-    Load session data from MongoDB:
-    1. Check if session exists in DB
-    2. Load all chunks with pre-computed embeddings
-    3. Create SessionRAG instance and rebuild indices from existing embeddings
-    """
+    """Load session data from MongoDB with better error handling"""
     session_logger = create_session_logger(session_id)
     session_logger.info(f"Loading session from MongoDB: {session_id}")
     
-    if not DB:
+    if DB is None:  # ✅ Fixed
         raise ValueError("Database not connected")
 
-    # 1. Load session metadata
-    session_doc = DB.sessions.find_one({"session_id": session_id})
-    if not session_doc:
-        raise ValueError(f"Session {session_id} not found in database")
-    
-    # Check session status
-    if session_doc.get("status") != "completed":
-        raise ValueError(f"Session not ready - status: {session_doc.get('status')}")
-
-    # 2. Load chunks with embeddings from MongoDB
-    session_logger.info(f"Loading chunks for: {session_doc.get('filename', 'unknown')}")
-    chunks_cursor = DB.chunks.find({"session_id": session_id}).sort("created_at", 1)
-    chunks_list = list(chunks_cursor)
-    
-    if not chunks_list:
-        raise ValueError(f"No chunks found for session {session_id}")
+    try:
+        # 1. Load session metadata
+        session_doc = DB.sessions.find_one({"session_id": session_id})
+        if not session_doc:
+            raise ValueError(f"Session {session_id} not found in database")
+        
+        # Check session status
+        if session_doc.get("status") != "completed":
+            raise ValueError(f"Session not ready - status: {session_doc.get('status')}")
+
+        # 2. Load chunks with embeddings from MongoDB
+        session_logger.info(f"Loading chunks for: {session_doc.get('filename', 'unknown')}")
+        chunks_cursor = DB.chunks.find({"session_id": session_id}).sort("created_at", 1)
+        chunks_list = list(chunks_cursor)
+        
+        if not chunks_list:
+            raise ValueError(f"No chunks found for session {session_id}")
 
-    session_logger.info(f"Found {len(chunks_list)} chunks with pre-computed embeddings")
+        session_logger.info(f"Found {len(chunks_list)} chunks with pre-computed embeddings")
 
-    # 3. Create SessionRAG instance
-    groq_api_key = os.getenv("GROQ_API_KEY")
-    session_rag = SessionRAG(session_id, groq_api_key)
-    
-    # 4. Load existing session data (rebuilds indices from stored embeddings)
-    session_logger.info(f"Rebuilding search indices from existing embeddings...")
-    session_rag.load_existing_session_data(chunks_list)
-
-    # 5. Create session store object
-    session_store = {
-        "session_rag": session_rag,
-        "indexed": True,
-        "metadata": {
-            "session_id": session_id,
-            "title": session_doc.get("filename", "Document"),
-            "chunk_count": len(chunks_list),
-            "loaded_at": datetime.utcnow(),
-            "document_info": {
-                "filename": session_doc.get("filename", "Unknown"),
-                "upload_date": session_doc.get("created_at")
+        # 3. Create SessionRAG instance
+        groq_api_key = os.getenv("GROQ_API_KEY")
+        
+        # Make sure to import from the correct module
+        from rag_optimized import OptimizedSessionRAG  # Or whatever your actual import is
+        session_rag = OptimizedSessionRAG(session_id, groq_api_key)
+        
+        # 4. Load existing session data (rebuilds indices from stored embeddings)
+        session_logger.info(f"Rebuilding search indices from existing embeddings...")
+        session_rag.load_existing_session_data(chunks_list)
+
+        # 5. Create session store object
+        session_store = {
+            "session_rag": session_rag,
+            "indexed": True,
+            "metadata": {
+                "session_id": session_id,
+                "title": session_doc.get("filename", "Document"),
+                "chunk_count": len(chunks_list),
+                "loaded_at": datetime.utcnow(),
+                "document_info": {
+                    "filename": session_doc.get("filename", "Unknown"),
+                    "upload_date": session_doc.get("created_at")
+                }
             }
         }
-    }
-    
-    session_logger.info("✓ Session loaded successfully with existing embeddings")
-    return session_store
+        
+        session_logger.info("✓ Session loaded successfully with existing embeddings")
+        return session_store
+        
+    except Exception as e:
+        session_logger.error(f"Failed to load session from MongoDB: {e}", exc_info=True)
+        raise ValueError(f"Failed to load session {session_id}: {str(e)}")
 
 def get_or_load_session(session_id: str) -> Dict[str, Any]:
     """
@@ -377,7 +380,7 @@ async def chat_with_document(session_id: str, request: ChatRequest):
 @app.get("/history/{session_id}")
 async def get_session_history(session_id: str):
     """Get chat history for a session from MongoDB"""
-    if not DB:
+    if DB is None:  # ✅ Correct way to check
         raise HTTPException(status_code=503, detail="Database not connected")
     
     history = await asyncio.to_thread(get_chat_history_safely, session_id)