feat: add backend

main.py

@@ -74,8 +74,7 @@ class MyCustomSyncHandler(BaseCallbackHandler):
 
 handlerToken = MyCustomSyncHandler()
 
-
-LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True, callbacks=[handlerToken])
+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.
@@ -97,6 +96,7 @@ QA = RetrievalQA.from_chain_type(
     return_source_documents=SHOW_SOURCES,
     chain_type_kwargs={
         "prompt": QA_CHAIN_PROMPT,
+        "callbacks": [handlerToken]
     },
 )
 
@@ -253,14 +253,9 @@ async def websocket_endpoint(websocket: WebSocket):
     while True:
         data = await websocket.receive_text()
 
-        QA(data)
+        res = QA(data)
+
+        await websocket.send_text(f"result: {res}")
 
-        finish = False
-        while finish == False:
-            finish = handlerToken.end
-            token = handlerToken.token
-            print(token)
-            print(finish)
-            await websocket.send_text(f"result: {token}")