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

langgenius · Nov 30, 2024 · 594666e · 594666e
1 parent e80f41a
commit 594666e
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/api/core/model_runtime/model_providers/google/llm/llm.py b/api/core/model_runtime/model_providers/google/llm/llm.py
@@ -254,8 +254,12 @@ def _handle_generate_response(
         assistant_prompt_message = AssistantPromptMessage(content=response.text)
 
         # calculate num tokens
-        prompt_tokens = self.get_num_tokens(model, credentials, prompt_messages)
-        completion_tokens = self.get_num_tokens(model, credentials, [assistant_prompt_message])
+        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])
 
         # transform usage
         usage = self._calc_response_usage(model, credentials, prompt_tokens, completion_tokens)