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"}