From 754c9c06e9e3a12833f00f1f5b4df9ad0af4e32d Mon Sep 17 00:00:00 2001
From: unknown <syedshahmeerali196@gmail.com>
Date: Thu, 26 Mar 2026 13:59:09 +0500
Subject: [PATCH] fix(vertex_ai): include toolUsePromptTokenCount in Gemini
 usage calculation

---
 .../vertex_and_google_ai_studio_gemini.py     | 10 ++++-
 litellm/types/llms/vertex_ai.py               |  2 +
 ...test_vertex_and_google_ai_studio_gemini.py | 42 +++++++++++++++++++
 3 files changed, 52 insertions(+), 2 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 36f51c5b2f5..825ebcec014 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
@@ -1666,7 +1666,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
 
         Addresses - https://github.com/BerriAI/litellm/pull/10141#discussion_r2052272035
         """
-        if usage_metadata.get("promptTokenCount", 0) + usage_metadata.get(
+        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
@@ -1848,8 +1851,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
         ):
             completion_tokens = reasoning_tokens + completion_tokens
         ## GET USAGE ##
+        prompt_tokens = usage_metadata.get("promptTokenCount", 0) + usage_metadata.get(
+            "toolUsePromptTokenCount", 0
+        )
         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,
diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py
index c49fc96a65b..7e5f4704828 100644
--- a/litellm/types/llms/vertex_ai.py
+++ b/litellm/types/llms/vertex_ai.py
@@ -286,6 +286,8 @@ class UsageMetadata(TypedDict, total=False):
     candidatesTokensDetails: List[
         PromptTokensDetails
     ]  # Alternative key name used in some responses
+    toolUsePromptTokenCount: int
+    toolUsePromptTokensDetails: List[PromptTokensDetails]
 
 
 class TokenCountDetailsResponse(TypedDict):
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 3102a695961..c85ef3ae808 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
@@ -1338,6 +1338,48 @@ def test_vertex_ai_usage_metadata_missing_token_count():
     )  # Default value for missing tokenCount
 
 
+def test_calculate_usage_with_tool_use_prompt_token_count():
+    """
+    Regression test for https://github.com/BerriAI/litellm/issues/23731.
+
+    When Gemini uses built-in tools (e.g. code execution), `usageMetadata` contains
+    `toolUsePromptTokenCount` for the tool system-prompt tokens.  These tokens must be:
+      1. Added to `prompt_tokens` in the returned Usage object.
+      2. Included in `is_candidate_token_count_inclusive` so that the inclusive/exclusive
+         check isn't thrown off by the extra token bucket, which previously caused
+         `thoughtsTokenCount` to be double-counted into `completion_tokens`.
+    """
+    from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
+        VertexGeminiConfig,
+    )
+    from litellm.types.llms.vertex_ai import UsageMetadata
+
+    v = VertexGeminiConfig()
+
+    # Gemini 2.5 Flash with code-execution tool enabled.
+    # candidatesTokenCount=28 is exclusive of thoughtsTokenCount=255, so total is:
+    # 123 (prompt) + 248 (tool prompt) + 255 (thoughts) + 28 (actual output) = 654
+    usage_metadata = UsageMetadata(
+        promptTokenCount=123,
+        toolUsePromptTokenCount=248,
+        candidatesTokenCount=28,
+        thoughtsTokenCount=255,
+        totalTokenCount=654,
+    )
+
+    # is_candidate_token_count_inclusive must return False (thoughts are separate)
+    assert VertexGeminiConfig.is_candidate_token_count_inclusive(usage_metadata) is False
+
+    result = v._calculate_usage(completion_response={"usageMetadata": usage_metadata})
+
+    # prompt_tokens must include toolUsePromptTokenCount
+    assert result.prompt_tokens == 123 + 248
+    # completion_tokens = candidatesTokenCount + thoughtsTokenCount (exclusive case)
+    assert result.completion_tokens == 28 + 255
+    assert result.total_tokens == 654
+    assert result.completion_tokens_details.reasoning_tokens == 255
+
+
 def test_vertex_ai_process_candidates_with_grounding_metadata():
     from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
         VertexGeminiConfig,