fix: add websocket in handlerToken

main.py

@@ -64,6 +64,7 @@ class MyCustomSyncHandler(BaseCallbackHandler):
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
     ) -> None:
         self.end = False
+        self.callback = None
 
     def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         self.end = True
@@ -71,7 +72,8 @@ class MyCustomSyncHandler(BaseCallbackHandler):
     def on_llm_new_token(self, token: str, **kwargs) -> Any:
         print(token)
 
-        self.state(token)
+        if self.callback :
+            self.callback(token)
 
         print(token)
 
@@ -257,9 +259,8 @@ async def websocket_endpoint(websocket: WebSocket,  client_id: int):
         while True:
             prompt = await websocket.receive_text()
 
-            handlerToken.callback = websocket.send_text;
-
             if (oldReceiveText != prompt):
+                handlerToken.callback = websocket.send_text
                 oldReceiveText = prompt
                 asyncio.run(QA(prompt))