fix: add websocket in handlerToken

main.py

@@ -20,8 +20,6 @@ from langchain.memory import ConversationBufferMemory
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.schema import LLMResult
 
-from varstate import State
-
 # from langchain.embeddings import HuggingFaceEmbeddings
 from load_models import load_model
 
@@ -58,9 +56,9 @@ DB = Chroma(
 RETRIEVER = DB.as_retriever()
 
 class MyCustomSyncHandler(BaseCallbackHandler):
-    def __init__(self, state):
+    def __init__(self):
         self.end = False
-        self.state = state
+        self.callback = None
 
     def on_llm_start(
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
@@ -77,14 +75,8 @@ class MyCustomSyncHandler(BaseCallbackHandler):
 
         print(token)
 
-
 # Create State
-
-tokenMessageLLM = State()
-
-get, update = tokenMessageLLM.create('')
-
-handlerToken = MyCustomSyncHandler(update)
+handlerToken = MyCustomSyncHandler()
 
 LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True, callbacks=[handlerToken])
 
@@ -253,8 +245,8 @@ async def create_upload_file(file: UploadFile):
 
     return {"filename": file.filename}
 
-@api_app.websocket("/ws")
-async def websocket_endpoint(websocket: WebSocket):
+@api_app.websocket("/ws/{client_id}")
+async def websocket_endpoint(websocket: WebSocket,  client_id: int):
     global QA
 
     await websocket.accept()
@@ -265,16 +257,11 @@ async def websocket_endpoint(websocket: WebSocket):
         while True:
             prompt = await websocket.receive_text()
 
-            statusProcess = handlerToken.end
+            handlerToken.callback = websocket.send_text;
 
-            if (oldReceiveText != prompt) and (statusProcess == False) :
+            if (oldReceiveText != prompt):
                 oldReceiveText = prompt
-                res = QA(prompt)
-
-            print(statusProcess);
-
-            tokenState = get()
-            await websocket.send_text(f"token: {tokenState}")
+                asyncio.run(QA(prompt))
 
     except WebSocketDisconnect:
         print('disconnect')