mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
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.
This commit is contained in:
mateo-berri
parent
fe75e8908d
commit
36665d16fd
9 changed files with 600 additions and 95 deletions
|
|
@ -2,7 +2,7 @@ import importlib.util
|
|||
import os
|
||||
import re
|
||||
import sys
|
||||
from collections.abc import Collection, Iterator, Mapping
|
||||
from collections.abc import Collection, Iterator, Mapping, Sequence
|
||||
from functools import lru_cache
|
||||
from itertools import accumulate, chain, repeat
|
||||
from logging import Logger
|
||||
|
|
@ -23,6 +23,7 @@ from litellm.constants import (
|
|||
STANDARD_CUSTOMER_ID_HEADERS,
|
||||
)
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor, reject_api_key_secret_reference
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.litellm_core_utils.url_utils import (
|
||||
SSRFError,
|
||||
|
|
@ -40,12 +41,13 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
|||
from litellm.types.router import (
|
||||
CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS,
|
||||
Deployment,
|
||||
DeploymentTypedDict,
|
||||
LiteLLMParamsTypedDict,
|
||||
reject_server_owned_wif_fields,
|
||||
server_owned_wif_fields_present,
|
||||
)
|
||||
from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params
|
||||
from litellm.types.utils import CustomPricingLiteLLMParams, oauth_token_exchange_litellm_params
|
||||
from litellm.types.utils import CustomPricingLiteLLMParams, LlmProviders, oauth_token_exchange_litellm_params
|
||||
from litellm.types.videos.main import DecodedVideoId
|
||||
|
||||
|
||||
|
|
@ -217,23 +219,50 @@ def deployment_clientside_auth_params(litellm_params: LiteLLMParamsTypedDict) ->
|
|||
return litellm_params["configurable_clientside_auth_params"]
|
||||
|
||||
|
||||
def _model_group_deployments(model: str, llm_router: Router) -> Sequence[DeploymentTypedDict]:
|
||||
registry: Final = llm_router.registry()
|
||||
provider_group: Final = model.split("/", 1)[0]
|
||||
return (
|
||||
registry.deployments(model, source="routed_group")
|
||||
or (registry.deployments(provider_group) if provider_group in provider_list else None)
|
||||
or ()
|
||||
)
|
||||
|
||||
|
||||
_VERTEX_TARGET_READING_PROVIDERS: Final[frozenset[str]] = frozenset(
|
||||
(LlmProviders.VERTEX_AI.value, LlmProviders.VERTEX_AI_BETA.value, LlmProviders.GEMINI.value, LlmProviders.GDC.value)
|
||||
)
|
||||
|
||||
|
||||
def _deployment_provider(litellm_params: LiteLLMParamsTypedDict) -> str | None:
|
||||
match litellm_params:
|
||||
case {"custom_llm_provider": str(declared)} if declared:
|
||||
return declared
|
||||
case {"model": str(model)}:
|
||||
return inferred_provider(model)
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _model_group_reads_vertex_target(model: str, llm_router: Router | None) -> bool:
|
||||
if llm_router is None:
|
||||
return True
|
||||
deployments: Final = _model_group_deployments(model, llm_router)
|
||||
providers: Final = frozenset(_deployment_provider(deployment["litellm_params"]) for deployment in deployments)
|
||||
return not deployments or None in providers or bool(providers & _VERTEX_TARGET_READING_PROVIDERS)
|
||||
|
||||
|
||||
def _allow_model_level_clientside_configurable_parameters(
|
||||
model: str, param: str, request_body_value: object, llm_router: Router | None
|
||||
) -> bool:
|
||||
"""
|
||||
The router, not the caller, picks which deployment serves the request, so an
|
||||
opt-in counts only when every deployment the model can route to carries it
|
||||
opt-in counts only when every deployment of the model group, already alias
|
||||
resolved, carries it
|
||||
"""
|
||||
if llm_router is None:
|
||||
return False
|
||||
registry: Final = llm_router.registry()
|
||||
routed_group: Final = registry.resolve(model, kind="alias") or model
|
||||
provider_group: Final = routed_group.split("/", 1)[0]
|
||||
deployments: Final = (
|
||||
registry.deployments(routed_group, source="routed_group")
|
||||
or (registry.deployments(provider_group) if provider_group in provider_list else None)
|
||||
or ()
|
||||
)
|
||||
deployments: Final = _model_group_deployments(model, llm_router)
|
||||
return bool(deployments) and all(
|
||||
_is_param_allowed(
|
||||
param=param,
|
||||
|
|
@ -322,21 +351,15 @@ def route_manages_mcp_servers(route: str | None) -> bool:
|
|||
VERTEX_TARGET_PARAMS: Final[frozenset[str]] = frozenset(
|
||||
("vertex_project", "vertex_location", "vertex_ai_project", "vertex_ai_location")
|
||||
)
|
||||
_RAG_INGEST_VERTEX_TARGET_PARAMS: Final[frozenset[str]] = frozenset(("vertex_project", "vertex_location"))
|
||||
_ROUTE_NAMED_VERTEX_TARGET_PARAMS: Final[Mapping[str, frozenset[str]]] = MappingProxyType(
|
||||
{
|
||||
"/vertex_ai/live": VERTEX_TARGET_PARAMS,
|
||||
"/rag/ingest": _RAG_INGEST_VERTEX_TARGET_PARAMS,
|
||||
"/v1/rag/ingest": _RAG_INGEST_VERTEX_TARGET_PARAMS,
|
||||
}
|
||||
_ROUTES_THAT_NAME_THEIR_VERTEX_TARGET: Final[frozenset[str]] = frozenset(
|
||||
("/vertex_ai/live", "/rag/ingest", "/v1/rag/ingest")
|
||||
)
|
||||
|
||||
|
||||
def _without_route_named_vertex_target(request_body: Mapping[str, object], route: str | None) -> Mapping[str, object]:
|
||||
route_named: Final = _ROUTE_NAMED_VERTEX_TARGET_PARAMS.get(route) if route is not None else None
|
||||
if route_named is None:
|
||||
if route not in _ROUTES_THAT_NAME_THEIR_VERTEX_TARGET:
|
||||
return request_body
|
||||
return MappingProxyType({key: value for key, value in request_body.items() if key not in route_named})
|
||||
return MappingProxyType({key: value for key, value in request_body.items() if key not in VERTEX_TARGET_PARAMS})
|
||||
|
||||
|
||||
_NESTED_CONFIG_KEYS: Final[tuple[str, ...]] = ("litellm_embedding_config", "extra_body")
|
||||
|
|
@ -509,14 +532,17 @@ _BANNED_METADATA_PARAMS: Final[tuple[str, ...]] = tuple(
|
|||
|
||||
|
||||
def _opt_in_model_groups(
|
||||
model: str, request_fallback_models: tuple[str, ...], llm_router: Router | None
|
||||
model: str, also_requested_models: tuple[str, ...], llm_router: Router | None
|
||||
) -> frozenset[str]:
|
||||
requested: Final = frozenset((model, *request_fallback_models))
|
||||
requested: Final = frozenset((model, *also_requested_models))
|
||||
if llm_router is None:
|
||||
return requested
|
||||
registry: Final = llm_router.registry()
|
||||
routed: Final = frozenset(target for name in requested if (target := registry.resolve(name, kind="alias")))
|
||||
return (requested | routed).union(*(router_fallback_model_names(name, llm_router) for name in requested | routed))
|
||||
routed: Final = frozenset(registry.resolve(name, kind="alias") or name for name in requested)
|
||||
fallbacks: Final = frozenset[str]().union(
|
||||
*(router_fallback_model_names(name, llm_router, _EVERY_FALLBACK_TYPE) for name in requested | routed)
|
||||
)
|
||||
return routed | frozenset(registry.resolve(name, kind="alias") or name for name in fallbacks)
|
||||
|
||||
|
||||
def _check_banned_params(
|
||||
|
|
@ -528,7 +554,7 @@ def _check_banned_params(
|
|||
manages_deployments: bool = False,
|
||||
mcp_server_fields_allowed: bool = False,
|
||||
banned_params: tuple[str, ...] = _BANNED_REQUEST_BODY_PARAMS,
|
||||
request_fallback_models: tuple[str, ...] = (),
|
||||
also_requested_models: tuple[str, ...] = (),
|
||||
) -> None:
|
||||
"""Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in.
|
||||
|
||||
|
|
@ -558,7 +584,8 @@ def _check_banned_params(
|
|||
llm_router=llm_router,
|
||||
)
|
||||
is True
|
||||
for model_group in _opt_in_model_groups(model, request_fallback_models, llm_router)
|
||||
for model_group in _opt_in_model_groups(model, also_requested_models, llm_router)
|
||||
if param not in VERTEX_TARGET_PARAMS or _model_group_reads_vertex_target(model_group, llm_router)
|
||||
):
|
||||
# Per-param opt-in: only THIS param is permitted by the
|
||||
# deployment's ``configurable_clientside_auth_params``. Skip
|
||||
|
|
@ -649,27 +676,39 @@ def fallback_target_model_name(target: object) -> str | None:
|
|||
return None
|
||||
|
||||
|
||||
def _direct_fallback_model_names(model: str, llm_router: Router) -> frozenset[str]:
|
||||
_EVERY_FALLBACK_TYPE: Final[tuple[str, ...]] = ("general", "context_window", "content_policy")
|
||||
|
||||
|
||||
def _direct_fallback_model_names(model: str, llm_router: Router, fallback_types: tuple[str, ...]) -> frozenset[str]:
|
||||
from litellm.proxy.auth.model_checks import get_all_fallbacks
|
||||
|
||||
return frozenset(
|
||||
name for target in get_all_fallbacks(model, llm_router) if (name := fallback_target_model_name(target))
|
||||
targets: Final = chain.from_iterable(
|
||||
get_all_fallbacks(model, llm_router, fallback_type) for fallback_type in fallback_types
|
||||
)
|
||||
return frozenset(name for target in targets if (name := fallback_target_model_name(target)))
|
||||
|
||||
|
||||
def _fallback_hop(
|
||||
frontier_and_reached: tuple[frozenset[str], frozenset[str]], llm_router: Router
|
||||
frontier_and_reached: tuple[frozenset[str], frozenset[str]], router_and_types: tuple[Router, tuple[str, ...]]
|
||||
) -> tuple[frozenset[str], frozenset[str]]:
|
||||
frontier, reached = frontier_and_reached
|
||||
llm_router, fallback_types = router_and_types
|
||||
next_hop: Final = (
|
||||
frozenset(chain.from_iterable(_direct_fallback_model_names(model, llm_router) for model in frontier)) - reached
|
||||
frozenset(
|
||||
chain.from_iterable(_direct_fallback_model_names(model, llm_router, fallback_types) for model in frontier)
|
||||
)
|
||||
- reached
|
||||
)
|
||||
return next_hop, reached | next_hop
|
||||
|
||||
|
||||
def router_fallback_model_names(model: str, llm_router: Router) -> frozenset[str]:
|
||||
def router_fallback_model_names(
|
||||
model: str, llm_router: Router, fallback_types: tuple[str, ...] = ("general",)
|
||||
) -> frozenset[str]:
|
||||
"""The groups a failure on `model` can reach through the router's fallbacks, chained and generic (`*`) ones."""
|
||||
hops: Final = accumulate(repeat(llm_router), _fallback_hop, initial=(frozenset({model}), frozenset({model})))
|
||||
hops: Final = accumulate(
|
||||
repeat((llm_router, fallback_types)), _fallback_hop, initial=(frozenset({model}), frozenset({model}))
|
||||
)
|
||||
_, reached = next(hop for hop in hops if not hop[0])
|
||||
return reached - {model}
|
||||
|
||||
|
|
@ -735,7 +774,11 @@ def is_request_body_safe(
|
|||
raise ValueError("Rejected Request: model_list is not allowed in the request body.")
|
||||
manages_deployments: Final = route_manages_deployments(route)
|
||||
root_body: Final = _without_route_named_vertex_target(request_body, route)
|
||||
request_fallback_models: Final = request_fallback_model_names(root_body)
|
||||
route_model: Final = _model_from_route_path(route) if route is not None else None
|
||||
also_requested_models: Final = (
|
||||
*request_fallback_model_names(root_body),
|
||||
*((route_model,) if route_model else ()),
|
||||
)
|
||||
_check_banned_params(
|
||||
root_body,
|
||||
general_settings,
|
||||
|
|
@ -743,7 +786,7 @@ def is_request_body_safe(
|
|||
model,
|
||||
manages_deployments=manages_deployments,
|
||||
mcp_server_fields_allowed=route_manages_mcp_servers(route),
|
||||
request_fallback_models=request_fallback_models,
|
||||
also_requested_models=also_requested_models,
|
||||
)
|
||||
for nested_key in _NESTED_CONFIG_KEYS:
|
||||
nested = _coerce_metadata_to_dict(request_body.get(nested_key))
|
||||
|
|
@ -754,7 +797,7 @@ def is_request_body_safe(
|
|||
llm_router,
|
||||
model,
|
||||
manages_deployments=manages_deployments,
|
||||
request_fallback_models=request_fallback_models,
|
||||
also_requested_models=also_requested_models,
|
||||
)
|
||||
for metadata_key in _NESTED_METADATA_KEYS:
|
||||
metadata = _coerce_metadata_to_dict(request_body.get(metadata_key))
|
||||
|
|
@ -766,7 +809,7 @@ def is_request_body_safe(
|
|||
model,
|
||||
manages_deployments=manages_deployments,
|
||||
banned_params=_BANNED_METADATA_PARAMS,
|
||||
request_fallback_models=request_fallback_models,
|
||||
also_requested_models=also_requested_models,
|
||||
)
|
||||
if any(isinstance(key, str) and key.startswith(f"{metadata_key}[") for key in request_body):
|
||||
_check_banned_params(
|
||||
|
|
@ -776,7 +819,7 @@ def is_request_body_safe(
|
|||
model,
|
||||
manages_deployments=manages_deployments,
|
||||
banned_params=_BANNED_METADATA_PARAMS,
|
||||
request_fallback_models=request_fallback_models,
|
||||
also_requested_models=also_requested_models,
|
||||
)
|
||||
for target in iter_request_fallback_targets(request_body):
|
||||
if isinstance(target, dict):
|
||||
|
|
@ -786,7 +829,7 @@ def is_request_body_safe(
|
|||
llm_router,
|
||||
model,
|
||||
manages_deployments=manages_deployments,
|
||||
request_fallback_models=request_fallback_models,
|
||||
also_requested_models=also_requested_models,
|
||||
)
|
||||
target_model = target.get("model")
|
||||
if isinstance(target_model, str):
|
||||
|
|
@ -807,7 +850,7 @@ def is_request_body_safe(
|
|||
model,
|
||||
manages_deployments=manages_deployments,
|
||||
banned_params=_BANNED_METADATA_PARAMS,
|
||||
request_fallback_models=request_fallback_models,
|
||||
also_requested_models=also_requested_models,
|
||||
)
|
||||
for user_config_params in _user_config_litellm_params(request_body): # pyright: ignore[reportUnknownArgumentType] # request_body is an untyped dict
|
||||
reject_api_key_secret_reference(user_config_params)
|
||||
|
|
|
|||
|
|
@ -3857,6 +3857,10 @@ VERTEX_LIVE_UNCONFIGURED_CLOSE_REASON: Final = (
|
|||
"Vertex AI auth failed: set a use_in_pass_through vertex model, default_vertex_config, or DEFAULT_VERTEXAI_* env"
|
||||
)
|
||||
|
||||
VERTEX_LIVE_UNPINNED_PROJECT_CLOSE_REASON: Final = (
|
||||
"Rejected Request: vertex_project must come from the proxy's configuration, see allow_client_side_credentials"
|
||||
)
|
||||
|
||||
VERTEX_PUBLISHER_MODEL_PREFIX: Final = "publishers/google/models/"
|
||||
|
||||
VERTEX_PUBLISHERS_SEGMENT: Final = "publishers/"
|
||||
|
|
@ -3947,6 +3951,10 @@ def _live_target_override_reason(
|
|||
if allow_client_side_credentials:
|
||||
return None
|
||||
match target:
|
||||
case _ConfiguredVertexLiveTarget(credentials=credentials) if vertex_project and (
|
||||
credentials is None or credentials.vertex_project is None
|
||||
):
|
||||
return VERTEX_LIVE_UNPINNED_PROJECT_CLOSE_REASON
|
||||
case _ConfiguredVertexLiveTarget(credentials=VertexPassThroughCredentials() as credentials):
|
||||
configured_field: Final = next(
|
||||
(
|
||||
|
|
@ -3986,12 +3994,13 @@ def _live_target_override_reason(
|
|||
)
|
||||
|
||||
|
||||
_VERTEX_RESOURCE_PROJECT_AND_LOCATION: Final = re.compile(r"projects/[^/]+/locations/[^/]+/")
|
||||
_VERTEX_RESOURCE_PROJECT_AND_LOCATION: Final = re.compile(r"projects/[^/]+/locations/(?P<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:
|
||||
|
|
@ -4000,8 +4009,8 @@ def _build_vertex_live_setup_model_rewriter(
|
|||
|
||||
Clients address the gateway the way they address LiteLLM (bare id or model alias); Vertex reads anything
|
||||
that is not a ``projects/...`` path as a project name and closes the socket. A ``projects/...`` path
|
||||
picks the project the proxy's token is spent on, so it is moved onto the connection's project and region
|
||||
unless ``allow_client_side_credentials`` lets callers choose them
|
||||
picks the project the proxy's token is spent on, so it is moved onto the connection's project, and onto its
|
||||
region when one is configured, unless ``allow_client_side_credentials`` lets callers choose them
|
||||
"""
|
||||
if vertex_project is None or vertex_location is None:
|
||||
return None
|
||||
|
|
@ -4011,7 +4020,8 @@ def _build_vertex_live_setup_model_rewriter(
|
|||
resource_prefix: Final = _VERTEX_RESOURCE_PROJECT_AND_LOCATION.match(setup_model)
|
||||
if keeps_caller_resource_path or resource_prefix is None:
|
||||
return setup_model
|
||||
return f"projects/{vertex_project}/locations/{vertex_location}/{setup_model[resource_prefix.end() :]}"
|
||||
location: Final = configured_location or resource_prefix.group("location")
|
||||
return f"projects/{vertex_project}/locations/{location}/{setup_model[resource_prefix.end() :]}"
|
||||
aliased: Final = _resolve_alias_to_upstream_model(setup_model, llm_router)
|
||||
return f"projects/{vertex_project}/locations/{vertex_location}/{_vertex_publisher_model_suffix(aliased)}"
|
||||
|
||||
|
|
@ -4155,6 +4165,7 @@ async def vertex_ai_live_websocket_passthrough(
|
|||
setup_model_rewriter=_build_vertex_live_setup_model_rewriter(
|
||||
vertex_project=resolved_project,
|
||||
vertex_location=resolved_location,
|
||||
configured_location=configured_location,
|
||||
llm_router=_get_llm_router(),
|
||||
keeps_caller_resource_path=allow_client_side_credentials,
|
||||
),
|
||||
|
|
|
|||
|
|
@ -2426,6 +2426,21 @@ def _rewrite_vertex_live_setup_model(text_data: str, setup_model_rewriter: Calla
|
|||
return json.dumps({**message, "setup": {**setup, "model": rewritten_model}})
|
||||
|
||||
|
||||
VERTEX_LIVE_UNDECODABLE_BINARY_CLOSE_REASON: Final = "Rejected Request: binary messages must be UTF-8 encoded JSON"
|
||||
|
||||
|
||||
def _rewrite_vertex_live_binary_setup_model(
|
||||
bytes_data: bytes, setup_model_rewriter: Callable[[str], str] | None
|
||||
) -> bytes | None:
|
||||
if setup_model_rewriter is None:
|
||||
return bytes_data
|
||||
try:
|
||||
text_data: Final = bytes_data.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
return None
|
||||
return _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter).encode("utf-8")
|
||||
|
||||
|
||||
def _resolved_vertex_live_setup(
|
||||
setup_data: Mapping[str, object], setup_model_rewriter: Callable[[str], str] | None
|
||||
) -> Mapping[str, object]:
|
||||
|
|
@ -2726,7 +2741,12 @@ async def websocket_passthrough_request(
|
|||
|
||||
await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter))
|
||||
elif bytes_data is not None:
|
||||
await upstream_ws.send(bytes_data)
|
||||
rewritten_bytes = _rewrite_vertex_live_binary_setup_model(bytes_data, setup_model_rewriter)
|
||||
if rewritten_bytes is None:
|
||||
await websocket.close(code=1008, reason=VERTEX_LIVE_UNDECODABLE_BINARY_CLOSE_REASON)
|
||||
await upstream_ws.close()
|
||||
break
|
||||
await upstream_ws.send(rewritten_bytes)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ from starlette.datastructures import UploadFile
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
||||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH, LITELLM_PROXY_ADMIN_NAME
|
||||
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
|
||||
LiteLLM_ManagedVectorStore,
|
||||
)
|
||||
|
|
@ -64,6 +64,7 @@ from litellm.proxy.vector_store_endpoints.utils import (
|
|||
from litellm.rag.ingestion.vertex_ai_ingestion import DEFAULT_VERTEX_RAG_LOCATION
|
||||
from litellm.rag.main import get_ingestion_class
|
||||
from litellm.repositories.table_repositories import ManagedVectorStoresRepository
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
|
@ -183,41 +184,95 @@ def _ingest_provider_error(vector_store_config: Mapping[str, object]) -> str | N
|
|||
return None
|
||||
|
||||
|
||||
_VERTEX_RAG_CORPUS_NAME: Final = re.compile(r"projects/(?P<project>[^/]+)/locations/(?P<location>[^/]+)/ragCorpora/")
|
||||
_VERTEX_RAG_CORPUS_NAME: Final = re.compile(
|
||||
r"projects/(?P<project>[^/%\\]+)/locations/(?P<location>[^/%\\]+)/ragCorpora/(?P<corpus>[^/%\\]+)"
|
||||
)
|
||||
_DOT_SEGMENTS: Final = frozenset({".", ".."})
|
||||
|
||||
|
||||
def _vertex_target_error(
|
||||
vector_store_config: Mapping[str, object],
|
||||
managed_store: LiteLLM_ManagedVectorStore | None,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
general_settings: Mapping[str, object],
|
||||
) -> str | None:
|
||||
if (
|
||||
managed_store is not None
|
||||
or vector_store_config.get("custom_llm_provider") != "vertex_ai"
|
||||
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
or general_settings.get("allow_client_side_credentials") is True
|
||||
):
|
||||
return None
|
||||
configured_project: Final = litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT")
|
||||
configured_location: Final = VertexBase.safe_get_vertex_ai_location({}) or DEFAULT_VERTEX_RAG_LOCATION
|
||||
corpus_name: Final = vector_store_config.get("vector_store_id")
|
||||
corpus: Final = _VERTEX_RAG_CORPUS_NAME.match(corpus_name) if isinstance(corpus_name, str) else None
|
||||
requested: Final = (
|
||||
("vertex_project", vector_store_config.get("vertex_project"), configured_project),
|
||||
("vertex_location", vector_store_config.get("vertex_location"), configured_location),
|
||||
("vector_store_id", corpus["project"] if corpus is not None else None, configured_project),
|
||||
("vector_store_id", corpus["location"] if corpus is not None else None, configured_location),
|
||||
)
|
||||
field: Final = next((name for name, asked, configured in requested if asked and asked != configured), None)
|
||||
if field is None:
|
||||
def _clean_vertex_corpus_name(corpus_name: object) -> re.Match[str] | None:
|
||||
if not isinstance(corpus_name, str) or not _DOT_SEGMENTS.isdisjoint(corpus_name.split("/")):
|
||||
return None
|
||||
return _VERTEX_RAG_CORPUS_NAME.fullmatch(corpus_name)
|
||||
|
||||
|
||||
def _vertex_target_message(field: str) -> str:
|
||||
return (
|
||||
f"Rejected Request: {field} picks a Vertex project or region other than the proxy's. "
|
||||
"Use a vector store a proxy admin created, or set general_settings.allow_client_side_credentials"
|
||||
)
|
||||
|
||||
|
||||
def _vertex_target_is_checked(
|
||||
vector_store_config: Mapping[str, object],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
general_settings: Mapping[str, object],
|
||||
) -> bool:
|
||||
return (
|
||||
vector_store_config.get("custom_llm_provider") == "vertex_ai"
|
||||
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN
|
||||
and general_settings.get("allow_client_side_credentials") is not True
|
||||
)
|
||||
|
||||
|
||||
def _configured_vertex_project() -> str | None:
|
||||
return litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT")
|
||||
|
||||
|
||||
def _corpus_on_configured_project(vector_store_config: Mapping[str, object]) -> Mapping[str, object]:
|
||||
corpus: Final = _clean_vertex_corpus_name(vector_store_config.get("vector_store_id"))
|
||||
configured_project: Final = _configured_vertex_project()
|
||||
if corpus is None or configured_project is None or not corpus["project"].isdigit():
|
||||
return vector_store_config
|
||||
return MappingProxyType(
|
||||
{
|
||||
**vector_store_config,
|
||||
"vector_store_id": (
|
||||
f"projects/{configured_project}/locations/{corpus['location']}/ragCorpora/{corpus['corpus']}"
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _vertex_target_error(vector_store_config: Mapping[str, object]) -> str | None:
|
||||
corpus_name: Final = vector_store_config.get("vector_store_id")
|
||||
corpus: Final = _clean_vertex_corpus_name(corpus_name)
|
||||
if corpus_name and corpus is None:
|
||||
return _vertex_target_message("vector_store_id")
|
||||
configured_project: Final = _configured_vertex_project()
|
||||
configured_location: Final = VertexBase.safe_get_vertex_ai_location({}) or DEFAULT_VERTEX_RAG_LOCATION
|
||||
requested: Final = (
|
||||
*(
|
||||
(name, vector_store_config.get(name), configured_project)
|
||||
for name in ("vertex_project", "vertex_ai_project")
|
||||
),
|
||||
*(
|
||||
(name, vector_store_config.get(name), configured_location)
|
||||
for name in ("vertex_location", "vertex_ai_location")
|
||||
),
|
||||
("vector_store_id", corpus["project"] if corpus is not None else None, configured_project),
|
||||
("vector_store_id", corpus["location"] if corpus is not None else None, configured_location),
|
||||
)
|
||||
field: Final = next((name for name, asked, configured in requested if asked and asked != configured), None)
|
||||
return None if field is None else _vertex_target_message(field)
|
||||
|
||||
|
||||
async def _admin_set_store_target(
|
||||
managed_store: LiteLLM_ManagedVectorStore | None, prisma_client: "PrismaClient | None"
|
||||
) -> bool:
|
||||
if managed_store is None:
|
||||
return False
|
||||
if managed_store.get("is_config", False):
|
||||
return True
|
||||
creator: Final = managed_store.get("user_id")
|
||||
if creator == LITELLM_PROXY_ADMIN_NAME:
|
||||
return True
|
||||
if creator is None or prisma_client is None:
|
||||
return False
|
||||
creator_row: Final = await UserRepository(prisma_client).find_by_id(creator)
|
||||
return creator_row is not None and creator_row.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
|
||||
|
||||
_MANAGED_STORE_CALLER_OPTIONS: Final = frozenset(
|
||||
{
|
||||
"vector_store_id",
|
||||
|
|
@ -730,15 +785,17 @@ async def rag_ingest(
|
|||
raise HTTPException(status_code=400, detail={"error": str(e)})
|
||||
|
||||
managed_store: Final = resolved_stores.get(request_vector_store_config.get("vector_store_id"))
|
||||
vertex_target_error: Final = _vertex_target_error(
|
||||
request_vector_store_config, managed_store, user_api_key_dict, general_settings_view()
|
||||
)
|
||||
if vertex_target_error is not None:
|
||||
raise HTTPException(status_code=400, detail={"error": vertex_target_error})
|
||||
base: Final = {
|
||||
requested_base: Final = {
|
||||
**_caller_vector_store_options(request_vector_store_config, managed_store),
|
||||
**_managed_store_overrides(managed_store),
|
||||
}
|
||||
vertex_target_checked: Final = _vertex_target_is_checked(
|
||||
requested_base, user_api_key_dict, general_settings_view()
|
||||
)
|
||||
base: Final = _corpus_on_configured_project(requested_base) if vertex_target_checked else requested_base
|
||||
vertex_target_error: Final = _vertex_target_error(base) if vertex_target_checked else None
|
||||
if vertex_target_error is not None and not await _admin_set_store_target(managed_store, prisma_client):
|
||||
raise HTTPException(status_code=400, detail={"error": vertex_target_error})
|
||||
request_data: Final[dict[str, Any]] = await add_litellm_data_to_request(
|
||||
data={},
|
||||
request=request,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -5136,7 +5136,11 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides:
|
|||
is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions")
|
||||
|
||||
@staticmethod
|
||||
def _opted_in_and_pinned_router(router_fallbacks: list[dict[str, list[str]]]) -> Router:
|
||||
def _opted_in_and_pinned_router(
|
||||
router_fallbacks: list[dict[str, list[str]]],
|
||||
fallback_setting: str = "fallbacks",
|
||||
model_group_alias: dict[str, str] | None = None,
|
||||
) -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
|
|
@ -5153,7 +5157,8 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides:
|
|||
("gemini-pinned", []),
|
||||
)
|
||||
],
|
||||
fallbacks=router_fallbacks,
|
||||
model_group_alias=model_group_alias or {},
|
||||
**{fallback_setting: router_fallbacks},
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -5175,13 +5180,16 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides:
|
|||
with pytest.raises(ValueError, match="vertex_location"):
|
||||
is_request_body_safe(request_body, {}, router, "gemini-optin", route="/v1/chat/completions")
|
||||
|
||||
@pytest.mark.parametrize("fallback_setting", ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"])
|
||||
@pytest.mark.parametrize(
|
||||
"fallback_group, is_allowed",
|
||||
[("gemini-pinned", False), ("gemini-optin-other", True)],
|
||||
ids=["pinned", "opted_in"],
|
||||
)
|
||||
def test_router_fallback_needs_the_same_opt_in(self, fallback_group: str, is_allowed: bool) -> None:
|
||||
router: Final = self._opted_in_and_pinned_router([{"gemini-optin": [fallback_group]}])
|
||||
def test_router_fallback_needs_the_same_opt_in(
|
||||
self, fallback_setting: str, fallback_group: str, is_allowed: bool
|
||||
) -> None:
|
||||
router: Final = self._opted_in_and_pinned_router([{"gemini-optin": [fallback_group]}], fallback_setting)
|
||||
request_body: Final = {"model": "gemini-optin", "vertex_location": "us-central1"}
|
||||
if is_allowed:
|
||||
assert is_request_body_safe(request_body, {}, router, "gemini-optin", route="/v1/chat/completions") is True
|
||||
|
|
@ -5189,6 +5197,67 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides:
|
|||
with pytest.raises(ValueError, match="vertex_location"):
|
||||
is_request_body_safe(request_body, {}, router, "gemini-optin", route="/v1/chat/completions")
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model, fields, is_allowed",
|
||||
[
|
||||
("gemini-optin", {"vertex_location": "us-central1"}, True),
|
||||
("openai-group", {"vertex_location": "us-central1"}, True),
|
||||
("gemini-optin", {"api_base": "https://eu.example.com"}, False),
|
||||
("studio-group", {"vertex_location": "us-central1"}, False),
|
||||
],
|
||||
ids=["vertex_field_to_openai_fallback", "vertex_field_on_openai", "api_base_to_openai_fallback", "gemini"],
|
||||
)
|
||||
def test_only_google_groups_need_the_vertex_opt_in(
|
||||
self, model: str, fields: dict[str, object], is_allowed: bool
|
||||
) -> None:
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gemini-optin",
|
||||
"litellm_params": {
|
||||
"model": "vertex_ai/gemini-2.5-flash",
|
||||
"configurable_clientside_auth_params": ["vertex_location", {"api_base": r"https://eu\..*"}],
|
||||
},
|
||||
},
|
||||
{"model_name": "openai-group", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}},
|
||||
{"model_name": "studio-group", "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "k"}},
|
||||
],
|
||||
fallbacks=[{"gemini-optin": ["openai-group"]}, {"studio-group": ["openai-group"]}],
|
||||
)
|
||||
request_body: Final = {"model": model, **fields}
|
||||
if is_allowed:
|
||||
assert is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions") is True
|
||||
return
|
||||
with pytest.raises(ValueError, match=r"vertex_location|api_base"):
|
||||
is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions")
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model, is_allowed", [("public-alias", True), ("gemini-optin", False)], ids=["alias", "aliased_target"]
|
||||
)
|
||||
def test_alias_chain_is_checked_one_hop_like_the_router_serves_it(self, model: str, is_allowed: bool) -> None:
|
||||
router: Final = self._opted_in_and_pinned_router(
|
||||
[], model_group_alias={"public-alias": "gemini-optin", "gemini-optin": "gemini-pinned"}
|
||||
)
|
||||
request_body: Final = {"model": model, "vertex_location": "us-central1"}
|
||||
if is_allowed:
|
||||
assert is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions") is True
|
||||
return
|
||||
with pytest.raises(ValueError, match="vertex_location"):
|
||||
is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions")
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"route",
|
||||
[
|
||||
"/openai/deployments/gemini-pinned/chat/completions",
|
||||
"/v1beta/models/gemini-pinned:generateContent",
|
||||
],
|
||||
)
|
||||
def test_model_named_in_the_route_needs_the_opt_in_too(self, route: str) -> None:
|
||||
router: Final = self._opted_in_and_pinned_router([])
|
||||
request_body: Final = {"model": "gemini-optin", "vertex_location": "us-central1"}
|
||||
with pytest.raises(ValueError, match="vertex_location"):
|
||||
is_request_body_safe(request_body, {}, router, "gemini-optin", route=route)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("opted_in", [True, False])
|
||||
async def test_opt_in_of_the_model_named_in_a_google_route_decides(self, monkeypatch, opted_in: bool) -> None:
|
||||
|
|
@ -5208,8 +5277,8 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides:
|
|||
"route, kept_params",
|
||||
[
|
||||
("/vertex_ai/live", _VERTEX_TARGET_PARAMS),
|
||||
("/rag/ingest", ("vertex_project", "vertex_location")),
|
||||
("/v1/rag/ingest", ("vertex_project", "vertex_location")),
|
||||
("/rag/ingest", _VERTEX_TARGET_PARAMS),
|
||||
("/v1/rag/ingest", _VERTEX_TARGET_PARAMS),
|
||||
],
|
||||
)
|
||||
def test_routes_that_name_their_vertex_target_keep_it(self, route: str, kept_params: tuple[str, ...]) -> None:
|
||||
|
|
@ -5229,8 +5298,7 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides:
|
|||
[
|
||||
("/vertex_ai/live", "api_base"),
|
||||
("/rag/ingest", "api_base"),
|
||||
("/rag/ingest", "vertex_ai_project"),
|
||||
("/rag/ingest", "vertex_ai_location"),
|
||||
("/v1/rag/ingest", "api_base"),
|
||||
],
|
||||
)
|
||||
def test_routes_that_name_their_vertex_target_reject_every_other_banned_param(self, route: str, param: str) -> None:
|
||||
|
|
|
|||
|
|
@ -6338,6 +6338,7 @@ class TestVertexAILiveWebsocketPassthrough:
|
|||
rewriter = passthrough_module._build_vertex_live_setup_model_rewriter(
|
||||
vertex_project="proj-db",
|
||||
vertex_location="global",
|
||||
configured_location="global",
|
||||
llm_router=llm_router,
|
||||
keeps_caller_resource_path=False,
|
||||
)
|
||||
|
|
@ -6706,6 +6707,103 @@ class TestVertexAILiveWebsocketPassthrough:
|
|||
== expected
|
||||
)
|
||||
|
||||
def _pin_no_vertex_project(
|
||||
self,
|
||||
monkeypatch,
|
||||
general_settings: Mapping[str, object],
|
||||
default_vertex_config: VertexPassThroughCredentials | None,
|
||||
):
|
||||
from litellm.proxy.pass_through_endpoints import (
|
||||
llm_passthrough_endpoints as passthrough_module,
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", dict(general_settings))
|
||||
monkeypatch.setattr(
|
||||
passthrough_module.passthrough_endpoint_router, "default_vertex_config", default_vertex_config
|
||||
)
|
||||
monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "deployment_key_to_vertex_credentials", {})
|
||||
self._clear_vertex_env(monkeypatch)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"default_vertex_config",
|
||||
[None, VertexPassThroughCredentials(vertex_location="europe-west4")],
|
||||
ids=["no_default_vertex_config", "default_vertex_config_without_a_project"],
|
||||
)
|
||||
async def test_caller_named_project_closes_before_any_upstream_connect_when_nothing_pins_one(
|
||||
self, monkeypatch, default_vertex_config: VertexPassThroughCredentials | None
|
||||
):
|
||||
self._pin_no_vertex_project(monkeypatch, {}, default_vertex_config)
|
||||
|
||||
websocket, ensure_token, ws_passthrough = await self._open_live({"vertex_project": "other-project"})
|
||||
|
||||
ws_passthrough.assert_not_awaited()
|
||||
ensure_token.assert_not_awaited()
|
||||
close_kwargs = websocket.close.await_args.kwargs
|
||||
assert close_kwargs["code"] == 1008
|
||||
assert "vertex_project" in close_kwargs["reason"]
|
||||
assert len(close_kwargs["reason"].encode("utf-8")) <= 123
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"general_settings, query, expected_project, expected_location",
|
||||
[
|
||||
({}, {"vertex_location": "europe-west4"}, "proj-env", "europe-west4"),
|
||||
(
|
||||
{"allow_client_side_credentials": True},
|
||||
{"vertex_project": "other-project", "vertex_location": "europe-west4"},
|
||||
"other-project",
|
||||
"europe-west4",
|
||||
),
|
||||
],
|
||||
ids=["caller_region", "proxy_wide_opt_in"],
|
||||
)
|
||||
async def test_with_nothing_pinned_the_caller_keeps_the_region_and_the_opt_in_keeps_the_project(
|
||||
self,
|
||||
monkeypatch,
|
||||
general_settings: Mapping[str, object],
|
||||
query: Mapping[str, str],
|
||||
expected_project: str,
|
||||
expected_location: str,
|
||||
):
|
||||
self._pin_no_vertex_project(monkeypatch, general_settings, None)
|
||||
|
||||
websocket, _, ws_passthrough = await self._open_live(query)
|
||||
|
||||
websocket.close.assert_not_awaited()
|
||||
passthrough_kwargs = ws_passthrough.await_args.kwargs
|
||||
assert passthrough_kwargs["custom_headers"]["x-goog-user-project"] == expected_project
|
||||
assert passthrough_kwargs["setup_model_rewriter"]("gemini-2.0-flash") == (
|
||||
f"projects/{expected_project}/locations/{expected_location}/publishers/google/models/gemini-2.0-flash"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"setup_model, expected",
|
||||
[
|
||||
(
|
||||
"projects/proj-env/locations/europe-west4/publishers/google/models/gemini-2.0-flash",
|
||||
"projects/proj-env/locations/europe-west4/publishers/google/models/gemini-2.0-flash",
|
||||
),
|
||||
(
|
||||
"projects/caller-project/locations/europe-west4/publishers/google/models/gemini-2.0-flash",
|
||||
"projects/proj-env/locations/europe-west4/publishers/google/models/gemini-2.0-flash",
|
||||
),
|
||||
],
|
||||
ids=["own_project", "other_project"],
|
||||
)
|
||||
async def test_setup_frame_keeps_its_region_when_no_region_is_pinned(
|
||||
self, monkeypatch, setup_model: str, expected: str
|
||||
):
|
||||
self._use_default_vertex_config(monkeypatch, {}, vertex_location=None)
|
||||
|
||||
websocket, _, ws_passthrough = await self._open_live({})
|
||||
|
||||
websocket.close.assert_not_awaited()
|
||||
rewriter = ws_passthrough.await_args.kwargs["setup_model_rewriter"]
|
||||
assert rewriter(setup_model) == expected
|
||||
|
||||
|
||||
class TestAnthropicProxyRoute:
|
||||
"""The /anthropic passthrough route: custom auth headers must not clobber the
|
||||
|
|
|
|||
|
|
@ -7396,6 +7396,7 @@ async def _run_setup_rewrite_passthrough(setup_model: str, llm_router, keeps_cal
|
|||
setup_model_rewriter=_build_vertex_live_setup_model_rewriter(
|
||||
vertex_project="proj-db",
|
||||
vertex_location="global",
|
||||
configured_location="global",
|
||||
llm_router=llm_router,
|
||||
keeps_caller_resource_path=keeps_caller_resource_path,
|
||||
),
|
||||
|
|
@ -7435,6 +7436,67 @@ async def test_websocket_passthrough_rewrites_gateway_alias_setup_model():
|
|||
assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"
|
||||
|
||||
|
||||
async def _run_binary_client_frame_passthrough(frame: bytes):
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
_build_vertex_live_setup_model_rewriter,
|
||||
)
|
||||
|
||||
upstream_ws = RecordingUpstreamWebSocket()
|
||||
websocket = _client_websocket(
|
||||
AsyncMock(
|
||||
side_effect=[
|
||||
{"type": "websocket.receive", "bytes": frame},
|
||||
{"type": "websocket.disconnect"},
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
with _patched_websocket_passthrough_environment(upstream_ws):
|
||||
await websocket_passthrough_request(
|
||||
websocket=websocket,
|
||||
target="wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent",
|
||||
custom_headers={"Authorization": "Bearer token"},
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
forward_headers=False,
|
||||
endpoint="/vertex_ai/live",
|
||||
accept_websocket=False,
|
||||
setup_model_rewriter=_build_vertex_live_setup_model_rewriter(
|
||||
vertex_project="proj-a",
|
||||
vertex_location="europe-west4",
|
||||
configured_location="europe-west4",
|
||||
llm_router=None,
|
||||
keeps_caller_resource_path=False,
|
||||
),
|
||||
)
|
||||
|
||||
return websocket, upstream_ws
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_moves_a_binary_setup_frame_onto_the_pinned_project():
|
||||
setup_frame = json.dumps(
|
||||
{"setup": {"model": "projects/proj-b/locations/europe-west4/publishers/google/models/gemini-2.0-flash"}}
|
||||
).encode("utf-8")
|
||||
|
||||
_, upstream_ws = await _run_binary_client_frame_passthrough(setup_frame)
|
||||
|
||||
upstream_ws.send.assert_awaited_once()
|
||||
sent_frame = upstream_ws.send.await_args.args[0]
|
||||
assert isinstance(sent_frame, bytes), "a binary client message must stay binary upstream"
|
||||
assert json.loads(sent_frame)["setup"]["model"] == (
|
||||
"projects/proj-a/locations/europe-west4/publishers/google/models/gemini-2.0-flash"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_closes_on_a_binary_frame_that_is_not_utf8():
|
||||
websocket, upstream_ws = await _run_binary_client_frame_passthrough(b"\xff\xfe\x00projects/proj-b")
|
||||
|
||||
upstream_ws.send.assert_not_awaited()
|
||||
websocket.close.assert_awaited_once()
|
||||
assert websocket.close.await_args.kwargs["code"] == 1008
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"setup_model",
|
||||
["gemini-live-2.5-flash", "models/gemini-live-2.5-flash", "publishers/google/models/gemini-live-2.5-flash"],
|
||||
|
|
@ -7456,7 +7518,11 @@ def test_vertex_live_setup_model_resolves_before_extraction(setup_model):
|
|||
)
|
||||
|
||||
rewriter = _build_vertex_live_setup_model_rewriter(
|
||||
vertex_project="proj-db", vertex_location="global", llm_router=None, keeps_caller_resource_path=False
|
||||
vertex_project="proj-db",
|
||||
vertex_location="global",
|
||||
configured_location="global",
|
||||
llm_router=None,
|
||||
keeps_caller_resource_path=False,
|
||||
)
|
||||
setup_data = {"model": setup_model}
|
||||
|
||||
|
|
@ -7508,6 +7574,7 @@ async def test_websocket_passthrough_logs_a_bare_alias_setup_model():
|
|||
setup_model_rewriter=_build_vertex_live_setup_model_rewriter(
|
||||
vertex_project="proj-db",
|
||||
vertex_location="global",
|
||||
configured_location="global",
|
||||
llm_router=None,
|
||||
keeps_caller_resource_path=False,
|
||||
),
|
||||
|
|
|
|||
|
|
@ -15,7 +15,8 @@ from fastapi.testclient import TestClient
|
|||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.constants import LITELLM_PROXY_ADMIN_NAME
|
||||
from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.proxy_server import app
|
||||
from litellm.types.vector_stores import LiteLLM_ManagedVectorStore
|
||||
|
|
@ -284,6 +285,7 @@ VERTEX_REGISTRY_STORE = {
|
|||
"vector_store_id": "projects/registry-project/locations/us-central1/ragCorpora/42",
|
||||
"custom_llm_provider": "vertex_ai",
|
||||
"litellm_params": {"vertex_project": "registry-project", "vertex_location": "us-central1"},
|
||||
"is_config": True,
|
||||
}
|
||||
UNSUPPORTED_INGEST_PROVIDER_ERROR = (
|
||||
"Provider '{provider}' is not supported for RAG ingestion. "
|
||||
|
|
@ -1028,8 +1030,20 @@ def _vertex_ingest(client: TestClient, vector_store: Mapping[str, str]) -> tuple
|
|||
"vector_store_id",
|
||||
),
|
||||
({"vector_store_id": f"projects/{_PROXY_VERTEX_PROJECT}/locations/us-east5/ragCorpora/7"}, "vector_store_id"),
|
||||
({"vertex_ai_project": "other-project"}, "vertex_ai_project"),
|
||||
({"vertex_ai_location": "us-east5"}, "vertex_ai_location"),
|
||||
({"vector_store_id": "projects/123456789/locations/us-east5/ragCorpora/7"}, "vector_store_id"),
|
||||
],
|
||||
ids=[
|
||||
"project_field",
|
||||
"location_field",
|
||||
"corpus_project",
|
||||
"corpus_project_behind_matching_field",
|
||||
"corpus_region",
|
||||
"project_alias_field",
|
||||
"location_alias_field",
|
||||
"corpus_region_by_project_number",
|
||||
],
|
||||
ids=["project_field", "location_field", "corpus_project", "corpus_project_behind_matching_field", "corpus_region"],
|
||||
)
|
||||
@pytest.mark.usefixtures("proxy_vertex_target")
|
||||
def test_rag_ingest_refuses_a_key_pointing_the_proxys_vertex_credentials_elsewhere(
|
||||
|
|
@ -1060,8 +1074,9 @@ def test_rag_ingest_refuses_a_key_naming_a_vertex_project_when_the_proxy_names_n
|
|||
{},
|
||||
{"vertex_project": _PROXY_VERTEX_PROJECT, "vertex_location": _PROXY_VERTEX_LOCATION},
|
||||
{"vector_store_id": f"projects/{_PROXY_VERTEX_PROJECT}/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/7"},
|
||||
{"vertex_ai_project": _PROXY_VERTEX_PROJECT, "vertex_ai_location": _PROXY_VERTEX_LOCATION},
|
||||
],
|
||||
ids=["proxy_defaults", "same_fields", "corpus_in_the_proxy_project"],
|
||||
ids=["proxy_defaults", "same_fields", "corpus_in_the_proxy_project", "same_alias_fields"],
|
||||
)
|
||||
@pytest.mark.usefixtures("proxy_vertex_target")
|
||||
def test_rag_ingest_lets_a_key_use_the_proxys_own_vertex_target(client_internal_user, vector_store: Mapping[str, str]):
|
||||
|
|
@ -1074,6 +1089,11 @@ def test_rag_ingest_lets_a_key_use_the_proxys_own_vertex_target(client_internal_
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"project_field, location_field",
|
||||
[("vertex_project", "vertex_location"), ("vertex_ai_project", "vertex_ai_location")],
|
||||
ids=["fields", "alias_fields"],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"user_role, general_settings",
|
||||
[
|
||||
|
|
@ -1084,14 +1104,14 @@ def test_rag_ingest_lets_a_key_use_the_proxys_own_vertex_target(client_internal_
|
|||
)
|
||||
@pytest.mark.usefixtures("proxy_vertex_target")
|
||||
def test_rag_ingest_lets_the_admin_choose_another_vertex_target(
|
||||
user_role: LitellmUserRoles, general_settings: Mapping[str, object]
|
||||
user_role: LitellmUserRoles, general_settings: Mapping[str, object], project_field: str, location_field: str
|
||||
):
|
||||
caller: Final = UserAPIKeyAuth(user_id="vertex-caller", user_role=user_role)
|
||||
original_overrides: Final = app.dependency_overrides.copy()
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: caller
|
||||
vector_store: Final = {
|
||||
"vertex_project": "other-project",
|
||||
"vertex_location": "us-east5",
|
||||
project_field: "other-project",
|
||||
location_field: "us-east5",
|
||||
"vector_store_id": "projects/other-project/locations/us-east5/ragCorpora/7",
|
||||
}
|
||||
try:
|
||||
|
|
@ -1107,6 +1127,119 @@ def test_rag_ingest_lets_the_admin_choose_another_vertex_target(
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("proxy_vertex_target")
|
||||
def test_rag_ingest_reads_a_corpus_named_by_project_number_on_the_proxys_own_project(client_internal_user):
|
||||
response, mock_aingest = _vertex_ingest(
|
||||
client_internal_user, {"vector_store_id": f"projects/123456789/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/7"}
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.json()
|
||||
assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"]["vector_store_id"] == (
|
||||
f"projects/{_PROXY_VERTEX_PROJECT}/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/7"
|
||||
)
|
||||
|
||||
|
||||
_PROXY_CORPUS_PREFIX: Final = f"projects/{_PROXY_VERTEX_PROJECT}/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora"
|
||||
_FOREIGN_CORPUS: Final = f"projects/other-project/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/1"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"vector_store_id",
|
||||
[
|
||||
f"./{_FOREIGN_CORPUS}",
|
||||
f"x/../{_FOREIGN_CORPUS}",
|
||||
f"{_PROXY_CORPUS_PREFIX}/1{'/..' * 6}/{_FOREIGN_CORPUS}",
|
||||
f"{_PROXY_CORPUS_PREFIX}/..",
|
||||
f"{_PROXY_CORPUS_PREFIX}/1%2F..%2F..",
|
||||
f"{_PROXY_CORPUS_PREFIX}/1\\..\\..",
|
||||
],
|
||||
ids=["dot_prefix", "dot_dot_prefix", "own_corpus_then_dot_dots", "dot_dot_corpus", "percent", "backslash"],
|
||||
)
|
||||
@pytest.mark.usefixtures("proxy_vertex_target")
|
||||
def test_rag_ingest_refuses_a_vertex_corpus_name_with_path_tricks(client_internal_user, vector_store_id: str):
|
||||
response, mock_aingest = _vertex_ingest(client_internal_user, {"vector_store_id": vector_store_id})
|
||||
|
||||
assert response.status_code == 400, response.json()
|
||||
assert response.json()["detail"]["error"].startswith("Rejected Request: vector_store_id picks a Vertex project")
|
||||
mock_aingest.assert_not_awaited()
|
||||
|
||||
|
||||
def _foreign_vertex_store(created_by: str | None) -> LiteLLM_ManagedVectorStore:
|
||||
return LiteLLM_ManagedVectorStore(
|
||||
vector_store_id=_FOREIGN_CORPUS, custom_llm_provider="vertex_ai", litellm_params={}, user_id=created_by
|
||||
)
|
||||
|
||||
|
||||
def _prisma_with_user_row(user_row: LiteLLM_UserTable | None) -> MagicMock:
|
||||
prisma_client: Final = MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row)
|
||||
prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(
|
||||
return_value=MagicMock(vector_store_metadata={})
|
||||
)
|
||||
prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock()
|
||||
return prisma_client
|
||||
|
||||
|
||||
def _ingest_into_saved_store(
|
||||
store: LiteLLM_ManagedVectorStore, creator_row: LiteLLM_UserTable | None
|
||||
) -> tuple[httpx.Response, AsyncMock]:
|
||||
caller: Final = UserAPIKeyAuth(user_id="vertex-caller", user_role=LitellmUserRoles.INTERNAL_USER)
|
||||
original_overrides: Final = app.dependency_overrides.copy()
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: caller
|
||||
aingest_patch, registry_patch = _patched_ingest_boundary(
|
||||
store, {"vector_store_id": store["vector_store_id"], "file_id": "file_123"}
|
||||
)
|
||||
try:
|
||||
with aingest_patch as mock_aingest, registry_patch, _patched_prisma_client(_prisma_with_user_row(creator_row)):
|
||||
response = TestClient(app).post(
|
||||
"/v1/rag/ingest", **_ingest_form({"vector_store_id": store["vector_store_id"]})
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides = original_overrides
|
||||
return response, mock_aingest
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"creator_row",
|
||||
[LiteLLM_UserTable(user_id="vertex-caller", user_role=LitellmUserRoles.INTERNAL_USER.value), None],
|
||||
ids=["saved_by_an_internal_user", "saved_by_a_key_with_no_user"],
|
||||
)
|
||||
@pytest.mark.usefixtures("proxy_vertex_target")
|
||||
def test_rag_ingest_refuses_a_saved_vertex_store_a_non_admin_pointed_elsewhere(creator_row: LiteLLM_UserTable | None):
|
||||
store: Final = _foreign_vertex_store(creator_row.user_id if creator_row is not None else None)
|
||||
|
||||
response, mock_aingest = _ingest_into_saved_store(store, creator_row)
|
||||
|
||||
assert response.status_code == 400, response.json()
|
||||
assert response.json()["detail"]["error"].startswith("Rejected Request: vector_store_id picks a Vertex project")
|
||||
mock_aingest.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"store, creator_row",
|
||||
[
|
||||
(_foreign_vertex_store(LITELLM_PROXY_ADMIN_NAME), None),
|
||||
(
|
||||
_foreign_vertex_store("admin-user"),
|
||||
LiteLLM_UserTable(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN.value),
|
||||
),
|
||||
(LiteLLM_ManagedVectorStore(**_foreign_vertex_store(None), is_config=True), None),
|
||||
],
|
||||
ids=["master_key_created", "admin_user_created", "config_store"],
|
||||
)
|
||||
@pytest.mark.usefixtures("proxy_vertex_target")
|
||||
def test_rag_ingest_lets_a_key_use_a_vertex_store_an_admin_pointed_elsewhere(
|
||||
store: LiteLLM_ManagedVectorStore, creator_row: LiteLLM_UserTable | None
|
||||
):
|
||||
response, mock_aingest = _ingest_into_saved_store(store, creator_row)
|
||||
|
||||
assert response.status_code == 200, response.json()
|
||||
assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] == {
|
||||
"vector_store_id": _FOREIGN_CORPUS,
|
||||
"custom_llm_provider": "vertex_ai",
|
||||
}
|
||||
|
||||
|
||||
def test_rag_ingest_db_managed_store_drops_the_callers_credential_name(client_internal_user):
|
||||
aingest_patch, registry_patch = _patched_ingest_boundary(
|
||||
DB_MANAGED_STORE, {"vector_store_id": "db-store", "file_id": "file_123"}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue