fix: add callback

main.py

@@ -241,17 +241,17 @@ async def websocket_endpoint(websocket: WebSocket,  client_id: int):
             pubsub = redisClient.pubsub()
             pubsub.subscribe(f'{client_id}')
 
-            with concurrent.futures.ThreadPoolExecutor() as executor:
-               executor.submit(QA(inputs=prompt, return_only_outputs=True, tags=f'{client_id}', include_run_info=True, callbacks=[handleCallback]))
-
-            i  = 0
-            for item in pubsub.listen():
-                if item["type"] == "message":
-                    message = item["data"].decode('utf-8')
-                    if message == "end": pubsub.unsubscribe({client_id})
-                    if(i == 10): pubsub.unsubscribe({client_id})
-                    await websocket.send_text(f'{message}')
-                i+=1
+            response = QA(inputs=prompt, return_only_outputs=True, tags=f'{client_id}', include_run_info=True, callbacks=[handleCallback])
+            await websocket.send_text(f'{response}')
+
+            # with concurrent.futures.ThreadPoolExecutor() as executor:
+            #    executor.submit(QA(inputs=prompt, return_only_outputs=True, tags=f'{client_id}', include_run_info=True, callbacks=[handleCallback]))
+
+            # for item in pubsub.listen():
+            #     if item["type"] == "message":
+            #         message = item["data"].decode('utf-8')
+            #         if message == "end": pubsub.unsubscribe({client_id})
+            #         await websocket.send_text(f'{message}')
 
 
     except WebSocketDisconnect: