fix: add websocketClient

load_models.py

@@ -69,7 +69,7 @@ def load_quantized_model_gguf_ggml(model_id, model_basename, device_type, loggin
         kwargs["stream"] = stream
 
         if stream == True:
-            kwargs["callbacks"] = [StreamingStdOutCallbackHandler()]
+            kwargs["callbacks"] = callbacks
 
         return LlamaCpp(**kwargs)
     except:

main.py

@@ -31,15 +31,13 @@ class Predict(BaseModel):
 class Delete(BaseModel):
     filename: str
 
-tokenWS = ''
-
 class MyCustomHandler(BaseCallbackHandler):
     def on_llm_new_token(self, token: str, **kwargs) -> None:
         global websocketClient
 
         websocketClient.send_text(f"Message text was: {token}")
 
-        print(f" token: {tokenWS}")
+        print(f" token: {token}")
 
     async def on_llm_start(
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any