diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 6e801ed5756..5954f60aeef 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -1225,14 +1225,16 @@ def transform_request_body( if cached_content is not None: data["cachedContent"] = cached_content - if service_tier := optional_params.pop("service_tier", None): - if isinstance(service_tier, str): - if service_tier.lower() == "default": - data["serviceTier"] = "standard" - else: - data["serviceTier"] = service_tier.lower() - else: - data["serviceTier"] = service_tier + # Vertex AI skips serviceTier in the body. Tier routing is via headers + # set in validate_environment(). See: https://github.com/BerriAI/litellm/issues/34914 + fields: Final[Mapping[str, object]] = cast(Mapping[str, object], optional_params) # cast-ok: untyped dict + popped_service_tier: Final = fields.get("service_tier") + optional_params.pop("service_tier", None) # pyright: ignore[reportUnknownMemberType] # untyped dict + normalized_service_tier: Final = ( + popped_service_tier.lower().removeprefix("service_tier_") if isinstance(popped_service_tier, str) else "" + ) + if custom_llm_provider == LlmProviders.GEMINI and normalized_service_tier in ("priority", "flex", "standard"): + data["serviceTier"] = normalized_service_tier if continuation_token := optional_params.pop("continuation_token", None): data["continuationToken"] = continuation_token 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 80244cfc3eb..a8ab437ebaa 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 @@ -1062,15 +1062,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return cls._drop_search_tools_mixed_with_functions(optional_params) 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 lowercased. - """ - if value.lower() == "auto": - optional_params["service_tier"] = "priority" - else: - optional_params["service_tier"] = value.lower() + normalized: Final = value.lower().removeprefix("service_tier_") + if normalized in ("auto", "default", "scale", "unspecified"): + return + optional_params["service_tier"] = normalized def _transform_computer_use_config(self, computer_use_config: dict) -> dict: """ @@ -3405,6 +3400,17 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): default_headers.update(api_key) elif api_key is not None: default_headers["Authorization"] = f"Bearer {api_key}" + + # Vertex AI: route service_tier via headers for actual tier enforcement. + # See: https://github.com/BerriAI/litellm/issues/34914 + optional_fields: Final[Mapping[str, object]] = cast(Mapping[str, object], optional_params) # cast-ok: bare dict + litellm_fields: Final[Mapping[str, object]] = cast(Mapping[str, object], litellm_params) # cast-ok: bare dict + service_tier: Final = optional_fields.get("service_tier") + provider: Final = litellm_fields.get("custom_llm_provider") + if isinstance(service_tier, str) and service_tier in ("priority", "flex") and provider != "gemini": + default_headers["X-Vertex-AI-LLM-Request-Type"] = "shared" + default_headers["X-Vertex-AI-LLM-Shared-Request-Type"] = service_tier + if headers is not None: default_headers.update(headers) 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 0539a939bb5..afe34bcfd1f 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 @@ -356,23 +356,64 @@ 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", + ["flex", "FLEX", "priority", "standard", "default", "SERVICE_TIER_FLEX"], +) +def test_service_tier_not_in_vertex_ai_request_body(service_tier): + """ + Vertex AI: service_tier should NOT be in the request body. + Tier routing is handled via X-Vertex-AI-LLM-Shared-Request-Type headers + set in validate_environment(), not the serviceTier body field. + See: https://github.com/BerriAI/litellm/issues/34914 + """ result = transform_request_body( - messages=messages, + messages=[{"role": "user", "content": "test"}], model="gemini-2.5-pro", - optional_params=optional_params, + optional_params={"service_tier": service_tier}, custom_llm_provider="vertex_ai", - litellm_params=litellm_params, + litellm_params={}, cached_content=None, ) - assert "serviceTier" in result - assert result["serviceTier"] == "flex" + assert "serviceTier" not in result + + +@pytest.mark.parametrize( + "service_tier, expected", + [ + ("flex", "flex"), + ("PRIORITY", "priority"), + ("SERVICE_TIER_FLEX", "flex"), + ], +) +def test_service_tier_forwarded_to_google_ai_studio_as_lowercase(service_tier, expected): + """Google AI Studio's generativelanguage API only accepts the lowercase enum names""" + result = transform_request_body( + messages=[{"role": "user", "content": "test"}], + model="gemini-2.5-pro", + optional_params={"service_tier": service_tier}, + custom_llm_provider="gemini", + litellm_params={}, + cached_content=None, + ) + + assert result["serviceTier"] == expected + + +@pytest.mark.parametrize("service_tier", ["auto", "default", "scale"]) +def test_service_tier_omitted_from_gemini_body_for_non_routing_tiers(service_tier): + """Non-routing tiers (auto, default, scale) should be omitted, not sent to the body.""" + result = transform_request_body( + messages=[{"role": "user", "content": "test"}], + model="gemini-2.5-pro", + optional_params={"service_tier": service_tier}, + custom_llm_provider="gemini", + litellm_params={}, + cached_content=None, + ) + + assert "serviceTier" not in result def test_extra_body_cache_not_forwarded_to_vertex_ai(): diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 03bc308a89d..c1eb3c35b41 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -6416,7 +6416,7 @@ def test_vertex_ai_service_tier_in_map_openai_params(): v = VertexGeminiConfig() - # Test pass-through + # Test pass-through with normalization (FLEX -> flex) optional_params = {} non_default_params = {"service_tier": "FLEX"} @@ -6429,7 +6429,20 @@ def test_vertex_ai_service_tier_in_map_openai_params(): assert result["service_tier"] == "flex" - # Test auto -> priority + # Test prefixed enum normalization (SERVICE_TIER_PRIORITY -> priority) + optional_params_prefixed = {} + non_default_params_prefixed = {"service_tier": "SERVICE_TIER_PRIORITY"} + + result_prefixed = v.map_openai_params( + non_default_params=non_default_params_prefixed, + optional_params=optional_params_prefixed, + model="gemini-3-pro-preview", + drop_params=True, + ) + + assert result_prefixed["service_tier"] == "priority" + + # Test auto -> omitted (lets provider decide) optional_params_auto = {} non_default_params_auto = {"service_tier": "auto"} @@ -6440,9 +6453,9 @@ def test_vertex_ai_service_tier_in_map_openai_params(): drop_params=True, ) - assert result_auto["service_tier"] == "priority" + assert "service_tier" not in result_auto - # Test AUTO (uppercase) -> priority + # Test AUTO (uppercase) -> omitted optional_params_auto_upper = {} non_default_params_auto_upper = {"service_tier": "AUTO"} @@ -6453,7 +6466,46 @@ def test_vertex_ai_service_tier_in_map_openai_params(): drop_params=True, ) - assert result_auto_upper["service_tier"] == "priority" + assert "service_tier" not in result_auto_upper + + +@pytest.mark.parametrize( + "service_tier, provider, expect_headers", + [ + ("priority", "vertex_ai", True), + ("flex", "vertex_ai", True), + ("priority", "vertex_ai_beta", True), + ("priority", "gemini", False), + ("flex", "gemini", False), + ("auto", "vertex_ai", False), + ("standard", "vertex_ai", False), + ], +) +def test_vertex_ai_validate_environment_service_tier_headers( + service_tier, provider, expect_headers +): + """validate_environment should inject Vertex AI tier headers for priority/flex + on Vertex AI only, not Google AI Studio.""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + v = VertexGeminiConfig() + headers = v.validate_environment( + headers=None, + model="gemini-3.5-flash-lite", + messages=[], + optional_params={"service_tier": service_tier}, + litellm_params={"custom_llm_provider": provider}, + api_key="test-key", + ) + + if expect_headers: + assert headers["X-Vertex-AI-LLM-Request-Type"] == "shared" + assert headers["X-Vertex-AI-LLM-Shared-Request-Type"] == service_tier + else: + assert "X-Vertex-AI-LLM-Request-Type" not in headers + assert "X-Vertex-AI-LLM-Shared-Request-Type" not in headers def test_vertex_ai_usage_metadata_with_video_tokens_in_prompt():