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..c62a37b3031 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,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, diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 64a06825773..df23de9e2e7 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/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 diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6487a8aa33f..e0bea83b237 100644 --- a/litellm/types/utils.py +++ b/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): 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..b563b916810 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 @@ -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.