Update app/main.py

app/main.py

@@ -1,10 +1,11 @@
 import os
+import traceback
 import logging
 import pathlib
 import time
 import re
 from typing import List
-from fastapi import FastAPI, Request, Depends
+from fastapi import FastAPI, Request, Depends, HTTPException
 from fastapi.middleware import Middleware
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse
@@ -27,7 +28,7 @@ session_assistant = ChatPDF()
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-
+isBusy = False
 
 def astreamer(generator):
     t0 = time.time()
@@ -39,6 +40,10 @@ def astreamer(generator):
 
 @app.get("/query")
 async def process_input(text: str):
+    global isBusy
+    if isBusy:
+        raise HTTPException(status_code=503, detail="Server is busy")
+    isBusy = True
     generator = None
     if text and len(text.strip()) > 0:
         if session_assistant.pdf_count > 0:
@@ -51,6 +56,7 @@ async def process_input(text: str):
     else:
         message = "The provided query is empty."
         generator = re.split(r'(\s)', message)
+    isBusy = False
     return StreamingResponse(astreamer(generator), media_type='text/event-stream')
 
 async def parse_body(request: Request):
@@ -59,6 +65,11 @@ async def parse_body(request: Request):
 
 @app.post("/upload")
 def upload(filename: str, data: bytes = Depends(parse_body)):
+    global isBusy
+    if isBusy:
+        raise HTTPException(status_code=503, detail="Server is busy")
+    isBusy = True
+    session_assistant.clear()
     if data:
         try:
             print("Filename: " + filename)
@@ -72,14 +83,17 @@ def upload(filename: str, data: bytes = Depends(parse_body)):
 
     message = "Files inserted successfully."
     generator = re.split(r'(\s)', message)
+    isBusy = False
     return StreamingResponse(astreamer(generator), media_type='text/event-stream')
 
 
 @app.get("/clear")
-def ping():
+def clear():
+    global isBusy
     session_assistant.clear()
     message = "All files have been cleared. The first query may take a little longer."
     generator = re.split(r'(\s)', message)
+    isBusy = False
     return StreamingResponse(astreamer(generator), media_type='text/event-stream')