fix: add websocket in handlerToken

main.py

@@ -5,7 +5,7 @@ import subprocess
 
 from typing import Any, Dict, List
 
-from fastapi import FastAPI, HTTPException, UploadFile, WebSocket
+from fastapi import FastAPI, HTTPException, UploadFile, WebSocket, WebSocketDisconnect
 from fastapi.staticfiles import StaticFiles
 
 from pydantic import BaseModel
@@ -58,7 +58,7 @@ RETRIEVER = DB.as_retriever()
 class MyCustomSyncHandler(BaseCallbackHandler):
     def __init__(self):
         self.end = False
-        self.token = ""
+        self.websocket = None
 
     def on_llm_start(
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
@@ -69,6 +69,11 @@ class MyCustomSyncHandler(BaseCallbackHandler):
         self.end = True
 
     def on_llm_new_token(self, token: str, **kwargs) -> Any:
+        if self.websocket != None:
+            self.websocket.send_text(token)
+
+        print(token)
+
         self.token += token
 
 
@@ -76,11 +81,8 @@ handlerToken = MyCustomSyncHandler()
 
 LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True, callbacks=[])
 
-template = """You are a helpful, respectful and honest assistant. You should only use the documents provided in the context to answer the questions.
-You should only respond only topics that contains in documents use to training. Use the following pieces of context to answer the question at the end.
-Always answer in the most helpful and safe way possible.
-If you don't know the answer to a question, just say that you "I don't know", never try to make up an answer and don't share false information.
-Use 15 sentences maximum. The answer must be as concise as possible. Always say "thanks for asking!" at the end of the answer.
+template = """You are a helpful, respectful and honest assistant.
+Always answer in the most helpful and safe way possible without trying to make up an answer, if you don't know the answer just say "I don't know" and don't share false information or topics that were not provided in your training. Use a maximum of 15 sentences. Your answer should be as concise and clear as possible. Always say "thank you for asking!" at the end of your answer.
 Context: {context}
 Question: {question}
 """
@@ -250,12 +252,16 @@ async def websocket_endpoint(websocket: WebSocket):
     global QA
 
     await websocket.accept()
-    while True:
-        data = await websocket.receive_text()
+    try:
+        while True:
+            handlerToken.websocket = websocket
 
-        res = QA(data)
+            data = await websocket.receive_text()
+            res = QA(data)
+            print(res)
 
-        await websocket.send_text(f"result: {res}")
+    except WebSocketDisconnect:
+        await websocket.send_text(f"disconnect")