fix: anthropic completion error in blocking mode (#591)

eff11526 · John Wang · GitHub · 07cde4f8 · eff11526
Unverified Commit eff11526 authored Jul 18, 2023 by John Wang Committed by GitHub Jul 18, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 3 deletions

llm_callback_handler.py api/core/callback_handler/llm_callback_handler.py +2 -3

No files found.
--- 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)