fix: add websocket in handlerToken

main.py

@@ -54,8 +54,6 @@ class MyCustomSyncHandler(BaseCallbackHandler):
         self.message += token
         redisClient.publish(f'{kwargs["tags"][0]}',  self.message)
 
-        print(self.message)
-
 LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True)
 
 prompt, memory = get_prompt_template(promptTemplate_type="llama", history=True)
@@ -71,6 +69,14 @@ QA = RetrievalQA.from_chain_type(
     },
 )
 
+async def pubsub_listener(websocket, client_id):
+    pubsub = redisClient.pubsub()
+    pubsub.subscribe(f'{client_id}')
+
+    while True:
+        message = pubsub.listen()
+        await websocket.send_text(f'{message}')
+
 app = FastAPI(title="homepage-app")
 api_app = FastAPI(title="api app")
 
@@ -123,7 +129,6 @@ def run_ingest_route():
             },
         )
 
-
         return {"response": "The training was successfully completed"}
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error occurred: {str(e)}")
@@ -226,27 +231,10 @@ async def websocket_endpoint(websocket: WebSocket,  client_id: int):
         while True:
             prompt = await websocket.receive_text()
 
-            pubsub = redisClient.pubsub()
-            pubsub.subscribe(f'{client_id}')
-
-            for item in pubsub.listen():
-                if item['type'] == 'message':
-                    await websocket.send_text(f'{item["data"]}')
-                if item['type'] == 'subscribe':
-                    QA(
-                        inputs=prompt,
-                        return_only_outputs=True,
-                        callbacks=[MyCustomSyncHandler()],
-                        tags=f'{client_id}',
-                        include_run_info=True
-                    )
-
-
-
-
-
-
+            pubsub_thread = threading.Thread(target=pubsub_listener(websocket, client_id))
+            pubsub_thread.start()
 
+            QA(inputs=prompt, return_only_outputs=True, callbacks=[MyCustomSyncHandler()], tags=f'{client_id}', include_run_info=True)
 
     except WebSocketDisconnect:
         print('disconnect')