Update app.py

app.py

@@ -443,6 +443,251 @@ async def list_active_sessions():
             "sessions": sessions
         }
 
+# Add this enhanced logging to your chat endpoint in app.py
+
+@app.post("/chat/{session_id}", response_model=ChatResponse)
+async def chat_with_document(session_id: str, request: ChatRequest):
+    session_logger = create_session_logger(session_id)
+    start_time = time.time()
+    
+    try:
+        session_logger.info(f"=== CHAT REQUEST START ===")
+        session_logger.info(f"Session ID: {session_id}")
+        session_logger.info(f"Message length: {len(request.message)}")
+        session_logger.info(f"Message preview: {request.message[:100]}...")
+        session_logger.info(f"RAG models initialized: {RAG_MODELS_INITIALIZED}")
+        session_logger.info(f"MongoDB connected: {APP_STATE['mongodb_connected']}")
+        
+        # Check if session is already loaded in memory
+        with STORE_LOCK:
+            session_in_memory = session_id in SESSION_STORES
+            session_logger.info(f"Session in memory: {session_in_memory}")
+            
+            if not session_in_memory:
+                # Lazy load: Load session from MongoDB when first chat request comes
+                session_logger.info("=== LAZY LOADING FROM MONGODB ===")
+                session_logger.info("Loading session from MongoDB for first chat request")
+                
+                try:
+                    load_start = time.time()
+                    session_store = await asyncio.to_thread(load_session_from_mongodb, session_id)
+                    load_time = time.time() - load_start
+                    session_logger.info(f"Session loaded from MongoDB in {load_time:.2f}s")
+                    
+                    SESSION_STORES[session_id] = session_store
+                    session_logger.info(f"Session added to memory store")
+                    session_logger.info(f"Chunks loaded: {len(session_store.get('session_rag', {}).chunks_data) if hasattr(session_store.get('session_rag'), 'chunks_data') else 'unknown'}")
+                    
+                except Exception as load_error:
+                    session_logger.error(f"Failed to load session: {load_error}", exc_info=True)
+                    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"]
+            session_logger.info(f"Using session RAG instance: {type(session_rag)}")
+        
+        session_logger.info(f"=== PROCESSING QUERY ===")
+        
+        # Process the query using SessionRAG
+        query_start = time.time()
+        result = await asyncio.to_thread(session_rag.query_documents, request.message, top_k=5)
+        query_time = time.time() - query_start
+        
+        session_logger.info(f"Query processed in {query_time:.2f}s")
+        session_logger.info(f"Result keys: {list(result.keys()) if isinstance(result, dict) else 'not dict'}")
+        
+        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.')
+        sources = result.get('sources', [])
+        
+        session_logger.info(f"Generated answer length: {len(answer)}")
+        session_logger.info(f"Sources found: {len(sources)}")
+        
+        # Save chat messages asynchronously
+        asyncio.create_task(save_chat_message_safely(session_id, "user", request.message))
+        asyncio.create_task(save_chat_message_safely(session_id, "assistant", answer))
+        
+        processing_time = time.time() - start_time
+        session_logger.info(f"=== CHAT REQUEST COMPLETE ===")
+        session_logger.info(f"Total processing time: {processing_time:.2f}s")
+        
+        return ChatResponse(
+            success=True,
+            answer=answer,
+            sources=sources,
+            chat_history=[],
+            processing_time=processing_time,
+            session_id=session_id,
+            query_analysis=result.get('query_analysis'),
+            confidence=result.get('confidence')
+        )
+        
+    except HTTPException:
+        session_logger.error(f"HTTP Exception in chat processing")
+        raise
+    except Exception as e:
+        session_logger.error(f"Chat processing failed: {e}", exc_info=True)
+        APP_STATE["errors"].append(f"Chat error: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Chat processing error: {str(e)}")
+
+# Also add this endpoint for better debugging
+@app.get("/debug/{session_id}")
+async def debug_session(session_id: str):
+    """Debug endpoint to check session status"""
+    with STORE_LOCK:
+        session_in_memory = session_id in SESSION_STORES
+        session_info = {}
+        
+        if session_in_memory:
+            store = SESSION_STORES[session_id]
+            session_rag = store.get("session_rag")
+            session_info = {
+                "in_memory": True,
+                "indexed": store.get("indexed", False),
+                "metadata": store.get("metadata", {}),
+                "chunks_count": len(session_rag.chunks_data) if hasattr(session_rag, 'chunks_data') else 0,
+                "has_dense_index": hasattr(session_rag, 'dense_index') and session_rag.dense_index is not None,
+                "has_bm25_index": hasattr(session_rag, 'bm25_index') and session_rag.bm25_index is not None,
+            }
+        else:
+            session_info = {"in_memory": False}
+    
+    # Check MongoDB
+    mongodb_info = {"connected": False, "session_exists": False, "chunks_count": 0}
+    if DB:
+        mongodb_info["connected"] = True
+        session_doc = DB.sessions.find_one({"session_id": session_id})
+        if session_doc:
+            mongodb_info["session_exists"] = True
+            mongodb_info["session_status"] = session_doc.get("status")
+            mongodb_info["filename"] = session_doc.get("filename")
+            chunks_count = DB.chunks.count_documents({"session_id": session_id})
+            mongodb_info["chunks_count"] = chunks_count
+    
+    return {
+        "session_id": session_id,
+        "memory": session_info,
+        "mongodb": mongodb_info,
+        "app_state": {
+            "rag_models_ready": RAG_MODELS_INITIALIZED,
+            "mongodb_connected": APP_STATE["mongodb_connected"],
+            "active_sessions": len(SESSION_STORES)
+        }
+    }
+
+# Enhanced load_session_from_mongodb with better error handling
+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_logger.info(f"Session ID: {session_id}")
+    
+    if not DB:
+        session_logger.error("Database not connected")
+        raise ValueError("Database not connected")
+
+    # Load session metadata
+    session_logger.info("Loading session metadata...")
+    session_doc = DB.sessions.find_one({"session_id": session_id})
+    if not session_doc:
+        session_logger.error(f"Session {session_id} not found in database")
+        raise ValueError(f"Session {session_id} not found in database")
+    
+    session_logger.info(f"Found session: {session_doc.get('filename', 'unknown')} - Status: {session_doc.get('status')}")
+    
+    if session_doc.get("status") != "completed":
+        session_logger.error(f"Session not ready - status: {session_doc.get('status')}")
+        raise ValueError(f"Session not ready - status: {session_doc.get('status')}")
+
+    # Load chunks with embeddings from MongoDB
+    session_logger.info(f"Loading chunks with embeddings for: {session_doc.get('filename', 'unknown')}")
+    chunks_cursor = DB.chunks.find({"session_id": session_id}).sort("created_at", 1)
+    chunks_list = list(chunks_cursor)
+    
+    session_logger.info(f"Found {len(chunks_list)} chunks in database")
+    
+    if not chunks_list:
+        session_logger.error(f"No chunks found for session {session_id}")
+        raise ValueError(f"No chunks found for session {session_id}")
+
+    # Verify chunks have embeddings
+    chunks_with_embeddings = 0
+    for chunk in chunks_list:
+        if chunk.get('embedding') is not None:
+            chunks_with_embeddings += 1
+    
+    session_logger.info(f"Chunks with embeddings: {chunks_with_embeddings}/{len(chunks_list)}")
+    
+    if chunks_with_embeddings == 0:
+        session_logger.error("No chunks have embeddings - document may not be fully processed")
+        raise ValueError("No chunks have embeddings - document may not be fully processed")
+
+    # Create SessionRAG instance and load pre-existing data
+    try:
+        session_logger.info("Creating SessionRAG instance...")
+        groq_api_key = os.getenv("GROQ_API_KEY")
+        session_rag = SessionRAG(session_id, groq_api_key)
+        
+        session_logger.info(f"Loading existing session data with {len(chunks_list)} chunks")
+        load_start = time.time()
+        session_rag.load_existing_session_data(chunks_list)
+        load_duration = time.time() - load_start
+        session_logger.info(f"Session data loaded in {load_duration:.2f}s")
+
+        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(),
+                "load_time": load_duration,
+                "document_info": {"filename": session_doc.get("filename", "Unknown")}
+            }
+        }
+        session_logger.info("=== SESSION LOADING COMPLETE ===")
+        return session_store
+        
+    except Exception as rag_error:
+        session_logger.error(f"Failed to create/load SessionRAG: {rag_error}", exc_info=True)
+        raise ValueError(f"Failed to initialize RAG system: {str(rag_error)}")
+
+# Add a health check specifically for the chat endpoint
+@app.get("/health/chat")
+async def chat_health_check():
+    """Specific health check for chat functionality"""
+    issues = []
+    
+    if not RAG_MODELS_INITIALIZED:
+        issues.append("RAG models not initialized")
+    
+    if not APP_STATE["mongodb_connected"]:
+        issues.append("MongoDB not connected")
+    
+    if not FAISS_AVAILABLE:
+        issues.append("FAISS not available")
+    
+    groq_key = os.getenv("GROQ_API_KEY")
+    if not groq_key:
+        issues.append("GROQ API key not configured")
+    
+    return {
+        "status": "healthy" if not issues else "unhealthy",
+        "issues": issues,
+        "components": {
+            "rag_models": RAG_MODELS_INITIALIZED,
+            "mongodb": APP_STATE["mongodb_connected"], 
+            "faiss": FAISS_AVAILABLE,
+            "groq_key_configured": bool(groq_key)
+        }
+    }
+
 if __name__ == "__main__":
     import uvicorn
     port = int(os.getenv("PORT", 7861))