Explorar o código

fix: stop completion response not save to db (#351)

John Wang %!s(int64=2) %!d(string=hai) anos
pai
achega
a11f36ca60
Modificáronse 1 ficheiros con 5 adicións e 1 borrados
  1. 5 1
      api/core/callback_handler/llm_callback_handler.py

+ 5 - 1
api/core/callback_handler/llm_callback_handler.py

@@ -75,7 +75,11 @@ class LLMCallbackHandler(BaseCallbackHandler):
         self.conversation_message_task.save_message(self.llm_message)
 
     def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        self.conversation_message_task.append_message_text(token)
+        try:
+            self.conversation_message_task.append_message_text(token)
+        except ConversationTaskStoppedException as ex:
+            self.on_llm_error(error=ex)
+
         self.llm_message.completion += token
 
     def on_llm_error(