fix: add type filter message

main.py

@@ -55,6 +55,20 @@ class MyCustomSyncHandler(BaseCallbackHandler):
         self.message += token
         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)
+
+    def on_llm_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> Any:
+        self.message = "end"
+        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)
+
 LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True)
 
 prompt, memory = get_prompt_template(promptTemplate_type="llama", history=True)
@@ -233,7 +247,7 @@ async def websocket_endpoint(websocket: WebSocket,  client_id: int):
             i  = 0
             for item in pubsub.listen():
                 i+=1
-                itemJson = str(item)
+
                 print(item["type"] == "message")
 
                 if(i > 10 ): pubsub.unsubscribe({client_id})