fix: add websocketClient

main.py

@@ -2,6 +2,8 @@ import os
 import glob
 import shutil
 import subprocess
+import contextvars
+
 from typing import Any, Dict, List
 
 from fastapi import FastAPI, HTTPException, UploadFile, WebSocket
@@ -31,11 +33,12 @@ class Predict(BaseModel):
 class Delete(BaseModel):
     filename: str
 
-class MyCustomHandler(BaseCallbackHandler):
-    async def on_llm_new_token(self, token: str, **kwargs) -> None:
-        global websocketClient
+websocketClient = contextvars.ContextVar("websocketClient")
 
-        await websocketClient.send_text(f"Message text was: {token}")
+class MyCustomHandler(AsyncCallbackHandler):
+    async def on_llm_new_token(self, token: str, **kwargs) -> None:
+        ws = websocketClient.get()
+        await ws.send_text(f"Message text was: {token}")
 
         print(f" token: {token}")
 
@@ -246,12 +249,10 @@ async def create_upload_file(file: UploadFile):
 @api_app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
     global QA
-    global websocketClient
 
     await websocket.accept()
     while True:
-        global websocketClient
-        websocketClient = websocket;
+        websocketClient.set(websocket);
 
         data = await websocket.receive_text()