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>

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)

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):

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:

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")