From 840ad6191db43259ba4a3a8303f3f85d016c6531 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 16 Jul 2026 15:00:04 +0000 Subject: [PATCH] fix(vertex_ai): surface Gemini toolUsePromptTokenCount in usage Grounded/tool-use Gemini requests report toolUsePromptTokenCount as its own slice of totalTokenCount. LiteLLM dropped the field entirely, so total_tokens had a silent unexplained gap and is_candidate_token_count_inclusive could double-count reasoning tokens. Surface it as prompt_tokens_details.tool_use_prompt_tokens and fold it into the inclusivity check. prompt_tokens and cost are left unchanged so grounding stays billed via its flat fee. Closes #33198 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../vertex_and_google_ai_studio_gemini.py | 16 +++- litellm/types/llms/vertex_ai.py | 2 + litellm/types/utils.py | 5 + ...test_vertex_and_google_ai_studio_gemini.py | 95 +++++++++++++++++++ 4 files changed, 114 insertions(+), 4 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 8c4bb1aa0c5..21aef9c6573 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 @@ -1731,15 +1731,20 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): """ Check if the candidate token count is inclusive of the thinking token count - if prompttokencount + candidatesTokenCount == totalTokenCount, then the candidate token count is inclusive of the thinking token count + if promptTokenCount + toolUsePromptTokenCount + candidatesTokenCount == totalTokenCount, then the candidate token count is inclusive of the thinking token count else the candidate token count is exclusive of the thinking token count + toolUsePromptTokenCount is folded into the prompt side of the equality because grounded / tool-use requests report it as a separate slice of totalTokenCount; ignoring it made this always return False and double-count reasoning tokens (https://github.com/BerriAI/litellm/discussions/33198) + Addresses - https://github.com/BerriAI/litellm/pull/10141#discussion_r2052272035 """ - if usage_metadata.get("promptTokenCount", 0) + usage_metadata.get( - "candidatesTokenCount", 0 - ) == usage_metadata.get("totalTokenCount", 0): + effective_prompt_tokens = usage_metadata.get("promptTokenCount", 0) + usage_metadata.get( + "toolUsePromptTokenCount", 0 + ) + if effective_prompt_tokens + usage_metadata.get("candidatesTokenCount", 0) == usage_metadata.get( + "totalTokenCount", 0 + ): return True else: return False @@ -1888,12 +1893,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): response_tokens_details = CompletionTokensDetailsWrapper() response_tokens_details.reasoning_tokens = reasoning_tokens + tool_use_prompt_tokens = usage_metadata.get("toolUsePromptTokenCount") + prompt_tokens_details = PromptTokensDetailsWrapper( cached_tokens=cached_tokens, audio_tokens=prompt_audio_tokens, text_tokens=prompt_text_tokens, image_tokens=prompt_image_tokens, video_tokens=prompt_video_tokens, + tool_use_prompt_tokens=tool_use_prompt_tokens, ) completion_tokens = response_tokens or completion_response["usageMetadata"].get("candidatesTokenCount", 0) diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 64a06825773..f1cb8ee914d 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -304,6 +304,8 @@ class UsageMetadata(TypedDict, total=False): thoughtsTokenCount: int responseTokensDetails: List[PromptTokensDetails] candidatesTokensDetails: List[PromptTokensDetails] # Alternative key name used in some responses + toolUsePromptTokenCount: int + toolUsePromptTokensDetails: List[PromptTokensDetails] class TokenCountDetailsResponse(TypedDict): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6487a8aa33f..e5e78dd68d9 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1474,6 +1474,9 @@ class PromptTokensDetailsWrapper( web_search_requests: Optional[int] = None """Number of web search requests made by the tool call. Used for Anthropic to calculate web search cost.""" + tool_use_prompt_tokens: Optional[int] = None + """Input tokens consumed by server-side tool-use prompts (e.g. Gemini grounding / code execution). Surfaced separately so they don't silently inflate text prompt tokens or cost.""" + character_count: Optional[int] = None """Character count sent to the model. Used for Vertex AI multimodal embeddings.""" @@ -1504,6 +1507,8 @@ class PromptTokensDetailsWrapper( del self.audio_length_seconds if self.web_search_requests is None: del self.web_search_requests + if self.tool_use_prompt_tokens is None: + del self.tool_use_prompt_tokens if self.cache_creation_tokens is None: del self.cache_creation_tokens if self.cache_creation_token_details is None: diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 40f9f4e7910..6099aadc932 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -5254,3 +5254,98 @@ def test_process_candidates_merges_thought_signatures_and_server_side_tools(): fields = model_response.choices[-1].message.provider_specific_fields assert fields["thought_signatures"] == ["sig-text"] assert fields["server_side_tool_invocations"][0]["id"] == "tool-1" + + +def test_gemini_grounded_tool_use_prompt_tokens_surfaced(): + """Regression for https://github.com/BerriAI/litellm/discussions/33198 + + Grounded Gemini requests report toolUsePromptTokenCount as its own slice of + totalTokenCount. Before the fix this field was dropped entirely, so + total_tokens had a silent, unexplained gap and reasoning tokens were at risk + of being double counted. The exact numbers below are from the reported request. + """ + v = VertexGeminiConfig() + + usage_metadata_dict = { + "promptTokenCount": 4647, + "candidatesTokenCount": 1495, + "thoughtsTokenCount": 10785, + "toolUsePromptTokenCount": 12499, + "totalTokenCount": 29426, + } + + completion_response = {"usageMetadata": usage_metadata_dict} + result = v._calculate_usage(completion_response=completion_response) + + assert result.prompt_tokens == 4647 + # candidatesTokenCount (1495) is exclusive of thoughts, so reasoning is added once + assert result.completion_tokens == 1495 + 10785 + assert result.total_tokens == 29426 + assert result.completion_tokens_details.reasoning_tokens == 10785 + + # the previously missing slice is now visible and accounts for the whole gap + assert result.prompt_tokens_details is not None + assert result.prompt_tokens_details.tool_use_prompt_tokens == 12499 + gap = result.total_tokens - result.prompt_tokens - result.completion_tokens + assert gap == result.prompt_tokens_details.tool_use_prompt_tokens + + +def test_is_candidate_token_count_inclusive_with_tool_use_prompt_tokens(): + """toolUsePromptTokenCount must be folded into the prompt side of the equality. + + When candidatesTokenCount is inclusive of thoughts, totalTokenCount is + promptTokenCount + toolUsePromptTokenCount + candidatesTokenCount. Ignoring + the tool-use slice made this return False and double-count reasoning tokens. + """ + inclusive_with_tool_use: UsageMetadata = { + "promptTokenCount": 100, + "candidatesTokenCount": 60, # already includes the 40 thoughts + "toolUsePromptTokenCount": 500, + "totalTokenCount": 660, + } + assert VertexGeminiConfig.is_candidate_token_count_inclusive(inclusive_with_tool_use) is True + + exclusive_with_tool_use: UsageMetadata = { + "promptTokenCount": 100, + "candidatesTokenCount": 20, + "thoughtsTokenCount": 40, + "toolUsePromptTokenCount": 500, + "totalTokenCount": 660, + } + assert VertexGeminiConfig.is_candidate_token_count_inclusive(exclusive_with_tool_use) is False + + +def test_gemini_inclusive_candidates_with_tool_use_no_double_count(): + """When candidates already includes thoughts and tool-use tokens are present, + reasoning tokens must not be added a second time.""" + v = VertexGeminiConfig() + + usage_metadata_dict = { + "promptTokenCount": 100, + "candidatesTokenCount": 60, # inclusive of the 40 thoughts + "thoughtsTokenCount": 40, + "toolUsePromptTokenCount": 500, + "totalTokenCount": 660, + } + + completion_response = {"usageMetadata": usage_metadata_dict} + result = v._calculate_usage(completion_response=completion_response) + + assert result.completion_tokens == 60 + assert result.prompt_tokens_details.tool_use_prompt_tokens == 500 + + +def test_gemini_usage_without_tool_use_prompt_tokens_omits_field(): + """Non-grounded requests must not gain a tool_use_prompt_tokens field.""" + v = VertexGeminiConfig() + + usage_metadata_dict = { + "promptTokenCount": 100, + "candidatesTokenCount": 50, + "totalTokenCount": 150, + } + + completion_response = {"usageMetadata": usage_metadata_dict} + result = v._calculate_usage(completion_response=completion_response) + + assert not hasattr(result.prompt_tokens_details, "tool_use_prompt_tokens")