fix: add handle

main.py

@@ -34,15 +34,13 @@ class Predict(BaseModel):
 class Delete(BaseModel):
     filename: str
 
-websocketClient = contextvars.ContextVar("websocketClient")
+websocket_state = None
 
 class MyCustomSyncHandler(AsyncCallbackHandler):
     async def on_llm_new_token(self, token: str, **kwargs) -> None:
-        ws = websocketClient.get()
+        print(f"{websocket_state}")
 
-        print(f"{ws}")
-
-        await ws.send_text(f"token: {token}")
+        await websocket_state(f"token: {token}")
 
         print(f"token: {token}")
 
@@ -247,12 +245,12 @@ async def websocket_endpoint(websocket: WebSocket):
 
     await websocket.accept()
     while True:
-        websocketClient.set(websocket);
+        websocket_state = websocket;
 
-        data = await websocket.receive_text()
+        data = await websocket_state.receive_text()
 
         res = QA(data)
 
         print(f"${res}")
 
-        await websocket.send_text(f"Message text was:")
+        await websocket_state.send_text(f"Message text was:")