fix: add tokenWs

main.py

@@ -14,7 +14,7 @@ from langchain.chains import RetrievalQA
 from langchain.embeddings import HuggingFaceInstructEmbeddings
 from langchain.prompts import PromptTemplate
 from langchain.memory import ConversationBufferMemory
-from langchain.callbacks.base import BaseCallbackHandler
+from langchain.callbacks.base import BaseCallbackHandler, AsyncCallbackHandler
 from langchain.schema import LLMResult
 
 # from langchain.embeddings import HuggingFaceEmbeddings
@@ -35,9 +35,9 @@ tokenWS = ''
 
 class MyCustomHandler(BaseCallbackHandler):
     def on_llm_new_token(self, token: str, **kwargs) -> None:
-        global tokenWS
+        global websocketClient
 
-        tokenWS = token
+        websocketClient.send_text(f"Message text was: {tokenWS}")
 
         print(f" token: {tokenWS}")
 
@@ -73,11 +73,14 @@ RETRIEVER = DB.as_retriever()
 
 LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True, callbacks=[MyCustomHandler()])
 
-template = """you are a helpful, respectful and honest assistant. When answering questions, you should only use the documents provided.
-You should only answer the topics that appear in these documents.
-Always answer in the most helpful and reliable way possible, if you don't know the answer to a question, just say you don't know, don't try to make up an answer,
-don't share false information. you should use no more than 15 sentences and all your answers should be as concise as possible.
-Always say "Thank you for asking!" at the end of your answer.
+template = """you are a helpful, respectful and honest assistant.
+Your name is Katara llma. You should only use the source documents provided 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 don't know, don't try to make up an answer, don't share false information.
+Use 15 sentences maximum. Keep the answer as concise as possible.
+Always say "thanks for asking!" at the end of the answer.
 Context: {history} \n {context}
 Question: {question}
 """
@@ -245,11 +248,15 @@ 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;
+
         data = await websocket.receive_text()
 
-        res = QA(data)
+        QA(data)
 
-        await websocket.send_text(f"Message text was: {tokenWS}")
+        await websocket.send_text(f"Message text was:")