scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

fix(vertex_ai): preserve calculate_usage signature and type grounding check

Browse source
This commit is contained in:
ec 2026-10-20 17:59:06 +08:00
parent 901e741493
commit 9d40381e6a
2 changed files with 34 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

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

29
tests/unit/test_gemini_4_argon_model_metadata.py
View file

@ -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