fix: add tokenWs

main.py

@@ -31,9 +31,15 @@ class Predict(BaseModel):
 class Delete(BaseModel):
     filename: str
 
+tokenWS = ''
+
 class MyCustomHandler(BaseCallbackHandler):
     def on_llm_new_token(self, token: str, **kwargs) -> None:
-        print(f" token: {token}")
+        global tokenWS
+
+        tokenWS = token
+
+        print(f" token: {tokenWS}")
 
     async def on_llm_start(
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
@@ -87,7 +93,6 @@ QA = RetrievalQA.from_chain_type(
     return_source_documents=SHOW_SOURCES,
     chain_type_kwargs={
         "prompt": QA_CHAIN_PROMPT,
-        "memory": memory,
     },
 )
 
@@ -238,16 +243,12 @@ async def create_upload_file(file: UploadFile):
 
 @api_app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
+    global QA
+
     await websocket.accept()
     while True:
         data = await websocket.receive_text()
 
-        res = QA(data)
-
-        qa_chain_response = res.stream(
-            {"query": data},
-        )
-
-        print(f"{qa_chain_response} stream")
+        QA(data)
 
-        await websocket.send_text(f"Message text was: {data}")
+        await websocket.send_text(f"Message text was: {tokenWS}")