fix: add websocket in handlerToken

main.py

@@ -4,6 +4,7 @@ import shutil
 import subprocess
 import asyncio
 
+
 from typing import Any, Dict, List
 
 from fastapi import FastAPI, HTTPException, UploadFile, WebSocket, WebSocketDisconnect
@@ -17,10 +18,10 @@ from langchain.embeddings import HuggingFaceInstructEmbeddings
 from langchain.prompts import PromptTemplate
 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
 
@@ -35,8 +36,6 @@ class Predict(BaseModel):
 class Delete(BaseModel):
     filename: str
 
-
-
 # if torch.backends.mps.is_available():
 #     DEVICE_TYPE = "mps"
 # elif torch.cuda.is_available():
@@ -59,9 +58,9 @@ DB = Chroma(
 RETRIEVER = DB.as_retriever()
 
 class MyCustomSyncHandler(BaseCallbackHandler):
-    def __init__(self):
+    def __init__(self, state):
         self.end = False
-        self.websocket = None
+        self.state = state
 
     def on_llm_start(
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
@@ -80,8 +79,13 @@ class MyCustomSyncHandler(BaseCallbackHandler):
         print(token)
 
 
+# Create State
+
+tokenMessageLLM = State()
+
+get, update = tokenMessageLLM.create('')
 
-handlerToken = MyCustomSyncHandler()
+handlerToken = MyCustomSyncHandler(update)
 
 LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True, callbacks=[handlerToken])
 
@@ -258,9 +262,9 @@ async def websocket_endpoint(websocket: WebSocket):
 
     try:
         while True:
-            handlerToken.websocket = websocket
+            tokenMessageLLM.after_create(lambda now, old: print(f"{old} updated to {now}."))
 
-            print(handlerToken.websocket)
+            print(tokenMessageLLM)
 
             data = await websocket.receive_text()
             res = QA(data)

requirements.txt

@@ -29,6 +29,7 @@ uvicorn
 fastapi
 websockets
 pydantic
+varstate
 
 # Streamlit related
 streamlit