diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index bc9d03a5f03..53eb13ed502 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -2,7 +2,7 @@ import importlib.util import os import re import sys -from collections.abc import Collection, Iterator, Mapping +from collections.abc import Collection, Iterator, Mapping, Sequence from functools import lru_cache from itertools import accumulate, chain, repeat from logging import Logger @@ -23,6 +23,7 @@ from litellm.constants import ( STANDARD_CUSTOMER_ID_HEADERS, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor, reject_api_key_secret_reference +from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import ( SSRFError, @@ -40,12 +41,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 @@ -217,23 +219,50 @@ def deployment_clientside_auth_params(litellm_params: LiteLLMParamsTypedDict) -> return litellm_params["configurable_clientside_auth_params"] +def _model_group_deployments(model: str, llm_router: Router) -> Sequence[DeploymentTypedDict]: + registry: Final = llm_router.registry() + provider_group: Final = model.split("/", 1)[0] + return ( + registry.deployments(model, source="routed_group") + or (registry.deployments(provider_group) if provider_group in provider_list else None) + or () + ) + + +_VERTEX_TARGET_READING_PROVIDERS: Final[frozenset[str]] = frozenset( + (LlmProviders.VERTEX_AI.value, LlmProviders.VERTEX_AI_BETA.value, LlmProviders.GEMINI.value, LlmProviders.GDC.value) +) + + +def _deployment_provider(litellm_params: LiteLLMParamsTypedDict) -> str | None: + match litellm_params: + case {"custom_llm_provider": str(declared)} if declared: + return declared + case {"model": str(model)}: + return inferred_provider(model) + case _: + return None + + +def _model_group_reads_vertex_target(model: str, llm_router: Router | None) -> bool: + if llm_router is None: + return True + deployments: Final = _model_group_deployments(model, llm_router) + providers: Final = frozenset(_deployment_provider(deployment["litellm_params"]) for deployment in deployments) + return not deployments or None in providers or bool(providers & _VERTEX_TARGET_READING_PROVIDERS) + + def _allow_model_level_clientside_configurable_parameters( model: str, param: str, request_body_value: object, llm_router: Router | None ) -> bool: """ The router, not the caller, picks which deployment serves the request, so an - opt-in counts only when every deployment the model can route to carries it + opt-in counts only when every deployment of the model group, already alias + resolved, carries it """ if llm_router is None: return False - registry: Final = llm_router.registry() - routed_group: Final = registry.resolve(model, kind="alias") or model - provider_group: Final = routed_group.split("/", 1)[0] - deployments: Final = ( - registry.deployments(routed_group, source="routed_group") - or (registry.deployments(provider_group) if provider_group in provider_list else None) - or () - ) + deployments: Final = _model_group_deployments(model, llm_router) return bool(deployments) and all( _is_param_allowed( param=param, @@ -322,21 +351,15 @@ def route_manages_mcp_servers(route: str | None) -> bool: VERTEX_TARGET_PARAMS: Final[frozenset[str]] = frozenset( ("vertex_project", "vertex_location", "vertex_ai_project", "vertex_ai_location") ) -_RAG_INGEST_VERTEX_TARGET_PARAMS: Final[frozenset[str]] = frozenset(("vertex_project", "vertex_location")) -_ROUTE_NAMED_VERTEX_TARGET_PARAMS: Final[Mapping[str, frozenset[str]]] = MappingProxyType( - { - "/vertex_ai/live": VERTEX_TARGET_PARAMS, - "/rag/ingest": _RAG_INGEST_VERTEX_TARGET_PARAMS, - "/v1/rag/ingest": _RAG_INGEST_VERTEX_TARGET_PARAMS, - } +_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]: - route_named: Final = _ROUTE_NAMED_VERTEX_TARGET_PARAMS.get(route) if route is not None else None - if route_named is None: + 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 route_named}) + 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") @@ -509,14 +532,17 @@ _BANNED_METADATA_PARAMS: Final[tuple[str, ...]] = tuple( def _opt_in_model_groups( - model: str, request_fallback_models: tuple[str, ...], llm_router: Router | None + model: str, also_requested_models: tuple[str, ...], llm_router: Router | None ) -> frozenset[str]: - requested: Final = frozenset((model, *request_fallback_models)) + requested: Final = frozenset((model, *also_requested_models)) if llm_router is None: return requested registry: Final = llm_router.registry() - routed: Final = frozenset(target for name in requested if (target := registry.resolve(name, kind="alias"))) - return (requested | routed).union(*(router_fallback_model_names(name, llm_router) for name in requested | routed)) + routed: Final = frozenset(registry.resolve(name, kind="alias") or name for name in requested) + fallbacks: Final = frozenset[str]().union( + *(router_fallback_model_names(name, llm_router, _EVERY_FALLBACK_TYPE) for name in requested | routed) + ) + return routed | frozenset(registry.resolve(name, kind="alias") or name for name in fallbacks) def _check_banned_params( @@ -528,7 +554,7 @@ def _check_banned_params( manages_deployments: bool = False, mcp_server_fields_allowed: bool = False, banned_params: tuple[str, ...] = _BANNED_REQUEST_BODY_PARAMS, - request_fallback_models: tuple[str, ...] = (), + also_requested_models: tuple[str, ...] = (), ) -> None: """Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in. @@ -558,7 +584,8 @@ def _check_banned_params( llm_router=llm_router, ) is True - for model_group in _opt_in_model_groups(model, request_fallback_models, llm_router) + for model_group in _opt_in_model_groups(model, also_requested_models, llm_router) + if param not in VERTEX_TARGET_PARAMS or _model_group_reads_vertex_target(model_group, llm_router) ): # Per-param opt-in: only THIS param is permitted by the # deployment's ``configurable_clientside_auth_params``. Skip @@ -649,27 +676,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} @@ -735,7 +774,11 @@ def is_request_body_safe( raise ValueError("Rejected Request: model_list is not allowed in the request body.") manages_deployments: Final = route_manages_deployments(route) root_body: Final = _without_route_named_vertex_target(request_body, route) - request_fallback_models: Final = request_fallback_model_names(root_body) + 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 ()), + ) _check_banned_params( root_body, general_settings, @@ -743,7 +786,7 @@ def is_request_body_safe( model, manages_deployments=manages_deployments, mcp_server_fields_allowed=route_manages_mcp_servers(route), - request_fallback_models=request_fallback_models, + also_requested_models=also_requested_models, ) for nested_key in _NESTED_CONFIG_KEYS: nested = _coerce_metadata_to_dict(request_body.get(nested_key)) @@ -754,7 +797,7 @@ def is_request_body_safe( llm_router, model, manages_deployments=manages_deployments, - request_fallback_models=request_fallback_models, + also_requested_models=also_requested_models, ) for metadata_key in _NESTED_METADATA_KEYS: metadata = _coerce_metadata_to_dict(request_body.get(metadata_key)) @@ -766,7 +809,7 @@ def is_request_body_safe( model, manages_deployments=manages_deployments, banned_params=_BANNED_METADATA_PARAMS, - request_fallback_models=request_fallback_models, + also_requested_models=also_requested_models, ) if any(isinstance(key, str) and key.startswith(f"{metadata_key}[") for key in request_body): _check_banned_params( @@ -776,7 +819,7 @@ def is_request_body_safe( model, manages_deployments=manages_deployments, banned_params=_BANNED_METADATA_PARAMS, - request_fallback_models=request_fallback_models, + also_requested_models=also_requested_models, ) for target in iter_request_fallback_targets(request_body): if isinstance(target, dict): @@ -786,7 +829,7 @@ def is_request_body_safe( llm_router, model, manages_deployments=manages_deployments, - request_fallback_models=request_fallback_models, + also_requested_models=also_requested_models, ) target_model = target.get("model") if isinstance(target_model, str): @@ -807,7 +850,7 @@ def is_request_body_safe( model, manages_deployments=manages_deployments, banned_params=_BANNED_METADATA_PARAMS, - request_fallback_models=request_fallback_models, + also_requested_models=also_requested_models, ) for user_config_params in _user_config_litellm_params(request_body): # pyright: ignore[reportUnknownArgumentType] # request_body is an untyped dict reject_api_key_secret_reference(user_config_params) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index df374c1a154..302f41491f5 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -3857,6 +3857,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/" @@ -3947,6 +3951,10 @@ def _live_target_override_reason( if allow_client_side_credentials: return None match target: + case _ConfiguredVertexLiveTarget(credentials=credentials) if vertex_project and ( + credentials is None or credentials.vertex_project is None + ): + return VERTEX_LIVE_UNPINNED_PROJECT_CLOSE_REASON case _ConfiguredVertexLiveTarget(credentials=VertexPassThroughCredentials() as credentials): configured_field: Final = next( ( @@ -3986,12 +3994,13 @@ def _live_target_override_reason( ) -_VERTEX_RESOURCE_PROJECT_AND_LOCATION: Final = re.compile(r"projects/[^/]+/locations/[^/]+/") +_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: @@ -4000,8 +4009,8 @@ def _build_vertex_live_setup_model_rewriter( 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. A ``projects/...`` path - picks the project the proxy's token is spent on, so it is moved onto the connection's project and region - unless ``allow_client_side_credentials`` lets callers choose them + 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 @@ -4011,7 +4020,8 @@ def _build_vertex_live_setup_model_rewriter( resource_prefix: Final = _VERTEX_RESOURCE_PROJECT_AND_LOCATION.match(setup_model) if keeps_caller_resource_path or resource_prefix is None: return setup_model - return f"projects/{vertex_project}/locations/{vertex_location}/{setup_model[resource_prefix.end() :]}" + 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)}" @@ -4155,6 +4165,7 @@ 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 7b97bbb670e..9534d33bd91 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2426,6 +2426,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]: @@ -2726,7 +2741,12 @@ async def websocket_passthrough_request( await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter)) elif bytes_data is not None: - await upstream_ws.send(bytes_data) + rewritten_bytes = _rewrite_vertex_live_binary_setup_model(bytes_data, setup_model_rewriter) + if rewritten_bytes is None: + await websocket.close(code=1008, reason=VERTEX_LIVE_UNDECODABLE_BINARY_CLOSE_REASON) + await upstream_ws.close() + break + await upstream_ws.send(rewritten_bytes) except asyncio.CancelledError: raise except Exception: diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 19c4bf5ff17..35e7ae71fb7 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -21,7 +21,7 @@ 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, ) @@ -64,6 +64,7 @@ from litellm.proxy.vector_store_endpoints.utils import ( from litellm.rag.ingestion.vertex_ai_ingestion import DEFAULT_VERTEX_RAG_LOCATION from litellm.rag.main import get_ingestion_class from litellm.repositories.table_repositories import ManagedVectorStoresRepository +from litellm.repositories.user_repository import UserRepository from litellm.secret_managers.main import get_secret_str from litellm.types.utils import ModelResponse @@ -183,41 +184,95 @@ 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/") +_VERTEX_RAG_CORPUS_NAME: Final = re.compile( + r"projects/(?P[^/%\\]+)/locations/(?P[^/%\\]+)/ragCorpora/(?P[^/%\\]+)" +) +_DOT_SEGMENTS: Final = frozenset({".", ".."}) -def _vertex_target_error( - vector_store_config: Mapping[str, object], - managed_store: LiteLLM_ManagedVectorStore | None, - user_api_key_dict: UserAPIKeyAuth, - general_settings: Mapping[str, object], -) -> str | None: - if ( - managed_store is not None - or vector_store_config.get("custom_llm_provider") != "vertex_ai" - or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN - or general_settings.get("allow_client_side_credentials") is True - ): - return None - configured_project: Final = litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT") - configured_location: Final = VertexBase.safe_get_vertex_ai_location({}) or DEFAULT_VERTEX_RAG_LOCATION - corpus_name: Final = vector_store_config.get("vector_store_id") - corpus: Final = _VERTEX_RAG_CORPUS_NAME.match(corpus_name) if isinstance(corpus_name, str) else None - requested: Final = ( - ("vertex_project", vector_store_config.get("vertex_project"), configured_project), - ("vertex_location", vector_store_config.get("vertex_location"), configured_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) - if field is None: +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']}" + ), + } + ) + + +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", @@ -730,15 +785,17 @@ async def rag_ingest( raise HTTPException(status_code=400, detail={"error": str(e)}) managed_store: Final = resolved_stores.get(request_vector_store_config.get("vector_store_id")) - vertex_target_error: Final = _vertex_target_error( - request_vector_store_config, managed_store, user_api_key_dict, general_settings_view() - ) - if vertex_target_error is not None: - raise HTTPException(status_code=400, detail={"error": vertex_target_error}) - 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() + ) + base: Final = _corpus_on_configured_project(requested_base) if vertex_target_checked else requested_base + vertex_target_error: Final = _vertex_target_error(base) if vertex_target_checked else None + if vertex_target_error is not None and not await _admin_set_store_target(managed_store, prisma_client): + 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/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 f5d2ebf234f..54bdc8e5959 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -5136,7 +5136,11 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: 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]]]) -> Router: + 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=[ { @@ -5153,7 +5157,8 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: ("gemini-pinned", []), ) ], - fallbacks=router_fallbacks, + model_group_alias=model_group_alias or {}, + **{fallback_setting: router_fallbacks}, ) @pytest.mark.parametrize( @@ -5175,13 +5180,16 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: 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_group: str, is_allowed: bool) -> None: - router: Final = self._opted_in_and_pinned_router([{"gemini-optin": [fallback_group]}]) + 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 @@ -5189,6 +5197,67 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: 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") + + @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: @@ -5208,8 +5277,8 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: "route, kept_params", [ ("/vertex_ai/live", _VERTEX_TARGET_PARAMS), - ("/rag/ingest", ("vertex_project", "vertex_location")), - ("/v1/rag/ingest", ("vertex_project", "vertex_location")), + ("/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: @@ -5229,8 +5298,7 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: [ ("/vertex_ai/live", "api_base"), ("/rag/ingest", "api_base"), - ("/rag/ingest", "vertex_ai_project"), - ("/rag/ingest", "vertex_ai_location"), + ("/v1/rag/ingest", "api_base"), ], ) def test_routes_that_name_their_vertex_target_reject_every_other_banned_param(self, route: str, param: str) -> None: 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 beaa7bb380d..75efce38a07 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 @@ -6338,6 +6338,7 @@ 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, ) @@ -6706,6 +6707,103 @@ class TestVertexAILiveWebsocketPassthrough: == 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( + "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 + class TestAnthropicProxyRoute: """The /anthropic passthrough route: custom auth headers must not clobber the 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 228a3934596..53feb88eac8 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 @@ -7396,6 +7396,7 @@ async def _run_setup_rewrite_passthrough(setup_model: str, llm_router, keeps_cal 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, ), @@ -7435,6 +7436,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"], @@ -7456,7 +7518,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, keeps_caller_resource_path=False + vertex_project="proj-db", + vertex_location="global", + configured_location="global", + llm_router=None, + keeps_caller_resource_path=False, ) setup_data = {"model": setup_model} @@ -7508,6 +7574,7 @@ async def test_websocket_passthrough_logs_a_bare_alias_setup_model(): setup_model_rewriter=_build_vertex_live_setup_model_rewriter( vertex_project="proj-db", vertex_location="global", + configured_location="global", llm_router=None, keeps_caller_resource_path=False, ), diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index 664556a8a4f..098b2065ae1 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -15,7 +15,8 @@ from fastapi.testclient import TestClient import litellm from litellm import Router -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.constants import LITELLM_PROXY_ADMIN_NAME +from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app from litellm.types.vector_stores import LiteLLM_ManagedVectorStore @@ -284,6 +285,7 @@ VERTEX_REGISTRY_STORE = { "vector_store_id": "projects/registry-project/locations/us-central1/ragCorpora/42", "custom_llm_provider": "vertex_ai", "litellm_params": {"vertex_project": "registry-project", "vertex_location": "us-central1"}, + "is_config": True, } UNSUPPORTED_INGEST_PROVIDER_ERROR = ( "Provider '{provider}' is not supported for RAG ingestion. " @@ -1028,8 +1030,20 @@ def _vertex_ingest(client: TestClient, vector_store: Mapping[str, str]) -> tuple "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", ], - ids=["project_field", "location_field", "corpus_project", "corpus_project_behind_matching_field", "corpus_region"], ) @pytest.mark.usefixtures("proxy_vertex_target") def test_rag_ingest_refuses_a_key_pointing_the_proxys_vertex_credentials_elsewhere( @@ -1060,8 +1074,9 @@ def test_rag_ingest_refuses_a_key_naming_a_vertex_project_when_the_proxy_names_n {}, {"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"], + 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]): @@ -1074,6 +1089,11 @@ def test_rag_ingest_lets_a_key_use_the_proxys_own_vertex_target(client_internal_ } +@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", [ @@ -1084,14 +1104,14 @@ def test_rag_ingest_lets_a_key_use_the_proxys_own_vertex_target(client_internal_ ) @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] + 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 = { - "vertex_project": "other-project", - "vertex_location": "us-east5", + project_field: "other-project", + location_field: "us-east5", "vector_store_id": "projects/other-project/locations/us-east5/ragCorpora/7", } try: @@ -1107,6 +1127,119 @@ def test_rag_ingest_lets_the_admin_choose_another_vertex_target( } +@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.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() + + +def _foreign_vertex_store(created_by: str | None) -> LiteLLM_ManagedVectorStore: + return LiteLLM_ManagedVectorStore( + vector_store_id=_FOREIGN_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( + "store, creator_row", + [ + (_foreign_vertex_store(LITELLM_PROXY_ADMIN_NAME), None), + ( + _foreign_vertex_store("admin-user"), + LiteLLM_UserTable(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN.value), + ), + (LiteLLM_ManagedVectorStore(**_foreign_vertex_store(None), is_config=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( + store: LiteLLM_ManagedVectorStore, creator_row: LiteLLM_UserTable | None +): + 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": _FOREIGN_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"}