From 4f4c0ba0bfe751502518198494d593b0ee5ed429 Mon Sep 17 00:00:00 2001 From: TobiMayr Date: Tue, 28 Jul 2026 17:47:25 +0100 Subject: [PATCH 1/3] fix(vertex_ai): route service_tier via headers for actual tier enforcement The serviceTier body field is accepted by Vertex AI (no 400) but does not activate tier routing. The X-Vertex-AI-LLM-Shared-Request-Type header is what actually routes to shared rate-limit pools with priority/flex enforcement. Changes: - validate_environment(): inject X-Vertex-AI-LLM-Request-Type and X-Vertex-AI-LLM-Shared-Request-Type headers when service_tier is priority or flex (Vertex AI only, not Google AI Studio) - _transform_request_body(): skip serviceTier in body for Vertex AI (keep for Google AI Studio gemini/ where body field works) - _map_service_tier_param(): normalize prefixed values (SERVICE_TIER_FLEX -> flex), skip non-routing tiers (auto, default, scale, unspecified) See: https://github.com/BerriAI/litellm/issues/34914 --- .../llms/vertex_ai/gemini/transformation.py | 13 ++-- .../vertex_and_google_ai_studio_gemini.py | 22 ++++--- .../test_vertex_ai_gemini_transformation.py | 63 +++++++++++++++---- ...test_vertex_and_google_ai_studio_gemini.py | 62 ++++++++++++++++-- 4 files changed, 128 insertions(+), 32 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 15e5e440b91..b3c05016146 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -1220,14 +1220,13 @@ def transform_request_body( if cached_content is not None: data["cachedContent"] = cached_content + # Vertex AI: skip serviceTier in body — tier routing is via headers + # set in validate_environment(). See: https://github.com/BerriAI/litellm/issues/34914 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 + if custom_llm_provider == LlmProviders.GEMINI and isinstance(service_tier, str): + normalized: Final = service_tier.lower().removeprefix("service_tier_") + if normalized in ("priority", "flex", "standard"): + data["serviceTier"] = normalized # Only add labels for Vertex AI endpoints (not Google GenAI/AI Studio) and only if non-empty if labels and custom_llm_provider != LlmProviders.GEMINI: 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 517ac9e98b1..118155e13f5 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 @@ -489,15 +489,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: """ @@ -2733,6 +2728,15 @@ 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 + service_tier: Final = optional_params.get("service_tier") + provider: Final = litellm_params.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 047af5d0ab2..09a517d97b9 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,23 +208,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 3f49869b19f..fd29142a29c 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 @@ -5432,7 +5432,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"} @@ -5445,7 +5445,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"} @@ -5456,9 +5469,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"} @@ -5469,7 +5482,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(): From e1aed5199257a74163b93bbc849da2cd611c1c6a Mon Sep 17 00:00:00 2001 From: TobiMayr Date: Wed, 14 Oct 2026 12:26:22 +0100 Subject: [PATCH 2/3] fix(vertex_ai): collapse the service tier body check for the strict ruff gate Co-authored-by: Cursor --- litellm/llms/vertex_ai/gemini/transformation.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index b3c05016146..ff145f2925a 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -1220,13 +1220,14 @@ def transform_request_body( if cached_content is not None: data["cachedContent"] = cached_content - # Vertex AI: skip serviceTier in body — tier routing is via headers + # Vertex AI skips serviceTier in the body. Tier routing is via headers # set in validate_environment(). See: https://github.com/BerriAI/litellm/issues/34914 - if service_tier := optional_params.pop("service_tier", None): - if custom_llm_provider == LlmProviders.GEMINI and isinstance(service_tier, str): - normalized: Final = service_tier.lower().removeprefix("service_tier_") - if normalized in ("priority", "flex", "standard"): - data["serviceTier"] = normalized + popped_service_tier: Final = optional_params.pop("service_tier", None) + 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 # Only add labels for Vertex AI endpoints (not Google GenAI/AI Studio) and only if non-empty if labels and custom_llm_provider != LlmProviders.GEMINI: From 81756c0c985c80a1ac708297588f55af172a4297 Mon Sep 17 00:00:00 2001 From: TobiMayr Date: Wed, 14 Oct 2026 12:58:50 +0100 Subject: [PATCH 3/3] fix(vertex_ai): read service_tier through a typed mapping --- litellm/llms/vertex_ai/gemini/transformation.py | 4 +++- .../vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py | 6 ++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index ff145f2925a..04f122cf97f 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -1222,7 +1222,9 @@ def transform_request_body( # Vertex AI skips serviceTier in the body. Tier routing is via headers # set in validate_environment(). See: https://github.com/BerriAI/litellm/issues/34914 - popped_service_tier: Final = optional_params.pop("service_tier", None) + 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 "" ) 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 118155e13f5..f46ff373bd9 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 @@ -2731,8 +2731,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # Vertex AI: route service_tier via headers for actual tier enforcement. # See: https://github.com/BerriAI/litellm/issues/34914 - service_tier: Final = optional_params.get("service_tier") - provider: Final = litellm_params.get("custom_llm_provider") + 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