fix: add websocket in handlerToken

main.py

@@ -2,6 +2,7 @@ import os
 import glob
 import shutil
 import subprocess
+import sys
 
 from typing import Any, Dict, List
 
@@ -16,6 +17,9 @@ from langchain.embeddings import HuggingFaceInstructEmbeddings
 from langchain.prompts import PromptTemplate
 from langchain.memory import ConversationBufferMemory
 from langchain.callbacks.base import BaseCallbackHandler
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+
+
 from langchain.schema import LLMResult
 
 # from langchain.embeddings import HuggingFaceEmbeddings
@@ -55,7 +59,7 @@ DB = Chroma(
 
 RETRIEVER = DB.as_retriever()
 
-class MyCustomSyncHandler(BaseCallbackHandler):
+class MyCustomSyncHandler(StreamingStdOutCallbackHandler):
     def __init__(self):
         self.end = False
         self.websocket = None
@@ -74,7 +78,6 @@ class MyCustomSyncHandler(BaseCallbackHandler):
 
         print(token)
 
-        self.token += token
 
 
 handlerToken = MyCustomSyncHandler()
@@ -252,6 +255,7 @@ async def websocket_endpoint(websocket: WebSocket):
     global QA
 
     await websocket.accept()
+
     try:
         while True:
             handlerToken.websocket = websocket
@@ -261,7 +265,10 @@ async def websocket_endpoint(websocket: WebSocket):
             print(res)
 
     except WebSocketDisconnect:
-        await websocket.send_text(f"disconnect")
+        print('disconnect')
+    except RuntimeError as error:
+        print(error)
+