fix(proxy)!: reject client-supplied Vertex location and project unless the admin opts in (#47028)

* fix(proxy): reject client-supplied Vertex location and project unless the admin opts in

A request body could carry vertex_location, vertex_project, or their
vertex_ai_* aliases and override the region and project pinned on a
Vertex AI deployment, sending the admin's credentials to a caller-chosen
project or region. Add the four fields to the banned request params so
they need allow_client_side_credentials or configurable_clientside_auth_params.

/vertex_ai/live and /rag/ingest name their Vertex project and location
by design and have no deployment pin to override, so they keep those two
fields while every other banned param is still rejected there

* fix(proxy): close the remaining Vertex target overrides and keep inert fields working

/vertex_ai/live now closes with 1008 when the query names a project or location that differs from a use_in_pass_through deployment's pin, unless that deployment or the proxy opts in

A model group's opt-in now counts only when every deployment in it carries it, and Google-style routes look the opt-in up on the model named in the URL

Vertex target names inside metadata, the live route's vertex_ai_* aliases, and /health/test_connection inheritance go back to the base behavior, since nothing reads them as a target

* test(batches): give the batch FakeRouter the model-list lookup the opt-in check reads

* fix(proxy): pin default_vertex_config on /vertex_ai/live and the proxy's Vertex project on /rag/ingest

* fix(proxy): require the Vertex opt-in on alias targets and fallbacks, and pass the live rewriter flag in main's tests

* test(proxy): give the batch tests' fake router the alias and routed-group lookups

* fix(proxy): close the Vertex opt-in, RAG corpus, and live setup-frame gaps found in review

A Vertex opt-in now counts only when every group the request can reach opts in, router,
context window, and content policy fallbacks and one alias hop included, while a group
with no Google deployment never needs the fields. A model named in the route path is
checked next to the body model.

RAG ingest from a non-admin key refuses a bare or dot-segment corpus id, a
vertex_ai_project or vertex_ai_location naming another target, and a store a non-admin
saved before the upgrade, and reads a corpus named by project number on the proxy's
own project.

/vertex_ai/live closes 1008 on any caller-named project when the proxy pins none, moves
a setup frame sent as a binary message onto the session's target, and resolves the
region from default_vertex_config before the model path.

* fix(proxy): let a live caller name the proxy's configured Vertex project when nothing pins one

With default_vertex_config set but no project pinned, /vertex_ai/live closed every query that named a project,
including the one the admin set in VERTEXAI_PROJECT or DEFAULT_VERTEXAI_PROJECT. A query naming one of those
projects now connects, and any other project still closes with 1008 before a token is minted.

* fix(proxy): keep an admin-created vector store's corpus as saved on RAG ingest

A non-admin ingest into a store an admin created named by project number had its corpus moved onto
the proxy's own project before the admin-store exemption ran, so the upload went to a different
corpus than the admin saved. The exemption now runs first, and a non-admin store still gets the
project-number rewrite and the target check.

* fix(proxy): authorize a RAG corpus moved onto the proxy's project and check client-side opt-ins once per request

* fix(proxy): merge opt-in fallback groups only when their deployments share provider and opt-in list

* fix(proxy): check team public models before allowing target overrides

* fix(proxy): include team wildcard targets in opt-in checks

* fix(proxy): validate overrides against authenticated team routes

* fix(proxy): retain URL model during deferred credential checks
This commit is contained in:
Mateo Wang 2026-11-10 20:50:17 +00:00 • committed by GitHub
parent 88dd18f4ab
commit 4867a56e9a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
23 changed files with 2163 additions and 134 deletions

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -93,6 +93,7 @@ from litellm.proxy.auth.auth_checks import (
key_model_aliases_for_auth_check,
team_model_aliases_for_auth_check,
)
from litellm.proxy.auth.auth_utils import deployment_clientside_auth_params
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports
@ -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<location>[^/]+)/")
def _build_vertex_live_setup_model_rewriter(
vertex_project: str | None,
vertex_location: str | None,
configured_location: str | None,
llm_router: Router | None,
keeps_caller_resource_path: bool,
) -> Callable[[str], str] | None:
"""
Rewrite the ``setup`` frame's model into the full Vertex resource path the Live API requires.
Clients address the gateway the way they address LiteLLM (bare id or model alias); Vertex reads anything
that is not a ``projects/...`` path as a project name and closes the socket
that is not a ``projects/...`` path as a project name and closes the socket. A ``projects/...`` path
picks the project the proxy's token is spent on, so it is moved onto the connection's project, and onto its
region when one is configured, unless ``allow_client_side_credentials`` lets callers choose them
"""
if vertex_project is None or vertex_location is None:
return None
def rewrite(setup_model: str) -> str:
if setup_model.startswith("projects/"):
return setup_model
resource_prefix: Final = _VERTEX_RESOURCE_PROJECT_AND_LOCATION.match(setup_model)
if keeps_caller_resource_path or resource_prefix is None:
return setup_model
location: Final = configured_location or resource_prefix.group("location")
return f"projects/{vertex_project}/locations/{location}/{setup_model[resource_prefix.end() :]}"
aliased: Final = _resolve_alias_to_upstream_model(setup_model, llm_router)
return f"projects/{vertex_project}/locations/{vertex_location}/{_vertex_publisher_model_suffix(aliased)}"
@ -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,
),
)

View file

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

View file

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

View file

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

View file

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

View file

@ -8,6 +8,7 @@ Provides:
import base64
import json
import re
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
@ -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<project>[^/%\\]+)/locations/(?P<location>[^/%\\]+)/ragCorpora/(?P<corpus>[^/%\\]+)"
)
_DOT_SEGMENTS: Final = frozenset({".", ".."})
def _clean_vertex_corpus_name(corpus_name: object) -> re.Match[str] | None:
if not isinstance(corpus_name, str) or not _DOT_SEGMENTS.isdisjoint(corpus_name.split("/")):
return None
return _VERTEX_RAG_CORPUS_NAME.fullmatch(corpus_name)
def _vertex_target_message(field: str) -> str:
return (
f"Rejected Request: {field} picks a Vertex project or region other than the proxy's. "
"Use a vector store a proxy admin created, or set general_settings.allow_client_side_credentials"
)
def _vertex_target_is_checked(
vector_store_config: Mapping[str, object],
user_api_key_dict: UserAPIKeyAuth,
general_settings: Mapping[str, object],
) -> bool:
return (
vector_store_config.get("custom_llm_provider") == "vertex_ai"
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN
and general_settings.get("allow_client_side_credentials") is not True
)
def _configured_vertex_project() -> str | None:
return litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT")
def _corpus_on_configured_project(vector_store_config: Mapping[str, object]) -> Mapping[str, object]:
corpus: Final = _clean_vertex_corpus_name(vector_store_config.get("vector_store_id"))
configured_project: Final = _configured_vertex_project()
if corpus is None or configured_project is None or not corpus["project"].isdigit():
return vector_store_config
return MappingProxyType(
{
**vector_store_config,
"vector_store_id": (
f"projects/{configured_project}/locations/{corpus['location']}/ragCorpora/{corpus['corpus']}"
),
}
)
async def _authorize_rewritten_corpus(
requested: Mapping[str, object], rewritten: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
) -> None:
rewritten_id: Final = rewritten.get("vector_store_id")
if isinstance(rewritten_id, str) and rewritten_id != requested.get("vector_store_id"):
await assert_user_can_access_vector_store_id(vector_store_id=rewritten_id, user_api_key_dict=user_api_key_dict)
def _vertex_target_error(vector_store_config: Mapping[str, object]) -> str | None:
corpus_name: Final = vector_store_config.get("vector_store_id")
corpus: Final = _clean_vertex_corpus_name(corpus_name)
if corpus_name and corpus is None:
return _vertex_target_message("vector_store_id")
configured_project: Final = _configured_vertex_project()
configured_location: Final = VertexBase.safe_get_vertex_ai_location({}) or DEFAULT_VERTEX_RAG_LOCATION
requested: Final = (
*(
(name, vector_store_config.get(name), configured_project)
for name in ("vertex_project", "vertex_ai_project")
),
*(
(name, vector_store_config.get(name), configured_location)
for name in ("vertex_location", "vertex_ai_location")
),
("vector_store_id", corpus["project"] if corpus is not None else None, configured_project),
("vector_store_id", corpus["location"] if corpus is not None else None, configured_location),
)
field: Final = next((name for name, asked, configured in requested if asked and asked != configured), None)
return None if field is None else _vertex_target_message(field)
async def _admin_set_store_target(
managed_store: LiteLLM_ManagedVectorStore | None, prisma_client: "PrismaClient | None"
) -> bool:
if managed_store is None:
return False
if managed_store.get("is_config", False):
return True
creator: Final = managed_store.get("user_id")
if creator == LITELLM_PROXY_ADMIN_NAME:
return True
if creator is None or prisma_client is None:
return False
creator_row: Final = await UserRepository(prisma_client).find_by_id(creator)
return creator_row is not None and creator_row.user_role == LitellmUserRoles.PROXY_ADMIN
_MANAGED_STORE_CALLER_OPTIONS: Final = frozenset(
{
"vector_store_id",
@ -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,

View file

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

View file

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

View file

@ -4,6 +4,7 @@ Unit tests for auth_utils functions related to rate limiting and customer ID ext
import base64
import logging
import time
from collections.abc import Callable, 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},

View file

@ -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"),
[

View file

@ -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 = []

View file

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

View file

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

View file

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

View file

@ -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"]}}}
)

View file

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

View file

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