From eff115267f67185c2dbbc9c3b4857ab28d51a189 Mon Sep 17 00:00:00 2001 From: John Wang Date: Tue, 18 Jul 2023 15:12:52 +0800 Subject: [PATCH] fix: anthropic completion error in blocking mode (#591) --- api/core/callback_handler/llm_callback_handler.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/api/core/callback_handler/llm_callback_handler.py b/api/core/callback_handler/llm_callback_handler.py index 6b3daf76e7..1369be710c 100644 --- a/api/core/callback_handler/llm_callback_handler.py +++ b/api/core/callback_handler/llm_callback_handler.py @@ -69,9 +69,8 @@ class LLMCallbackHandler(BaseCallbackHandler): if not self.conversation_message_task.streaming: self.conversation_message_task.append_message_text(response.generations[0][0].text) self.llm_message.completion = response.generations[0][0].text - self.llm_message.completion_tokens = response.llm_output['token_usage']['completion_tokens'] - else: - self.llm_message.completion_tokens = self.llm.get_num_tokens(self.llm_message.completion) + + self.llm_message.completion_tokens = self.llm.get_num_tokens(self.llm_message.completion) self.conversation_message_task.save_message(self.llm_message)