Update app.py

app.py

@@ -1,19 +1,13 @@
-# app_optimized.py - Performance-Optimized FastAPI App
 import asyncio
-import concurrent.futures
 import logging
 import os
 import sys
 import threading
 import time
-import traceback
-from collections import defaultdict
 from contextlib import asynccontextmanager
-from datetime import datetime, timedelta
-from functools import partial
+from datetime import datetime
 from typing import Any, Dict, List, Optional
 
-import numpy as np
 import pymongo
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
@@ -26,8 +20,7 @@ except ImportError:
     FAISS_AVAILABLE = False
 
 try:
-    # Import the optimized SessionRAG class
-    from rag import OptimizedSessionRAG, initialize_models
+    from rag import SessionRAG, initialize_models
     RAG_AVAILABLE = True
 except ImportError:
     RAG_AVAILABLE = False
@@ -47,25 +40,17 @@ logger = logging.getLogger(__name__)
 MONGO_CLIENT = None
 DB = None
 RAG_MODELS_INITIALIZED = False
-SESSION_LAST_ACCESS = {}
+SESSION_STORES = {}  # In-memory cache: {session_id: {session_rag, metadata, indexed}}
+STORE_LOCK = threading.RLock()
+
 APP_STATE = {
     "startup_time": None,
     "mongodb_connected": False,
     "rag_models_ready": False,
-    "active_sessions": 0,
     "total_queries": 0,
     "errors": []
 }
 
-# --- 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)
@@ -90,25 +75,25 @@ 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]})
 
 def connect_mongodb():
+    """Connect to MongoDB Atlas"""
     global MONGO_CLIENT, DB
     try:
         mongodb_url = os.getenv("MONGODB_URL", "mongodb://localhost:27017/")
         logger.info(f"Connecting to MongoDB...")
         MONGO_CLIENT = pymongo.MongoClient(
             mongodb_url, 
-            serverSelectionTimeoutMS=5000,
-            maxPoolSize=50,
-            waitQueueTimeoutMS=2500
+            serverSelectionTimeoutMS=5000
         )
         MONGO_CLIENT.admin.command('ping')
         DB = MONGO_CLIENT["legal_rag_system"]
         
-        # Ensure indices exist
+        # Create indexes
         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)
@@ -119,13 +104,14 @@ def connect_mongodb():
         return True
     except Exception as e:
         logger.error(f"MongoDB connection failed: {e}")
+        APP_STATE["errors"].append(f"MongoDB error: {str(e)}")
         return False
 
 def init_rag_models():
-    """Initialize shared RAG models once at startup"""
+    """Initialize shared RAG models (embedding model, NLP model, etc.)"""
     global RAG_MODELS_INITIALIZED
     if not RAG_AVAILABLE or not FAISS_AVAILABLE:
-        logger.error("RAG module or FAISS not available - cannot initialize models.")
+        logger.error("RAG module or FAISS not available")
         return False
     try:
         model_id = os.getenv("EMBEDDING_MODEL_ID", "sentence-transformers/all-MiniLM-L6-v2")
@@ -141,162 +127,163 @@ def init_rag_models():
         APP_STATE["errors"].append(f"RAG init failed: {str(e)}")
         return False
 
-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}")
+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
+    """
+    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")
 
-        # 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)
+    # 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')}")
 
-        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")}
+    # 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")
+
+    # 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")
             }
         }
-        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)
-
-def cleanup_session_resources(session_id: str):
-    """Safely clean up session resources"""
+    session_logger.info("✓ Session loaded successfully with existing embeddings")
+    return session_store
+
+def get_or_load_session(session_id: str) -> Dict[str, Any]:
+    """
+    Get session from memory cache, or load from MongoDB if not in memory.
+    Thread-safe with locking.
+    """
     with STORE_LOCK:
+        # Check if already loaded in memory
         if session_id in SESSION_STORES:
-            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 recently"""
-    now = datetime.utcnow()
-    expired_ids = [
-        sid for sid, last_access in list(SESSION_LAST_ACCESS.items())
-        if (now - last_access).total_seconds() > STORE_TTL
-    ]
-    if expired_ids:
-        logger.info(f"Cleaning up {len(expired_ids)} expired sessions")
-        for sid in expired_ids:
-            cleanup_session_resources(sid)
-            SESSION_LAST_ACCESS.pop(sid, None)
-            
-async def periodic_cleanup():
-    while True:
-        await asyncio.sleep(CLEANUP_INTERVAL)
-        logger.info("Running periodic cleanup of expired sessions...")
-        cleanup_expired_sessions()
+            logger.info(f"Session {session_id[:8]} already in memory")
+            return SESSION_STORES[session_id]
+        
+        # Not in memory - load from MongoDB
+        logger.info(f"Session {session_id[:8]} not in memory, loading from MongoDB...")
+        session_store = load_session_from_mongodb(session_id)
+        SESSION_STORES[session_id] = session_store
+        logger.info(f"Session {session_id[:8]} loaded and cached in memory")
+        return session_store
 
 async def save_chat_message_safely(session_id: str, role: str, message: str):
-    """Save chat messages asynchronously"""
-    if not DB: 
+    """Save chat messages to MongoDB asynchronously"""
+    if not DB:
         return
     try:
-        def _save_message():
-            DB.chats.insert_one({
-                "session_id": session_id, 
-                "role": role, 
-                "message": message, 
+        await asyncio.to_thread(
+            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}")
 
+def get_chat_history_safely(session_id: str, limit: int = 50) -> List[Dict[str, Any]]:
+    """Get chat history from MongoDB with error handling"""
+    if not DB:
+        return []
+    try:
+        chats_cursor = DB.chats.find({"session_id": session_id}).sort("created_at", -1).limit(limit)
+        return list(chats_cursor)[::-1]  # Reverse for chronological order
+    except Exception as e:
+        logger.error(f"Failed to get chat history for session {session_id}: {e}")
+        return []
+
+
 # --- Application Lifespan ---
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    cleanup_task = None
+    """Application startup and shutdown"""
     APP_STATE["startup_time"] = datetime.utcnow()
-    logger.info("Starting Optimized RAG Chat Service...")
+    logger.info("Starting RAG Chat Service...")
 
+    # Initialize MongoDB and models
     connect_mongodb()
     init_rag_models()
 
-    cleanup_task = asyncio.create_task(periodic_cleanup())
-    logger.info("Background session cleanup task started.")
+    logger.info("✓ Service ready")
     
     yield
     
+    # Cleanup on shutdown
     logger.info("Shutting down...")
-    if cleanup_task:
-        cleanup_task.cancel()
     if MONGO_CLIENT:
         MONGO_CLIENT.close()
-    EXECUTOR.shutdown(wait=True)
-    logger.info("Shutdown complete.")
+    logger.info("✓ Shutdown complete")
+
 
 # --- FastAPI App ---
 app = FastAPI(
-    title="Optimized RAG Chat Service",
-    description="High-performance, session-isolated RAG chat service with pre-computed embeddings.",
+    title="Session-based RAG Chat Service",
+    description="RAG system with MongoDB session persistence",
     version="4.0.0",
     lifespan=lifespan
 )
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"], 
-    allow_credentials=True, 
-    allow_methods=["*"], 
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
     allow_headers=["*"],
 )
 
 @app.get("/")
 async def root():
-    return {"service": "Optimized RAG Chat Service", "version": "4.0.0"}
+    return {
+        "service": "Session-based RAG Chat Service",
+        "version": "4.0.0",
+        "description": "Embeddings stored in MongoDB, lazy-loaded on demand"
+    }
 
 @app.get("/health", response_model=HealthResponse)
 async def health_check():
+    """Health check endpoint"""
     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.get("indexed", False))
@@ -312,7 +299,7 @@ async def health_check():
         faiss_available=FAISS_AVAILABLE,
         active_sessions=active_sessions,
         memory_usage={
-            "loaded_sessions": active_sessions, 
+            "loaded_sessions": active_sessions,
             "indexed_sessions": indexed_sessions
         },
         uptime_seconds=uptime,
@@ -321,50 +308,51 @@ async def health_check():
 
 @app.post("/chat/{session_id}", response_model=ChatResponse)
 async def chat_with_document(session_id: str, request: ChatRequest):
+    """
+    Main chat endpoint:
+    1. Load session from MongoDB if not in memory (lazy loading)
+    2. Process query using RAG pipeline
+    3. Save chat messages to MongoDB
+    4. Return answer with sources
+    """
     session_logger = create_session_logger(session_id)
     start_time = time.time()
     
     try:
-        session_logger.info(f"Received chat request: {request.message[:100]}...")
-        
-        # Check if session is loaded in memory
-        with STORE_LOCK:
-            if session_id not in SESSION_STORES:
-                # Lazy load: Load session from MongoDB asynchronously
-                session_logger.info("Loading session from MongoDB for first chat request")
-                try:
-                    session_store = await load_session_from_mongodb_async(session_id)
-                    SESSION_STORES[session_id] = session_store
-                    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)}")
-            
-            # Update last access time
-            SESSION_LAST_ACCESS[session_id] = datetime.utcnow()
-            session_rag = SESSION_STORES[session_id]["session_rag"]
+        session_logger.info(f"Chat request: {request.message[:100]}...")
         
-        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)
-        
-        result = await asyncio.get_event_loop().run_in_executor(EXECUTOR, _process_query)
+        # Get or load session (lazy loading from MongoDB)
+        try:
+            session_store = await asyncio.to_thread(get_or_load_session, session_id)
+            session_rag = session_store["session_rag"]
+        except Exception as load_error:
+            session_logger.error(f"Failed to load session: {load_error}")
+            raise HTTPException(
+                status_code=404,
+                detail=f"Session not found or failed to load: {str(load_error)}"
+            )
+        
+        # Process query using RAG pipeline
+        session_logger.info(f"Processing query with RAG...")
+        result = await asyncio.to_thread(
+            session_rag.query_documents,
+            request.message,
+            top_k=5
+        )
         
         if 'error' in result:
-            session_logger.error(f"Query processing error: {result['error']}")
+            session_logger.error(f"Query 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 (non-blocking)
+        # Save chat messages asynchronously to MongoDB
         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"Query processed successfully in {processing_time:.2f}s")
+        session_logger.info(f"✓ Query processed in {processing_time:.2f}s")
         
         return ChatResponse(
             success=True,
@@ -381,315 +369,71 @@ async def chat_with_document(session_id: str, request: ChatRequest):
     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)}")
+        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):
-    """Retrieve chat history for a session"""
+    """Get chat history for a session from MongoDB"""
     if not DB:
         raise HTTPException(status_code=503, detail="Database not connected")
     
-    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 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
-        }
-
-# 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
-    
+    history = await asyncio.to_thread(get_chat_history_safely, session_id)
     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)
-        }
+        "chat_history": history,
+        "count": len(history)
     }
 
-# 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}")
-    
+@app.get("/session/{session_id}/info")
+async def get_session_info(session_id: str):
+    """Get session metadata from MongoDB"""
     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)
+        raise HTTPException(status_code=503, detail="Database not connected")
     
-    session_logger.info(f"Found {len(chunks_list)} chunks in database")
+    session_doc = await asyncio.to_thread(DB.sessions.find_one, {"session_id": session_id})
+    if not session_doc:
+        raise HTTPException(status_code=404, detail="Session not found")
     
-    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
+    # Convert ObjectId to string for JSON serialization
+    session_doc['_id'] = str(session_doc['_id'])
     
-    session_logger.info(f"Chunks with embeddings: {chunks_with_embeddings}/{len(chunks_list)}")
+    # Check if loaded in memory
+    with STORE_LOCK:
+        in_memory = session_id in SESSION_STORES
     
-    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")
+    return {
+        "session_id": session_id,
+        "metadata": session_doc,
+        "in_memory": in_memory
+    }
 
-    # 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")}
+@app.delete("/session/{session_id}/cache")
+async def clear_session_cache(session_id: str):
+    """Remove session from memory cache (data remains in MongoDB)"""
+    with STORE_LOCK:
+        if session_id in SESSION_STORES:
+            store = SESSION_STORES.pop(session_id)
+            session_rag = store.get("session_rag")
+            if hasattr(session_rag, 'cleanup'):
+                session_rag.cleanup()
+            logger.info(f"Session {session_id[:8]} removed from memory cache")
+            return {
+                "success": True,
+                "message": f"Session removed from memory cache",
+                "note": "Data remains in MongoDB"
             }
-        }
-        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)
-        }
+        "success": False,
+        "message": "Session not found in memory cache"
     }
 
+
 if __name__ == "__main__":
     import uvicorn
     port = int(os.getenv("PORT", 7861))
-    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)
+    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)