fix: add type Union

main.py

@@ -43,33 +43,34 @@ EMBEDDINGS = HuggingFaceInstructEmbeddings(model_name=EMBEDDING_MODEL_NAME, mode
 DB = Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=EMBEDDINGS, client_settings=CHROMA_SETTINGS)
 RETRIEVER = DB.as_retriever()
 
-
 redisClient = redis.Redis(host='localhost', port=6379, db=0)
 
-
 class MyCustomSyncHandler(BaseCallbackHandler):
-    def __init__(self):
+    def __init__(self, redisClient):
         self.message = ''
+        self.redisClient = redisClient
 
     def on_llm_new_token(self, token: str, **kwargs) -> Any:
         self.message += token
-        redisClient.publish(f'{kwargs["tags"][0]}',  self.message)
+        self.redisClient.publish(f'{kwargs["tags"][0]}',  self.message)
 
     def on_llm_end(self, response: LLMResult, **kwargs: Any) -> Any:
         self.message = "end"
-        redisClient.publish(f'{kwargs["tags"][0]}',  self.message)
+        self.redisClient.publish(f'{kwargs["tags"][0]}',  self.message)
 
     def on_llm_error(
         self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
     ) -> Any:
         self.message = "end"
-        redisClient.publish(f'{kwargs["tags"][0]}',  self.message)
+        self.redisClient.publish(f'{kwargs["tags"][0]}',  self.message)
 
     def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> Any:
         self.message = "end"
-        redisClient.publish(f'{kwargs["tags"][0]}',  self.message)
+        self.redisClient.publish(f'{kwargs["tags"][0]}',  self.message)
 
-LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True)
+handleCallback = MyCustomSyncHandler(redisClient)
+
+LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True, callbacks=[handleCallback])
 
 prompt, memory = get_prompt_template(promptTemplate_type="llama", history=True)
 
@@ -84,7 +85,6 @@ QA = RetrievalQA.from_chain_type(
     },
 )
 
-
 app = FastAPI(title="homepage-app")
 api_app = FastAPI(title="api app")
 
@@ -242,7 +242,7 @@ async def websocket_endpoint(websocket: WebSocket,  client_id: int):
             pubsub.subscribe(f'{client_id}')
 
             with concurrent.futures.ThreadPoolExecutor() as executor:
-               executor.submit(QA(inputs=prompt, return_only_outputs=True, callbacks=[MyCustomSyncHandler()], tags=f'{client_id}', include_run_info=True))
+               executor.submit(QA(inputs=prompt, return_only_outputs=True, tags=f'{client_id}', include_run_info=True))
 
             i  = 0
             for item in pubsub.listen():