Update app/main.py

app/main.py

@@ -25,31 +25,53 @@ session_assistant = ChatPDF()
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-isBusy = False
+
+class Session:
+    isBusy = False # Processing upload or query response
+    curUserID = ""
+    prevUserID = ""
+    lastQueryTimestamp = 0
+
+session = Session()
 
 @app.middleware("http")
 async def resolve_availability(request: Request, call_next):
-    global isBusy
-    if isBusy:
+    if session.isBusy:
         return PlainTextResponse("Server is busy", status_code=503)
-    response = await call_next(request)
-    return response
+    
+    request_args = dict(request.query_params)
+    if not "id" in request_args:
+        return PlainTextResponse("Bad request", status_code=400)
+    
+    now = time.time()
+    if session.curUserID == request_args["id"]:
+        session.lastQueryTimestamp = now
+        return await call_next(request)
+    
+    if session.prevUserID == request_args["id"]:
+        return PlainTextResponse("Session has expired", status_code=419)
+    
+    if now - session.lastQueryTimestamp >= 300:
+        session.lastQueryTimestamp = now
+        session.prevUserID = session.curUserID
+        session.curUserID = request_args["id"]
+        return await call_next(request)
+    
+    return PlainTextResponse("Server is busy", status_code=503)
 
 
 def astreamer(generator):
-    global isBusy
     t0 = time.time()
     for i in generator:
         logger.info(f"Chunk being yielded (time {int((time.time()-t0)*1000)}ms) - {i}")
         yield i
     logger.info(f"X-Process-Time: {int((time.time()-t0)*1000)}ms")
-    isBusy = False
+    session.isBusy = False
 
 
 @app.get("/query")
 async def process_input(text: str):
-    global isBusy
-    isBusy = True
+    session.isBusy = True
     generator = None
     if text and len(text.strip()) > 0:
         if session_assistant.pdf_count > 0:
@@ -67,8 +89,7 @@ async def process_input(text: str):
 
 @app.post("/upload")
 def upload(files: list[UploadFile]):
-    global isBusy
-    isBusy = True
+    session.isBusy = True
     try:
         os.makedirs(files_dir)
         for file in files:
@@ -89,8 +110,7 @@ def upload(files: list[UploadFile]):
 
 @app.get("/clear")
 def clear():
-    global isBusy
-    isBusy = True
+    session.isBusy = True
     session_assistant.clear()
     message = "Your files have been cleared successfully."
     generator = re.split(r'(\s)', message)