fix

app/main.py

@@ -53,11 +53,19 @@ async def websocket_endpoint(websocket: WebSocket, client_id: int):
             data = await websocket.receive_text()
             data = data.strip()
             if len(data) > 0:
-                generator = session_assistant.ask(data)
-                for text in generator:
-                    message = {"time":current_time,"clientId":client_id,"message":text}
+                if not session_assistant.pdf_count > 0:
+                    message = {"time":current_time,"clientId":client_id,"message":"Please, add a PDF document first."}
                     # await manager.broadcast(json.dumps(message))
                     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"}

app/rag.py

@@ -107,11 +107,10 @@ class ChatPDF:
 
     def ask(self, query: str):
         logger.info("retrieving the response to the query")
-        if not self.pdf_count > 0:
-            return "Please, add a PDF document first."
-
+        print("GENERATING STREAM")
         streaming_response = self.query_engine.query(query)
-        return streaming_response.response_gen
+        print("PASSING STREAM")
+        return streaming_response
 
     def clear(self):
         self.pdf_count = 0