Update main.py

app/main.py

@@ -73,33 +73,21 @@ async def websocket_endpoint(websocket: WebSocket, client_id: int):
         await manager.broadcast(json.dumps(message))
 
 
-def fake_json_streamer():
+def astreamer(generator):
     t0 = time.time()
-    for i in range(10):
+    for i in generator:
         print(f"Chunk being yielded (time {int((time.time()-t0)*1000)}ms)", flush=True)
-        yield json.dumps( {"message": "Hello World"}) + '\n'
+        yield i
         time.sleep(0.5)
     print(f"Over (time {int((time.time()-t0)*1000)}ms)", flush=True)
 
 
-async def astreamer(generator):
-    try:
-        print("streaming........")
-        for i in generator:
-            print(i)
-            yield (i)
-            await asyncio.sleep(.1)
-    except asyncio.CancelledError as e:
-        yield ('cancelled')
-
-
 @app.get("/query")
 async def process_input(text: str):
     if text and len(text.strip()) > 0:
         text = text.strip()
         streaming_response = session_assistant.ask(text)
-        return StreamingResponse(fake_json_streamer(), media_type='text/event-stream')
-        # return StreamingResponse(astreamer(streaming_response.response_gen), media_type='text/event-stream')
+        return StreamingResponse(astreamer(streaming_response.response_gen), media_type='text/event-stream')
 
 
 @app.post("/upload")