From 9d40381e6a5c346a9e4ae029c8607b3ee11febe4 Mon Sep 17 00:00:00 2001 From: ec Date: Tue, 20 Oct 2026 17:59:06 +0800 Subject: [PATCH] fix(vertex_ai): preserve calculate_usage signature and type grounding check --- .../vertex_and_google_ai_studio_gemini.py | 18 ++++-------- .../test_gemini_4_argon_model_metadata.py | 29 +++++++++++++++++++ 2 files changed, 34 insertions(+), 13 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 2a2be812591..037a9542f56 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 @@ -2275,8 +2275,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): @staticmethod def _response_bills_grounding_tool_use_tokens( - completion_response: GenerateContentResponseBody | BidiGenerateContentServerMessage, - model: str | None = None, + completion_response: Mapping[str, object], ) -> bool: """ Gemini 4+ bills every token the model processes, including internal tool system @@ -2284,16 +2283,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): per-query search grounding fee). Source: https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/guides/gemini-4-argon#migrate-to-4-argon-from-31-pro """ - resolved_model: Final = _served_model_name(completion_response.get("modelVersion")) or model - if not resolved_model: - return False - return VertexGeminiConfig._is_gemini_4_or_newer(resolved_model) + resolved_model: Final = _served_model_name(completion_response.get("modelVersion")) + return bool(resolved_model and VertexGeminiConfig._is_gemini_4_or_newer(resolved_model)) @staticmethod def _calculate_usage( completion_response: GenerateContentResponseBody | BidiGenerateContentServerMessage, cache_creation_tokens: int = 0, - model: str | None = None, ) -> Usage: if completion_response is not None and "usageMetadata" not in completion_response: raise ValueError(f"usageMetadata not found in completion_response. Got={completion_response}") @@ -2441,7 +2437,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): 0 if ( VertexGeminiConfig._response_has_search_grounding(completion_response) - and not VertexGeminiConfig._response_bills_grounding_tool_use_tokens(completion_response, model=model) + and not VertexGeminiConfig._response_bills_grounding_tool_use_tokens(completion_response) ) else (tool_use_prompt_tokens or 0) ) @@ -2482,9 +2478,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): cls, completion_response: GenerateContentResponseBody | BidiGenerateContentServerMessage, cache_creation_tokens: int = 0, - model: str | None = None, ) -> Usage: - return cls._calculate_usage(completion_response, cache_creation_tokens, model=model) + return cls._calculate_usage(completion_response, cache_creation_tokens) @staticmethod def _check_finish_reason( @@ -3087,7 +3082,6 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): usage: Final = VertexGeminiConfig._calculate_usage( completion_response=completion_response, cache_creation_tokens=_cache_creation_tokens(logging_obj), - model=model_response.model, ) VertexGeminiConfig._set_grounding_usage_counters(usage, grounding_metadata) @@ -4067,11 +4061,9 @@ class ModelResponseIterator: if "usageMetadata" not in processed_chunk: return None - logging_model: Final[object] = getattr(self.logging_obj, "model", None) usage: Final = VertexGeminiConfig.calculate_usage( completion_response=processed_chunk, cache_creation_tokens=_cache_creation_tokens(self.logging_obj), - model=model_response.model or (logging_model if isinstance(logging_model, str) else None), ) VertexGeminiConfig.set_grounding_usage_counters(usage, self._grounding_through_this_hop(grounding_metadata)) diff --git a/tests/unit/test_gemini_4_argon_model_metadata.py b/tests/unit/test_gemini_4_argon_model_metadata.py index abc779de267..0e388e04af5 100644 --- a/tests/unit/test_gemini_4_argon_model_metadata.py +++ b/tests/unit/test_gemini_4_argon_model_metadata.py @@ -176,3 +176,32 @@ def test_gemini_4_argon_grounding_bills_tool_use_prompt_tokens_and_search_querie assert usage.prompt_tokens_details.text_tokens == 50 + 300 assert usage.prompt_tokens_details.tool_use_tokens == 300 assert usage.prompt_tokens_details.web_search_requests == 1 + + pre_gemini_4_response = VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response( + completion_response={**completion_response, "modelVersion": "gemini-3.1-pro-preview"}, + model_response=ModelResponse(), + model="gemini-3.1-pro-preview", + logging_obj=MagicMock(optional_params={}), + raw_response=raw_response, + ) + assert pre_gemini_4_response.usage.prompt_tokens == 50 + assert pre_gemini_4_response.usage.prompt_tokens_details.text_tokens == 50 + assert pre_gemini_4_response.usage.prompt_tokens_details.tool_use_tokens == 300 + assert pre_gemini_4_response.usage.prompt_tokens_details.web_search_requests == 1 + + +@pytest.mark.parametrize( + ("completion_response", "expected"), + [ + ({"modelVersion": "gemini-4-argon"}, True), + ({"modelVersion": "gemini-4-argon-cyber-permissive"}, True), + ({"modelVersion": "gemini-3.1-pro-preview"}, False), + ({}, False), + ], +) +def test_response_bills_grounding_tool_use_tokens( + completion_response: dict[str, object], + expected: bool, +) -> None: + assert VertexGeminiConfig._response_bills_grounding_tool_use_tokens(completion_response) is expected +