Skip to content

Commit 594666e

Browse files
authored
fix: use Gemini response metadata for token counting (#11226)
1 parent e80f41a commit 594666e

File tree

1 file changed

+6
-2
lines changed
  • api/core/model_runtime/model_providers/google/llm

1 file changed

+6
-2
lines changed

api/core/model_runtime/model_providers/google/llm/llm.py

+6-2
Original file line numberDiff line numberDiff line change
@@ -254,8 +254,12 @@ def _handle_generate_response(
254254
assistant_prompt_message = AssistantPromptMessage(content=response.text)
255255

256256
# calculate num tokens
257-
prompt_tokens = self.get_num_tokens(model, credentials, prompt_messages)
258-
completion_tokens = self.get_num_tokens(model, credentials, [assistant_prompt_message])
257+
if response.usage_metadata:
258+
prompt_tokens = response.usage_metadata.prompt_token_count
259+
completion_tokens = response.usage_metadata.candidates_token_count
260+
else:
261+
prompt_tokens = self.get_num_tokens(model, credentials, prompt_messages)
262+
completion_tokens = self.get_num_tokens(model, credentials, [assistant_prompt_message])
259263

260264
# transform usage
261265
usage = self._calc_response_usage(model, credentials, prompt_tokens, completion_tokens)

0 commit comments

Comments
 (0)