Update app.py

app.py

@@ -46,7 +46,7 @@ RAG_INITIALIZED = False
 SESSION_STORES = {}
 STORE_LOCK = threading.RLock()
 CLEANUP_INTERVAL = 3600  # 1 hour cleanup interval
-STORE_TTL = 24 * 3600  # 24 hours TTL for in-memory stores
+STORE_TTL = 30 * 60  # 24 hours TTL for in-memory stores
 
 # Request/Response models
 class ChatRequest(BaseModel):
@@ -322,37 +322,56 @@ def get_chat_history(session_id: str, limit: int = 50) -> List[Dict[str, Any]]:
         logger.error(f"Failed to get chat history for session {session_id}: {e}")
         return []
 
+import asyncio
+from contextlib import asynccontextmanager
+
+# Global cleanup task
+cleanup_task = None
+
 def cleanup_old_stores():
-    """Background cleanup of old in-memory stores"""
-    while True:
-        try:
-            current_time = datetime.utcnow()
-            expired_sessions = []
-            
-            with STORE_LOCK:
-                for session_id, store in SESSION_STORES.items():
-                    loaded_at = store["metadata"]["loaded_at"]
-                    if (current_time - loaded_at).total_seconds() > STORE_TTL:
-                        expired_sessions.append(session_id)
-                
-                for session_id in expired_sessions:
-                    # Clean up FAISS index and other resources
-                    if SESSION_STORES[session_id].get("faiss_index"):
-                        del SESSION_STORES[session_id]["faiss_index"]
-                    del SESSION_STORES[session_id]
-                    logger.info(f"Cleaned up expired store for session: {session_id}")
-            
-            if expired_sessions:
-                logger.info(f"Cleaned up {len(expired_sessions)} expired session stores")
+    """Background cleanup of old in-memory stores - single run"""
+    try:
+        current_time = datetime.utcnow()
+        expired_sessions = []
+        
+        with STORE_LOCK:
+            for session_id, store in SESSION_STORES.items():
+                loaded_at = store["metadata"]["loaded_at"]
+                if (current_time - loaded_at).total_seconds() > STORE_TTL:
+                    expired_sessions.append(session_id)
             
-        except Exception as e:
-            logger.error(f"Cleanup error: {e}")
+            for session_id in expired_sessions:
+                # Clean up FAISS index and other resources
+                if SESSION_STORES[session_id].get("faiss_index"):
+                    del SESSION_STORES[session_id]["faiss_index"]
+                del SESSION_STORES[session_id]
+                logger.info(f"Cleaned up expired store for session: {session_id}")
         
-        time.sleep(CLEANUP_INTERVAL)
+        if expired_sessions:
+            logger.info(f"Cleaned up {len(expired_sessions)} expired session stores")
+        
+    except Exception as e:
+        logger.error(f"Cleanup error: {e}")
 
-@app.on_event("startup")
-async def startup_event():
-    """Initialize connections on startup"""
+async def periodic_cleanup():
+    """Async periodic cleanup task"""
+    global cleanup_task
+    try:
+        while True:
+            cleanup_old_stores()
+            await asyncio.sleep(CLEANUP_INTERVAL)
+    except asyncio.CancelledError:
+        logger.info("Cleanup task cancelled")
+        raise
+    except Exception as e:
+        logger.error(f"Periodic cleanup error: {e}")
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Application lifespan manager"""
+    global cleanup_task
+    
+    # Startup
     logger.info("Starting up Advanced RAG Chat Service...")
     
     # Connect to MongoDB
@@ -365,12 +384,37 @@ async def startup_event():
         logger.error("Failed to initialize RAG system")
         raise Exception("RAG initialization failed")
     
-    # Start background cleanup thread
-    cleanup_thread = threading.Thread(target=cleanup_old_stores, daemon=True)
-    cleanup_thread.start()
-    logger.info("Background cleanup thread started")
+    # Start background cleanup task
+    cleanup_task = asyncio.create_task(periodic_cleanup())
+    logger.info("Background cleanup task started")
     
     logger.info("Startup completed successfully")
+    
+    yield
+    
+    # Shutdown
+    logger.info("Shutting down Advanced RAG Chat Service...")
+    
+    if cleanup_task:
+        cleanup_task.cancel()
+        try:
+            await cleanup_task
+        except asyncio.CancelledError:
+            pass
+    
+    if MONGO_CLIENT:
+        MONGO_CLIENT.close()
+    
+    logger.info("Shutdown completed")
+
+# Replace the FastAPI app initialization
+app = FastAPI(
+    title="Advanced RAG Chat Service", 
+    version="1.0.0",
+    lifespan=lifespan
+)
+
+
 
 @app.get("/health", response_model=HealthResponse)
 async def health_check():