fix: use Gemini response metadata for token counting (#11226)

1 year ago · 594666eb61
parent e80f41a701
commit 594666eb61
1 changed files with 6 additions and 2 deletions
--- a/api/core/model_runtime/model_providers/google/llm/llm.py
+++ b/api/core/model_runtime/model_providers/google/llm/llm.py
@ -254,6 +254,10 @@ class GoogleLargeLanguageModel(LargeLanguageModel):
        assistant_prompt_message = AssistantPromptMessage(content=response.text)

        # calculate num tokens
+        if response.usage_metadata:
+            prompt_tokens = response.usage_metadata.prompt_token_count
+            completion_tokens = response.usage_metadata.candidates_token_count
+        else:
            prompt_tokens = self.get_num_tokens(model, credentials, prompt_messages)
            completion_tokens = self.get_num_tokens(model, credentials, [assistant_prompt_message])