From 4867a56e9aef1310c741b90ee9c1a9e1bbd49e0d Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 10 Nov 2026 20:50:17 +0000 Subject: [PATCH] fix(proxy)!: reject client-supplied Vertex location and project unless the admin opts in (#47028) * fix(proxy): reject client-supplied Vertex location and project unless the admin opts in A request body could carry vertex_location, vertex_project, or their vertex_ai_* aliases and override the region and project pinned on a Vertex AI deployment, sending the admin's credentials to a caller-chosen project or region. Add the four fields to the banned request params so they need allow_client_side_credentials or configurable_clientside_auth_params. /vertex_ai/live and /rag/ingest name their Vertex project and location by design and have no deployment pin to override, so they keep those two fields while every other banned param is still rejected there * fix(proxy): close the remaining Vertex target overrides and keep inert fields working /vertex_ai/live now closes with 1008 when the query names a project or location that differs from a use_in_pass_through deployment's pin, unless that deployment or the proxy opts in A model group's opt-in now counts only when every deployment in it carries it, and Google-style routes look the opt-in up on the model named in the URL Vertex target names inside metadata, the live route's vertex_ai_* aliases, and /health/test_connection inheritance go back to the base behavior, since nothing reads them as a target * test(batches): give the batch FakeRouter the model-list lookup the opt-in check reads * fix(proxy): pin default_vertex_config on /vertex_ai/live and the proxy's Vertex project on /rag/ingest * fix(proxy): require the Vertex opt-in on alias targets and fallbacks, and pass the live rewriter flag in main's tests * test(proxy): give the batch tests' fake router the alias and routed-group lookups * fix(proxy): close the Vertex opt-in, RAG corpus, and live setup-frame gaps found in review A Vertex opt-in now counts only when every group the request can reach opts in, router, context window, and content policy fallbacks and one alias hop included, while a group with no Google deployment never needs the fields. A model named in the route path is checked next to the body model. RAG ingest from a non-admin key refuses a bare or dot-segment corpus id, a vertex_ai_project or vertex_ai_location naming another target, and a store a non-admin saved before the upgrade, and reads a corpus named by project number on the proxy's own project. /vertex_ai/live closes 1008 on any caller-named project when the proxy pins none, moves a setup frame sent as a binary message onto the session's target, and resolves the region from default_vertex_config before the model path. * fix(proxy): let a live caller name the proxy's configured Vertex project when nothing pins one With default_vertex_config set but no project pinned, /vertex_ai/live closed every query that named a project, including the one the admin set in VERTEXAI_PROJECT or DEFAULT_VERTEXAI_PROJECT. A query naming one of those projects now connects, and any other project still closes with 1008 before a token is minted. * fix(proxy): keep an admin-created vector store's corpus as saved on RAG ingest A non-admin ingest into a store an admin created named by project number had its corpus moved onto the proxy's own project before the admin-store exemption ran, so the upload went to a different corpus than the admin saved. The exemption now runs first, and a non-admin store still gets the project-number rewrite and the target check. * fix(proxy): authorize a RAG corpus moved onto the proxy's project and check client-side opt-ins once per request * fix(proxy): merge opt-in fallback groups only when their deployments share provider and opt-in list * fix(proxy): check team public models before allowing target overrides * fix(proxy): include team wildcard targets in opt-in checks * fix(proxy): validate overrides against authenticated team routes * fix(proxy): retain URL model during deferred credential checks --- litellm/proxy/auth/auth_utils.py | 401 +++++++++--- litellm/proxy/auth/user_api_key_auth.py | 19 +- litellm/proxy/batches_endpoints/endpoints.py | 1 + .../litellm_executed_batches.py | 5 +- .../health_endpoints/_health_endpoints.py | 8 +- .../openai_files_endpoints/files_endpoints.py | 1 + .../llm_passthrough_endpoints.py | 153 ++++- .../pass_through_endpoints.py | 46 +- .../passthrough_endpoint_router.py | 10 +- litellm/proxy/prompts/prompt_endpoints.py | 1 + litellm/proxy/proxy_server.py | 5 +- litellm/proxy/rag_endpoints/endpoints.py | 119 +++- litellm/rag/ingestion/vertex_ai_ingestion.py | 4 +- .../test_vertex_image_regional_uplift.py | 10 +- tests/unit/proxy/auth/test_auth_utils.py | 591 +++++++++++++++++- .../test_user_api_key_auth_request_flow.py | 92 ++- .../proxy/batches_endpoints/test_endpoints.py | 8 +- .../test_litellm_executed_batches.py | 22 + .../health_endpoints/test_health_endpoints.py | 6 + .../test_llm_pass_through_endpoints.py | 404 +++++++++++- .../test_pass_through_endpoints.py | 97 ++- .../public_v1/test_model_hub.py | 4 + .../proxy/rag_endpoints/test_rag_endpoints.py | 290 ++++++++- 23 files changed, 2163 insertions(+), 134 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 715b98a441a..7413490e44f 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -2,8 +2,9 @@ import importlib.util import os import re import sys -from collections.abc import Collection, Iterator, Mapping, Sequence -from functools import lru_cache, partial +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 @@ -24,6 +25,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.cloud_storage_security import REQUEST_S3_BUCKET_PARAMS, STORED_S3_BUCKET_PARAM 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, @@ -42,11 +44,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 @@ -212,35 +216,74 @@ 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"] + + +_DEPLOYMENT_LIST: Final = TypeAdapter(tuple[DeploymentTypedDict, ...], config=ConfigDict(arbitrary_types_allowed=True)) + + +def _model_group_deployments( + model: str, llm_router: Router, team_id: str | None = None +) -> Sequence[DeploymentTypedDict]: + registry: Final = llm_router.registry() + provider_group: Final = model.split("/", 1)[0] + if team_id is not None and registry.routing_group(model) is None: + team_deployments: Final = registry.deployments( + model, source="all", team_id=team_id + ) or _DEPLOYMENT_LIST.validate_python(registry.deployments(model, source="wildcard", team_id=team_id)) + if team_deployments: + return team_deployments + return ( + registry.deployments(model, source="routed_group", team_id=team_id) + or (registry.deployments(provider_group, team_id=team_id) 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, team_id: str | None = None) -> bool: + if llm_router is None: + return True + deployments: Final = _model_group_deployments(model, llm_router, team_id) + 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 + model: str, param: str, request_body_value: object, llm_router: Router | None, team_id: str | None = 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, team_id) + 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 ) @@ -319,6 +362,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 @@ -411,6 +468,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 @@ -483,6 +544,105 @@ _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 _routed_opt_in_hop( + frontier_and_reached: tuple[frozenset[str], frozenset[str]], router_and_team: tuple[Router, str | None] +) -> tuple[frozenset[str], frozenset[str]]: + frontier, reached = frontier_and_reached + llm_router, team_id = router_and_team + deployments: Final = chain.from_iterable(_model_group_deployments(model, llm_router, team_id) for model in frontier) + routed_groups: Final = frozenset(deployment["model_name"] for deployment in deployments) - reached + fallbacks: Final = frozenset[str]().union( + *(router_fallback_model_names(name, llm_router, _EVERY_FALLBACK_TYPE) for name in routed_groups) + ) + registry: Final = llm_router.registry() + next_hop: Final = ( + routed_groups | frozenset(registry.resolve(name, kind="alias") or name for name in fallbacks) + ) - reached + return next_hop, reached | next_hop + + +def _opt_in_model_groups( + model: str, also_requested_models: tuple[str, ...], llm_router: Router | None, team_id: str | None = 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) + ) + groups: Final = routed | frozenset(registry.resolve(name, kind="alias") or name for name in fallbacks) + hops: Final = accumulate(repeat((llm_router, team_id)), _routed_opt_in_hop, initial=(groups, groups)) + _, reached = next(hop for hop in hops if not hop[0]) + return reached + + +_MAX_OPT_IN_FALLBACK_MODELS: Final = 100 + + +@dataclass(frozen=True, slots=True) +class _OptInGroups: + every: tuple[str, ...] + vertex_target: tuple[str, ...] + team_id: str | None = None + + +_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, team_id: str | None +) -> frozenset[tuple[str | None, _OptInKey]] | str: + deployments: Final = _model_group_deployments(model_group, llm_router, team_id) + 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, team_id: str | None +) -> tuple[str, ...]: + if llm_router is None: + return tuple(sorted(model_groups)) + return tuple( + { + _same_answer_key(model_group, llm_router, team_id): 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, team_id: str | None = 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, team_id), llm_router, team_id + ) + return _OptInGroups( + every=every, + vertex_target=tuple(group for group in every if _model_group_reads_vertex_target(group, llm_router, team_id)), + team_id=team_id, + ) + _ClientSideAllowlistSetting: TypeAlias = Literal[ "allowed_client_side_s3_buckets", "allowed_client_side_s3_endpoint_urls" @@ -538,12 +698,14 @@ def _banned_param_error(param: str) -> ValueError: def _check_banned_params( body: Mapping[str, object], - general_settings: dict, + general_settings: Mapping[str, object], 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], + defer_model_checks: bool = False, defer_s3_buckets: bool = False, ) -> None: """Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in. @@ -560,7 +722,7 @@ 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: @@ -573,14 +735,20 @@ def _check_banned_params( continue if defer_s3_buckets and param in REQUEST_S3_BUCKET_PARAMS: continue - if ( + if defer_model_checks: + continue + if all( _allow_model_level_clientside_configurable_parameters( - model=model, + model=model_group, param=param, request_body_value=body[param], llm_router=llm_router, + team_id=opt_in_groups().team_id, ) 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 @@ -664,27 +832,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} @@ -714,12 +894,14 @@ def _reject_url_valued_fallback_target(value: str) -> None: def is_request_body_safe( - request_body: dict, - general_settings: dict, + request_body: Mapping[str, object], + general_settings: Mapping[str, object], llm_router: Router | None, model: str, *, route: str | None = None, + team_id: str | None = None, + defer_model_checks: bool = False, defer_s3_buckets: bool = False, ) -> bool: """ @@ -750,28 +932,72 @@ 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) - check_banned_params: Final = partial( - _check_banned_params, - general_settings=general_settings, - llm_router=llm_router, - model=model, + 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 or route_model or "", also_requested_models, llm_router, team_id) + ) + _check_banned_params( + root_body, + general_settings, + llm_router, manages_deployments=manages_deployments, + mcp_server_fields_allowed=route_manages_mcp_servers(route), + opt_in_groups=opt_in_groups, + defer_model_checks=defer_model_checks, defer_s3_buckets=defer_s3_buckets, ) - check_banned_params(request_body, mcp_server_fields_allowed=route_manages_mcp_servers(route)) 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) + _check_banned_params( + nested, + general_settings, + llm_router, + manages_deployments=manages_deployments, + opt_in_groups=opt_in_groups, + defer_model_checks=defer_model_checks, + defer_s3_buckets=defer_s3_buckets, + ) 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) - 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}[")) + _check_banned_params( + metadata, + general_settings, + llm_router, + manages_deployments=manages_deployments, + banned_params=_BANNED_METADATA_PARAMS, + opt_in_groups=opt_in_groups, + defer_model_checks=defer_model_checks, + defer_s3_buckets=defer_s3_buckets, + ) + if any(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, + manages_deployments=manages_deployments, + banned_params=_BANNED_METADATA_PARAMS, + opt_in_groups=opt_in_groups, + defer_model_checks=defer_model_checks, + defer_s3_buckets=defer_s3_buckets, + ) for target in iter_request_fallback_targets(request_body): if isinstance(target, dict): - check_banned_params(target) + _check_banned_params( + target, + general_settings, + llm_router, + manages_deployments=manages_deployments, + opt_in_groups=opt_in_groups, + defer_model_checks=defer_model_checks, + defer_s3_buckets=defer_s3_buckets, + ) target_model = target.get("model") if isinstance(target_model, str): _reject_url_valued_fallback_target(target_model) @@ -784,8 +1010,17 @@ def is_request_body_safe( _reject_request_credential_params(litellm_params, manages_deployments=manages_deployments) 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) - for user_config_params in _user_config_litellm_params(request_body): # pyright: ignore[reportUnknownArgumentType] # request_body is an untyped dict + _check_banned_params( + litellm_params_metadata, + general_settings, + llm_router, + manages_deployments=manages_deployments, + banned_params=_BANNED_METADATA_PARAMS, + opt_in_groups=opt_in_groups, + defer_model_checks=defer_model_checks, + defer_s3_buckets=defer_s3_buckets, + ) + for user_config_params in _user_config_litellm_params(request_body): reject_api_key_secret_reference(user_config_params) return True @@ -880,6 +1115,8 @@ def enforce_team_s3_bucket_allowlist( general_settings: Mapping[str, object], llm_router: Router | None, team_metadata: Mapping[str, object] | None, + *, + team_id: str | None = None, ) -> None: """Re-run the request body check that auth relaxed for bucket params, now admitting the caller team's buckets.""" proxy_buckets: Final = _ClientSideStorageAllowlist.model_validate(general_settings).allowed_client_side_s3_buckets @@ -893,6 +1130,7 @@ def enforce_team_s3_bucket_allowlist( }, llm_router=llm_router, model=model if isinstance(model, str) else "", + team_id=team_id, ) except ValueError as rejection: raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail=f"Authentication Error, {rejection}") @@ -902,6 +1140,8 @@ async def pre_db_read_auth_checks( request: Request, request_data: dict, route: str, + *, + defer_team_checks: bool = False, ): """ 1. Checks if request size is under max_request_size_mb (if set) @@ -921,12 +1161,17 @@ async def pre_db_read_auth_checks( await check_if_request_size_is_safe(request=request) # Check 2. Request body is safe + defer_model_checks: Final = ( + defer_team_checks and llm_router is not None and bool(llm_router.registry().index("team_model")) + ) + request.state.deferred_model_body_check = defer_model_checks # rebind-ok: hand state to auth is_request_body_safe( 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, + defer_model_checks=defer_model_checks, defer_s3_buckets=checks_team_s3_buckets_in_handler(request), ) @@ -2454,37 +2699,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) @@ -2503,6 +2719,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 diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 1dca6075a3d..5bf4db6a725 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -109,12 +109,14 @@ from litellm.proxy.auth.auth_object_prefetch import ( ) from litellm.proxy.auth.auth_utils import ( abbreviate_api_key, + checks_team_s3_buckets_in_handler, fallback_target_model_name, get_end_user_id_from_request_body, get_model_from_request, get_request_route, get_request_route_template, is_invalid_virtual_key_error, + is_request_body_safe, iter_request_fallback_targets, log_model_access_denial, normalize_request_route, @@ -2276,6 +2278,7 @@ async def user_api_key_auth_builder( azure_apim_header: str | None, request_data: dict[str, object], custom_litellm_key_header: str | None = None, + defer_team_checks: bool = False, ) -> UserAPIKeyAuth: from litellm.proxy.proxy_server import ( general_settings, @@ -2310,6 +2313,7 @@ async def user_api_key_auth_builder( request_data=request_data, request=request, route=route, + defer_team_checks=defer_team_checks, ) pass_through_endpoints: Final[list[dict] | None] = general_settings.get("pass_through_endpoints", None) ## CHECK IF X-LITELM-API-KEY IS PASSED IN - supercedes Authorization header @@ -3684,6 +3688,18 @@ async def run_centralized_common_checks( user_custom_auth, ) + if not sub_call and getattr(request.state, "deferred_model_body_check", False) is True: + body_model: Final = request_data.get("model") + is_request_body_safe( + request_body=request_data, + general_settings=_GENERAL_SETTINGS_ADAPTER.validate_python(general_settings), + llm_router=llm_router, + model=body_model if isinstance(body_model, str) else "", + route=route, + team_id=user_api_key_auth_obj.team_id, + defer_s3_buckets=checks_team_s3_buckets_in_handler(request), + ) + # Public routes (e.g. /health/liveness) are exempt from # auth in the builder — the wrapper must not retroactively apply # authz on top, or k8s readiness probes and other unauthenticated @@ -4520,6 +4536,7 @@ async def user_api_key_auth( azure_apim_header=azure_apim_header, request_data=request_data, custom_litellm_key_header=custom_litellm_key_header, + defer_team_checks=True, ) except Exception: # The body was read first, so a caller who sent both a malformed body and @@ -4999,7 +5016,7 @@ async def authorize_internal_virtual_key( ) ) route: Final = get_request_route(request=request) - await pre_db_read_auth_checks(request_data=request_data, request=request, route=route) + await pre_db_read_auth_checks(request_data=request_data, request=request, route=route, defer_team_checks=True) auth: Final = await validate_resolved_virtual_key( request=request, request_data=request_data, diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 00fa7f3c33e..0b443572fe6 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -303,6 +303,7 @@ async def create_batch( general_settings=general_settings_view(), llm_router=llm_router, team_metadata=user_api_key_dict.team_metadata, + team_id=user_api_key_dict.team_id, ) _raise_when_routing_headers_disagree(request, data.get("model")) verbose_proxy_logger.debug( diff --git a/litellm/proxy/batches_endpoints/litellm_executed_batches.py b/litellm/proxy/batches_endpoints/litellm_executed_batches.py index 4dd6496812e..8c78e650339 100644 --- a/litellm/proxy/batches_endpoints/litellm_executed_batches.py +++ b/litellm/proxy/batches_endpoints/litellm_executed_batches.py @@ -452,7 +452,7 @@ class LiteLLMExecutedBatchRunner: ) -> LiteLLMBatch: endpoint: Final = _validate_endpoint(create_request.get("endpoint")) content: Final = await self._download_input(unified_input_file_id, user_api_key_dict) - parsed: Final = parse_batch_input(content, endpoint, self._body_rejection(model)) + parsed: Final = parse_batch_input(content, endpoint, self._body_rejection(model, user_api_key_dict.team_id)) if isinstance(parsed, InvalidBatchInput): raise batch_error(400, f"Invalid batch input file: {parsed.describe()}") llm_batch_id: Final = f"{LITELLM_EXECUTED_BATCH_ID_PREFIX}{uuid_module.uuid4().hex}" @@ -533,7 +533,7 @@ class LiteLLMExecutedBatchRunner: return failed return await self.batches.load_batch(batch.id) or batch - def _body_rejection(self, model: str) -> BodyRejection: + def _body_rejection(self, model: str, team_id: str | None) -> BodyRejection: def reject(body: Mapping[str, object]) -> str | None: try: is_request_body_safe( @@ -541,6 +541,7 @@ class LiteLLMExecutedBatchRunner: general_settings=dict(self.general_settings), llm_router=self.llm_router, model=model, + team_id=team_id, ) except ValueError as e: return str(e) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 333687b07cd..1211c56975b 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -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( diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index ca1470bc94e..8b8d9ec81be 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -639,6 +639,7 @@ async def create_file( general_settings=general_settings_view(), llm_router=llm_router, team_metadata=user_api_key_dict.team_metadata, + team_id=user_api_key_dict.team_id, ) file_params: Final = await extract_file_creation_params( request=request, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 5bdb2219194..6ab31e7d4b0 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -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 @@ -4241,6 +4242,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/" @@ -4265,11 +4270,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 @@ -4283,36 +4302,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[^/]+)/") + + 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)}" @@ -4366,11 +4482,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 @@ -4452,7 +4579,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, ), ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 66634688b9a..42a7c458ef2 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2613,6 +2613,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]: @@ -3042,6 +3057,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 @@ -3409,8 +3446,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, ): diff --git a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py index e793dac4d06..877c4ece6c7 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py @@ -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 diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index b7775b99419..04d6c78f08a 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -1111,6 +1111,7 @@ async def test_prompt( general_settings=general_settings, llm_router=llm_router, model=data.get("model", ""), + team_id=user_api_key_dict.team_id, ) # Use ProxyBaseLLMRequestProcessing to go through all proxy logic diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3fddca6f4d8..f69d9bd4724 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -16496,7 +16496,9 @@ async def model_info_lookup(model: str, custom_llm_provider: str | None = None): dependencies=[Depends(user_api_key_auth)], response_model=RawRequestTypedDict, ) -async def transform_request(request: TransformRequestBody): +async def transform_request( + request: TransformRequestBody, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)] +): from litellm.utils import return_raw_request try: @@ -16505,6 +16507,7 @@ async def transform_request(request: TransformRequestBody): general_settings=general_settings, llm_router=llm_router, model=request.request_body.get("model", ""), + team_id=user_api_key_dict.team_id, ) except ValueError as e: raise HTTPException(status_code=400, detail={"error": str(e)}) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index f4710fa30e3..47050f01644 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -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 @@ -21,12 +22,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, @@ -70,9 +72,12 @@ from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, with_managed_search_file_ids, ) +from litellm.rag.ingestion.vertex_ai_ingestion import DEFAULT_VERTEX_RAG_LOCATION from litellm.rag.main import get_ingestion_class from litellm.rag.rag_query import RAGQuery 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.rag import RAGIngestResponse from litellm.types.utils import Choices, ModelResponse @@ -192,6 +197,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[^/%\\]+)/locations/(?P[^/%\\]+)/ragCorpora/(?P[^/%\\]+)" +) +_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", @@ -757,6 +859,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, @@ -782,15 +885,25 @@ async def rag_ingest( general_settings=general_settings, llm_router=llm_router, model="", + route=get_request_route(request=request), + team_id=user_api_key_dict.team_id, ) 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, diff --git a/litellm/rag/ingestion/vertex_ai_ingestion.py b/litellm/rag/ingestion/vertex_ai_ingestion.py index cceddf31117..8157c0c6405 100644 --- a/litellm/rag/ingestion/vertex_ai_ingestion.py +++ b/litellm/rag/ingestion/vertex_ai_ingestion.py @@ -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( diff --git a/tests/integration/spend/test_vertex_image_regional_uplift.py b/tests/integration/spend/test_vertex_image_regional_uplift.py index 6900a6d489a..dfd52d28f1f 100644 --- a/tests/integration/spend/test_vertex_image_regional_uplift.py +++ b/tests/integration/spend/test_vertex_image_regional_uplift.py @@ -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) diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 82b80ad8da9..73a2c534628 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -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, Mapping from types import SimpleNamespace from typing import Final, Optional @@ -371,7 +372,11 @@ _SERVICE_ACCOUNT_TOKEN_REFERENCE: Final = "oidc/file//var/run/secrets/kubernetes def _user_config_with(litellm_params: dict[str, str]) -> dict[str, object]: - return {"model_list": [{"model_name": "devstral", "litellm_params": {"model": "mistral/devstral-latest", **litellm_params}}]} + return { + "model_list": [ + {"model_name": "devstral", "litellm_params": {"model": "mistral/devstral-latest", **litellm_params}} + ] + } def _user_config_defaulting_to(litellm_params: dict[str, str]) -> dict[str, object]: @@ -391,7 +396,11 @@ def _user_config_search_tool_with(litellm_params: dict[str, str]) -> dict[str, o "search_tools": [ { "search_tool_name": "web", - "litellm_params": {"search_provider": "tavily", "api_base": "https://attacker.example", **litellm_params}, + "litellm_params": { + "search_provider": "tavily", + "api_base": "https://attacker.example", + **litellm_params, + }, } ] } @@ -512,7 +521,9 @@ def test_a_call_cannot_carry_an_oidc_api_key_in_litellm_params(route: str | None ) -@pytest.mark.parametrize("route", ["/model/new", "/model/update", "/model/deployment-1/update", "/health/test_connection"]) +@pytest.mark.parametrize( + "route", ["/model/new", "/model/update", "/model/deployment-1/update", "/health/test_connection"] +) def test_a_deployment_write_may_carry_an_oidc_api_key_in_litellm_params(route: str): """A proxy admin configures an oidc/ api_key on a deployment through these routes, and the model management role gate refuses anyone else with a 403.""" @@ -3675,6 +3686,7 @@ def test_request_cannot_select_chatgpt_auth_file(allow_client_side_credentials: model="chatgpt-account-a", ) + class TestIsRequestBodySafeBlocksBedrockProjectOverride: """``aws_bedrock_project_id`` pins a deployment to a Bedrock project so that project's data-retention policy applies to its requests. A @@ -3934,7 +3946,7 @@ class TestIsRequestBodySafeBlocksNVCFFunctionOverride: monkeypatch.setattr( auth_utils, "_allow_model_level_clientside_configurable_parameters", - lambda model, param, request_body_value, llm_router: param == "nvcf_function_id", + lambda model, param, request_body_value, llm_router, team_id=None: param == "nvcf_function_id", ) assert ( @@ -4019,7 +4031,7 @@ class TestIsRequestBodySafeBlocksRivaUseSsl: monkeypatch.setattr( auth_utils, "_allow_model_level_clientside_configurable_parameters", - lambda model, param, request_body_value, llm_router: param == "use_ssl", + lambda model, param, request_body_value, llm_router, team_id=None: param == "use_ssl", ) assert ( @@ -4457,7 +4469,7 @@ def test_model_level_allow_does_not_skip_subsequent_banned_params(monkeypatch): monkeypatch.setattr( auth_utils, "_allow_model_level_clientside_configurable_parameters", - lambda model, param, request_body_value, llm_router: param == "api_base", + lambda model, param, request_body_value, llm_router, team_id=None: param == "api_base", ) with pytest.raises(ValueError, match="Rejected Request: langfuse_host is not allowed in request") as exc: @@ -4981,6 +4993,538 @@ 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", + ) + + @pytest.mark.parametrize("model", ["chat", "chat-alias", "entry"]) + @pytest.mark.parametrize("opted_in", [False, True]) + def test_team_public_model_requires_the_team_deployment_opt_in(self, model: str, opted_in: bool) -> None: + router: Final = Router( + model_list=[ + {"model_name": "chat", "litellm_params": {"model": "openai/test-model", "api_key": "test-key"}}, + {"model_name": "entry", "litellm_params": {"model": "openai/test-model", "api_key": "test-key"}}, + { + "model_name": "team-chat", + "model_info": {"team_id": "team-one", "team_public_model_name": "chat"}, + "litellm_params": { + "model": "vertex_ai/test-model", + "vertex_location": "europe-west4", + "configurable_clientside_auth_params": ["vertex_location"] if opted_in else [], + }, + }, + ], + model_group_alias={"chat-alias": "chat"}, + fallbacks=[{"entry": ["chat"]}], + ) + request_body: Final = {"model": model, "vertex_location": "us-central1"} + if opted_in: + assert ( + is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions", team_id="team-one") + is True + ) + return + with pytest.raises(ValueError, match="vertex_location"): + is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions", team_id="team-one") + + @pytest.mark.parametrize( + "model, opted_in, global_exact, is_allowed", + [ + ("chat-test-model", False, False, False), + ("chat-test-model", True, False, True), + ("other-test-model", False, False, True), + ("chat-test-model", False, True, True), + ], + ) + def test_team_public_wildcard_keeps_its_own_opt_in( + self, model: str, opted_in: bool, global_exact: bool, is_allowed: bool + ) -> None: + router: Final = Router( + model_list=[ + {"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "test-key"}}, + { + "model_name": "team-chat", + "model_info": {"team_id": "team-one", "team_public_model_name": "chat-*"}, + "litellm_params": { + "model": "vertex_ai/*", + "configurable_clientside_auth_params": ["vertex_location"] if opted_in else [], + }, + }, + ] + + ( + [ + { + "model_name": "chat-test-model", + "litellm_params": {"model": "openai/test-model", "api_key": "test-key"}, + } + ] + if global_exact + else [] + ) + ) + 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", team_id="team-one") + is True + ) + return + with pytest.raises(ValueError, match="vertex_location"): + is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions", team_id="team-one") + + @pytest.mark.parametrize("fallback_setting", ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"]) + def test_team_internal_fallback_to_another_public_name_keeps_its_pin(self, fallback_setting: str) -> None: + router: Final = Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/test-model", "api_key": "test-key"}} + for name in ("chat", "backup") + ] + + [ + { + "model_name": internal, + "model_info": {"team_id": "team-one", "team_public_model_name": public}, + "litellm_params": { + "model": "vertex_ai/test-model", + "configurable_clientside_auth_params": opt_in, + }, + } + for internal, public, opt_in in ( + ("team-chat", "chat", ["vertex_location"]), + ("team-backup", "backup", []), + ) + ], + **{fallback_setting: [{"team-chat": ["backup"]}]}, + ) + with pytest.raises(ValueError, match="vertex_location"): + is_request_body_safe( + {"model": "chat", "vertex_location": "us-central1"}, + {}, + router, + "chat", + route="/v1/chat/completions", + team_id="team-one", + ) + + 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", [ @@ -5446,6 +5990,29 @@ def _enforce_for_team( ) +@pytest.mark.parametrize("opted_in", [False, True]) +def test_team_bucket_validation_preserves_the_authenticated_model_opt_in(opted_in: bool) -> None: + router: Final = Router( + model_list=[ + {"model_name": "chat", "litellm_params": {"model": "openai/test-model"}}, + { + "model_name": "team-chat", + "model_info": {"team_id": "team-one", "team_public_model_name": "chat"}, + "litellm_params": { + "model": "vertex_ai/test-model", + "configurable_clientside_auth_params": ["vertex_location"] if opted_in else [], + }, + }, + ] + ) + body: Final = {"model": "chat", "s3_bucket_name": "team-a-bucket", "vertex_location": "us-central1"} + if opted_in: + enforce_team_s3_bucket_allowlist(body, {}, router, _TEAM_A_METADATA, team_id="team-one") + return + with pytest.raises(HTTPException, match="vertex_location"): + enforce_team_s3_bucket_allowlist(body, {}, router, _TEAM_A_METADATA, team_id="team-one") + + def _rejection_for_team(body: Mapping[str, object], team_metadata: Mapping[str, object] | None) -> HTTPException: with pytest.raises(HTTPException) as rejection: _enforce_for_team(body, team_metadata) @@ -5499,7 +6066,13 @@ def test_team_allowlist_admits_no_other_blocked_field() -> None: def _routed_request(method: str, route_template: str | None, path: str) -> Request: - scope: Final[dict[str, object]] = {"type": "http", "method": method, "path": path, "headers": [], "query_string": b""} + scope: Final[dict[str, object]] = { + "type": "http", + "method": method, + "path": path, + "headers": [], + "query_string": b"", + } return Request(scope if route_template is None else {**scope, "route": SimpleNamespace(path=route_template)}) @@ -5661,7 +6234,9 @@ async def test_a_non_numeric_max_request_size_mb_is_ignored_without_premium(): def test_chatgpt_model_configuration_reaches_role_authorization( allow_client_side_credentials: bool, route: str ) -> None: - body: Final = {"litellm_params": {"model": "chatgpt/gpt-6-luna", "chatgpt_auth_file": "/tokens/account-a/auth.json"}} + body: Final = { + "litellm_params": {"model": "chatgpt/gpt-6-luna", "chatgpt_auth_file": "/tokens/account-a/auth.json"} + } assert is_request_body_safe( request_body=body, general_settings={"allow_client_side_credentials": allow_client_side_credentials}, diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 1489afc02af..8a5d4f7969b 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -3279,7 +3279,9 @@ class TestJWTOAuth2Coexistence: async def test_auto_register_passes_validated_org_context_to_generated_key(self): jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" general_settings = {"enable_jwt_auth": True} - user_api_key_cache: Final = _validated_jwt_admission_cache(org_id="validated-org", end_user_id="validated-end-user") + user_api_key_cache: Final = _validated_jwt_admission_cache( + org_id="validated-org", end_user_id="validated-end-user" + ) prisma_client = MagicMock() jwt_handler = MagicMock() jwt_handler.is_jwt.return_value = True @@ -4851,6 +4853,92 @@ def _proxy_attrs_for_centralized_checks(user_custom_auth=None, flag=False, maste } +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route, include_model, path_params", + [ + ("/v1/files", True, {}), + ("/v1/batches", True, {}), + ("/v1/chat/completions", True, {}), + ("/v1beta/models/chat:generateContent", False, {"model_name": "chat"}), + ("/openai/deployments/chat/chat/completions", False, {"model": "chat"}), + ], +) +@pytest.mark.parametrize("team_id, is_allowed", [(None, True), ("other-team", True), ("pinned-team", False)]) +async def test_deferred_body_check_uses_only_the_authenticated_teams_deployments( + monkeypatch: pytest.MonkeyPatch, + team_id: str | None, + is_allowed: bool, + route: str, + include_model: bool, + path_params: Mapping[str, str], +) -> None: + from litellm.proxy.utils import ProxyLogging + + async def authenticate(request: Request, api_key: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + team_id=team_id, api_key=api_key, user_role=LitellmUserRoles.INTERNAL_USER, models=["chat"] + ) + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "chat", + "litellm_params": { + "model": "vertex_ai/test-model", + "configurable_clientside_auth_params": ["vertex_location"], + }, + }, + { + "model_name": "private-chat", + "model_info": {"team_id": "pinned-team", "team_public_model_name": "chat"}, + "litellm_params": {"model": "vertex_ai/test-model", "vertex_location": "europe-west4"}, + }, + ] + ) + key_cache: Final = UserApiKeyCache() + attributes: Final = { + **_proxy_attrs_for_centralized_checks(user_custom_auth=authenticate), + "llm_router": router, + "llm_model_list": None, + "user_api_key_cache": key_cache, + "proxy_logging_obj": ProxyLogging(user_api_key_cache=key_cache), + } + for name, value in attributes.items(): + monkeypatch.setattr(litellm.proxy.proxy_server, name, value) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": route, + "headers": [], + "path_params": path_params, + "route": SimpleNamespace(path=route), + }, + receive=AsyncMock( + return_value={ + "type": "http.request", + "more_body": False, + "body": json.dumps( + { + "vertex_location": "us-central1", + **({"model": "chat"} if include_model else {}), + **({"s3_bucket_name": "team-bucket"} if route in {"/v1/files", "/v1/batches"} else {}), + } + ).encode(), + } + ), + ) + if is_allowed: + identity: Final = await user_api_key_auth(request=request, api_key="Bearer external-token") + assert identity.team_id == team_id + return + with pytest.raises(ProxyException) as rejection: + await user_api_key_auth(request=request, api_key="Bearer external-token") + assert int(rejection.value.code) == 401 + assert "vertex_location" in str(rejection.value) + + @pytest.mark.asyncio async def test_centralized_common_checks_runs_for_standard_auth(): """Regardless of which _user_api_key_auth_builder path returned, the @@ -5556,8 +5644,6 @@ async def _run_centralized_checks_for_end_user( setattr(_proxy_server_mod, attribute, value) - - @pytest.mark.parametrize( ("make_token", "make_end_user", "gpt_4o_spend", "model", "route", "must_raise"), [ diff --git a/tests/unit/proxy/batches_endpoints/test_endpoints.py b/tests/unit/proxy/batches_endpoints/test_endpoints.py index cee8a129172..127cf85710b 100644 --- a/tests/unit/proxy/batches_endpoints/test_endpoints.py +++ b/tests/unit/proxy/batches_endpoints/test_endpoints.py @@ -221,6 +221,7 @@ def harness(monkeypatch: pytest.MonkeyPatch): registry.access_groups = MagicMock(return_value={}) registry.resolve = MagicMock(side_effect=lambda name, **_: name) registry.all = MagicMock(return_value=[]) + registry.deployments = MagicMock(return_value=[]) creds_resolver = MagicMock(side_effect=_registry_creds_lookup) registry.credentials = _route_credentials({"with_provider": creds_resolver}) router.model_list = [] @@ -474,9 +475,7 @@ async def test_create__team_with_an_allowlist_and_no_request_bucket_leaves_the_b @pytest.mark.asyncio -@pytest.mark.parametrize( - "request_bucket,expected_bucket", [("team-a-bucket", "team-a-bucket"), (None, "model-bucket")] -) +@pytest.mark.parametrize("request_bucket,expected_bucket", [("team-a-bucket", "team-a-bucket"), (None, "model-bucket")]) @pytest.mark.parametrize( "routing", [ @@ -1564,6 +1563,7 @@ def retrieve_harness(): registry.access_groups = MagicMock(return_value={}) registry.resolve = MagicMock(side_effect=lambda name, **_: name) registry.all = MagicMock(return_value=[]) + registry.deployments = MagicMock(return_value=[]) creds_resolver = MagicMock(side_effect=_registry_creds_lookup) router.model_list = [] router.aretrieve_batch = AsyncMock(return_value=make_batch()) @@ -2151,6 +2151,7 @@ def list_harness(): registry.access_groups = MagicMock(return_value={}) registry.resolve = MagicMock(side_effect=lambda name, **_: name) registry.all = MagicMock(return_value=[]) + registry.deployments = MagicMock(return_value=[]) creds_resolver = MagicMock(side_effect=_registry_creds_lookup) registry.credentials = _route_credentials({"with_provider": creds_resolver}) router.model_list = [] @@ -2576,6 +2577,7 @@ def cancel_harness(): registry.access_groups = MagicMock(return_value={}) registry.resolve = MagicMock(side_effect=lambda name, **_: name) registry.all = MagicMock(return_value=[]) + registry.deployments = MagicMock(return_value=[]) creds_resolver = MagicMock(side_effect=_registry_creds_lookup) registry.credentials = _route_credentials({"with_provider": creds_resolver}) router.model_list = [] diff --git a/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py b/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py index f9918d5a1ee..ce2e740c39c 100644 --- a/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py +++ b/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py @@ -252,11 +252,20 @@ 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 class FakeRegistry: + def routing_group(self, name: str) -> None: + return None + + def index(self, name: Literal["team_model"]) -> Mapping[tuple[str, str], Sequence[int]]: + return MappingProxyType({}) + def names(self, model_name: str | None = None, *, kind: str = "groups") -> list[str]: assert kind == "ids" return [DEPLOYMENT_ID] if model_name == BATCH_MODEL else [] @@ -264,6 +273,19 @@ 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", team_id: str | None = None + ) -> list[object]: + return [] + + +class FakeSettings: + def is_set(self, name: str) -> bool: + return False + class FakeStorageBackend: def __init__(self, contents: Mapping[str, bytes]) -> None: diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py index 2258eb95fb8..4aff6adadd3 100644 --- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py @@ -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.""" diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 70f5e855e16..6e551f8ff0a 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -6399,7 +6399,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", ), ], ) @@ -6425,7 +6429,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 @@ -6517,6 +6523,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( diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 4f73e8935e8..8222ca09b24 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7712,7 +7712,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, ) @@ -7740,7 +7740,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, ), ) @@ -7778,6 +7780,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"], @@ -7799,7 +7862,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} @@ -7849,7 +7916,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, ), ) @@ -7914,14 +7985,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"]}}} ) diff --git a/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py b/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py index 143f5bfb83b..7facd6a62d6 100644 --- a/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py +++ b/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py @@ -2,6 +2,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from types import MappingProxyType +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -26,6 +27,7 @@ class _FakeRouter: infos: Mapping[str, ModelGroupInfo] model_group_alias: Mapping[str, str] = MappingProxyType({}) + team_model_to_deployment_indices: Mapping[tuple[str, str], Sequence[int]] = MappingProxyType({}) def get_model_group_info(self, model_group: str, team_id: str | None = None) -> ModelGroupInfo | None: return self.infos.get(model_group) @@ -351,6 +353,8 @@ def test_no_router_answers_with_a_problem_rather_than_the_openai_error_shape(mon def test_an_unexpected_router_failure_answers_as_a_problem_not_the_openai_error_shape(monkeypatch): class _Exploding: + team_model_to_deployment_indices: Final[Mapping[tuple[str, str], Sequence[int]]] = MappingProxyType({}) + def get_model_group_info(self, model_group: str, team_id: str | None = None) -> ModelGroupInfo: raise RuntimeError("router blew up") diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index 6eff950dddf..1d5a1f8fa30 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -18,7 +18,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.openai_files_endpoints.common_utils import ( TeamProviderDeployment, @@ -292,6 +293,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. " @@ -1501,6 +1503,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"}