Update main.py

app/main.py

@@ -1,6 +1,7 @@
 import os
 import shutil
 import time
+import re
 from typing import List
 from fastapi import FastAPI, UploadFile
 from fastapi.middleware import Middleware
@@ -22,21 +23,33 @@ app = FastAPI(middleware=middleware)
 files_dir = os.path.expanduser("~/wtp_be_files/")
 session_assistant = ChatPDF()
 
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 
 def astreamer(generator):
     t0 = time.time()
     for i in generator:
-        print(f"Chunk being yielded (time {int((time.time()-t0)*1000)}ms)", flush=True)
+        logger.info(f"Chunk being yielded (time {int((time.time()-t0)*1000)}ms)", flush=True)
         yield i
-    print(f"Over (time {int((time.time()-t0)*1000)}ms)", flush=True)
+    logger.info(f"Over (time {int((time.time()-t0)*1000)}ms)", flush=True)
 
 
 @app.get("/query")
 async def process_input(text: str):
+    generator = None
     if text and len(text.strip()) > 0:
-        text = text.strip()
-        streaming_response = session_assistant.ask(text)
-        return StreamingResponse(astreamer(streaming_response.response_gen), media_type='text/event-stream')
+        if session_assistant.pdf_count > 0:
+            text = text.strip()
+            streaming_response = session_assistant.ask(text)
+            generator = streaming_response.response_gen
+        else:
+            message = "Please add a PDF document first."
+            generator = re.split(r'(\s)', message)
+    else:
+        message = "The provided query is empty."
+        generator = re.split(r'(\s)', message)
+    return StreamingResponse(astreamer(generator), media_type='text/event-stream')
 
 
 @app.post("/upload")
@@ -55,13 +68,17 @@ def upload(files: list[UploadFile]):
         session_assistant.ingest(files_dir)
         shutil.rmtree(files_dir)
 
-    return "Files inserted!"
+    message = "Files inserted successfully."
+    generator = re.split(r'(\s)', message)
+    return StreamingResponse(astreamer(generator), media_type='text/event-stream')
 
 
 @app.get("/clear")
 def ping():
     session_assistant.clear()
-    return "All files have been cleared."
+    message = "All files have been cleared."
+    generator = re.split(r'(\s)', message)
+    return StreamingResponse(astreamer(generator), media_type='text/event-stream')
 
 
 @app.get("/")

requirements.txt

@@ -6,6 +6,5 @@ llama-index-readers-file
 llama-index-llms-llama-cpp
 llama-index-embeddings-fastembed
 qdrant-client
-websockets
 unstructured
 fastembed==0.2.7