fix: add types

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 AsyncCallbackHandler, BaseCallbackHandler
+from langchain.callbacks.base import BaseCallbackHandler
 from langchain.schema import LLMResult
 
 # from langchain.embeddings import HuggingFaceEmbeddings
@@ -31,7 +31,7 @@ class Predict(BaseModel):
 class Delete(BaseModel):
     filename: str
 
-class MyCustomAsyncHandler(AsyncCallbackHandler):
+class MyCustomHandler(BaseCallbackHandler):
     def on_llm_new_token(self, token: str, **kwargs) -> None:
         print(f" token: {token}")
 
@@ -44,6 +44,19 @@ class MyCustomAsyncHandler(AsyncCallbackHandler):
     async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         print("finish")
 
+class CustomHandler(BaseCallbackHandler):
+    def on_llm_new_token(self, token: str, **kwargs) -> None:
+        print(f" CustomHandler: {token}")
+
+    async def on_llm_start(
+        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+    ) -> None:
+        class_name = serialized["name"]
+        print("CustomHandler start")
+
+    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        print("CustomHandler finish")
+
 # if torch.backends.mps.is_available():
 #     DEVICE_TYPE = "mps"
 # elif torch.cuda.is_available():
@@ -65,7 +78,7 @@ DB = Chroma(
 
 RETRIEVER = DB.as_retriever()
 
-LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True, callbacks=[MyCustomAsyncHandler()])
+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.
@@ -87,7 +100,8 @@ QA = RetrievalQA.from_chain_type(
     return_source_documents=SHOW_SOURCES,
     chain_type_kwargs={
         "prompt": QA_CHAIN_PROMPT,
-        "memory": memory
+        "memory": memory,
+        "callbacks": [CustomHandler()]
     },
 )
 
@@ -179,8 +193,6 @@ async def predict(data: Predict):
     if user_prompt:
         res = QA(user_prompt)
 
-        print(res)
-
         answer, docs = res["result"], res["source_documents"]
 
         prompt_response_dict = {
@@ -194,17 +206,6 @@ async def predict(data: Predict):
                 (os.path.basename(str(document.metadata["source"])), str(document.page_content))
             )
 
-        qa_chain_response = res.stream(
-            {"query": user_prompt},
-        )
-
-        print(f"{qa_chain_response} stream")
-
-        # generated_text = ""
-        # for new_text in STREAMER:
-        #     generated_text += new_text
-        #     print(generated_text)
-
         return {"response": prompt_response_dict}
     else:
         raise HTTPException(status_code=400, detail="Prompt Incorrect")
@@ -254,4 +255,13 @@ async def websocket_endpoint(websocket: WebSocket):
     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")
+
         await websocket.send_text(f"Message text was: {data}")