chore: add a service_tier field mapping from openai to gemini

This commit is contained in:
Mark McDonald 2026-03-24 15:59:42 +08:00 committed by Sameer Kankute
parent 25b266fe2c
commit 7617d4e6db
No known key found for this signature in database
2 changed files with 28 additions and 2 deletions

View file

@ -363,6 +363,17 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
"""
return Tools(googleSearch={})
def _map_service_tier_param(self, value: str, optional_params: dict) -> None:
"""
Map OpenAI service_tier (string) to Gemini serviceTier.
'auto' maps to 'priority'.
Other values are passed blindly.
"""
if value == "auto":
optional_params["service_tier"] = "priority"
else:
optional_params["service_tier"] = value
def _transform_computer_use_config(self, computer_use_config: dict) -> dict:
"""
Transform Computer Use configuration to Gemini API format.
@ -1123,7 +1134,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
optional_params, [_tools]
)
elif param == "service_tier" and isinstance(value, str):
optional_params["service_tier"] = value
self._map_service_tier_param(value, optional_params)
elif param == "include_server_side_tool_invocations" and value is True:
optional_params["include_server_side_tool_invocations"] = True
if litellm.vertex_ai_safety_settings is not None:

View file

@ -3674,18 +3674,33 @@ def test_vertex_ai_service_tier_in_map_openai_params():
)
v = VertexGeminiConfig()
# Test pass-through
optional_params = {}
non_default_params = {"service_tier": "FLEX"}
result = v.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model="gemini-pro",
model="gemini-3-pro-preview",
drop_params=True,
)
assert result["service_tier"] == "FLEX"
# Test auto -> priority
optional_params_auto = {}
non_default_params_auto = {"service_tier": "auto"}
result_auto = v.map_openai_params(
non_default_params=non_default_params_auto,
optional_params=optional_params_auto,
model="gemini-3-pro-preview",
drop_params=True,
)
assert result_auto["service_tier"] == "priority"
def test_vertex_ai_usage_metadata_with_video_tokens_in_prompt():
"""Test promptTokensDetails with VIDEO modality for video inputs.