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

Merge 9d40381e6a into 3b8ad37296

Browse source
This commit is contained in:
Caffeine Coder 2026-10-27 02:21:48 +00:00 • committed by GitHub
commit e6a4db47b9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 212 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/llms/gemini/chat/transformation.py
View file

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

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

@ -672,6 +672,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:
"""
@ -746,7 +767,6 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
"functions",
"tool_choice",
"response_format",
"n",
"stop",
"extra_headers",
"seed",
@ -760,6 +780,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"])
@ -1642,7 +1665,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":
@ -2254,6 +2278,19 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
return True
return False
@staticmethod
def _response_bills_grounding_tool_use_tokens(
completion_response: Mapping[str, object],
) -> 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"))
return bool(resolved_model and VertexGeminiConfig._is_gemini_4_or_newer(resolved_model))
@staticmethod
def _calculate_usage(
completion_response: GenerateContentResponseBody | BidiGenerateContentServerMessage,
@ -2390,7 +2427,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)
)
else (tool_use_prompt_tokens or 0)
)
billable_text_tokens: Final = (

165
tests/unit/test_gemini_4_argon_model_metadata.py
View file

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