feat: add backend

main.py

@@ -32,9 +32,7 @@ class Predict(BaseModel):
 class Delete(BaseModel):
     filename: str
 
-class MyCustomSyncHandler(BaseCallbackHandler):
-    def on_llm_new_token(self, token: str, **kwargs) -> None:
-        print(f"token: {token}")
+
 
 # if torch.backends.mps.is_available():
 #     DEVICE_TYPE = "mps"
@@ -57,7 +55,31 @@ DB = Chroma(
 
 RETRIEVER = DB.as_retriever()
 
-LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True, callbacks=[MyCustomSyncHandler()])
+
+machine = Machine(lump)
+solid = State('solid')
+liquid = State('liquid')
+
+class MyCustomSyncHandler(BaseCallbackHandler):
+    def __init__(self):
+        self.end = False
+
+    async def on_llm_start(
+        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+    ) -> None:
+        self.end = False
+
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        self.end = True
+
+    def on_llm_new_token(self, token: str, **kwargs) -> Any:
+        return token
+
+
+handlerToken = MyCustomSyncHandler()
+
+
+LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True, callbacks=[handlerToken])
 
 template = """you are a helpful, respectful and honest assistant. 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.
@@ -233,12 +255,14 @@ async def websocket_endpoint(websocket: WebSocket):
 
     await websocket.accept()
     while True:
-        websocket_state = websocket;
+        data = await websocket.receive_text()
 
-        data = await websocket_state.receive_text()
+        QA(data)
 
-        res = QA(data)
+        finish = False
+        while finish == False:
+            finish = handlerToken.end
+            token = handlerToken.on_llm_new_token()
+            await websocket.send_text(f"result: {token}")
 
-        print(f"${res}")
 
-        await websocket_state.send_text(f"Message text was:")