Migrate to stream generator

app/main.py

@@ -22,22 +22,26 @@ app = FastAPI(middleware=middleware)
 
 files_dir = os.path.expanduser("~/wtp_be_files/")
 session_assistant = ChatPDF()
-session_messages = []
+# session_messages = []
 
 @app.get("/query")
 def process_input(text: str):
     if text and len(text.strip()) > 0:
         text = text.strip()
-        agent_text = session_assistant.ask(text)
-        session_messages.append((text, True))
-        session_messages.append((agent_text, False))
-        return agent_text
+        print("PRINTING STREAM")
+        agent_text_stream = session_assistant.ask(text)
+        print(stream_gen)
+        for text in agent_text_stream:
+            print(text)
+        # session_messages.append((text, True))
+        # session_messages.append((agent_text, False))
+        return "Query resolved!"
 
 
 @app.post("/upload")
 def upload(files: list[UploadFile]):
     session_assistant.clear()
-    session_messages = []
+    # session_messages = []
     
     try:
         os.makedirs(files_dir)

app/rag.py

@@ -151,7 +151,10 @@ class ChatPDF:
         #     response_synthesizer=response_synthesizer,
         # )
 
-        self.query_engine = index.as_query_engine()
+        self.query_engine = index.as_query_engine(
+            streaming=True,
+            # similarity_top_k=6,
+        )
 
     def ask(self, query: str):
         if not self.query_engine:
@@ -159,9 +162,8 @@ class ChatPDF:
 
         logger.info("retrieving the response to the query")
         # response = self.query_engine.query(str_or_query_bundle=query)
-        response = self.query_engine.query(query)
-        print(response)
-        return response
+        streaming_response = self.query_engine.query(query)
+        return streaming_response.response_gen
 
     def clear(self):
         self.query_engine = None