From e4a99d774ca325459bb104db67367d729c120e5c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 19:18:54 +0000 Subject: [PATCH] fix(vertex_ai): route service_tier through shared request headers on Vertex Gemini Co-authored-by: TobiMayr Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/gemini/transformation.py | 22 +----- .../vertex_and_google_ai_studio_gemini.py | 69 ++++++++++++----- .../test_vertex_ai_gemini_transformation.py | 37 ++++++---- ...test_vertex_and_google_ai_studio_gemini.py | 74 +++++++++++++++++++ 4 files changed, 153 insertions(+), 49 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index d9738c38e43..fda781d85af 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -88,17 +88,6 @@ 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 @@ -1236,14 +1225,11 @@ def transform_request_body( if cached_content is not None: data["cachedContent"] = cached_content - if service_tier := optional_params.pop("service_tier", None): + if ( + service_tier := optional_params.pop("service_tier", None) + ) and custom_llm_provider == LlmProviders.GEMINI.value: if isinstance(service_tier, str): - 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": + if service_tier.lower() == "default": data["serviceTier"] = "standard" else: data["serviceTier"] = service_tier.lower() 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 257fafbb6b8..fd0369c6f04 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 @@ -248,6 +248,26 @@ class VertexAIBaseConfig: ] +def _vertex_service_tier_headers( + service_tier: str | None, + custom_llm_provider: str | None, +) -> tuple[tuple[str, str], ...]: + if custom_llm_provider not in ( + litellm.LlmProviders.VERTEX_AI.value, + litellm.LlmProviders.VERTEX_AI_BETA.value, + ): + return () + if service_tier is None: + return () + normalized_service_tier: Final = service_tier.lower() + if normalized_service_tier not in ("flex", "priority"): + return () + return ( + ("X-Vertex-AI-LLM-Request-Type", "shared"), + ("X-Vertex-AI-LLM-Shared-Request-Type", normalized_service_tier), + ) + + class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): """ Reference: https://cloud.google.com/vertex-ai/docs/generative-ai/chat/test-chat-prompts @@ -2779,14 +2799,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): def validate_environment( self, - headers: dict | None, + headers: Mapping[str, str] | None, model: str, messages: list[AllMessageValues], - optional_params: dict, - litellm_params: dict, - api_key: str | dict | None = None, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | Mapping[str, str] | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict[str, str]: default_headers: Final = { "Content-Type": "application/json", } @@ -2794,9 +2814,22 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): default_headers.update(api_key) elif api_key is not None: default_headers["Authorization"] = f"Bearer {api_key}" + service_tier_value: Final = optional_params.get("service_tier") + service_tier: Final[str | None] = service_tier_value if isinstance(service_tier_value, str) else None + custom_llm_provider_value: Final = litellm_params.get("custom_llm_provider") + custom_llm_provider: Final[str | None] = ( + custom_llm_provider_value if isinstance(custom_llm_provider_value, str) else None + ) + service_tier_headers: Final = dict( + _vertex_service_tier_headers( + service_tier=service_tier, + custom_llm_provider=custom_llm_provider, + ) + ) + + default_headers.update(service_tier_headers) if headers is not None: default_headers.update(headers) - return default_headers @@ -3025,12 +3058,12 @@ class VertexLLM(VertexBase): ) headers: Final = VertexGeminiConfig().validate_environment( - api_key=auth_header, - headers=extra_headers, + api_key=cast(str | Mapping[str, str] | None, auth_header), # cast-ok: Vertex auth headers contain strings + headers=cast(Mapping[str, str] | None, extra_headers), # cast-ok: request headers contain strings model=model, messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, + optional_params=cast(Mapping[str, object], optional_params), # cast-ok: request parameter keys are strings + litellm_params=cast(Mapping[str, object], litellm_params), # cast-ok: request parameter keys are strings ) request_body: Final = await async_transform_request_body( @@ -3128,12 +3161,12 @@ class VertexLLM(VertexBase): ) headers: Final = VertexGeminiConfig().validate_environment( - api_key=auth_header, - headers=extra_headers, + api_key=cast(str | Mapping[str, str] | None, auth_header), # cast-ok: Vertex auth headers contain strings + headers=cast(Mapping[str, str] | None, extra_headers), # cast-ok: request headers contain strings model=model, messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, + optional_params=cast(Mapping[str, object], optional_params), # cast-ok: request parameter keys are strings + litellm_params=cast(Mapping[str, object], litellm_params), # cast-ok: request parameter keys are strings ) request_body: Final = await async_transform_request_body( @@ -3311,12 +3344,12 @@ class VertexLLM(VertexBase): use_psc_endpoint_format=use_psc_endpoint_format, ) headers: Final = VertexGeminiConfig().validate_environment( - api_key=auth_header, - headers=extra_headers, + api_key=cast(str | Mapping[str, str] | None, auth_header), # cast-ok: Vertex auth headers contain strings + headers=cast(Mapping[str, str] | None, extra_headers), # cast-ok: request headers contain strings model=model, messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, + optional_params=cast(Mapping[str, object], optional_params), # cast-ok: request parameter keys are strings + litellm_params=cast(Mapping[str, object], litellm_params), # cast-ok: request parameter keys are strings ) ## TRANSFORMATION ## 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 67d0ee487c2..c1f98c63d1a 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 @@ -209,17 +209,22 @@ def test_vertex_ai_includes_labels(): @pytest.mark.parametrize( - ("service_tier", "expected_service_tier"), + ("custom_llm_provider", "service_tier"), [ - ("flex", "SERVICE_TIER_FLEX"), - ("FLEX", "SERVICE_TIER_FLEX"), - ("priority", "SERVICE_TIER_PRIORITY"), - ("default", "SERVICE_TIER_STANDARD"), - ("SERVICE_TIER_FLEX", "SERVICE_TIER_FLEX"), + ("vertex_ai", "flex"), + ("vertex_ai", "FLEX"), + ("vertex_ai", "priority"), + ("vertex_ai", "default"), + ("vertex_ai", "SERVICE_TIER_FLEX"), + ("vertex_ai_beta", "flex"), + ("vertex_ai_beta", "FLEX"), + ("vertex_ai_beta", "priority"), + ("vertex_ai_beta", "default"), + ("vertex_ai_beta", "SERVICE_TIER_FLEX"), ], ) -def test_service_tier_forwarded_to_vertex_ai( - service_tier: str, expected_service_tier: str +def test_service_tier_not_forwarded_to_vertex_body( + custom_llm_provider: str, service_tier: str ) -> None: messages: Final = [{"role": "user", "content": "test"}] optional_params: Final = {"service_tier": service_tier} @@ -229,25 +234,31 @@ def test_service_tier_forwarded_to_vertex_ai( messages=messages, model="gemini-2.5-pro", optional_params=optional_params, - custom_llm_provider="vertex_ai", + custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, cached_content=None, ) - assert result["serviceTier"] == expected_service_tier + assert "serviceTier" not in result -def test_service_tier_forwarded_to_gemini() -> None: +@pytest.mark.parametrize( + ("service_tier", "expected_service_tier"), + [("flex", "flex"), ("PRIORITY", "priority"), ("default", "standard")], +) +def test_service_tier_forwarded_to_gemini( + service_tier: str, expected_service_tier: str +) -> None: result: Final = transform_request_body( messages=[{"role": "user", "content": "test"}], model="gemini-2.5-pro", - optional_params={"service_tier": "flex"}, + optional_params={"service_tier": service_tier}, custom_llm_provider="gemini", litellm_params={}, cached_content=None, ) - assert result["serviceTier"] == "flex" + assert result["serviceTier"] == expected_service_tier 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 5c0e9f1a305..fefe49d5bda 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 @@ -5286,6 +5286,80 @@ def test_vertex_ai_web_search_options_in_map_openai_params(): assert "web_search_options" not in optional_params, "web_search_options should be removed after transformation" +@pytest.mark.parametrize("custom_llm_provider", ["vertex_ai", "vertex_ai_beta"]) +@pytest.mark.parametrize( + ("service_tier", "expected_service_tier"), + [ + ("flex", "flex"), + ("FLEX", "flex"), + ("priority", "priority"), + ("PRIORITY", "priority"), + ], +) +def test_validate_environment_adds_shared_service_tier_headers( + custom_llm_provider: str, service_tier: str, expected_service_tier: str +) -> None: + result: Final = VertexGeminiConfig().validate_environment( + headers=None, + model="gemini-2.5-pro", + messages=[], + optional_params={"service_tier": service_tier}, + litellm_params={"custom_llm_provider": custom_llm_provider}, + ) + + assert result["X-Vertex-AI-LLM-Request-Type"] == "shared" + assert result["X-Vertex-AI-LLM-Shared-Request-Type"] == expected_service_tier + + +@pytest.mark.parametrize("service_tier", [None, "default", "DEFAULT"]) +def test_validate_environment_omits_shared_service_tier_headers_for_default_and_unset( + service_tier: str | None, +) -> None: + optional_params: Final = {} if service_tier is None else {"service_tier": service_tier} + result: Final = VertexGeminiConfig().validate_environment( + headers=None, + model="gemini-2.5-pro", + messages=[], + optional_params=optional_params, + litellm_params={"custom_llm_provider": "vertex_ai"}, + ) + + assert "X-Vertex-AI-LLM-Request-Type" not in result + assert "X-Vertex-AI-LLM-Shared-Request-Type" not in result + + +@pytest.mark.parametrize("service_tier", ["flex", "priority", "FLEX", "PRIORITY"]) +def test_validate_environment_omits_shared_service_tier_headers_for_gemini( + service_tier: str, +) -> None: + result: Final = VertexGeminiConfig().validate_environment( + headers=None, + model="gemini-2.5-pro", + messages=[], + optional_params={"service_tier": service_tier}, + litellm_params={"custom_llm_provider": "gemini"}, + ) + + assert "X-Vertex-AI-LLM-Request-Type" not in result + assert "X-Vertex-AI-LLM-Shared-Request-Type" not in result + + +def test_validate_environment_preserves_caller_service_tier_header_overrides() -> None: + result: Final = VertexGeminiConfig().validate_environment( + headers={ + "X-Vertex-AI-LLM-Request-Type": "caller-request-type", + "X-Vertex-AI-LLM-Shared-Request-Type": "caller-tier", + }, + model="gemini-2.5-pro", + messages=[], + optional_params={"service_tier": "flex"}, + litellm_params={"custom_llm_provider": "vertex_ai"}, + ) + + assert result["X-Vertex-AI-LLM-Request-Type"] == "caller-request-type" + assert result["X-Vertex-AI-LLM-Shared-Request-Type"] == "caller-tier" + + def test_vertex_ai_service_tier_in_map_openai_params(): """Test that service_tier is correctly mapped to optional_params.""" from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (