fix: add websocket in handlerToken

main.py

@@ -41,23 +41,19 @@ EMBEDDINGS = HuggingFaceInstructEmbeddings(model_name=EMBEDDING_MODEL_NAME, mode
 DB = Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=EMBEDDINGS, client_settings=CHROMA_SETTINGS)
 RETRIEVER = DB.as_retriever()
 
+
+redisClient = redis.Redis(host='localhost', port=6379, db=0)
+
+
 class MyCustomSyncHandler(BaseCallbackHandler):
-    def on_llm_start(
-        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
-    ) -> None:
-        print(f'on_llm_start self {self}')
-        print(f'on_llm_start kwargs {prompts}')
-        print(f'on_llm_start token {kwargs}')
-
-    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
-        print(f'on_llm_end self {self}')
-        print(f'on_llm_end kwargs {response}')
-        print(f'on_llm_end token {kwargs}')
+    def __init__(self):
+        self.message = ''
 
     def on_llm_new_token(self, token: str, **kwargs) -> Any:
-        print(f'on_llm_new_token self {self}')
-        print(f'on_llm_new_token kwargs {kwargs}')
-        print(f'on_llm_new_token token {token}')
+        message += token
+        redisClient.publish(f'{kwargs["tags"][0]}',  message)
+
+        print(message)
 
 LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True)
 
@@ -74,10 +70,6 @@ QA = RetrievalQA.from_chain_type(
     },
 )
 
-redisClient = redis.Redis(host='localhost', port=6379, db=0)
-
-redisClient.set('foo', 'bar')
-
 app = FastAPI(title="homepage-app")
 api_app = FastAPI(title="api app")
 
@@ -232,6 +224,14 @@ async def websocket_endpoint(websocket: WebSocket,  client_id: int):
     try:
         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(item["data"])
+
             QA(
                 inputs=prompt,
                 return_only_outputs=True,
@@ -240,9 +240,8 @@ async def websocket_endpoint(websocket: WebSocket,  client_id: int):
                 include_run_info=True
             )
 
-            response = redisClient.get('foo')
 
-            await websocket.send_text(response)
+
 
     except WebSocketDisconnect:
         print('disconnect')