From f07e37ce6973664d9f9f93c678bfc0cc33840ec9 Mon Sep 17 00:00:00 2001 From: unknown Date: Thu, 26 Mar 2026 14:28:38 +0500 Subject: [PATCH] fix(vertex_ai): fix toolUsePromptTokenCount in blocked/content-policy response paths --- .../vertex_and_google_ai_studio_gemini.py | 24 ++++++++----------- litellm/types/llms/vertex_ai.py | 4 +++- 2 files changed, 13 insertions(+), 15 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 291d88890f1..a701d4ac30d 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1603,14 +1603,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): model_response.choices = [choice] ## GET USAGE ## + _usage_metadata = completion_response["usageMetadata"] usage = Usage( - prompt_tokens=completion_response["usageMetadata"].get( - "promptTokenCount", 0 - ), - completion_tokens=completion_response["usageMetadata"].get( - "candidatesTokenCount", 0 - ), - total_tokens=completion_response["usageMetadata"].get("totalTokenCount", 0), + prompt_tokens=_usage_metadata.get("promptTokenCount", 0) + + _usage_metadata.get("toolUsePromptTokenCount", 0), + completion_tokens=_usage_metadata.get("candidatesTokenCount", 0), + total_tokens=_usage_metadata.get("totalTokenCount", 0), ) setattr(model_response, "usage", usage) @@ -1641,14 +1639,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): model_response.choices = [choice] ## GET USAGE ## + _usage_metadata = completion_response["usageMetadata"] usage = Usage( - prompt_tokens=completion_response["usageMetadata"].get( - "promptTokenCount", 0 - ), - completion_tokens=completion_response["usageMetadata"].get( - "candidatesTokenCount", 0 - ), - total_tokens=completion_response["usageMetadata"].get("totalTokenCount", 0), + prompt_tokens=_usage_metadata.get("promptTokenCount", 0) + + _usage_metadata.get("toolUsePromptTokenCount", 0), + completion_tokens=_usage_metadata.get("candidatesTokenCount", 0), + total_tokens=_usage_metadata.get("totalTokenCount", 0), ) setattr(model_response, "usage", usage) diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 7e5f4704828..cd523f784ec 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -287,7 +287,9 @@ class UsageMetadata(TypedDict, total=False): PromptTokensDetails ] # Alternative key name used in some responses toolUsePromptTokenCount: int - toolUsePromptTokensDetails: List[PromptTokensDetails] + toolUsePromptTokensDetails: List[ + PromptTokensDetails + ] # TODO: parse into prompt_tokens_details per modality class TokenCountDetailsResponse(TypedDict):