diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 6e801ed5756..3b4d70f1c49 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -1225,7 +1225,9 @@ 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 service_tier.lower() == "default": data["serviceTier"] = "standard" 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..9e3bd596028 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 @@ -730,6 +730,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 @@ -3390,12 +3410,12 @@ 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[str, str]: default_headers: Final = { @@ -3405,9 +3425,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 @@ -3663,12 +3696,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 ) transform: Final = partial( @@ -3776,12 +3809,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 ) transform: Final = partial( @@ -3965,12 +3998,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 0539a939bb5..760a6c3382d 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,57 @@ 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( + ("custom_llm_provider", "service_tier"), + [ + ("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_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} + litellm_params: Final = {} - result = transform_request_body( + result: Final = transform_request_body( 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 "serviceTier" in result - assert result["serviceTier"] == "flex" + assert "serviceTier" not in result + + +@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": service_tier}, + custom_llm_provider="gemini", + litellm_params={}, + cached_content=None, + ) + + 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 03bc308a89d..5d31c2e72d6 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 @@ -6408,6 +6408,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 (