Update main.py

main.py

@@ -1,5 +1,8 @@
-from fastapi import FastAPI, File, UploadFile, Depends, HTTPException
+from fastapi import FastAPI, File, UploadFile, HTTPException, Response, Request
 from fastapi.middleware.cors import CORSMiddleware
+import uuid
+from datetime import datetime, timedelta
+import asyncio
 
 from typing import List, Dict, Any
 from io import BytesIO, StringIO
@@ -182,6 +185,29 @@ class Conversational_Chain:
         return ChatPromptTemplate.from_messages(messages)
 
 
+class UserSessionManager:
+    def __init__(self):
+        self.sessions = {}
+        self.last_request_time = {}
+
+    def get_session(self, user_id: str):
+        if user_id not in self.sessions:
+            self.sessions[user_id] = None
+            self.last_request_time = datetime.now()
+        return self.sessions[user_id]
+
+    def set_session(self, user_id: str, conversational_chain):
+        self.sessions[user_id] = conversational_chain
+        self.last_request_time[user_id] = datetime.now()
+
+    def delete_inactive_sessions(self, inactive_period: timedelta):
+        current_time = datetime.now()
+        for user_id, last_request_time in list(self.last_request_time.items()):
+            if current_time - last_request_time > inactive_period:
+                del self.sessions[user_id]
+                del self.last_request_time[user_id]
+
+
 app = FastAPI()
 
 origins = ["https://viboognesh-react-chat.static.hf.space"]
@@ -195,12 +221,30 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+user_session_manager = UserSessionManager()
+
+
+@app.middleware("http")
+async def update_last_request_time(request: Request, call_next):
+    user_id = request.cookies.get("user_id")
+    if user_id:
+        user_session_manager.last_request_time[user_id] = datetime.now()
+    response = await call_next(request)
+    return response
 
-app.state.conversation_chain = None
+
+async def check_inactivity():
+    inactive_period = timedelta(hours=2)
+    while True:
+        await asyncio.sleep(600)
+        user_session_manager.delete_inactive_sessions(inactive_period)
+
+
+app.add_task(check_inactivity())
 
 
 @app.post("/upload_files/")
-async def upload_files(files: List[UploadFile] = File(...)):
+async def upload_files(response: Response, files: List[UploadFile] = File(...)):
     file_details = []
     try:
         for file in files:
@@ -211,10 +255,18 @@ async def upload_files(files: List[UploadFile] = File(...)):
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e))
 
+    user_id = response.cookies.get("user_id")
+    if not user_id:
+        user_id = str(uuid.uuid4())
+        response.set_cookie(key="user_id", value=user_id)
+
     try:
-        app.state.conversational_chain = Conversational_Chain(
+        conversational_chain = Conversational_Chain(
             file_details
         ).create_conversational_chain()
+        user_session_manager.set_session(
+            user_id=user_id, conversational_chain=conversational_chain
+        )
         print("conversational_chain_manager created")
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -224,17 +276,23 @@ async def upload_files(files: List[UploadFile] = File(...)):
 
 @app.get("/predict/")
 async def predict(query: str):
+    user_id = response.cookies.get("user_id")
+    if not user_id:
+        user_id = str(uuid.uuid4())
+        response.set_cookie(key="user_id", value=user_id)
+
     try:
-        if app.state.conversation_chain is None:
+        conversational_chain = user_session_manager.get_session(user_id=user_id)
+        if conversational_chain is None:
             system_prompt = "Answer the question and also ask the user to upload files to ask questions from the files.\n"
             llm_model = ChatOpenAI()
             response = llm_model.invoke(system_prompt + query)
             answer = response.content
         else:
-            response = app.state.conversation_chain.invoke(query)
+            response = conversational_chain.invoke(query)
             answer = response["answer"]
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
     print("predict called")
-    return {"answer": answer}
+    return {"answer": answer}