Update app/main.py

app/main.py

@@ -1,14 +1,12 @@
 import os
-import traceback
 import logging
-import pathlib
+import shutil
 import time
 import re
-from typing import List
-from fastapi import FastAPI, Request, Depends, HTTPException
+from fastapi import FastAPI, Request, UploadFile
 from fastapi.middleware import Middleware
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import StreamingResponse
+from fastapi.responses import PlainTextResponse, StreamingResponse
 from .rag import ChatPDF
 
 middleware = [
@@ -23,26 +21,34 @@ middleware = [
 app = FastAPI(middleware=middleware)
 
 files_dir = os.path.expanduser("~/wtp_be_files/")
-os.makedirs(files_dir)
 session_assistant = ChatPDF()
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 isBusy = False
 
+@app.middleware("http")
+async def resolve_availability(request: Request, call_next):
+    global isBusy
+    if isBusy:
+        return PlainTextResponse("Server is busy", status_code=503)
+    response = await call_next(request)
+    return response
+
+
 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)")
+        logger.info(f"Chunk being yielded (time {int((time.time()-t0)*1000)}ms) - {i}")
         yield i
-    logger.info(f"Over (time {int((time.time()-t0)*1000)}ms)")
+    logger.info(f"X-Process-Time: {int((time.time()-t0)*1000)}ms")
+    isBusy = False
 
 
 @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:
@@ -51,49 +57,43 @@ async def process_input(text: str):
             streaming_response = session_assistant.ask(text)
             generator = streaming_response.response_gen
         else:
-            message = "Please add a PDF document first."
+            message = "Please provide the PDF document you'd like to add."
             generator = re.split(r'(\s)', message)
     else:
-        message = "The provided query is empty."
+        message = "Your query is empty. Please provide a query for me to process."
         generator = re.split(r'(\s)', message)
-    isBusy = False
     return StreamingResponse(astreamer(generator), media_type='text/event-stream')
 
-async def parse_body(request: Request):
-    data: bytes = await request.body()
-    return data
 
 @app.post("/upload")
-def upload(filename: str, data: bytes = Depends(parse_body)):
+def upload(files: list[UploadFile]):
     global isBusy
-    if isBusy:
-        raise HTTPException(status_code=503, detail="Server is busy")
     isBusy = True
-    session_assistant.clear()
-    if data:
-        try:
-            print("Filename: " + filename)
-            path = f"{files_dir}/{filename}"
-            with open(path, "wb") as f:
-                f.write(data)
-            session_assistant.ingest(files_dir)
-            pathlib.Path(path).unlink()
-        except Exception as e:
-            logging.error(traceback.format_exc())
-
-    message = "Files inserted successfully."
+    try:
+        os.makedirs(files_dir)
+        for file in files:
+            try:
+                path = f"{files_dir}/{file.filename}"
+                file.file.seek(0)
+                with open(path, 'wb') as destination:
+                    shutil.copyfileobj(file.file, destination)
+            finally:
+                file.file.close()
+    finally:
+        session_assistant.ingest(files_dir)
+        shutil.rmtree(files_dir)
+    message = "All files have been added successfully to your account. Your first query may take a little longer as the system indexes your documents. Please be patient while we process your request."
     generator = re.split(r'(\s)', message)
-    isBusy = False
     return StreamingResponse(astreamer(generator), media_type='text/event-stream')
 
 
 @app.get("/clear")
 def clear():
     global isBusy
+    isBusy = True
     session_assistant.clear()
-    message = "All files have been cleared. The first query may take a little longer."
+    message = "Your files have been cleared successfully."
     generator = re.split(r'(\s)', message)
-    isBusy = False
     return StreamingResponse(astreamer(generator), media_type='text/event-stream')