From 901e74149368edd596529ce2ecf41257a60792e9 Mon Sep 17 00:00:00 2001 From: ec Date: Tue, 20 Oct 2026 13:40:28 +0800 Subject: [PATCH] fix(gemini,vertex_ai): gate candidate_count/penalties and bill grounding tokens on gemini-4-argon --- litellm/llms/gemini/chat/transformation.py | 7 +- .../vertex_and_google_ai_studio_gemini.py | 56 +++++++- .../test_gemini_4_argon_model_metadata.py | 136 ++++++++++++++++++ 3 files changed, 192 insertions(+), 7 deletions(-) diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py index 936919af83d..bf7a1349365 100644 --- a/litellm/llms/gemini/chat/transformation.py +++ b/litellm/llms/gemini/chat/transformation.py @@ -92,11 +92,8 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): "tool_choice", "functions", "response_format", - "n", "stop", "logprobs", - "frequency_penalty", - "presence_penalty", "seed", "modalities", "parallel_tool_calls", @@ -104,6 +101,10 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): "include_server_side_tool_invocations", "service_tier", ] + if self._supports_candidate_count(model): + supported_params.append("n") + if self._supports_penalty_parameters(model): + supported_params.extend(["frequency_penalty", "presence_penalty"]) if supports_reasoning(model, custom_llm_provider="gemini") or self._is_gemini_3_or_newer(model): supported_params.append("reasoning_effort") supported_params.append("thinking") 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 ce632bc0eef..2a2be812591 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 @@ -671,6 +671,27 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) -> bool: return cls._is_gemini_3_or_newer(model) + @staticmethod + def _is_gemini_4_or_newer(model: str) -> bool: + """ + Check if the model is Gemini 4 or newer. + """ + if not VertexGeminiConfig._is_gemini_3_or_newer(model): + return False + model_name: Final = model.split("/")[-1].lower() + return bool(re.match(r"^gemini-(?:[4-9]|\d{2,})(?:\.\d+)?(?:-|$)", model_name)) + + @classmethod + def is_gemini_4_or_newer( + cls, + model: str, + ) -> bool: + return cls._is_gemini_4_or_newer(model) + + def _supports_candidate_count(self, model: str) -> bool: + # Gemini 4+ models reject candidate_count (`n`) with HTTP 400 + return not VertexGeminiConfig._is_gemini_4_or_newer(model) + @staticmethod def _forward_gemini_function_call_id(model: str) -> bool: """ @@ -741,7 +762,6 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "functions", "tool_choice", "response_format", - "n", "stop", "extra_headers", "seed", @@ -755,6 +775,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "service_tier", ] + if self._supports_candidate_count(model): + supported_params.append("n") + # Add penalty parameters only for non-preview models if self._supports_penalty_parameters(model): supported_params.extend(["frequency_penalty", "presence_penalty"]) @@ -1637,7 +1660,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ): # sending stream = False, can cause it to get passed unchecked and raise issues optional_params["stream"] = value elif param == "n": - optional_params["candidate_count"] = value + if self._supports_candidate_count(model): + optional_params["candidate_count"] = value elif param == "audio" and isinstance(value, dict): optional_params["speechConfig"] = self._map_chat_audio_params(model, value, drop_params) elif param == "stop": @@ -2249,10 +2273,27 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return True return False + @staticmethod + def _response_bills_grounding_tool_use_tokens( + completion_response: GenerateContentResponseBody | BidiGenerateContentServerMessage, + model: str | None = None, + ) -> bool: + """ + Gemini 4+ bills every token the model processes, including internal tool system + instructions and function declarations when using grounding (in addition to the + 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) + @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}") @@ -2398,7 +2439,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): billable_tool_use_prompt_tokens: Final = ( 0 - if VertexGeminiConfig._response_has_search_grounding(completion_response) + if ( + VertexGeminiConfig._response_has_search_grounding(completion_response) + and not VertexGeminiConfig._response_bills_grounding_tool_use_tokens(completion_response, model=model) + ) else (tool_use_prompt_tokens or 0) ) billable_text_tokens: Final = ( @@ -2438,8 +2482,9 @@ 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) + return cls._calculate_usage(completion_response, cache_creation_tokens, model=model) @staticmethod def _check_finish_reason( @@ -3042,6 +3087,7 @@ 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) @@ -4021,9 +4067,11 @@ 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 d4c3a69ca20..abc779de267 100644 --- a/tests/unit/test_gemini_4_argon_model_metadata.py +++ b/tests/unit/test_gemini_4_argon_model_metadata.py @@ -1,12 +1,18 @@ from pathlib import Path from typing import Final +from unittest.mock import MagicMock import pytest from pydantic import TypeAdapter +import litellm from litellm import get_model_info +from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig +from litellm.types.llms.vertex_ai import UsageMetadata +from litellm.types.utils import ModelResponse REPO_ROOT: Final = Path(__file__).parents[2] MODELS: Final = ( @@ -32,6 +38,7 @@ def test_gemini_4_argon_routes_to_its_provider_on_the_thinking_level_path( assert info["litellm_provider"] == expected_provider assert info["max_output_tokens"] == info["max_input_tokens"] assert VertexGeminiConfig.is_gemini_3_or_newer(routed_model) is True + assert VertexGeminiConfig.is_gemini_4_or_newer(routed_model) is True @pytest.mark.parametrize("model", [row[0] for row in MODELS]) @@ -40,3 +47,132 @@ def test_gemini_4_argon_entry_source_and_backup_match(model: str) -> None: backup_entry = _cost_map_entry(REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json", model) assert backup_entry == main_entry + + +@pytest.mark.parametrize( + ("model", "expected"), + [ + ("gemini-4-argon", True), + ("gemini/gemini-4-argon", True), + ("vertex_ai/gemini-4-argon", True), + ("vertex_ai/gemini-4-argon-cyber-permissive", True), + ("gemini-3.1-pro-preview", False), + ("gemini-3.8-flash", False), + ("gemini-2.5-flash", False), + ("gemma-4-31b-it", False), + ], +) +def test_is_gemini_4_or_newer_classification(model: str, expected: bool) -> None: + assert VertexGeminiConfig.is_gemini_4_or_newer(model) is expected + + +@pytest.mark.parametrize( + ("config_cls", "model", "provider"), + [ + (VertexGeminiConfig, "gemini-4-argon", "vertex_ai"), + (VertexGeminiConfig, "gemini-4-argon-cyber-permissive", "vertex_ai"), + (GoogleAIStudioGeminiConfig, "gemini-4-argon", "gemini"), + ], +) +def test_gemini_4_argon_rejects_candidate_count_and_penalty_parameters( + config_cls: type[VertexGeminiConfig], + model: str, + provider: str, +) -> None: + config = config_cls() + supported = config.get_supported_openai_params(model) + assert "n" not in supported + assert "frequency_penalty" not in supported + assert "presence_penalty" not in supported + + mapped = config.map_openai_params( + non_default_params={ + "n": 2, + "frequency_penalty": 0.5, + "presence_penalty": 0.3, + "max_tokens": 256, + }, + optional_params={}, + model=model, + drop_params=False, + ) + assert "candidate_count" not in mapped + assert "frequency_penalty" not in mapped + assert "presence_penalty" not in mapped + assert mapped["max_output_tokens"] == 256 + + for unsupported_kwarg in ({"n": 2}, {"frequency_penalty": 0.5}, {"presence_penalty": 0.3}): + with pytest.raises(UnsupportedParamsError): + litellm.utils.get_optional_params( + model=model, + custom_llm_provider=provider, + drop_params=False, + **unsupported_kwarg, + ) + + +@pytest.mark.parametrize("config_cls", [VertexGeminiConfig, GoogleAIStudioGeminiConfig]) +def test_pre_gemini_4_keeps_candidate_count_support(config_cls: type[VertexGeminiConfig]) -> None: + config = config_cls() + for model in ("gemini-3.1-pro-preview", "gemini-2.5-flash"): + supported = config.get_supported_openai_params(model) + assert "n" in supported + mapped = config.map_openai_params( + non_default_params={"n": 3}, + optional_params={}, + model=model, + drop_params=False, + ) + assert mapped["candidate_count"] == 3 + + +def test_google_ai_studio_gemini_3_excludes_penalty_parameters() -> None: + config = GoogleAIStudioGeminiConfig() + supported = config.get_supported_openai_params("gemini-3.1-pro-preview") + assert "frequency_penalty" not in supported + assert "presence_penalty" not in supported + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_gemini_4_argon_grounding_bills_tool_use_prompt_tokens_and_search_queries() -> None: + """ + Gemini 4 Argon bills every token the model processes, including internal tool system + instructions and function declarations when using grounding, in addition to the per-query + Google Search grounding fee. + """ + completion_response = { + "modelVersion": "gemini-4-argon", + "candidates": [ + { + "content": {"parts": [{"text": "Grounded answer"}], "role": "model"}, + "finishReason": "STOP", + "groundingMetadata": {"webSearchQueries": ["gemini 4 argon"]}, + } + ], + "usageMetadata": UsageMetadata( + promptTokenCount=50, + candidatesTokenCount=100, + toolUsePromptTokenCount=300, + totalTokenCount=450, + promptTokensDetails=[{"modality": "TEXT", "tokenCount": 50}], + toolUsePromptTokensDetails=[{"modality": "TEXT", "tokenCount": 300}], + ), + } + + raw_response = MagicMock() + raw_response.text = "{}" + raw_response.headers = {} + + response = VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response( + completion_response=completion_response, + model_response=ModelResponse(), + model="gemini-4-argon", + logging_obj=MagicMock(optional_params={}), + raw_response=raw_response, + ) + + usage = response.usage + assert usage.prompt_tokens == 50 + 300 + 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