fix(vertex_ai): surface Gemini toolUsePromptTokenCount so grounded token counts reconcile

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py

@@ -1731,7 +1731,7 @@ 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 + candidatesTokenCount + toolUsePromptTokenCount == 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
 
@@ -1739,7 +1739,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
         """
         if usage_metadata.get("promptTokenCount", 0) + usage_metadata.get(
             "candidatesTokenCount", 0
-        ) == usage_metadata.get("totalTokenCount", 0):
+        ) + usage_metadata.get("toolUsePromptTokenCount", 0) == usage_metadata.get("totalTokenCount", 0):
             return True
         else:
             return False
@@ -1888,20 +1888,24 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
                 response_tokens_details = CompletionTokensDetailsWrapper()
             response_tokens_details.reasoning_tokens = reasoning_tokens
 
+        tool_use_prompt_tokens: Optional[int] = 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)
         if not VertexGeminiConfig.is_candidate_token_count_inclusive(usage_metadata) and reasoning_tokens:
             completion_tokens = reasoning_tokens + completion_tokens
+        prompt_tokens = usage_metadata.get("promptTokenCount", 0) + (tool_use_prompt_tokens or 0)
         ## GET USAGE ##
         usage = Usage(
-            prompt_tokens=usage_metadata.get("promptTokenCount", 0),
+            prompt_tokens=prompt_tokens,
             completion_tokens=completion_tokens,
             total_tokens=usage_metadata.get("totalTokenCount", 0),
             prompt_tokens_details=prompt_tokens_details,

litellm/types/llms/vertex_ai.py

@@ -299,8 +299,10 @@ class UsageMetadata(TypedDict, total=False):
     candidatesTokenCount: int
     responseTokenCount: int
     cachedContentTokenCount: int
+    toolUsePromptTokenCount: int
     promptTokensDetails: List[PromptTokensDetails]
     cacheTokensDetails: List[PromptTokensDetails]
+    toolUsePromptTokensDetails: List[PromptTokensDetails]
     thoughtsTokenCount: int
     responseTokensDetails: List[PromptTokensDetails]
     candidatesTokensDetails: List[PromptTokensDetails]  # Alternative key name used in some responses

litellm/types/utils.py

@@ -1492,6 +1492,9 @@ class PromptTokensDetailsWrapper(
     cache_creation_token_details: Optional[CacheCreationTokenDetails] = None
     """Details of cache creation tokens sent to the model. Used for tracking 5m/1h cache creation tokens for Anthropic prompt caching."""
 
+    tool_use_prompt_tokens: Optional[int] = None
+    """Number of tokens present in tool-use prompt(s). Used for Gemini grounded requests so prompt_tokens + completion_tokens reconciles with total_tokens."""
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         if self.character_count is None:
@@ -1508,6 +1511,8 @@ class PromptTokensDetailsWrapper(
             del self.cache_creation_tokens
         if self.cache_creation_token_details is None:
             del self.cache_creation_token_details
+        if self.tool_use_prompt_tokens is None:
+            del self.tool_use_prompt_tokens
 
 
 class ServerToolUse(BaseModel):

tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py

@@ -3508,6 +3508,61 @@ def test_gemini_token_usage_standard_response():
     assert result.completion_tokens_details.image_tokens == 10
 
 
+def test_gemini_grounded_request_tool_use_prompt_tokens():
+    """
+    Regression test for https://github.com/BerriAI/litellm/issues/33530
+
+    Grounded Gemini requests (googleSearch) report toolUsePromptTokenCount, which
+    Gemini folds into totalTokenCount but LiteLLM previously dropped. This caused
+    prompt_tokens + completion_tokens != total_tokens. The grounding tokens must be
+    surfaced on prompt_tokens_details.tool_use_prompt_tokens and folded into
+    prompt_tokens so the accounting reconciles.
+    """
+    v = VertexGeminiConfig()
+
+    usage_metadata_dict = {
+        "promptTokenCount": 100,
+        "candidatesTokenCount": 50,
+        "toolUsePromptTokenCount": 30,
+        "totalTokenCount": 180,
+    }
+
+    completion_response = {"usageMetadata": usage_metadata_dict}
+    result = v._calculate_usage(completion_response=completion_response)
+
+    assert result.prompt_tokens == 130
+    assert result.completion_tokens == 50
+    assert result.total_tokens == 180
+    assert result.prompt_tokens + result.completion_tokens == result.total_tokens
+    assert result.prompt_tokens_details is not None
+    assert result.prompt_tokens_details.tool_use_prompt_tokens == 30
+
+
+def test_gemini_grounded_request_with_thinking_no_double_count():
+    """
+    With grounding active, is_candidate_token_count_inclusive must account for
+    toolUsePromptTokenCount; otherwise it wrongly treats an inclusive candidate
+    count as exclusive and double-counts thoughtsTokenCount into completion_tokens.
+    """
+    v = VertexGeminiConfig()
+
+    usage_metadata_dict = {
+        "promptTokenCount": 100,
+        "candidatesTokenCount": 50,
+        "thoughtsTokenCount": 20,
+        "toolUsePromptTokenCount": 30,
+        "totalTokenCount": 180,
+    }
+
+    completion_response = {"usageMetadata": usage_metadata_dict}
+    result = v._calculate_usage(completion_response=completion_response)
+
+    assert result.prompt_tokens == 130
+    assert result.completion_tokens == 50
+    assert result.total_tokens == 180
+    assert result.prompt_tokens + result.completion_tokens == result.total_tokens
+
+
 def test_gemini_image_gen_usage_metadata_prompt_vs_completion_separation():
     """
     Test that image generation models correctly separate prompt and completion token details.