minor fix

app/main.py

@@ -1,11 +1,8 @@
 import os
 import shutil
-import json
-import asyncio
 import time
-from datetime import datetime
 from typing import List
-from fastapi import FastAPI, UploadFile, WebSocket, WebSocketDisconnect
+from fastapi import FastAPI, UploadFile
 from fastapi.middleware import Middleware
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse
@@ -25,60 +22,12 @@ app = FastAPI(middleware=middleware)
 files_dir = os.path.expanduser("~/wtp_be_files/")
 session_assistant = ChatPDF()
 
-class ConnectionManager:
-    def __init__(self):
-        self.active_connections: List[WebSocket] = []
-
-    async def connect(self, websocket: WebSocket):
-        await websocket.accept()
-        self.active_connections.append(websocket)
-
-    def disconnect(self, websocket: WebSocket):
-        self.active_connections.remove(websocket)
-
-    async def send_personal_message(self, message: str, websocket: WebSocket):
-        await websocket.send_text(message)
-
-    async def broadcast(self, message: str):
-        for connection in self.active_connections:
-            await connection.send_text(message)
-
-manager = ConnectionManager()
-
-@app.websocket("/ws/{client_id}")
-async def websocket_endpoint(websocket: WebSocket, client_id: int):
-    await manager.connect(websocket)
-    now = datetime.now()
-    current_time = now.strftime("%H:%M")
-    try:
-        while True:
-            data = await websocket.receive_text()
-            data = data.strip()
-            if len(data) > 0:
-                if not session_assistant.pdf_count > 0:
-                    message = {"time":current_time,"clientId":client_id,"message":"Please, add a PDF document first."}
-                    await manager.send_personal_message(json.dumps(message), websocket)
-                else:
-                    print("FETCHING STREAM")
-                    streaming_response = session_assistant.ask(data)
-                    print("STARTING STREAM")
-                    for text in streaming_response.response_gen:
-                        message = {"time":current_time,"clientId":client_id,"message":text}
-                        # await manager.broadcast(json.dumps(message))
-                        await manager.send_personal_message(json.dumps(message), websocket)
-                    print("ENDING STREAM")
-    except WebSocketDisconnect:
-        manager.disconnect(websocket)
-        message = {"time":current_time,"clientId":client_id,"message":"Offline"}
-        await manager.broadcast(json.dumps(message))
-
 
 def astreamer(generator):
     t0 = time.time()
     for i in generator:
         print(f"Chunk being yielded (time {int((time.time()-t0)*1000)}ms)", flush=True)
         yield i
-        # time.sleep(0.5)
     print(f"Over (time {int((time.time()-t0)*1000)}ms)", flush=True)
 
 
@@ -87,19 +36,7 @@ async def process_input(text: str):
     if text and len(text.strip()) > 0:
         text = text.strip()
         streaming_response = session_assistant.ask(text)
-        headers = {
-            'X-Content-Type-Options': 'nosniff',
-            'Content-Type': 'text/event-stream',
-            'Cache-Control': 'no-cache',          
-            'Content-Encoding': 'none',
-            'Connection': 'keep-alive',
-            'Access-Control-Allow-Origin': '*'
-            }
-        return StreamingResponse(
-            astreamer(streaming_response.response_gen),
-            headers=headers,
-            media_type='text/event-stream'
-        )
+        return StreamingResponse(astreamer(streaming_response.response_gen), media_type='text/event-stream')
 
 
 @app.post("/upload")

app/rag.py

@@ -6,7 +6,6 @@ from llama_index.core import (
     StorageContext,
     Settings,
     get_response_synthesizer)
-# from llama_index.legacy.readers.file.base import SimpleDirectoryReader
 from llama_index.core.node_parser import SentenceSplitter
 from llama_index.core.schema import TextNode, MetadataMode
 from llama_index.core.vector_stores import VectorStoreQuery
@@ -109,9 +108,7 @@ class ChatPDF:
 
     def ask(self, query: str):
         logger.info("retrieving the response to the query")
-        print("GENERATING STREAM")
         streaming_response = self.query_engine.query(query)
-        print("PASSING STREAM")
         return streaming_response
 
     def clear(self):