diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 7b832b321da..d9738c38e43 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -88,6 +88,17 @@ def _apply_gemini_mime_type_aliases(mime_type: str) -> str: return _GEMINI_MIME_TYPE_ALIASES.get(normalized, normalized) +def _map_vertex_service_tier(service_tier: str) -> str: + normalized_service_tier: Final = service_tier.lower() + if normalized_service_tier == "flex": + return "SERVICE_TIER_FLEX" + if normalized_service_tier == "priority": + return "SERVICE_TIER_PRIORITY" + if normalized_service_tier == "default": + return "SERVICE_TIER_STANDARD" + return service_tier + + def _get_vertex_base() -> "VertexBase": """Lazily return the shared VertexBase instance to avoid a module-load-time cyclic import.""" global _GCS_METADATA_VERTEX_BASE @@ -1227,7 +1238,12 @@ def transform_request_body( if service_tier := optional_params.pop("service_tier", None): if isinstance(service_tier, str): - if service_tier.lower() == "default": + if custom_llm_provider in ( + LlmProviders.VERTEX_AI.value, + LlmProviders.VERTEX_AI_BETA.value, + ): + data["serviceTier"] = _map_vertex_service_tier(service_tier) + elif service_tier.lower() == "default": data["serviceTier"] = "standard" else: data["serviceTier"] = service_tier.lower() diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index b97724319c7..67d0ee487c2 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -208,13 +208,24 @@ def test_vertex_ai_includes_labels(): assert result["labels"] == {"project": "test", "team": "ai"} -def test_service_tier_forwarded_to_vertex_ai(): - """Test that service_tier in optional_params is mapped to serviceTier in request body.""" - messages = [{"role": "user", "content": "test"}] - optional_params = {"service_tier": "flex"} - litellm_params = {} +@pytest.mark.parametrize( + ("service_tier", "expected_service_tier"), + [ + ("flex", "SERVICE_TIER_FLEX"), + ("FLEX", "SERVICE_TIER_FLEX"), + ("priority", "SERVICE_TIER_PRIORITY"), + ("default", "SERVICE_TIER_STANDARD"), + ("SERVICE_TIER_FLEX", "SERVICE_TIER_FLEX"), + ], +) +def test_service_tier_forwarded_to_vertex_ai( + service_tier: str, expected_service_tier: str +) -> None: + messages: Final = [{"role": "user", "content": "test"}] + optional_params: Final = {"service_tier": service_tier} + litellm_params: Final = {} - result = transform_request_body( + result: Final = transform_request_body( messages=messages, model="gemini-2.5-pro", optional_params=optional_params, @@ -223,7 +234,19 @@ def test_service_tier_forwarded_to_vertex_ai(): cached_content=None, ) - assert "serviceTier" in result + assert result["serviceTier"] == expected_service_tier + + +def test_service_tier_forwarded_to_gemini() -> None: + result: Final = transform_request_body( + messages=[{"role": "user", "content": "test"}], + model="gemini-2.5-pro", + optional_params={"service_tier": "flex"}, + custom_llm_provider="gemini", + litellm_params={}, + cached_content=None, + ) + assert result["serviceTier"] == "flex"