fix: add handle

main.py

@@ -38,26 +38,12 @@ websocketClient = contextvars.ContextVar("websocketClient")
 
 class MyCustomSyncHandler(BaseCallbackHandler):
     def on_llm_new_token(self, token: str, **kwargs) -> None:
-        print(f"Sync handler being called in a `thread_pool_executor`: token: {token}")
-
-
-class MyCustomAsyncHandler(AsyncCallbackHandler):
-    """Async callback handler that can be used to handle callbacks from langchain."""
-
-    async def on_llm_start(
-        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
-    ) -> None:
-        """Run when chain starts running."""
-        print("zzzz....")
-        await asyncio.sleep(0.3)
-        class_name = serialized["name"]
-        print("Hi! I just woke up. Your llm is starting")
-
-    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
-        """Run when chain ends running."""
-        print("zzzz....")
-        await asyncio.sleep(0.3)
-        print("Hi! I just woke up. Your llm is ending")
+        ws = websocketClient.get()
+
+        ws.receive_text(token)
+
+        print(f"token: {token}")
+
 # if torch.backends.mps.is_available():
 #     DEVICE_TYPE = "mps"
 # elif torch.cuda.is_available():
@@ -79,7 +65,7 @@ DB = Chroma(
 
 RETRIEVER = DB.as_retriever()
 
-LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True, callbacks=[MyCustomSyncHandler(), MyCustomAsyncHandler()])
+LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True, callbacks=[MyCustomSyncHandler()])
 
 template = """you are a helpful, respectful and honest assistant.
 Your name is Katara llma. You should only use the source documents provided to answer the questions.
@@ -263,6 +249,6 @@ async def websocket_endpoint(websocket: WebSocket):
 
         data = await websocket.receive_text()
 
-        QA(data)
+        res = QA(data)
 
         await websocket.send_text(f"Message text was:")