Merge 81756c0c98 into e9b1c11098

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

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)
 

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():

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():