Update app.py

app.py

@@ -1,3 +1,4 @@
+# app_optimized.py - Performance-Optimized FastAPI App
 import asyncio
 import concurrent.futures
 import logging
@@ -25,13 +26,13 @@ except ImportError:
     FAISS_AVAILABLE = False
 
 try:
-    # We import the SessionRAG class and the model initializer.
-    from rag import SessionRAG, initialize_models
+    # Import the optimized SessionRAG class
+    from rag_optimized import OptimizedSessionRAG, initialize_models
     RAG_AVAILABLE = True
 except ImportError:
     RAG_AVAILABLE = False
 
-# Configure comprehensive logging
+# Configure logging
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - [%(funcName)s:%(lineno)d] - %(message)s',
@@ -46,7 +47,7 @@ logger = logging.getLogger(__name__)
 MONGO_CLIENT = None
 DB = None
 RAG_MODELS_INITIALIZED = False
-SESSION_LAST_ACCESS = {}  # Track last access time for each session
+SESSION_LAST_ACCESS = {}
 APP_STATE = {
     "startup_time": None,
     "mongodb_connected": False,
@@ -56,12 +57,15 @@ APP_STATE = {
     "errors": []
 }
 
-# --- In-memory Session Management & Threading ---
+# --- Optimized Session Management ---
 SESSION_STORES = {}
 STORE_LOCK = threading.RLock()
 CLEANUP_INTERVAL = 1800  # 30 minutes
 STORE_TTL = 1800  # 30 minutes
 
+# Thread pool for async operations
+EXECUTOR = concurrent.futures.ThreadPoolExecutor(max_workers=4)
+
 # --- Pydantic Models ---
 class ChatRequest(BaseModel):
     message: str = Field(..., min_length=1, max_length=5000)
@@ -70,25 +74,12 @@ class ChatResponse(BaseModel):
     success: bool
     answer: str
     sources: List[Dict[str, Any]] = Field(default_factory=list)
-    chat_history: List[Dict[str, Any]] = Field(default_factory=list)
     processing_time: float
     session_id: str
     query_analysis: Optional[Dict[str, Any]] = None
     confidence: Optional[float] = None
     error_details: Optional[str] = None
 
-class InitRequest(BaseModel):
-    force_reload: bool = Field(default=False)
-
-class InitResponse(BaseModel):
-    success: bool
-    session_id: str
-    message: str
-    chunk_count: int = Field(default=0)
-    title: str = Field(default="Unknown Document")
-    document_info: Optional[Dict[str, Any]] = None
-    error_details: Optional[str] = None
-
 class HealthResponse(BaseModel):
     status: str
     mongodb_connected: bool
@@ -99,7 +90,6 @@ class HealthResponse(BaseModel):
     uptime_seconds: float
     last_error: Optional[str] = None
 
-
 # --- Helper Functions ---
 def create_session_logger(session_id: str):
     return logging.LoggerAdapter(logger, {'session_id': session_id[:8]})
@@ -110,12 +100,20 @@ def connect_mongodb():
         mongodb_url = os.getenv("MONGODB_URL", "mongodb://localhost:27017/")
         logger.info(f"Connecting to MongoDB...")
         MONGO_CLIENT = pymongo.MongoClient(
-            mongodb_url, serverSelectionTimeoutMS=5000
+            mongodb_url, 
+            serverSelectionTimeoutMS=5000,
+            maxPoolSize=50,
+            waitQueueTimeoutMS=2500
         )
         MONGO_CLIENT.admin.command('ping')
         DB = MONGO_CLIENT["legal_rag_system"]
+        
+        # Ensure indices exist
         DB.chats.create_index("session_id", background=True)
         DB.chats.create_index("created_at", expireAfterSeconds=24 * 60 * 60, background=True)
+        DB.sessions.create_index("session_id", unique=True, background=True)
+        DB.chunks.create_index("session_id", background=True)
+        
         APP_STATE["mongodb_connected"] = True
         logger.info("MongoDB connected successfully")
         return True
@@ -124,7 +122,7 @@ def connect_mongodb():
         return False
 
 def init_rag_models():
-    """Initializes the shared, stateless RAG models once at startup."""
+    """Initialize shared RAG models once at startup"""
     global RAG_MODELS_INITIALIZED
     if not RAG_AVAILABLE or not FAISS_AVAILABLE:
         logger.error("RAG module or FAISS not available - cannot initialize models.")
@@ -143,75 +141,82 @@ def init_rag_models():
         APP_STATE["errors"].append(f"RAG init failed: {str(e)}")
         return False
 
-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")
-
-    # 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")
-    if session_doc.get("status") != "completed":
-        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)
-    
-    if not chunks_list:
-        raise ValueError(f"No chunks found for session {session_id}")
-
-    # Create SessionRAG instance and load pre-existing data
-    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")
-    session_rag.load_existing_session_data(chunks_list)
+async def load_session_from_mongodb_async(session_id: str) -> Dict[str, Any]:
+    """OPTIMIZED: Load session from MongoDB in async way"""
+    def _load_session_sync():
+        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")
+
+        # 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")
+        if session_doc.get("status") != "completed":
+            raise ValueError(f"Session not ready - status: {session_doc.get('status')}")
+
+        # Load chunks with embeddings efficiently
+        session_logger.info(f"Loading chunks with embeddings for: {session_doc.get('filename', 'unknown')}")
+        
+        # Use projection to only load needed fields
+        chunks_cursor = DB.chunks.find(
+            {"session_id": session_id},
+            {
+                "chunk_id": 1, 
+                "content": 1, 
+                "title": 1, 
+                "section_type": 1,
+                "importance_score": 1, 
+                "entities": 1, 
+                "embedding": 1
+            }
+        ).sort("created_at", 1)
+        
+        chunks_list = list(chunks_cursor)
+        
+        if not chunks_list:
+            raise ValueError(f"No chunks found for session {session_id}")
 
-    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"),
-            "chunk_count": len(chunks_list),
-            "loaded_at": datetime.utcnow(),
-            "document_info": {"filename": session_doc.get("filename", "Unknown")}
+        # Create OptimizedSessionRAG instance
+        groq_api_key = os.getenv("GROQ_API_KEY")
+        session_rag = OptimizedSessionRAG(session_id, groq_api_key)
+        
+        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,
+            "metadata": {
+                "session_id": session_id,
+                "title": session_doc.get("filename", "Document"),
+                "chunk_count": len(chunks_list),
+                "loaded_at": datetime.utcnow(),
+                "load_time": getattr(session_rag, 'load_time', 0),
+                "index_build_time": getattr(session_rag, 'index_build_time', 0),
+                "document_info": {"filename": session_doc.get("filename", "Unknown")}
+            }
         }
-    }
-    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 []
-    try:
-        chats_cursor = DB.chats.find({"session_id": session_id}).sort("created_at", -1).limit(limit)
-        # Reverse to get chronological order [oldest -> newest]
-        return list(chats_cursor)[::-1]
-    except Exception as e:
-        logger.error(f"Failed to get chat history for session {session_id}: {e}")
-        return []
+        session_logger.info(f"Session loaded from MongoDB in {session_rag.load_time:.2f}s")
+        return session_store
+    
+    # Run in thread pool to avoid blocking
+    return await asyncio.get_event_loop().run_in_executor(EXECUTOR, _load_session_sync)
 
-# --- Session Cleanup ---
 def cleanup_session_resources(session_id: str):
-    """Safely cleans up a session's resources."""
+    """Safely clean up session resources"""
     with STORE_LOCK:
         if session_id in SESSION_STORES:
-            store = SESSION_STORES.pop(session_id) # Atomically get and remove
+            store = SESSION_STORES.pop(session_id)
             session_rag = store.get("session_rag")
             if hasattr(session_rag, 'cleanup'):
                 session_rag.cleanup()
             logger.info(f"Cleaned up session from memory: {session_id[:8]}")
 
 def cleanup_expired_sessions():
-    """Clean up sessions that haven't been accessed in 30 minutes"""
+    """Clean up sessions that haven't been accessed recently"""
     now = datetime.utcnow()
     expired_ids = [
         sid for sid, last_access in list(SESSION_LAST_ACCESS.items())
@@ -230,14 +235,19 @@ async def periodic_cleanup():
         cleanup_expired_sessions()
 
 async def save_chat_message_safely(session_id: str, role: str, message: str):
-    """Saves chat messages in a non-blocking way."""
+    """Save chat messages asynchronously"""
     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()}
-        )
+        def _save_message():
+            DB.chats.insert_one({
+                "session_id": session_id, 
+                "role": role, 
+                "message": message, 
+                "created_at": datetime.utcnow()
+            })
+        
+        await asyncio.get_event_loop().run_in_executor(EXECUTOR, _save_message)
     except Exception as e:
         logger.error(f"Failed to save chat message for session {session_id}: {e}")
 
@@ -246,7 +256,7 @@ async def save_chat_message_safely(session_id: str, role: str, message: str):
 async def lifespan(app: FastAPI):
     cleanup_task = None
     APP_STATE["startup_time"] = datetime.utcnow()
-    logger.info("Starting Advanced RAG Chat Service...")
+    logger.info("Starting Optimized RAG Chat Service...")
 
     connect_mongodb()
     init_rag_models()
@@ -261,13 +271,14 @@ async def lifespan(app: FastAPI):
         cleanup_task.cancel()
     if MONGO_CLIENT:
         MONGO_CLIENT.close()
+    EXECUTOR.shutdown(wait=True)
     logger.info("Shutdown complete.")
 
-# --- FastAPI App and Endpoints ---
+# --- FastAPI App ---
 app = FastAPI(
-    title="Advanced RAG Chat Service",
-    description="A robust, session-isolated RAG chat service.",
-    version="3.1.0",
+    title="Optimized RAG Chat Service",
+    description="High-performance, session-isolated RAG chat service with pre-computed embeddings.",
+    version="4.0.0",
     lifespan=lifespan
 )
 
@@ -281,7 +292,7 @@ app.add_middleware(
 
 @app.get("/")
 async def root():
-    return {"service": "Advanced RAG Chat Service", "version": "3.1.0"}
+    return {"service": "Optimized RAG Chat Service", "version": "4.0.0"}
 
 @app.get("/health", response_model=HealthResponse)
 async def health_check():
@@ -300,7 +311,10 @@ async def health_check():
         rag_models_initialized=RAG_MODELS_INITIALIZED,
         faiss_available=FAISS_AVAILABLE,
         active_sessions=active_sessions,
-        memory_usage={"loaded_sessions": active_sessions, "indexed_sessions": indexed_sessions},
+        memory_usage={
+            "loaded_sessions": active_sessions, 
+            "indexed_sessions": indexed_sessions
+        },
         uptime_seconds=uptime,
         last_error=APP_STATE["errors"][-1] if APP_STATE["errors"] else None
     )
@@ -313,15 +327,15 @@ async def chat_with_document(session_id: str, request: ChatRequest):
     try:
         session_logger.info(f"Received chat request: {request.message[:100]}...")
         
-        # Check if session is already loaded in memory
+        # Check if session is loaded in memory
         with STORE_LOCK:
             if session_id not in SESSION_STORES:
-                # Lazy load: Load session from MongoDB when first chat request comes
+                # Lazy load: Load session from MongoDB asynchronously
                 session_logger.info("Loading session from MongoDB for first chat request")
                 try:
-                    session_store = await asyncio.to_thread(load_session_from_mongodb, session_id)
+                    session_store = await load_session_from_mongodb_async(session_id)
                     SESSION_STORES[session_id] = session_store
-                    session_logger.info("Session loaded successfully from MongoDB")
+                    session_logger.info(f"Session loaded successfully from MongoDB in {session_store['metadata'].get('load_time', 0):.2f}s")
                 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)}")
@@ -330,10 +344,13 @@ async def chat_with_document(session_id: str, request: ChatRequest):
             SESSION_LAST_ACCESS[session_id] = datetime.utcnow()
             session_rag = SESSION_STORES[session_id]["session_rag"]
         
-        session_logger.info(f"Processing query with SessionRAG...")
+        session_logger.info(f"Processing query with OptimizedSessionRAG...")
+        
+        # Process the query using OptimizedSessionRAG - this is FAST now
+        def _process_query():
+            return session_rag.query_documents(request.message, top_k=5)
         
-        # Process the query using SessionRAG
-        result = await asyncio.to_thread(session_rag.query_documents, request.message, top_k=5)
+        result = await asyncio.get_event_loop().run_in_executor(EXECUTOR, _process_query)
         
         if 'error' in result:
             session_logger.error(f"Query processing error: {result['error']}")
@@ -342,7 +359,7 @@ async def chat_with_document(session_id: str, request: ChatRequest):
         APP_STATE["total_queries"] += 1
         answer = result.get('answer', 'Unable to generate an answer.')
         
-        # Save chat messages asynchronously
+        # Save chat messages asynchronously (non-blocking)
         asyncio.create_task(save_chat_message_safely(session_id, "user", request.message))
         asyncio.create_task(save_chat_message_safely(session_id, "assistant", answer))
         
@@ -353,7 +370,6 @@ async def chat_with_document(session_id: str, request: ChatRequest):
             success=True,
             answer=answer,
             sources=result.get('sources', []),
-            chat_history=[],
             processing_time=processing_time,
             session_id=session_id,
             query_analysis=result.get('query_analysis'),
@@ -369,21 +385,66 @@ async def chat_with_document(session_id: str, request: ChatRequest):
 
 @app.get("/history/{session_id}")
 async def get_session_history(session_id: str):
-    """Retrieves chat history for a session."""
+    """Retrieve chat history for a session"""
     if not DB:
         raise HTTPException(status_code=503, detail="Database not connected")
     
-    history = await asyncio.to_thread(get_chat_history_safely, session_id)
+    def _get_history():
+        try:
+            chats_cursor = DB.chats.find({"session_id": session_id}).sort("created_at", -1).limit(50)
+            return list(chats_cursor)[::-1]  # Reverse to get chronological order
+        except Exception as e:
+            logger.error(f"Failed to get chat history for session {session_id}: {e}")
+            return []
+    
+    history = await asyncio.get_event_loop().run_in_executor(EXECUTOR, _get_history)
     return {"session_id": session_id, "chat_history": history}
 
 @app.delete("/session/{session_id}")
 async def cleanup_session(session_id: str):
-    """Manually cleans up a specific session from memory."""
+    """Manually clean up a specific session from memory"""
     cleanup_session_resources(session_id)
+    SESSION_LAST_ACCESS.pop(session_id, None)
     return {"success": True, "message": f"Session {session_id} cleaned up."}
 
+@app.get("/session/{session_id}/info")
+async def get_session_info(session_id: str):
+    """Get information about a loaded session"""
+    with STORE_LOCK:
+        if session_id not in SESSION_STORES:
+            raise HTTPException(status_code=404, detail="Session not loaded in memory")
+        
+        metadata = SESSION_STORES[session_id]["metadata"]
+        return {
+            "session_id": session_id,
+            "loaded": True,
+            "metadata": metadata,
+            "last_access": SESSION_LAST_ACCESS.get(session_id)
+        }
+
+@app.get("/sessions")
+async def list_active_sessions():
+    """List all currently active sessions in memory"""
+    with STORE_LOCK:
+        sessions = []
+        for session_id, store in SESSION_STORES.items():
+            metadata = store.get("metadata", {})
+            sessions.append({
+                "session_id": session_id,
+                "title": metadata.get("title", "Unknown"),
+                "chunk_count": metadata.get("chunk_count", 0),
+                "loaded_at": metadata.get("loaded_at"),
+                "load_time": metadata.get("load_time", 0),
+                "last_access": SESSION_LAST_ACCESS.get(session_id)
+            })
+        
+        return {
+            "active_sessions": len(sessions),
+            "sessions": sessions
+        }
+
 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)
+    logger.info(f"Starting optimized server on http://0.0.0.0:{port}")
+    uvicorn.run("app_optimized:app", host="0.0.0.0", port=port, reload=False, workers=1)