fix: stop completion response not save to db

2025-12-06 19:42:42 +08:00 · 2023-06-13 12:52:52 +08:00
1 changed files with 5 additions and 1 deletions
--- a/api/core/callback_handler/llm_callback_handler.py
+++ b/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(