scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 199

fix(vertex_ai): route service_tier through shared request headers on Vertex Gemini

Browse source 
Co-authored-by: TobiMayr <tobi.mayr@proton.me>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-10-17 19:18:54 +00:00
parent f2a169a258
commit e4a99d774c
4 changed files with 153 additions and 49 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
litellm/llms/vertex_ai/gemini/transformation.py
View file

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

69
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -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 ##

37
tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py
View file

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

74
tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

@ -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 (