This commit is contained in:
Mateo Wang 2026-11-05 08:47:12 -08:00 • committed by GitHub
commit 42dc70485d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 1779 additions and 106 deletions

View file

@ -2,8 +2,9 @@ import importlib.util
import os
import re
import sys
from collections.abc import Collection, Iterator, Mapping
from functools import lru_cache
from collections.abc import Callable, Collection, Iterator, Mapping, Sequence
from dataclasses import dataclass
from functools import cache, lru_cache
from itertools import accumulate, chain, repeat
from logging import Logger
from types import MappingProxyType
@ -23,6 +24,7 @@ from litellm.constants import (
STANDARD_CUSTOMER_ID_HEADERS,
)
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor, reject_api_key_secret_reference
from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.url_utils import (
SSRFError,
@ -41,11 +43,13 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
from litellm.types.router import (
CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS,
Deployment,
DeploymentTypedDict,
LiteLLMParamsTypedDict,
reject_server_owned_wif_fields,
server_owned_wif_fields_present,
)
from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params
from litellm.types.utils import CustomPricingLiteLLMParams, oauth_token_exchange_litellm_params
from litellm.types.utils import CustomPricingLiteLLMParams, LlmProviders, oauth_token_exchange_litellm_params
from litellm.types.videos.main import DecodedVideoId
@ -211,35 +215,63 @@ def _is_param_allowed(
return False
def deployment_clientside_auth_params(litellm_params: LiteLLMParamsTypedDict) -> CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS:
if "configurable_clientside_auth_params" not in litellm_params:
return None
return litellm_params["configurable_clientside_auth_params"]
def _model_group_deployments(model: str, llm_router: Router) -> Sequence[DeploymentTypedDict]:
registry: Final = llm_router.registry()
provider_group: Final = model.split("/", 1)[0]
return (
registry.deployments(model, source="routed_group")
or (registry.deployments(provider_group) if provider_group in provider_list else None)
or ()
)
_VERTEX_TARGET_READING_PROVIDERS: Final[frozenset[str]] = frozenset(
(LlmProviders.VERTEX_AI.value, LlmProviders.VERTEX_AI_BETA.value, LlmProviders.GEMINI.value, LlmProviders.GDC.value)
)
def _deployment_provider(litellm_params: LiteLLMParamsTypedDict) -> str | None:
match litellm_params:
case {"custom_llm_provider": str(declared)} if declared:
return declared
case {"model": str(model)}:
return inferred_provider(model)
case _:
return None
def _model_group_reads_vertex_target(model: str, llm_router: Router | None) -> bool:
if llm_router is None:
return True
deployments: Final = _model_group_deployments(model, llm_router)
providers: Final = frozenset(_deployment_provider(deployment["litellm_params"]) for deployment in deployments)
return not deployments or None in providers or bool(providers & _VERTEX_TARGET_READING_PROVIDERS)
def _allow_model_level_clientside_configurable_parameters(
model: str, param: str, request_body_value: object, llm_router: Router | None
) -> bool:
"""
Check if model is allowed to use configurable client-side params
- get matching model
- check if 'clientside_configurable_parameters' is set for model
-
The router, not the caller, picks which deployment serves the request, so an
opt-in counts only when every deployment of the model group, already alias
resolved, carries it
"""
if llm_router is None:
return False
# check if model is set
registry: Final = llm_router.registry()
model_info = registry.group(model)
if model_info is None:
# check if wildcard model is set
if model.split("/", 1)[0] in provider_list:
model_info = registry.group(model.split("/", 1)[0])
if model_info is None:
return False
if model_info is None or model_info.configurable_clientside_auth_params is None:
return False
return _is_param_allowed(
param=param,
request_body_value=request_body_value,
configurable_clientside_auth_params=model_info.configurable_clientside_auth_params,
deployments: Final = _model_group_deployments(model, llm_router)
return bool(deployments) and all(
_is_param_allowed(
param=param,
request_body_value=request_body_value,
configurable_clientside_auth_params=deployment_clientside_auth_params(deployment["litellm_params"]),
)
for deployment in deployments
)
@ -318,6 +350,20 @@ def route_manages_mcp_servers(route: str | None) -> bool:
return route is not None and _MCP_SERVER_MANAGEMENT_ROUTE.fullmatch(route) is not None
VERTEX_TARGET_PARAMS: Final[frozenset[str]] = frozenset(
("vertex_project", "vertex_location", "vertex_ai_project", "vertex_ai_location")
)
_ROUTES_THAT_NAME_THEIR_VERTEX_TARGET: Final[frozenset[str]] = frozenset(
("/vertex_ai/live", "/rag/ingest", "/v1/rag/ingest")
)
def _without_route_named_vertex_target(request_body: Mapping[str, object], route: str | None) -> Mapping[str, object]:
if route not in _ROUTES_THAT_NAME_THEIR_VERTEX_TARGET:
return request_body
return MappingProxyType({key: value for key, value in request_body.items() if key not in VERTEX_TARGET_PARAMS})
_NESTED_CONFIG_KEYS: Final[tuple[str, ...]] = ("litellm_embedding_config", "extra_body")
# Metadata containers that carry per-request configuration consumed by the
@ -410,6 +456,10 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = (
"aws_external_id",
"aws_session_tags",
"vertex_credentials",
"vertex_location",
"vertex_project",
"vertex_ai_location",
"vertex_ai_project",
# Azure managed-identity / federated-auth token. The Azure provider
# transformer reads ``azure_ad_token`` (top-level or via
# ``extra_body``) and resolves it through ``get_secret`` before
@ -478,15 +528,85 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = (
*sorted(CustomPricingLiteLLMParams.model_fields.keys()),
)
_BANNED_METADATA_PARAMS: Final[tuple[str, ...]] = tuple(
param for param in _BANNED_REQUEST_BODY_PARAMS if param not in VERTEX_TARGET_PARAMS
)
def _opt_in_model_groups(
model: str, also_requested_models: tuple[str, ...], llm_router: Router | None
) -> frozenset[str]:
requested: Final = frozenset((model, *also_requested_models))
if llm_router is None:
return requested
registry: Final = llm_router.registry()
routed: Final = frozenset(registry.resolve(name, kind="alias") or name for name in requested)
fallbacks: Final = frozenset[str]().union(
*(router_fallback_model_names(name, llm_router, _EVERY_FALLBACK_TYPE) for name in requested | routed)
)
return routed | frozenset(registry.resolve(name, kind="alias") or name for name in fallbacks)
_MAX_OPT_IN_FALLBACK_MODELS: Final = 100
@dataclass(frozen=True, slots=True)
class _OptInGroups:
every: tuple[str, ...]
vertex_target: tuple[str, ...]
_OptInKey = tuple[str | tuple[str, str | None], ...]
def _opt_in_key(litellm_params: LiteLLMParamsTypedDict) -> _OptInKey:
return tuple(
item if isinstance(item, str) else ("api_base", item.get("api_base"))
for item in deployment_clientside_auth_params(litellm_params) or ()
)
def _same_answer_key(model_group: str, llm_router: Router) -> frozenset[tuple[str | None, _OptInKey]] | str:
deployments: Final = _model_group_deployments(model_group, llm_router)
key: Final = frozenset(
(_deployment_provider(deployment["litellm_params"]), _opt_in_key(deployment["litellm_params"]))
for deployment in deployments
)
return key or model_group
def _distinct_opt_in_groups(model_groups: frozenset[str], llm_router: Router | None) -> tuple[str, ...]:
if llm_router is None:
return tuple(sorted(model_groups))
return tuple(
{_same_answer_key(model_group, llm_router): model_group for model_group in sorted(model_groups)}.values()
)
def _resolve_opt_in_groups(
model: str, also_requested_models: tuple[str, ...], llm_router: Router | None
) -> _OptInGroups:
if len(frozenset(also_requested_models)) > _MAX_OPT_IN_FALLBACK_MODELS:
raise ValueError(
f"Rejected Request: the request names more than {_MAX_OPT_IN_FALLBACK_MODELS} fallback models, "
"too many to check against configurable_clientside_auth_params."
)
every: Final = _distinct_opt_in_groups(_opt_in_model_groups(model, also_requested_models, llm_router), llm_router)
return _OptInGroups(
every=every,
vertex_target=tuple(group for group in every if _model_group_reads_vertex_target(group, llm_router)),
)
def _check_banned_params(
body: Mapping[str, object],
general_settings: dict,
llm_router: Router | None,
model: str,
*,
manages_deployments: bool = False,
mcp_server_fields_allowed: bool = False,
banned_params: tuple[str, ...] = _BANNED_REQUEST_BODY_PARAMS,
opt_in_groups: Callable[[], _OptInGroups],
) -> None:
"""Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in.
@ -501,21 +621,24 @@ def _check_banned_params(
reject_server_owned_wif_params(wif_body)
if not manages_deployments:
reject_federated_credential_reference(body)
for param in _BANNED_REQUEST_BODY_PARAMS:
for param in banned_params:
if param not in body:
continue
if general_settings.get("allow_client_side_credentials") is True:
# Proxy-wide opt-in: every banned param is permitted, exit
# entirely so the rest of the loop doesn't waste work.
return
if (
if all(
_allow_model_level_clientside_configurable_parameters(
model=model,
model=model_group,
param=param,
request_body_value=body[param],
llm_router=llm_router,
)
is True
for model_group in (
opt_in_groups().vertex_target if param in VERTEX_TARGET_PARAMS else opt_in_groups().every
)
):
# Per-param opt-in: only THIS param is permitted by the
# deployment's ``configurable_clientside_auth_params``. Skip
@ -606,27 +729,39 @@ def fallback_target_model_name(target: object) -> str | None:
return None
def _direct_fallback_model_names(model: str, llm_router: Router) -> frozenset[str]:
_EVERY_FALLBACK_TYPE: Final[tuple[str, ...]] = ("general", "context_window", "content_policy")
def _direct_fallback_model_names(model: str, llm_router: Router, fallback_types: tuple[str, ...]) -> frozenset[str]:
from litellm.proxy.auth.model_checks import get_all_fallbacks
return frozenset(
name for target in get_all_fallbacks(model, llm_router) if (name := fallback_target_model_name(target))
targets: Final = chain.from_iterable(
get_all_fallbacks(model, llm_router, fallback_type) for fallback_type in fallback_types
)
return frozenset(name for target in targets if (name := fallback_target_model_name(target)))
def _fallback_hop(
frontier_and_reached: tuple[frozenset[str], frozenset[str]], llm_router: Router
frontier_and_reached: tuple[frozenset[str], frozenset[str]], router_and_types: tuple[Router, tuple[str, ...]]
) -> tuple[frozenset[str], frozenset[str]]:
frontier, reached = frontier_and_reached
llm_router, fallback_types = router_and_types
next_hop: Final = (
frozenset(chain.from_iterable(_direct_fallback_model_names(model, llm_router) for model in frontier)) - reached
frozenset(
chain.from_iterable(_direct_fallback_model_names(model, llm_router, fallback_types) for model in frontier)
)
- reached
)
return next_hop, reached | next_hop
def router_fallback_model_names(model: str, llm_router: Router) -> frozenset[str]:
def router_fallback_model_names(
model: str, llm_router: Router, fallback_types: tuple[str, ...] = ("general",)
) -> frozenset[str]:
"""The groups a failure on `model` can reach through the router's fallbacks, chained and generic (`*`) ones."""
hops: Final = accumulate(repeat(llm_router), _fallback_hop, initial=(frozenset({model}), frozenset({model})))
hops: Final = accumulate(
repeat((llm_router, fallback_types)), _fallback_hop, initial=(frozenset({model}), frozenset({model}))
)
_, reached = next(hop for hop in hops if not hop[0])
return reached - {model}
@ -691,33 +826,60 @@ def is_request_body_safe(
if "model_list" in request_body:
raise ValueError("Rejected Request: model_list is not allowed in the request body.")
manages_deployments: Final = route_manages_deployments(route)
root_body: Final = _without_route_named_vertex_target(request_body, route)
route_model: Final = _model_from_route_path(route) if route is not None else None
also_requested_models: Final = (
*request_fallback_model_names(root_body),
*((route_model,) if route_model else ()),
)
opt_in_groups: Final = cache(lambda: _resolve_opt_in_groups(model, also_requested_models, llm_router))
_check_banned_params(
request_body,
root_body,
general_settings,
llm_router,
model,
manages_deployments=manages_deployments,
mcp_server_fields_allowed=route_manages_mcp_servers(route),
opt_in_groups=opt_in_groups,
)
for nested_key in _NESTED_CONFIG_KEYS:
nested = _coerce_metadata_to_dict(request_body.get(nested_key))
if nested is not None:
_check_banned_params(nested, general_settings, llm_router, model, manages_deployments=manages_deployments)
_check_banned_params(
nested,
general_settings,
llm_router,
manages_deployments=manages_deployments,
opt_in_groups=opt_in_groups,
)
for metadata_key in _NESTED_METADATA_KEYS:
metadata = _coerce_metadata_to_dict(request_body.get(metadata_key))
if metadata is not None:
_check_banned_params(metadata, general_settings, llm_router, model, manages_deployments=manages_deployments)
_check_banned_params(
metadata,
general_settings,
llm_router,
manages_deployments=manages_deployments,
banned_params=_BANNED_METADATA_PARAMS,
opt_in_groups=opt_in_groups,
)
if any(isinstance(key, str) and key.startswith(f"{metadata_key}[") for key in request_body):
_check_banned_params(
extract_nested_form_metadata(form_data=request_body, prefix=f"{metadata_key}["),
general_settings,
llm_router,
model,
manages_deployments=manages_deployments,
banned_params=_BANNED_METADATA_PARAMS,
opt_in_groups=opt_in_groups,
)
for target in iter_request_fallback_targets(request_body):
if isinstance(target, dict):
_check_banned_params(target, general_settings, llm_router, model, manages_deployments=manages_deployments)
_check_banned_params(
target,
general_settings,
llm_router,
manages_deployments=manages_deployments,
opt_in_groups=opt_in_groups,
)
target_model = target.get("model")
if isinstance(target_model, str):
_reject_url_valued_fallback_target(target_model)
@ -731,7 +893,12 @@ def is_request_body_safe(
litellm_params_metadata: Final = _coerce_metadata_to_dict(litellm_params.get("metadata"))
if litellm_params_metadata is not None:
_check_banned_params(
litellm_params_metadata, general_settings, llm_router, model, manages_deployments=manages_deployments
litellm_params_metadata,
general_settings,
llm_router,
manages_deployments=manages_deployments,
banned_params=_BANNED_METADATA_PARAMS,
opt_in_groups=opt_in_groups,
)
for user_config_params in _user_config_litellm_params(request_body): # pyright: ignore[reportUnknownArgumentType] # request_body is an untyped dict
reject_api_key_secret_reference(user_config_params)
@ -831,7 +998,7 @@ async def pre_db_read_auth_checks(
request_body=request_data,
general_settings=general_settings,
llm_router=llm_router,
model=request_data.get("model", ""), # [TODO] use model passed in url as well (azure openai routes)
model=request_data.get("model") or _model_from_route_path(route) or "",
route=route,
)
@ -2314,37 +2481,8 @@ def get_model_from_request(
team_id=team_id,
)
model = _format_model_candidates(candidates)
# If no explicit model was found, try to extract from route
if model is None:
# Parse model from route that follows the pattern /openai/deployments/{model}/*
match: Final = re.match(r"/openai/deployments/([^/]+)", route)
if match:
model = match.group(1)
# If still not found, extract model from Google generateContent-style routes.
# These routes put the model in the path and allow "/" inside the model id.
# Examples:
# - /v1beta/models/gemini-2.0-flash:generateContent
# - /v1beta/models/bedrock/claude-sonnet-3.7:generateContent
# - /models/custom/ns/model:streamGenerateContent
if model is None and not route.lower().startswith("/vertex"):
google_match = re.search(r"/(?:v1beta|beta)/models/([^:]+):", route)
if google_match:
model = google_match.group(1)
if model is None and not route.lower().startswith("/vertex"):
google_match = re.search(r"^/models/([^:]+):", route)
if google_match:
model = google_match.group(1)
# If still not found, extract from Vertex AI passthrough route
# Pattern: /vertex_ai/.../models/{model_id}:*
# Example: /vertex_ai/v1/.../models/gemini-1.5-pro:generateContent
if model is None and route.lower().startswith("/vertex"):
vertex_match: Final = re.search(r"/models/([^:]+)", route)
if vertex_match:
model = vertex_match.group(1)
model = _model_from_route_path(route)
if route.lower().startswith("/bedrock"):
bedrock_model: Final = _model_from_bedrock_route(route)
@ -2363,6 +2501,23 @@ def get_model_from_request(
return model
def _model_from_route_path(route: str) -> str | None:
"""
Google generateContent-style routes allow "/" inside the model id, e.g.
/v1beta/models/bedrock/claude-sonnet-3.7:generateContent
"""
azure_deployment: Final = re.match(r"/openai/deployments/([^/]+)", route)
if azure_deployment:
return azure_deployment.group(1)
if route.lower().startswith("/vertex"):
vertex_model: Final = re.search(r"/models/([^:]+)", route)
return vertex_model.group(1) if vertex_model else None
google_model: Final = re.search(r"/(?:v1beta|beta)/models/([^:]+):", route) or re.search(
r"^/models/([^:]+):", route
)
return google_model.group(1) if google_model else None
def _router_model_from_azure_route(route: str, llm_router: Router | None) -> str | None:
if llm_router is None:
return None

View file

@ -62,6 +62,7 @@ from litellm.proxy.auth.auth_exception_handler import (
)
from litellm.proxy.auth.auth_utils import (
_BANNED_REQUEST_BODY_PARAMS, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the request-body check
VERTEX_TARGET_PARAMS,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils import http_parsing_utils
@ -156,6 +157,11 @@ _CONFIG_CONNECTION_FIELDS: Final[frozenset[str]] = frozenset(
)
_BANNED_CONNECTION_PARAMS: Final[tuple[str, ...]] = tuple(
param for param in _BANNED_REQUEST_BODY_PARAMS if param not in VERTEX_TARGET_PARAMS
)
def _request_inherits_config_credentials(
config_params: Mapping[str, object],
request_params: Mapping[str, object],
@ -181,7 +187,7 @@ def _request_inherits_config_credentials(
return False
if allow_client_side_credentials:
return True
return not any(param in request_params for param in _BANNED_REQUEST_BODY_PARAMS)
return not any(param in request_params for param in _BANNED_CONNECTION_PARAMS)
def _admin_api_key_from_request(

View file

@ -93,6 +93,7 @@ from litellm.proxy.auth.auth_checks import (
key_model_aliases_for_auth_check,
team_model_aliases_for_auth_check,
)
from litellm.proxy.auth.auth_utils import deployment_clientside_auth_params
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports
@ -4103,6 +4104,10 @@ VERTEX_LIVE_UNCONFIGURED_CLOSE_REASON: Final = (
"Vertex AI auth failed: set a use_in_pass_through vertex model, default_vertex_config, or DEFAULT_VERTEXAI_* env"
)
VERTEX_LIVE_UNPINNED_PROJECT_CLOSE_REASON: Final = (
"Rejected Request: vertex_project must come from the proxy's configuration, see allow_client_side_credentials"
)
VERTEX_PUBLISHER_MODEL_PREFIX: Final = "publishers/google/models/"
VERTEX_PUBLISHERS_SEGMENT: Final = "publishers/"
@ -4127,11 +4132,25 @@ def _get_llm_router() -> Router | None:
return llm_router
def _resolve_vertex_live_credentials(
@dataclass(frozen=True, slots=True)
class _ConfiguredVertexLiveTarget:
credentials: VertexPassThroughCredentials | None
@dataclass(frozen=True, slots=True)
class _DeploymentVertexLiveTarget:
credentials: VertexPassThroughCredentials
deployment_params: LiteLLMParamsTypedDict
_VertexLiveTarget = _ConfiguredVertexLiveTarget | _DeploymentVertexLiveTarget
def _resolve_vertex_live_target(
vertex_project: str | None,
vertex_location: str | None,
model: str | None,
) -> VertexPassThroughCredentials | None:
) -> _VertexLiveTarget:
"""
Resolution order: an explicit project/location registration, then ``default_vertex_config`` (which the proxy
fills from the ``DEFAULT_VERTEXAI_*`` env vars whenever the yaml leaves it out), then any DB model entry
@ -4145,36 +4164,133 @@ def _resolve_vertex_live_credentials(
location=vertex_location,
)
if keyed is not None and keyed.vertex_project is not None:
return keyed
from_deployments: Final = passthrough_endpoint_router.get_vertex_credentials_from_router_deployments(model=model)
return _ConfiguredVertexLiveTarget(credentials=keyed)
from_deployments: Final = passthrough_endpoint_router.get_vertex_deployment_from_router(model=model)
if from_deployments is not None:
return from_deployments
deployment, credentials = from_deployments
return _DeploymentVertexLiveTarget(credentials=credentials, deployment_params=deployment["litellm_params"])
if keyed is not None:
return keyed
return _ConfiguredVertexLiveTarget(credentials=keyed)
passthrough_endpoint_router.set_default_vertex_config()
return passthrough_endpoint_router.get_vertex_credentials(
project_id=vertex_project,
location=vertex_location,
return _ConfiguredVertexLiveTarget(
credentials=passthrough_endpoint_router.get_vertex_credentials(
project_id=vertex_project,
location=vertex_location,
)
)
def _requested_vertex_live_target(
credentials: VertexPassThroughCredentials, vertex_project: str | None, vertex_location: str | None
) -> tuple[tuple[str, str | None, str | None], ...]:
return (
("vertex_project", vertex_project, credentials.vertex_project),
("vertex_location", vertex_location, credentials.vertex_location),
)
def _proxy_configured_vertex_projects() -> frozenset[str]:
return frozenset(
project
for project in (
litellm.vertex_project,
get_secret_str("VERTEXAI_PROJECT"),
get_secret_str("DEFAULT_VERTEXAI_PROJECT"),
)
if project
)
def _names_an_unconfigured_project(
credentials: VertexPassThroughCredentials | None, vertex_project: str | None
) -> bool:
return (
bool(vertex_project)
and (credentials is None or credentials.vertex_project is None)
and vertex_project not in _proxy_configured_vertex_projects()
)
def _live_target_override_reason(
target: _VertexLiveTarget,
vertex_project: str | None,
vertex_location: str | None,
allow_client_side_credentials: bool,
) -> str | None:
if allow_client_side_credentials:
return None
match target:
case _ConfiguredVertexLiveTarget(credentials=credentials) if _names_an_unconfigured_project(
credentials, vertex_project
):
return VERTEX_LIVE_UNPINNED_PROJECT_CLOSE_REASON
case _ConfiguredVertexLiveTarget(credentials=VertexPassThroughCredentials() as credentials):
configured_field: Final = next(
(
name
for name, asked, pinned in _requested_vertex_live_target(
credentials, vertex_project, vertex_location
)
if asked and pinned and asked != pinned
),
None,
)
if configured_field is None:
return None
return (
f"Rejected Request: {configured_field} is pinned by default_vertex_config, "
"see allow_client_side_credentials"
)
case _ConfiguredVertexLiveTarget():
return None
case _DeploymentVertexLiveTarget(credentials=credentials, deployment_params=deployment_params):
opted_in: Final = deployment_clientside_auth_params(deployment_params) or ()
deployment_field: Final = next(
(
name
for name, asked, pinned in _requested_vertex_live_target(
credentials, vertex_project, vertex_location
)
if asked and asked != pinned and name not in opted_in
),
None,
)
if deployment_field is None:
return None
return (
f"Rejected Request: {deployment_field} is pinned by the deployment, "
"see configurable_clientside_auth_params"
)
_VERTEX_RESOURCE_PROJECT_AND_LOCATION: Final = re.compile(r"projects/[^/]+/locations/(?P<location>[^/]+)/")
def _build_vertex_live_setup_model_rewriter(
vertex_project: str | None,
vertex_location: str | None,
configured_location: str | None,
llm_router: Router | None,
keeps_caller_resource_path: bool,
) -> Callable[[str], str] | None:
"""
Rewrite the ``setup`` frame's model into the full Vertex resource path the Live API requires.
Clients address the gateway the way they address LiteLLM (bare id or model alias); Vertex reads anything
that is not a ``projects/...`` path as a project name and closes the socket
that is not a ``projects/...`` path as a project name and closes the socket. A ``projects/...`` path
picks the project the proxy's token is spent on, so it is moved onto the connection's project, and onto its
region when one is configured, unless ``allow_client_side_credentials`` lets callers choose them
"""
if vertex_project is None or vertex_location is None:
return None
def rewrite(setup_model: str) -> str:
if setup_model.startswith("projects/"):
return setup_model
resource_prefix: Final = _VERTEX_RESOURCE_PROJECT_AND_LOCATION.match(setup_model)
if keeps_caller_resource_path or resource_prefix is None:
return setup_model
location: Final = configured_location or resource_prefix.group("location")
return f"projects/{vertex_project}/locations/{location}/{setup_model[resource_prefix.end() :]}"
aliased: Final = _resolve_alias_to_upstream_model(setup_model, llm_router)
return f"projects/{vertex_project}/locations/{vertex_location}/{_vertex_publisher_model_suffix(aliased)}"
@ -4228,11 +4344,22 @@ async def vertex_ai_live_websocket_passthrough(
await websocket.accept()
incoming_headers: Final = dict(websocket.headers)
vertex_credentials_config: Final = _resolve_vertex_live_credentials(
live_target: Final = _resolve_vertex_live_target(
vertex_project=vertex_project,
vertex_location=vertex_location,
model=model,
)
allow_client_side_credentials: Final = _proxy_general_settings().get("allow_client_side_credentials") is True
override_reason: Final = _live_target_override_reason(
live_target,
vertex_project=vertex_project,
vertex_location=vertex_location,
allow_client_side_credentials=allow_client_side_credentials,
)
if override_reason is not None:
await websocket.close(code=1008, reason=override_reason)
return
vertex_credentials_config: Final = live_target.credentials
configured_project: Final = vertex_project or (
vertex_credentials_config.vertex_project if vertex_credentials_config is not None else None
@ -4314,7 +4441,9 @@ async def vertex_ai_live_websocket_passthrough(
setup_model_rewriter=_build_vertex_live_setup_model_rewriter(
vertex_project=resolved_project,
vertex_location=resolved_location,
configured_location=configured_location,
llm_router=_get_llm_router(),
keeps_caller_resource_path=allow_client_side_credentials,
),
)

View file

@ -2565,6 +2565,21 @@ def _rewrite_vertex_live_setup_model(text_data: str, setup_model_rewriter: Calla
return json.dumps({**message, "setup": {**setup, "model": rewritten_model}})
VERTEX_LIVE_UNDECODABLE_BINARY_CLOSE_REASON: Final = "Rejected Request: binary messages must be UTF-8 encoded JSON"
def _rewrite_vertex_live_binary_setup_model(
bytes_data: bytes, setup_model_rewriter: Callable[[str], str] | None
) -> bytes | None:
if setup_model_rewriter is None:
return bytes_data
try:
text_data: Final = bytes_data.decode("utf-8")
except UnicodeDecodeError:
return None
return _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter).encode("utf-8")
def _resolved_vertex_live_setup(
setup_data: Mapping[str, object], setup_model_rewriter: Callable[[str], str] | None
) -> Mapping[str, object]:
@ -2994,6 +3009,28 @@ async def _relay_client_frame(
return verdict
_UNDECODABLE_BINARY_CLIENT_FRAME: Final = _RefuseClientFrame(
refusal=WebsocketFrameModelRefused(
message=VERTEX_LIVE_UNDECODABLE_BINARY_CLOSE_REASON, error_type="invalid_request_error"
)
)
async def _relay_client_binary_frame(
websocket: WebSocket,
upstream_ws: ClientConnection,
bytes_data: bytes,
setup_model_rewriter: Callable[[str], str] | None,
gate: WebsocketFrameModelGate | None,
ledger: _WebsocketTurnLedger,
) -> _ForwardClientFrame | _RefuseClientFrame:
rewritten_bytes: Final = _rewrite_vertex_live_binary_setup_model(bytes_data, setup_model_rewriter)
if rewritten_bytes is None:
await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=VERTEX_LIVE_UNDECODABLE_BINARY_CLOSE_REASON)
return _UNDECODABLE_BINARY_CLIENT_FRAME
return await _relay_client_frame(websocket, upstream_ws, bytes_data, rewritten_bytes, gate, ledger)
_EXPLICIT_SUCCESS_KWARGS: Final = frozenset({"logging_obj", "passthrough_logging_payload", "custom_llm_provider"})
WALKOUT_TURN_DRAIN_SECONDS: Final = 600.0
@ -3361,8 +3398,13 @@ async def websocket_passthrough_request(
metered_total += _metered_client_text(meter, text_data)
elif bytes_data is not None:
if isinstance(
await _relay_client_frame(
websocket, upstream_ws, bytes_data, bytes_data, client_frame_model_gate, ledger
await _relay_client_binary_frame(
websocket,
upstream_ws,
bytes_data,
setup_model_rewriter,
client_frame_model_gate,
ledger,
),
_RefuseClientFrame,
):

View file

@ -223,6 +223,12 @@ class PassthroughEndpointRouter:
return provider
def get_vertex_credentials_from_router_deployments(self, model: str | None) -> VertexPassThroughCredentials | None:
found: Final = self.get_vertex_deployment_from_router(model)
return found[1] if found is not None else None
def get_vertex_deployment_from_router(
self, model: str | None
) -> tuple[DeploymentTypedDict, VertexPassThroughCredentials] | None:
"""
Resolve vertex pass-through credentials from the live router deployments flagged ``use_in_pass_through``.
@ -242,7 +248,7 @@ class PassthroughEndpointRouter:
)
matched: Final = next(
(
credentials
(deployment, credentials)
for deployment, credentials in resolved
if model is not None and self._deployment_matches_model(deployment, model)
),
@ -260,7 +266,7 @@ class PassthroughEndpointRouter:
)
if len(targets) != 1:
return None
return resolved[0][1]
return resolved[0]
def _resolve_vertex_deployment_credentials(
self, litellm_params: LiteLLMParamsTypedDict

View file

@ -8,6 +8,7 @@ Provides:
import base64
import json
import re
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
@ -20,12 +21,13 @@ from starlette.datastructures import UploadFile
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH, LITELLM_PROXY_ADMIN_NAME
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
LiteLLM_ManagedVectorStore,
)
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.proxy._types import *
from litellm.proxy.auth.auth_utils import is_request_body_safe
from litellm.proxy.auth.auth_utils import get_request_route, is_request_body_safe
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
@ -59,8 +61,11 @@ from litellm.proxy.vector_store_endpoints.endpoints import (
from litellm.proxy.vector_store_endpoints.utils import (
assert_user_can_access_vector_store_id,
)
from litellm.rag.ingestion.vertex_ai_ingestion import DEFAULT_VERTEX_RAG_LOCATION
from litellm.rag.main import get_ingestion_class
from litellm.repositories.table_repositories import ManagedVectorStoresRepository
from litellm.repositories.user_repository import UserRepository
from litellm.secret_managers.main import get_secret_str
from litellm.types.utils import ModelResponse
if TYPE_CHECKING:
@ -179,6 +184,103 @@ def _ingest_provider_error(vector_store_config: Mapping[str, object]) -> str | N
return None
_VERTEX_RAG_CORPUS_NAME: Final = re.compile(
r"projects/(?P<project>[^/%\\]+)/locations/(?P<location>[^/%\\]+)/ragCorpora/(?P<corpus>[^/%\\]+)"
)
_DOT_SEGMENTS: Final = frozenset({".", ".."})
def _clean_vertex_corpus_name(corpus_name: object) -> re.Match[str] | None:
if not isinstance(corpus_name, str) or not _DOT_SEGMENTS.isdisjoint(corpus_name.split("/")):
return None
return _VERTEX_RAG_CORPUS_NAME.fullmatch(corpus_name)
def _vertex_target_message(field: str) -> str:
return (
f"Rejected Request: {field} picks a Vertex project or region other than the proxy's. "
"Use a vector store a proxy admin created, or set general_settings.allow_client_side_credentials"
)
def _vertex_target_is_checked(
vector_store_config: Mapping[str, object],
user_api_key_dict: UserAPIKeyAuth,
general_settings: Mapping[str, object],
) -> bool:
return (
vector_store_config.get("custom_llm_provider") == "vertex_ai"
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN
and general_settings.get("allow_client_side_credentials") is not True
)
def _configured_vertex_project() -> str | None:
return litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT")
def _corpus_on_configured_project(vector_store_config: Mapping[str, object]) -> Mapping[str, object]:
corpus: Final = _clean_vertex_corpus_name(vector_store_config.get("vector_store_id"))
configured_project: Final = _configured_vertex_project()
if corpus is None or configured_project is None or not corpus["project"].isdigit():
return vector_store_config
return MappingProxyType(
{
**vector_store_config,
"vector_store_id": (
f"projects/{configured_project}/locations/{corpus['location']}/ragCorpora/{corpus['corpus']}"
),
}
)
async def _authorize_rewritten_corpus(
requested: Mapping[str, object], rewritten: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
) -> None:
rewritten_id: Final = rewritten.get("vector_store_id")
if isinstance(rewritten_id, str) and rewritten_id != requested.get("vector_store_id"):
await assert_user_can_access_vector_store_id(vector_store_id=rewritten_id, user_api_key_dict=user_api_key_dict)
def _vertex_target_error(vector_store_config: Mapping[str, object]) -> str | None:
corpus_name: Final = vector_store_config.get("vector_store_id")
corpus: Final = _clean_vertex_corpus_name(corpus_name)
if corpus_name and corpus is None:
return _vertex_target_message("vector_store_id")
configured_project: Final = _configured_vertex_project()
configured_location: Final = VertexBase.safe_get_vertex_ai_location({}) or DEFAULT_VERTEX_RAG_LOCATION
requested: Final = (
*(
(name, vector_store_config.get(name), configured_project)
for name in ("vertex_project", "vertex_ai_project")
),
*(
(name, vector_store_config.get(name), configured_location)
for name in ("vertex_location", "vertex_ai_location")
),
("vector_store_id", corpus["project"] if corpus is not None else None, configured_project),
("vector_store_id", corpus["location"] if corpus is not None else None, configured_location),
)
field: Final = next((name for name, asked, configured in requested if asked and asked != configured), None)
return None if field is None else _vertex_target_message(field)
async def _admin_set_store_target(
managed_store: LiteLLM_ManagedVectorStore | None, prisma_client: "PrismaClient | None"
) -> bool:
if managed_store is None:
return False
if managed_store.get("is_config", False):
return True
creator: Final = managed_store.get("user_id")
if creator == LITELLM_PROXY_ADMIN_NAME:
return True
if creator is None or prisma_client is None:
return False
creator_row: Final = await UserRepository(prisma_client).find_by_id(creator)
return creator_row is not None and creator_row.user_role == LitellmUserRoles.PROXY_ADMIN
_MANAGED_STORE_CALLER_OPTIONS: Final = frozenset(
{
"vector_store_id",
@ -660,6 +762,7 @@ async def rag_ingest(
from litellm.proxy.proxy_server import (
add_litellm_data_to_request,
general_settings,
general_settings_view,
llm_router,
prisma_client,
proxy_config,
@ -684,15 +787,24 @@ async def rag_ingest(
general_settings=general_settings,
llm_router=llm_router,
model="",
route=get_request_route(request=request),
)
except ValueError as e:
raise HTTPException(status_code=400, detail={"error": str(e)})
managed_store: Final = resolved_stores.get(request_vector_store_config.get("vector_store_id"))
base: Final = {
requested_base: Final = {
**_caller_vector_store_options(request_vector_store_config, managed_store),
**_managed_store_overrides(managed_store),
}
vertex_target_checked: Final = _vertex_target_is_checked(
requested_base, user_api_key_dict, general_settings_view()
) and not await _admin_set_store_target(managed_store, prisma_client)
base: Final = _corpus_on_configured_project(requested_base) if vertex_target_checked else requested_base
await _authorize_rewritten_corpus(requested_base, base, user_api_key_dict)
vertex_target_error: Final = _vertex_target_error(base) if vertex_target_checked else None
if vertex_target_error is not None:
raise HTTPException(status_code=400, detail={"error": vertex_target_error})
request_data: Final[dict[str, Any]] = await add_litellm_data_to_request(
data={},
request=request,

View file

@ -28,6 +28,8 @@ if TYPE_CHECKING:
from litellm import Router
from litellm.types.rag import RAGIngestOptions
DEFAULT_VERTEX_RAG_LOCATION: Final = "us-central1"
def _present_fields(fields: tuple[tuple[str, object], ...]) -> Mapping[str, object]:
return {name: value for name, value in fields if value}
@ -90,7 +92,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase):
# Get project, location, and credentials using VertexBase methods
self.project_id = self.safe_get_vertex_ai_project(litellm_params)
self.location = self.get_vertex_ai_location(litellm_params) or "us-central1"
self.location = self.get_vertex_ai_location(litellm_params) or DEFAULT_VERTEX_RAG_LOCATION
self.vertex_credentials = self.safe_get_vertex_ai_credentials(litellm_params)
async def embed(

View file

@ -100,6 +100,7 @@ def _deployment(
location: Mapping[str, JsonValue],
multiplier: JsonValue = 1.1,
backend: str = _IMAGE_BACKEND,
clientside_params: tuple[str, ...] = (),
) -> str:
return scenario.model(
model=f"vertex_ai/{backend}",
@ -111,6 +112,7 @@ def _deployment(
output_cost_per_token=4e-5,
output_cost_per_image=0.05,
**({} if multiplier is None else {"regional_endpoint_uplift_multiplier": multiplier}),
**({"configurable_clientside_auth_params": list(clientside_params)} if clientside_params else {}),
**location,
)
@ -380,7 +382,13 @@ def test_request_body_location_overrides_the_deployment(
gateway: Gateway, deployment_location: str, extra: Mapping[str, JsonValue], expected: float
) -> None:
with wire_server(_scripted()) as wire, gateway.scenario() as scenario:
model: Final = _deployment(gateway, scenario, wire, location={"vertex_location": deployment_location})
model: Final = _deployment(
gateway,
scenario,
wire,
location={"vertex_location": deployment_location},
clientside_params=("vertex_location", "vertex_ai_location"),
)
call_id, response = _generate(gateway, model, extra=extra)
_assert_images(response, call_id, expected)
_assert_one_upstream_image_call(wire)

View file

@ -4,6 +4,7 @@ Unit tests for auth_utils functions related to rate limiting and customer ID ext
import base64
import logging
import time
from collections.abc import Callable
from typing import Final, Optional
from unittest.mock import MagicMock, patch
@ -34,6 +35,7 @@ from litellm.proxy.auth.auth_utils import (
get_project_model_tpm_limit,
get_request_route_template,
is_request_body_safe,
pre_db_read_auth_checks,
)
from litellm.router import Router
from litellm.types.utils import oauth_token_exchange_litellm_params, server_owned_wif_litellm_params
@ -4978,6 +4980,426 @@ class TestIsRequestBodySafeBlocksAwsIdentitySelectors:
)
_VERTEX_TARGET_PARAMS: Final = ("vertex_location", "vertex_project", "vertex_ai_location", "vertex_ai_project")
def _vertex_pinned_router(configurable_clientside_auth_params: list[str]) -> Router:
return Router(
model_list=[
{
"model_name": "gemini-2.5-flash",
"litellm_params": {
"model": "vertex_ai/gemini-2.5-flash",
"vertex_project": "admin-project",
"vertex_location": "europe-west4",
"configurable_clientside_auth_params": configurable_clientside_auth_params,
},
}
]
)
class TestIsRequestBodySafeBlocksVertexTargetOverrides:
@pytest.mark.parametrize("param", _VERTEX_TARGET_PARAMS)
@pytest.mark.parametrize(
"placement",
[
lambda param: {param: "us-central1"},
lambda param: {"extra_body": {param: "us-central1"}},
],
ids=["root", "extra_body"],
)
def test_vertex_target_is_rejected_without_admin_opt_in(
self, param: str, placement: Callable[[str], dict[str, object]]
) -> None:
with pytest.raises(ValueError, match=param):
is_request_body_safe(
request_body={"model": "gemini-2.5-flash", **placement(param)},
general_settings={},
llm_router=_vertex_pinned_router([]),
model="gemini-2.5-flash",
route="/v1/chat/completions",
)
@pytest.mark.parametrize("param", _VERTEX_TARGET_PARAMS)
@pytest.mark.parametrize(
"placement",
[
lambda param: {"metadata": {param: "us-central1"}},
lambda param: {"litellm_metadata": {param: "us-central1"}},
lambda param: {f"metadata[{param}]": "us-central1"},
lambda param: {"litellm_params": {"metadata": {param: "us-central1"}}},
],
ids=["metadata", "litellm_metadata", "form_metadata", "litellm_params_metadata"],
)
def test_vertex_target_in_metadata_never_reaches_the_provider_and_is_allowed(
self, param: str, placement: Callable[[str], dict[str, object]]
) -> None:
assert (
is_request_body_safe(
request_body={"model": "gemini-2.5-flash", **placement(param)},
general_settings={},
llm_router=_vertex_pinned_router([]),
model="gemini-2.5-flash",
route="/v1/chat/completions",
)
is True
)
def test_vertex_target_is_allowed_under_proxy_wide_opt_in(self) -> None:
assert (
is_request_body_safe(
request_body={"model": "gemini-2.5-flash", **{param: "us-central1" for param in _VERTEX_TARGET_PARAMS}},
general_settings={"allow_client_side_credentials": True},
llm_router=None,
model="gemini-2.5-flash",
route="/v1/chat/completions",
)
is True
)
def test_deployment_opt_in_allows_only_the_listed_vertex_target(self) -> None:
router: Final = _vertex_pinned_router(["vertex_location"])
assert (
is_request_body_safe(
request_body={"model": "gemini-2.5-flash", "vertex_location": "us-central1"},
general_settings={},
llm_router=router,
model="gemini-2.5-flash",
route="/v1/chat/completions",
)
is True
)
with pytest.raises(ValueError, match="vertex_project"):
is_request_body_safe(
request_body={
"model": "gemini-2.5-flash",
"vertex_location": "us-central1",
"vertex_project": "caller-project",
},
general_settings={},
llm_router=router,
model="gemini-2.5-flash",
route="/v1/chat/completions",
)
@pytest.mark.parametrize("opted_in_first", [True, False])
def test_group_opt_in_does_not_count_while_any_deployment_lacks_it(self, opted_in_first: bool) -> None:
opted_in: Final = {
"model_name": "gemini-group",
"litellm_params": {
"model": "vertex_ai/gemini-2.5-flash",
"vertex_location": "us-central1",
"configurable_clientside_auth_params": ["vertex_location"],
},
}
pinned: Final = {
"model_name": "gemini-group",
"litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "vertex_location": "europe-west4"},
}
router: Final = Router(model_list=[opted_in, pinned] if opted_in_first else [pinned, opted_in])
with pytest.raises(ValueError, match="vertex_location"):
is_request_body_safe(
request_body={"model": "gemini-group", "vertex_location": "us-central1"},
general_settings={},
llm_router=router,
model="gemini-group",
route="/v1/chat/completions",
)
def test_model_with_no_deployment_has_no_opt_in(self) -> None:
with pytest.raises(ValueError, match="vertex_location"):
is_request_body_safe(
request_body={"model": "not-configured", "vertex_location": "us-central1"},
general_settings={},
llm_router=_vertex_pinned_router(["vertex_location"]),
model="not-configured",
route="/v1/chat/completions",
)
def test_group_opt_in_counts_when_every_deployment_carries_it(self) -> None:
router: Final = Router(
model_list=[
{
"model_name": "gemini-group",
"litellm_params": {
"model": "vertex_ai/gemini-2.5-flash",
"vertex_location": location,
"configurable_clientside_auth_params": ["vertex_location"],
},
}
for location in ("us-central1", "europe-west4")
]
)
assert (
is_request_body_safe(
request_body={"model": "gemini-group", "vertex_location": "us-central1"},
general_settings={},
llm_router=router,
model="gemini-group",
route="/v1/chat/completions",
)
is True
)
@pytest.mark.parametrize(
"model, is_allowed", [("public-alias", False), ("public-gemini-2.5-flash", True)], ids=["alias", "direct"]
)
def test_alias_opt_in_comes_from_the_group_the_alias_routes_to(self, model: str, is_allowed: bool) -> None:
router: Final = Router(
model_list=[
{
"model_name": "gemini-eu/*",
"litellm_params": {"model": "vertex_ai/*", "vertex_location": "europe-west4"},
},
{
"model_name": "public-*",
"litellm_params": {
"model": "vertex_ai/*",
"vertex_location": "europe-west4",
"configurable_clientside_auth_params": ["vertex_location"],
},
},
],
model_group_alias={"public-alias": "gemini-eu/gemini-2.5-flash"},
)
request_body: Final = {"model": model, "vertex_location": "us-central1"}
if is_allowed:
assert is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions") is True
return
with pytest.raises(ValueError, match="vertex_location"):
is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions")
@staticmethod
def _opted_in_and_pinned_router(
router_fallbacks: list[dict[str, list[str]]],
fallback_setting: str = "fallbacks",
model_group_alias: dict[str, str] | None = None,
) -> Router:
return Router(
model_list=[
{
"model_name": name,
"litellm_params": {
"model": "vertex_ai/gemini-2.5-flash",
"vertex_location": "europe-west4",
"configurable_clientside_auth_params": opt_in,
},
}
for name, opt_in in (
("gemini-optin", ["vertex_location"]),
("gemini-optin-other", ["vertex_location"]),
("gemini-pinned", []),
)
],
model_group_alias=model_group_alias or {},
**{fallback_setting: router_fallbacks},
)
@pytest.mark.parametrize(
"fields, is_allowed",
[
({"vertex_location": "us-central1", "fallbacks": ["gemini-pinned"]}, False),
({"fallbacks": [{"model": "gemini-pinned", "vertex_location": "us-central1"}]}, False),
({"vertex_location": "us-central1", "fallbacks": ["gemini-optin-other"]}, True),
({"fallbacks": [{"model": "gemini-optin-other", "vertex_location": "us-central1"}]}, True),
],
ids=["pinned", "pinned_entry_with_field", "opted_in", "opted_in_entry_with_field"],
)
def test_request_fallback_needs_the_same_opt_in(self, fields: dict[str, object], is_allowed: bool) -> None:
router: Final = self._opted_in_and_pinned_router([])
request_body: Final = {"model": "gemini-optin", **fields}
if is_allowed:
assert is_request_body_safe(request_body, {}, router, "gemini-optin", route="/v1/chat/completions") is True
return
with pytest.raises(ValueError, match="vertex_location"):
is_request_body_safe(request_body, {}, router, "gemini-optin", route="/v1/chat/completions")
@pytest.mark.parametrize("fallback_setting", ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"])
@pytest.mark.parametrize(
"fallback_group, is_allowed",
[("gemini-pinned", False), ("gemini-optin-other", True)],
ids=["pinned", "opted_in"],
)
def test_router_fallback_needs_the_same_opt_in(
self, fallback_setting: str, fallback_group: str, is_allowed: bool
) -> None:
router: Final = self._opted_in_and_pinned_router([{"gemini-optin": [fallback_group]}], fallback_setting)
request_body: Final = {"model": "gemini-optin", "vertex_location": "us-central1"}
if is_allowed:
assert is_request_body_safe(request_body, {}, router, "gemini-optin", route="/v1/chat/completions") is True
return
with pytest.raises(ValueError, match="vertex_location"):
is_request_body_safe(request_body, {}, router, "gemini-optin", route="/v1/chat/completions")
@pytest.mark.parametrize(
"model, fields, is_allowed",
[
("gemini-optin", {"vertex_location": "us-central1"}, True),
("openai-group", {"vertex_location": "us-central1"}, True),
("gemini-optin", {"api_base": "https://eu.example.com"}, False),
("studio-group", {"vertex_location": "us-central1"}, False),
],
ids=["vertex_field_to_openai_fallback", "vertex_field_on_openai", "api_base_to_openai_fallback", "gemini"],
)
def test_only_google_groups_need_the_vertex_opt_in(
self, model: str, fields: dict[str, object], is_allowed: bool
) -> None:
router: Final = Router(
model_list=[
{
"model_name": "gemini-optin",
"litellm_params": {
"model": "vertex_ai/gemini-2.5-flash",
"configurable_clientside_auth_params": ["vertex_location", {"api_base": r"https://eu\..*"}],
},
},
{"model_name": "openai-group", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}},
{"model_name": "studio-group", "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "k"}},
],
fallbacks=[{"gemini-optin": ["openai-group"]}, {"studio-group": ["openai-group"]}],
)
request_body: Final = {"model": model, **fields}
if is_allowed:
assert is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions") is True
return
with pytest.raises(ValueError, match=r"vertex_location|api_base"):
is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions")
def test_many_fallback_entries_on_one_wildcard_are_checked_in_linear_time(self) -> None:
router: Final = Router(model_list=[{"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "k"}}])
request_body: Final = {
"model": "gpt-4o-mini",
"fallbacks": [{"model": f"model-{index % 100}", "vertex_location": f"r{index}"} for index in range(5000)],
}
started: Final = time.perf_counter()
assert is_request_body_safe(request_body, {}, router, "gpt-4o-mini", route="/v1/chat/completions") is True
assert time.perf_counter() - started < 10
def test_too_many_fallback_models_for_an_opt_in_check_are_rejected(self) -> None:
router: Final = Router(model_list=[{"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "k"}}])
request_body: Final = {
"model": "gpt-4o-mini",
"vertex_location": "us-central1",
"fallbacks": [f"model-{index}" for index in range(101)],
}
with pytest.raises(ValueError, match="more than 100 fallback models"):
is_request_body_safe(request_body, {}, router, "gpt-4o-mini", route="/v1/chat/completions")
def test_one_wildcard_serving_two_providers_checks_the_google_one(self) -> None:
router: Final = Router(model_list=[{"model_name": "*", "litellm_params": {"model": "*"}}])
request_body: Final = {
"model": "vertex_ai/gemini-2.5-flash",
"vertex_location": "us-central1",
"fallbacks": ["xai/grok-3"],
}
with pytest.raises(ValueError, match="vertex_location"):
is_request_body_safe(request_body, {}, router, "vertex_ai/gemini-2.5-flash", route="/v1/chat/completions")
def test_a_pinned_group_sharing_a_deployment_id_with_an_opted_in_group_is_still_checked(self) -> None:
router: Final = Router(
model_list=[
{
"model_name": name,
"litellm_params": {
"model": "vertex_ai/gemini-2.5-flash",
"vertex_location": "europe-west4",
"configurable_clientside_auth_params": opt_in,
},
"model_info": {"id": "shared-id"},
}
for name, opt_in in (("a-pinned", []), ("z-optin", ["vertex_location"]))
]
)
request_body: Final = {"model": "a-pinned", "vertex_location": "us-central1", "fallbacks": ["z-optin"]}
with pytest.raises(ValueError, match="vertex_location"):
is_request_body_safe(request_body, {}, router, "a-pinned", route="/v1/chat/completions")
@pytest.mark.parametrize(
"model, is_allowed", [("public-alias", True), ("gemini-optin", False)], ids=["alias", "aliased_target"]
)
def test_alias_chain_is_checked_one_hop_like_the_router_serves_it(self, model: str, is_allowed: bool) -> None:
router: Final = self._opted_in_and_pinned_router(
[], model_group_alias={"public-alias": "gemini-optin", "gemini-optin": "gemini-pinned"}
)
request_body: Final = {"model": model, "vertex_location": "us-central1"}
if is_allowed:
assert is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions") is True
return
with pytest.raises(ValueError, match="vertex_location"):
is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions")
@pytest.mark.parametrize(
"route",
[
"/openai/deployments/gemini-pinned/chat/completions",
"/v1beta/models/gemini-pinned:generateContent",
],
)
def test_model_named_in_the_route_needs_the_opt_in_too(self, route: str) -> None:
router: Final = self._opted_in_and_pinned_router([])
request_body: Final = {"model": "gemini-optin", "vertex_location": "us-central1"}
with pytest.raises(ValueError, match="vertex_location"):
is_request_body_safe(request_body, {}, router, "gemini-optin", route=route)
@pytest.mark.asyncio
@pytest.mark.parametrize("opted_in", [True, False])
async def test_opt_in_of_the_model_named_in_a_google_route_decides(self, monkeypatch, opted_in: bool) -> None:
route: Final = "/v1beta/models/gemini-2.5-flash:generateContent"
monkeypatch.setattr(
"litellm.proxy.proxy_server.llm_router", _vertex_pinned_router(["vertex_location"] if opted_in else [])
)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
request: Final = Request({"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""})
if opted_in:
await pre_db_read_auth_checks(request=request, request_data={"vertex_location": "us-central1"}, route=route)
return
with pytest.raises(ValueError, match="vertex_location"):
await pre_db_read_auth_checks(request=request, request_data={"vertex_location": "us-central1"}, route=route)
@pytest.mark.parametrize(
"route, kept_params",
[
("/vertex_ai/live", _VERTEX_TARGET_PARAMS),
("/rag/ingest", _VERTEX_TARGET_PARAMS),
("/v1/rag/ingest", _VERTEX_TARGET_PARAMS),
],
)
def test_routes_that_name_their_vertex_target_keep_it(self, route: str, kept_params: tuple[str, ...]) -> None:
assert (
is_request_body_safe(
request_body={param: "caller-chosen" for param in kept_params},
general_settings={},
llm_router=None,
model="",
route=route,
)
is True
)
@pytest.mark.parametrize(
"route, param",
[
("/vertex_ai/live", "api_base"),
("/rag/ingest", "api_base"),
("/v1/rag/ingest", "api_base"),
],
)
def test_routes_that_name_their_vertex_target_reject_every_other_banned_param(self, route: str, param: str) -> None:
with pytest.raises(ValueError, match=param):
is_request_body_safe(
request_body={"vertex_location": "us-central1", param: "caller-chosen"},
general_settings={},
llm_router=None,
model="",
route=route,
)
@pytest.mark.parametrize(
"allowed_param, input_value, should_return_true",
[

View file

@ -252,6 +252,9 @@ class FakeRouter:
def registry(self) -> "FakeRegistry":
return FakeRegistry()
def settings(self) -> "FakeSettings":
return FakeSettings()
def get_model_group_info(self, model_group: str) -> None:
return None
@ -264,6 +267,17 @@ class FakeRegistry:
def group(self, name: str, *, field: str = "info") -> None:
return None
def resolve(self, name: str, *, kind: str = "group") -> None:
return None
def deployments(self, model_name: str | None = None, *, source: str = "group") -> list[object]:
return []
class FakeSettings:
def is_set(self, name: str) -> bool:
return False
class FakeStorageBackend:
def __init__(self, contents: Mapping[str, bytes]) -> None:

View file

@ -5271,6 +5271,12 @@ class TestConfigBaseForHealthCheck:
assert "vertex_credentials" not in base
assert base["rpm"] == 100
@pytest.mark.parametrize("param", ["vertex_project", "vertex_location", "vertex_ai_project", "vertex_ai_location"])
def test_request_naming_a_vertex_target_still_inherits_config_credentials(self, param: str) -> None:
base = self._base(self.CONFIG, {"model": "openai/gpt-4o", param: "us-central1"}, selected_by_id=False)
assert base["vertex_credentials"] == "configured-creds"
assert base["api_key"] == "sk-configured"
def test_add_model_flow_keeps_its_own_credentials(self):
"""Adding a second deployment for an already-configured name sends a
complete connection; it is tested as sent, not as configured."""

View file

@ -6374,7 +6374,11 @@ class TestVertexAILiveWebsocketPassthrough:
),
(
"projects/other/locations/us-central1/publishers/google/models/gemini-2.0-flash",
"projects/other/locations/us-central1/publishers/google/models/gemini-2.0-flash",
"projects/proj-db/locations/global/publishers/google/models/gemini-2.0-flash",
),
(
"projects/other/locations/us-central1/endpoints/1234",
"projects/proj-db/locations/global/endpoints/1234",
),
],
)
@ -6400,7 +6404,9 @@ class TestVertexAILiveWebsocketPassthrough:
rewriter = passthrough_module._build_vertex_live_setup_model_rewriter(
vertex_project="proj-db",
vertex_location="global",
configured_location="global",
llm_router=llm_router,
keeps_caller_resource_path=False,
)
assert rewriter is not None
@ -6492,6 +6498,402 @@ class TestVertexAILiveWebsocketPassthrough:
assert "default_vertex_config" in close_kwargs["reason"]
assert len(close_kwargs["reason"].encode("utf-8")) <= 123
def _pinned_live_router(self, extra_params: Mapping[str, object]) -> litellm.Router:
return litellm.Router(
model_list=[
{
"model_name": "gemini-live",
"litellm_params": {
"model": "vertex_ai/gemini-live-2.5-flash",
"use_in_pass_through": True,
"vertex_project": "proj-db",
"vertex_location": "global",
"vertex_credentials": '{"type": "service_account"}',
**extra_params,
},
}
]
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"query, field",
[
({"vertex_location": "us-central1"}, "vertex_location"),
({"vertex_project": "caller-project"}, "vertex_project"),
({"vertex_project": "proj-db", "vertex_location": "us-central1"}, "vertex_location"),
],
)
async def test_query_overriding_a_pinned_deployment_closes_before_minting_a_token(
self, monkeypatch, query: Mapping[str, str], field: str
):
from litellm.proxy.pass_through_endpoints import (
llm_passthrough_endpoints as passthrough_module,
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", self._pinned_live_router({}))
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "default_vertex_config", None)
self._clear_vertex_env(monkeypatch)
websocket = self._websocket()
ensure_token = AsyncMock(return_value=("token-abc", "proj-db"))
ws_passthrough = AsyncMock()
with (
patch.object(passthrough_module.vertex_llm_base, "ensure_access_token_async", ensure_token),
patch.object(passthrough_module, "websocket_passthrough_request", ws_passthrough),
):
await passthrough_module.vertex_ai_live_websocket_passthrough(
websocket=websocket,
user_api_key_dict=UserAPIKeyAuth(),
**query,
)
close_kwargs = websocket.close.await_args.kwargs
assert close_kwargs["code"] == 1008
assert field in close_kwargs["reason"]
assert len(close_kwargs["reason"].encode("utf-8")) <= 123
ensure_token.assert_not_awaited()
ws_passthrough.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"extra_params, general_settings, query, expected_location",
[
({}, {}, {"vertex_project": "proj-db", "vertex_location": "global"}, "global"),
(
{"configurable_clientside_auth_params": ["vertex_location"]},
{},
{"vertex_location": "us-central1"},
"us-central1",
),
({}, {"allow_client_side_credentials": True}, {"vertex_location": "us-central1"}, "us-central1"),
],
ids=["matches_the_pin", "deployment_opt_in", "proxy_wide_opt_in"],
)
async def test_query_the_admin_allows_reaches_vertex(
self,
monkeypatch,
extra_params: Mapping[str, object],
general_settings: Mapping[str, object],
query: Mapping[str, str],
expected_location: str,
):
from litellm.proxy.pass_through_endpoints import (
llm_passthrough_endpoints as passthrough_module,
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", self._pinned_live_router(extra_params))
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", dict(general_settings))
monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "default_vertex_config", None)
self._clear_vertex_env(monkeypatch)
websocket = self._websocket()
ensure_token = AsyncMock(return_value=("token-abc", "proj-db"))
ws_passthrough = AsyncMock()
with (
patch.object(passthrough_module.vertex_llm_base, "ensure_access_token_async", ensure_token),
patch.object(passthrough_module, "websocket_passthrough_request", ws_passthrough),
):
await passthrough_module.vertex_ai_live_websocket_passthrough(
websocket=websocket,
user_api_key_dict=UserAPIKeyAuth(),
**query,
)
websocket.close.assert_not_awaited()
rewriter = ws_passthrough.await_args.kwargs["setup_model_rewriter"]
assert rewriter("gemini-live") == (
f"projects/proj-db/locations/{expected_location}/publishers/google/models/gemini-live-2.5-flash"
)
def _use_default_vertex_config(self, monkeypatch, general_settings: Mapping[str, object], vertex_location):
from litellm.proxy.pass_through_endpoints import (
llm_passthrough_endpoints as passthrough_module,
)
from litellm.types.passthrough_endpoints.vertex_ai import (
VertexPassThroughCredentials,
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", dict(general_settings))
monkeypatch.setattr(
passthrough_module.passthrough_endpoint_router,
"default_vertex_config",
VertexPassThroughCredentials(
vertex_project="proj-env",
vertex_location=vertex_location,
vertex_credentials='{"type": "env_account"}',
),
)
monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "deployment_key_to_vertex_credentials", {})
self._clear_vertex_env(monkeypatch)
async def _open_live(self, query: Mapping[str, str]):
from litellm.proxy.pass_through_endpoints import (
llm_passthrough_endpoints as passthrough_module,
)
websocket = self._websocket()
ensure_token = AsyncMock(
side_effect=lambda credentials, project_id, custom_llm_provider: ("token-abc", project_id or "proj-env")
)
ws_passthrough = AsyncMock()
with (
patch.object(passthrough_module.vertex_llm_base, "ensure_access_token_async", ensure_token),
patch.object(passthrough_module, "websocket_passthrough_request", ws_passthrough),
):
await passthrough_module.vertex_ai_live_websocket_passthrough(
websocket=websocket,
user_api_key_dict=UserAPIKeyAuth(),
**query,
)
return websocket, ensure_token, ws_passthrough
@pytest.mark.asyncio
@pytest.mark.parametrize(
"query, field",
[
({"vertex_project": "caller-project"}, "vertex_project"),
({"vertex_location": "us-central1"}, "vertex_location"),
({"vertex_project": "proj-env", "vertex_location": "us-central1"}, "vertex_location"),
],
)
async def test_query_moving_default_vertex_config_closes_before_minting_a_token(
self, monkeypatch, query: Mapping[str, str], field: str
):
self._use_default_vertex_config(monkeypatch, {}, vertex_location="europe-west4")
websocket, ensure_token, ws_passthrough = await self._open_live(query)
close_kwargs = websocket.close.await_args.kwargs
assert close_kwargs["code"] == 1008
assert f"{field} is pinned by default_vertex_config" in close_kwargs["reason"]
assert len(close_kwargs["reason"].encode("utf-8")) <= 123
ensure_token.assert_not_awaited()
ws_passthrough.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"general_settings, default_location, query, expected_project, expected_location",
[
(
{},
"europe-west4",
{"vertex_project": "proj-env", "vertex_location": "europe-west4"},
"proj-env",
"europe-west4",
),
({}, None, {"vertex_location": "us-central1"}, "proj-env", "us-central1"),
(
{"allow_client_side_credentials": True},
"europe-west4",
{"vertex_project": "caller-project", "vertex_location": "us-central1"},
"caller-project",
"us-central1",
),
],
ids=["matches_the_pin", "location_left_unset", "proxy_wide_opt_in"],
)
async def test_default_vertex_config_serves_the_queries_it_allows(
self,
monkeypatch,
general_settings: Mapping[str, object],
default_location,
query: Mapping[str, str],
expected_project: str,
expected_location: str,
):
self._use_default_vertex_config(monkeypatch, general_settings, vertex_location=default_location)
websocket, ensure_token, ws_passthrough = await self._open_live(query)
websocket.close.assert_not_awaited()
assert ensure_token.await_args.kwargs["project_id"] == expected_project
rewriter = ws_passthrough.await_args.kwargs["setup_model_rewriter"]
assert rewriter("gemini-live-2.5-flash") == (
f"projects/{expected_project}/locations/{expected_location}/publishers/google/models/gemini-live-2.5-flash"
)
@pytest.mark.asyncio
async def test_a_registered_project_and_location_pair_still_serves_its_query(self, monkeypatch):
from litellm.proxy.pass_through_endpoints import (
llm_passthrough_endpoints as passthrough_module,
)
from litellm.types.passthrough_endpoints.vertex_ai import (
VertexPassThroughCredentials,
)
self._use_default_vertex_config(monkeypatch, {}, vertex_location="europe-west4")
monkeypatch.setattr(
passthrough_module.passthrough_endpoint_router,
"deployment_key_to_vertex_credentials",
{
"caller-project-us-central1": VertexPassThroughCredentials(
vertex_project="caller-project",
vertex_location="us-central1",
vertex_credentials='{"type": "registered_account"}',
)
},
)
websocket, ensure_token, _ = await self._open_live(
{"vertex_project": "caller-project", "vertex_location": "us-central1"}
)
websocket.close.assert_not_awaited()
ensure_token.assert_awaited_once_with(
credentials='{"type": "registered_account"}',
project_id="caller-project",
custom_llm_provider="vertex_ai_beta",
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"general_settings, expected",
[
({}, "projects/proj-env/locations/europe-west4/publishers/google/models/gemini-2.0-flash"),
(
{"allow_client_side_credentials": True},
"projects/caller-project/locations/us-central1/publishers/google/models/gemini-2.0-flash",
),
],
ids=["pinned", "proxy_wide_opt_in"],
)
async def test_setup_frame_resource_path_cannot_move_the_session_to_another_project(
self, monkeypatch, general_settings: Mapping[str, object], expected: str
):
self._use_default_vertex_config(monkeypatch, general_settings, vertex_location="europe-west4")
_, _, ws_passthrough = await self._open_live({})
rewriter = ws_passthrough.await_args.kwargs["setup_model_rewriter"]
assert (
rewriter("projects/caller-project/locations/us-central1/publishers/google/models/gemini-2.0-flash")
== expected
)
def _pin_no_vertex_project(
self,
monkeypatch,
general_settings: Mapping[str, object],
default_vertex_config: VertexPassThroughCredentials | None,
):
from litellm.proxy.pass_through_endpoints import (
llm_passthrough_endpoints as passthrough_module,
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", dict(general_settings))
monkeypatch.setattr(
passthrough_module.passthrough_endpoint_router, "default_vertex_config", default_vertex_config
)
monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "deployment_key_to_vertex_credentials", {})
self._clear_vertex_env(monkeypatch)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"default_vertex_config",
[None, VertexPassThroughCredentials(vertex_location="europe-west4")],
ids=["no_default_vertex_config", "default_vertex_config_without_a_project"],
)
async def test_caller_named_project_closes_before_any_upstream_connect_when_nothing_pins_one(
self, monkeypatch, default_vertex_config: VertexPassThroughCredentials | None
):
self._pin_no_vertex_project(monkeypatch, {}, default_vertex_config)
websocket, ensure_token, ws_passthrough = await self._open_live({"vertex_project": "other-project"})
ws_passthrough.assert_not_awaited()
ensure_token.assert_not_awaited()
close_kwargs = websocket.close.await_args.kwargs
assert close_kwargs["code"] == 1008
assert "vertex_project" in close_kwargs["reason"]
assert len(close_kwargs["reason"].encode("utf-8")) <= 123
@pytest.mark.asyncio
@pytest.mark.parametrize(
"configured_by, default_vertex_config",
[
("VERTEXAI_PROJECT", None),
("DEFAULT_VERTEXAI_PROJECT", VertexPassThroughCredentials(vertex_location="europe-west4")),
],
ids=["vertexai_project_env", "default_vertexai_project_env"],
)
async def test_with_nothing_pinned_the_caller_may_name_the_project_the_admin_configured(
self, monkeypatch, configured_by: str, default_vertex_config: VertexPassThroughCredentials | None
):
self._pin_no_vertex_project(monkeypatch, {}, default_vertex_config)
monkeypatch.setattr(litellm, "vertex_project", None)
monkeypatch.delenv("VERTEXAI_PROJECT", raising=False)
monkeypatch.setenv(configured_by, "proj-configured")
websocket, _, ws_passthrough = await self._open_live(
{"vertex_project": "proj-configured", "vertex_location": "europe-west4"}
)
websocket.close.assert_not_awaited()
assert ws_passthrough.await_args.kwargs["custom_headers"]["x-goog-user-project"] == "proj-configured"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"general_settings, query, expected_project, expected_location",
[
({}, {"vertex_location": "europe-west4"}, "proj-env", "europe-west4"),
(
{"allow_client_side_credentials": True},
{"vertex_project": "other-project", "vertex_location": "europe-west4"},
"other-project",
"europe-west4",
),
],
ids=["caller_region", "proxy_wide_opt_in"],
)
async def test_with_nothing_pinned_the_caller_keeps_the_region_and_the_opt_in_keeps_the_project(
self,
monkeypatch,
general_settings: Mapping[str, object],
query: Mapping[str, str],
expected_project: str,
expected_location: str,
):
self._pin_no_vertex_project(monkeypatch, general_settings, None)
websocket, _, ws_passthrough = await self._open_live(query)
websocket.close.assert_not_awaited()
passthrough_kwargs = ws_passthrough.await_args.kwargs
assert passthrough_kwargs["custom_headers"]["x-goog-user-project"] == expected_project
assert passthrough_kwargs["setup_model_rewriter"]("gemini-2.0-flash") == (
f"projects/{expected_project}/locations/{expected_location}/publishers/google/models/gemini-2.0-flash"
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"setup_model, expected",
[
(
"projects/proj-env/locations/europe-west4/publishers/google/models/gemini-2.0-flash",
"projects/proj-env/locations/europe-west4/publishers/google/models/gemini-2.0-flash",
),
(
"projects/caller-project/locations/europe-west4/publishers/google/models/gemini-2.0-flash",
"projects/proj-env/locations/europe-west4/publishers/google/models/gemini-2.0-flash",
),
],
ids=["own_project", "other_project"],
)
async def test_setup_frame_keeps_its_region_when_no_region_is_pinned(
self, monkeypatch, setup_model: str, expected: str
):
self._use_default_vertex_config(monkeypatch, {}, vertex_location=None)
websocket, _, ws_passthrough = await self._open_live({})
websocket.close.assert_not_awaited()
rewriter = ws_passthrough.await_args.kwargs["setup_model_rewriter"]
assert rewriter(setup_model) == expected
@pytest.mark.asyncio
@pytest.mark.parametrize("model, expected_model", [("gemini-3.8-live", "gemini-3.8-live"), (None, "")])
async def test_credential_failure_logs_the_named_model_under_vertex_ai(

View file

@ -7674,7 +7674,7 @@ async def test_websocket_passthrough_keeps_normal_upstream_close_normal():
websocket.close.assert_awaited_once_with()
async def _run_setup_rewrite_passthrough(setup_model: str, llm_router) -> str:
async def _run_setup_rewrite_passthrough(setup_model: str, llm_router, keeps_caller_resource_path: bool = False) -> str:
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
_build_vertex_live_setup_model_rewriter,
)
@ -7702,7 +7702,9 @@ async def _run_setup_rewrite_passthrough(setup_model: str, llm_router) -> str:
setup_model_rewriter=_build_vertex_live_setup_model_rewriter(
vertex_project="proj-db",
vertex_location="global",
configured_location="global",
llm_router=llm_router,
keeps_caller_resource_path=keeps_caller_resource_path,
),
)
@ -7740,6 +7742,67 @@ async def test_websocket_passthrough_rewrites_gateway_alias_setup_model():
assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"
async def _run_binary_client_frame_passthrough(frame: bytes):
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
_build_vertex_live_setup_model_rewriter,
)
upstream_ws = RecordingUpstreamWebSocket()
websocket = _client_websocket(
AsyncMock(
side_effect=[
{"type": "websocket.receive", "bytes": frame},
{"type": "websocket.disconnect"},
]
)
)
with _patched_websocket_passthrough_environment(upstream_ws):
await websocket_passthrough_request(
websocket=websocket,
target="wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent",
custom_headers={"Authorization": "Bearer token"},
user_api_key_dict=UserAPIKeyAuth(),
forward_headers=False,
endpoint="/vertex_ai/live",
accept_websocket=False,
setup_model_rewriter=_build_vertex_live_setup_model_rewriter(
vertex_project="proj-a",
vertex_location="europe-west4",
configured_location="europe-west4",
llm_router=None,
keeps_caller_resource_path=False,
),
)
return websocket, upstream_ws
@pytest.mark.asyncio
async def test_websocket_passthrough_moves_a_binary_setup_frame_onto_the_pinned_project():
setup_frame = json.dumps(
{"setup": {"model": "projects/proj-b/locations/europe-west4/publishers/google/models/gemini-2.0-flash"}}
).encode("utf-8")
_, upstream_ws = await _run_binary_client_frame_passthrough(setup_frame)
upstream_ws.send.assert_awaited_once()
sent_frame = upstream_ws.send.await_args.args[0]
assert isinstance(sent_frame, bytes), "a binary client message must stay binary upstream"
assert json.loads(sent_frame)["setup"]["model"] == (
"projects/proj-a/locations/europe-west4/publishers/google/models/gemini-2.0-flash"
)
@pytest.mark.asyncio
async def test_websocket_passthrough_closes_on_a_binary_frame_that_is_not_utf8():
websocket, upstream_ws = await _run_binary_client_frame_passthrough(b"\xff\xfe\x00projects/proj-b")
upstream_ws.send.assert_not_awaited()
websocket.close.assert_awaited_once()
assert websocket.close.await_args.kwargs["code"] == 1008
@pytest.mark.parametrize(
"setup_model",
["gemini-live-2.5-flash", "models/gemini-live-2.5-flash", "publishers/google/models/gemini-live-2.5-flash"],
@ -7761,7 +7824,11 @@ def test_vertex_live_setup_model_resolves_before_extraction(setup_model):
)
rewriter = _build_vertex_live_setup_model_rewriter(
vertex_project="proj-db", vertex_location="global", llm_router=None
vertex_project="proj-db",
vertex_location="global",
configured_location="global",
llm_router=None,
keeps_caller_resource_path=False,
)
setup_data = {"model": setup_model}
@ -7811,7 +7878,11 @@ async def test_websocket_passthrough_logs_a_bare_alias_setup_model():
endpoint="/vertex_ai/live",
accept_websocket=False,
setup_model_rewriter=_build_vertex_live_setup_model_rewriter(
vertex_project="proj-db", vertex_location="global", llm_router=None
vertex_project="proj-db",
vertex_location="global",
configured_location="global",
llm_router=None,
keeps_caller_resource_path=False,
),
)
@ -7876,14 +7947,24 @@ async def test_websocket_passthrough_rewrites_alias_of_unrecognised_upstream_mod
@pytest.mark.asyncio
async def test_websocket_passthrough_leaves_full_resource_setup_model_untouched():
full_resource = "projects/other/locations/us-central1/publishers/google/models/gemini-2.0-flash-live-preview-04-09"
@pytest.mark.parametrize(
"keeps_caller_resource_path, expected_prefix",
[(True, "projects/other/locations/us-central1/"), (False, "projects/proj-db/locations/global/")],
ids=["caller_path_kept", "moved_onto_session_target"],
)
async def test_websocket_passthrough_full_resource_setup_model(
keeps_caller_resource_path: bool, expected_prefix: str
) -> None:
model_suffix = "publishers/google/models/gemini-2.0-flash-live-preview-04-09"
sent_frame = await _run_setup_rewrite_passthrough(full_resource, llm_router=None)
sent_frame = await _run_setup_rewrite_passthrough(
f"projects/other/locations/us-central1/{model_suffix}",
llm_router=None,
keeps_caller_resource_path=keeps_caller_resource_path,
)
assert json.loads(sent_frame)["setup"]["model"] == full_resource
assert sent_frame == json.dumps(
{"setup": {"model": full_resource, "generationConfig": {"responseModalities": ["TEXT"]}}}
{"setup": {"model": f"{expected_prefix}{model_suffix}", "generationConfig": {"responseModalities": ["TEXT"]}}}
)

View file

@ -15,7 +15,8 @@ from fastapi.testclient import TestClient
import litellm
from litellm import Router
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.constants import LITELLM_PROXY_ADMIN_NAME
from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.proxy_server import app
from litellm.types.vector_stores import LiteLLM_ManagedVectorStore
@ -284,6 +285,7 @@ VERTEX_REGISTRY_STORE = {
"vector_store_id": "projects/registry-project/locations/us-central1/ragCorpora/42",
"custom_llm_provider": "vertex_ai",
"litellm_params": {"vertex_project": "registry-project", "vertex_location": "us-central1"},
"is_config": True,
}
UNSUPPORTED_INGEST_PROVIDER_ERROR = (
"Provider '{provider}' is not supported for RAG ingestion. "
@ -990,6 +992,292 @@ def test_rag_ingest_unmanaged_store_keeps_the_callers_full_config(client_interna
assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] == caller_config
_PROXY_VERTEX_PROJECT: Final = "proxy-project"
_PROXY_VERTEX_LOCATION: Final = "europe-west4"
@pytest.fixture
def proxy_vertex_target(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "vertex_project", None)
monkeypatch.setattr(litellm, "vertex_location", None)
monkeypatch.setenv("VERTEXAI_PROJECT", _PROXY_VERTEX_PROJECT)
monkeypatch.setenv("VERTEXAI_LOCATION", _PROXY_VERTEX_LOCATION)
def _vertex_ingest(client: TestClient, vector_store: Mapping[str, str]) -> tuple[httpx.Response, AsyncMock]:
aingest_patch, registry_patch = _patched_ingest_boundary(
None, {"vector_store_id": "projects/p/locations/l/ragCorpora/7", "file_id": "file_123"}
)
with aingest_patch as mock_aingest, registry_patch, _patched_prisma_client(None):
response = client.post("/v1/rag/ingest", **_ingest_form({"custom_llm_provider": "vertex_ai", **vector_store}))
return response, mock_aingest
@pytest.mark.parametrize(
"vector_store, field",
[
({"vertex_project": "other-project"}, "vertex_project"),
({"vertex_location": "us-east5"}, "vertex_location"),
(
{"vector_store_id": f"projects/other-project/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/7"},
"vector_store_id",
),
(
{
"vertex_project": _PROXY_VERTEX_PROJECT,
"vector_store_id": f"projects/other-project/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/7",
},
"vector_store_id",
),
({"vector_store_id": f"projects/{_PROXY_VERTEX_PROJECT}/locations/us-east5/ragCorpora/7"}, "vector_store_id"),
({"vertex_ai_project": "other-project"}, "vertex_ai_project"),
({"vertex_ai_location": "us-east5"}, "vertex_ai_location"),
({"vector_store_id": "projects/123456789/locations/us-east5/ragCorpora/7"}, "vector_store_id"),
],
ids=[
"project_field",
"location_field",
"corpus_project",
"corpus_project_behind_matching_field",
"corpus_region",
"project_alias_field",
"location_alias_field",
"corpus_region_by_project_number",
],
)
@pytest.mark.usefixtures("proxy_vertex_target")
def test_rag_ingest_refuses_a_key_pointing_the_proxys_vertex_credentials_elsewhere(
client_internal_user, vector_store: Mapping[str, str], field: str
):
response, mock_aingest = _vertex_ingest(client_internal_user, vector_store)
assert response.status_code == 400, response.json()
assert response.json()["detail"]["error"].startswith(f"Rejected Request: {field} picks a Vertex project")
mock_aingest.assert_not_awaited()
def test_rag_ingest_refuses_a_key_naming_a_vertex_project_when_the_proxy_names_none(
client_internal_user, monkeypatch: pytest.MonkeyPatch
):
monkeypatch.setattr(litellm, "vertex_project", None)
monkeypatch.delenv("VERTEXAI_PROJECT", raising=False)
response, mock_aingest = _vertex_ingest(client_internal_user, {"vertex_project": "other-project"})
assert response.status_code == 400, response.json()
mock_aingest.assert_not_awaited()
@pytest.mark.parametrize(
"vector_store",
[
{},
{"vertex_project": _PROXY_VERTEX_PROJECT, "vertex_location": _PROXY_VERTEX_LOCATION},
{"vector_store_id": f"projects/{_PROXY_VERTEX_PROJECT}/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/7"},
{"vertex_ai_project": _PROXY_VERTEX_PROJECT, "vertex_ai_location": _PROXY_VERTEX_LOCATION},
],
ids=["proxy_defaults", "same_fields", "corpus_in_the_proxy_project", "same_alias_fields"],
)
@pytest.mark.usefixtures("proxy_vertex_target")
def test_rag_ingest_lets_a_key_use_the_proxys_own_vertex_target(client_internal_user, vector_store: Mapping[str, str]):
response, mock_aingest = _vertex_ingest(client_internal_user, vector_store)
assert response.status_code == 200, response.json()
assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] == {
"custom_llm_provider": "vertex_ai",
**vector_store,
}
@pytest.mark.parametrize(
"project_field, location_field",
[("vertex_project", "vertex_location"), ("vertex_ai_project", "vertex_ai_location")],
ids=["fields", "alias_fields"],
)
@pytest.mark.parametrize(
"user_role, general_settings",
[
(LitellmUserRoles.PROXY_ADMIN, {}),
(LitellmUserRoles.INTERNAL_USER, {"allow_client_side_credentials": True}),
],
ids=["proxy_admin", "proxy_wide_opt_in"],
)
@pytest.mark.usefixtures("proxy_vertex_target")
def test_rag_ingest_lets_the_admin_choose_another_vertex_target(
user_role: LitellmUserRoles, general_settings: Mapping[str, object], project_field: str, location_field: str
):
caller: Final = UserAPIKeyAuth(user_id="vertex-caller", user_role=user_role)
original_overrides: Final = app.dependency_overrides.copy()
app.dependency_overrides[user_api_key_auth] = lambda: caller
vector_store: Final = {
project_field: "other-project",
location_field: "us-east5",
"vector_store_id": "projects/other-project/locations/us-east5/ragCorpora/7",
}
try:
with patch.dict("litellm.proxy.proxy_server.general_settings", general_settings):
response, mock_aingest = _vertex_ingest(TestClient(app), vector_store)
finally:
app.dependency_overrides = original_overrides
assert response.status_code == 200, response.json()
assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] == {
"custom_llm_provider": "vertex_ai",
**vector_store,
}
@pytest.mark.usefixtures("proxy_vertex_target")
def test_rag_ingest_reads_a_corpus_named_by_project_number_on_the_proxys_own_project(client_internal_user):
response, mock_aingest = _vertex_ingest(
client_internal_user, {"vector_store_id": f"projects/123456789/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/7"}
)
assert response.status_code == 200, response.json()
assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"]["vector_store_id"] == (
f"projects/{_PROXY_VERTEX_PROJECT}/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/7"
)
_PROXY_CORPUS_PREFIX: Final = f"projects/{_PROXY_VERTEX_PROJECT}/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora"
_FOREIGN_CORPUS: Final = f"projects/other-project/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/1"
@pytest.mark.usefixtures("proxy_vertex_target")
def test_rag_ingest_refuses_another_teams_store_named_by_a_project_number(client_internal_user):
other_teams_store: Final = LiteLLM_ManagedVectorStore(
vector_store_id=f"{_PROXY_CORPUS_PREFIX}/1",
custom_llm_provider="vertex_ai",
litellm_params={},
user_id="other-user",
team_id="other-team",
)
registry: Final = MagicMock()
registry.get_litellm_managed_vector_store_from_registry.side_effect = lambda vector_store_id: (
other_teams_store if vector_store_id == other_teams_store["vector_store_id"] else None
)
with (
patch( # test-quality-ok: aingest is the endpoint's downstream boundary; the test asserts it is never called
"litellm.proxy.rag_endpoints.endpoints.litellm.aingest",
new=AsyncMock(return_value={"vector_store_id": f"{_PROXY_CORPUS_PREFIX}/1", "file_id": "file_123"}),
) as mock_aingest,
patch.object(litellm, "vector_store_registry", registry), # test-quality-ok: seeds the saved stores by id
_patched_prisma_client(None),
):
response = client_internal_user.post(
"/v1/rag/ingest",
**_ingest_form(
{
"custom_llm_provider": "vertex_ai",
"vector_store_id": f"projects/0/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/1",
}
),
)
assert response.status_code == 403, response.json()
mock_aingest.assert_not_awaited()
@pytest.mark.parametrize(
"vector_store_id",
[
f"./{_FOREIGN_CORPUS}",
f"x/../{_FOREIGN_CORPUS}",
f"{_PROXY_CORPUS_PREFIX}/1{'/..' * 6}/{_FOREIGN_CORPUS}",
f"{_PROXY_CORPUS_PREFIX}/..",
f"{_PROXY_CORPUS_PREFIX}/1%2F..%2F..",
f"{_PROXY_CORPUS_PREFIX}/1\\..\\..",
],
ids=["dot_prefix", "dot_dot_prefix", "own_corpus_then_dot_dots", "dot_dot_corpus", "percent", "backslash"],
)
@pytest.mark.usefixtures("proxy_vertex_target")
def test_rag_ingest_refuses_a_vertex_corpus_name_with_path_tricks(client_internal_user, vector_store_id: str):
response, mock_aingest = _vertex_ingest(client_internal_user, {"vector_store_id": vector_store_id})
assert response.status_code == 400, response.json()
assert response.json()["detail"]["error"].startswith("Rejected Request: vector_store_id picks a Vertex project")
mock_aingest.assert_not_awaited()
_FOREIGN_CORPUS_BY_NUMBER: Final = f"projects/987654321/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/1"
def _foreign_vertex_store(created_by: str | None, corpus: str = _FOREIGN_CORPUS) -> LiteLLM_ManagedVectorStore:
return LiteLLM_ManagedVectorStore(
vector_store_id=corpus, custom_llm_provider="vertex_ai", litellm_params={}, user_id=created_by
)
def _prisma_with_user_row(user_row: LiteLLM_UserTable | None) -> MagicMock:
prisma_client: Final = MagicMock()
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row)
prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(
return_value=MagicMock(vector_store_metadata={})
)
prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock()
return prisma_client
def _ingest_into_saved_store(
store: LiteLLM_ManagedVectorStore, creator_row: LiteLLM_UserTable | None
) -> tuple[httpx.Response, AsyncMock]:
caller: Final = UserAPIKeyAuth(user_id="vertex-caller", user_role=LitellmUserRoles.INTERNAL_USER)
original_overrides: Final = app.dependency_overrides.copy()
app.dependency_overrides[user_api_key_auth] = lambda: caller
aingest_patch, registry_patch = _patched_ingest_boundary(
store, {"vector_store_id": store["vector_store_id"], "file_id": "file_123"}
)
try:
with aingest_patch as mock_aingest, registry_patch, _patched_prisma_client(_prisma_with_user_row(creator_row)):
response = TestClient(app).post(
"/v1/rag/ingest", **_ingest_form({"vector_store_id": store["vector_store_id"]})
)
finally:
app.dependency_overrides = original_overrides
return response, mock_aingest
@pytest.mark.parametrize(
"creator_row",
[LiteLLM_UserTable(user_id="vertex-caller", user_role=LitellmUserRoles.INTERNAL_USER.value), None],
ids=["saved_by_an_internal_user", "saved_by_a_key_with_no_user"],
)
@pytest.mark.usefixtures("proxy_vertex_target")
def test_rag_ingest_refuses_a_saved_vertex_store_a_non_admin_pointed_elsewhere(creator_row: LiteLLM_UserTable | None):
store: Final = _foreign_vertex_store(creator_row.user_id if creator_row is not None else None)
response, mock_aingest = _ingest_into_saved_store(store, creator_row)
assert response.status_code == 400, response.json()
assert response.json()["detail"]["error"].startswith("Rejected Request: vector_store_id picks a Vertex project")
mock_aingest.assert_not_awaited()
@pytest.mark.parametrize("corpus", [_FOREIGN_CORPUS, _FOREIGN_CORPUS_BY_NUMBER], ids=["by_id", "by_number"])
@pytest.mark.parametrize(
"created_by, is_config, creator_row",
[
(LITELLM_PROXY_ADMIN_NAME, False, None),
("admin-user", False, LiteLLM_UserTable(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN.value)),
(None, True, None),
],
ids=["master_key_created", "admin_user_created", "config_store"],
)
@pytest.mark.usefixtures("proxy_vertex_target")
def test_rag_ingest_lets_a_key_use_a_vertex_store_an_admin_pointed_elsewhere(
corpus: str, created_by: str | None, is_config: bool, creator_row: LiteLLM_UserTable | None
):
store: Final = LiteLLM_ManagedVectorStore(**_foreign_vertex_store(created_by, corpus), is_config=is_config)
response, mock_aingest = _ingest_into_saved_store(store, creator_row)
assert response.status_code == 200, response.json()
assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] == {
"vector_store_id": corpus,
"custom_llm_provider": "vertex_ai",
}
def test_rag_ingest_db_managed_store_drops_the_callers_credential_name(client_internal_user):
aingest_patch, registry_patch = _patched_ingest_boundary(
DB_MANAGED_STORE, {"vector_store_id": "db-store", "file_id": "file_123"}