stream response

app/main.py

@@ -45,38 +45,23 @@ manager = ConnectionManager()
 
 @app.websocket("/ws/{client_id}")
 async def websocket_endpoint(websocket: WebSocket, client_id: int):
-    print("CONNECTING")
     await manager.connect(websocket)
-    print("1")
     now = datetime.now()
     current_time = now.strftime("%H:%M")
-    print("2")
     try:
         while True:
-            print("3")
-            data = await websocket.receive_text()
-            # print("4")
-            # await manager.send_personal_message(f"You wrote: {data}", websocket)
-            print("5")
-            message = {"time":current_time,"clientId":client_id,"message":data}
-            print("6")
-            await manager.broadcast(json.dumps(message))
-            print("7")
-            
+            data = await websocket.receive_text().strip()
+            if len(data) > 0:
+                generator = session_assistant.ask(data)
+                for text in generator
+                    message = {"time":current_time,"clientId":client_id,"message":text}
+                    # await manager.broadcast(json.dumps(message))
+                    await manager.send_personal_message(json.dumps(message), websocket)
     except WebSocketDisconnect:
-        print("DISCONNECTING")
         manager.disconnect(websocket)
-        print("1")
         message = {"time":current_time,"clientId":client_id,"message":"Offline"}
-        print("2")
         await manager.broadcast(json.dumps(message))
-        print("3")
-
-@app.get("/query")
-async def process_input(text: str):
-    if text and len(text.strip()) > 0:
-        text = text.strip()
-        return StreamingResponse(session_assistant.ask(text), media_type='text/event-stream')
+    # return StreamingResponse(, media_type='text/event-stream')
 
 
 @app.post("/upload")

app/rag.py

@@ -102,15 +102,10 @@ class ChatPDF:
             similarity_top_k=6,
         )
 
-    async def ask(self, query: str):
+    def ask(self, query: str):
         logger.info("retrieving the response to the query")
         streaming_response = self.query_engine.query(query)
-        print(streaming_response)
-        print("PRINTING STREAM")
-        generator = streaming_response.response_gen
-        for text in generator:
-            print(text)
-            yield text
+        return streaming_response.response_gen
 
     def clear(self):
         pass