mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge ddb4c37fff into dfabf32cd5
This commit is contained in:
commit
42dc70485d
14 changed files with 1779 additions and 106 deletions
|
|
@ -2,8 +2,9 @@ import importlib.util
|
|||
import os
|
||||
import re
|
||||
import sys
|
||||
from collections.abc import Collection, Iterator, Mapping
|
||||
from functools import lru_cache
|
||||
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
|
||||
|
|
@ -23,6 +24,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,
|
||||
|
|
@ -41,11 +43,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
|
||||
|
||||
|
||||
|
|
@ -211,35 +215,63 @@ 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"]
|
||||
|
||||
|
||||
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:
|
||||
"""
|
||||
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)
|
||||
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
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -318,6 +350,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
|
||||
|
|
@ -410,6 +456,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
|
||||
|
|
@ -478,15 +528,85 @@ _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 _opt_in_model_groups(
|
||||
model: str, also_requested_models: tuple[str, ...], llm_router: Router | 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)
|
||||
)
|
||||
return routed | frozenset(registry.resolve(name, kind="alias") or name for name in fallbacks)
|
||||
|
||||
|
||||
_MAX_OPT_IN_FALLBACK_MODELS: Final = 100
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _OptInGroups:
|
||||
every: tuple[str, ...]
|
||||
vertex_target: tuple[str, ...]
|
||||
|
||||
|
||||
_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) -> frozenset[tuple[str | None, _OptInKey]] | str:
|
||||
deployments: Final = _model_group_deployments(model_group, llm_router)
|
||||
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) -> tuple[str, ...]:
|
||||
if llm_router is None:
|
||||
return tuple(sorted(model_groups))
|
||||
return tuple(
|
||||
{_same_answer_key(model_group, llm_router): 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
|
||||
) -> _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), llm_router)
|
||||
return _OptInGroups(
|
||||
every=every,
|
||||
vertex_target=tuple(group for group in every if _model_group_reads_vertex_target(group, llm_router)),
|
||||
)
|
||||
|
||||
|
||||
def _check_banned_params(
|
||||
body: Mapping[str, object],
|
||||
general_settings: dict,
|
||||
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],
|
||||
) -> None:
|
||||
"""Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in.
|
||||
|
||||
|
|
@ -501,21 +621,24 @@ 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:
|
||||
# Proxy-wide opt-in: every banned param is permitted, exit
|
||||
# entirely so the rest of the loop doesn't waste work.
|
||||
return
|
||||
if (
|
||||
if all(
|
||||
_allow_model_level_clientside_configurable_parameters(
|
||||
model=model,
|
||||
model=model_group,
|
||||
param=param,
|
||||
request_body_value=body[param],
|
||||
llm_router=llm_router,
|
||||
)
|
||||
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
|
||||
|
|
@ -606,27 +729,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}
|
||||
|
||||
|
|
@ -691,33 +826,60 @@ 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)
|
||||
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, also_requested_models, llm_router))
|
||||
_check_banned_params(
|
||||
request_body,
|
||||
root_body,
|
||||
general_settings,
|
||||
llm_router,
|
||||
model,
|
||||
manages_deployments=manages_deployments,
|
||||
mcp_server_fields_allowed=route_manages_mcp_servers(route),
|
||||
opt_in_groups=opt_in_groups,
|
||||
)
|
||||
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, general_settings, llm_router, model, manages_deployments=manages_deployments)
|
||||
_check_banned_params(
|
||||
nested,
|
||||
general_settings,
|
||||
llm_router,
|
||||
manages_deployments=manages_deployments,
|
||||
opt_in_groups=opt_in_groups,
|
||||
)
|
||||
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, general_settings, llm_router, model, manages_deployments=manages_deployments)
|
||||
_check_banned_params(
|
||||
metadata,
|
||||
general_settings,
|
||||
llm_router,
|
||||
manages_deployments=manages_deployments,
|
||||
banned_params=_BANNED_METADATA_PARAMS,
|
||||
opt_in_groups=opt_in_groups,
|
||||
)
|
||||
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}["),
|
||||
general_settings,
|
||||
llm_router,
|
||||
model,
|
||||
manages_deployments=manages_deployments,
|
||||
banned_params=_BANNED_METADATA_PARAMS,
|
||||
opt_in_groups=opt_in_groups,
|
||||
)
|
||||
for target in iter_request_fallback_targets(request_body):
|
||||
if isinstance(target, dict):
|
||||
_check_banned_params(target, general_settings, llm_router, model, manages_deployments=manages_deployments)
|
||||
_check_banned_params(
|
||||
target,
|
||||
general_settings,
|
||||
llm_router,
|
||||
manages_deployments=manages_deployments,
|
||||
opt_in_groups=opt_in_groups,
|
||||
)
|
||||
target_model = target.get("model")
|
||||
if isinstance(target_model, str):
|
||||
_reject_url_valued_fallback_target(target_model)
|
||||
|
|
@ -731,7 +893,12 @@ def is_request_body_safe(
|
|||
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, general_settings, llm_router, model, manages_deployments=manages_deployments
|
||||
litellm_params_metadata,
|
||||
general_settings,
|
||||
llm_router,
|
||||
manages_deployments=manages_deployments,
|
||||
banned_params=_BANNED_METADATA_PARAMS,
|
||||
opt_in_groups=opt_in_groups,
|
||||
)
|
||||
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)
|
||||
|
|
@ -831,7 +998,7 @@ async def pre_db_read_auth_checks(
|
|||
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,
|
||||
)
|
||||
|
||||
|
|
@ -2314,37 +2481,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)
|
||||
|
|
@ -2363,6 +2501,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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -4103,6 +4104,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/"
|
||||
|
|
@ -4127,11 +4132,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
|
||||
|
|
@ -4145,36 +4164,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)}"
|
||||
|
||||
|
|
@ -4228,11 +4344,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
|
||||
|
|
@ -4314,7 +4441,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,
|
||||
),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -2565,6 +2565,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]:
|
||||
|
|
@ -2994,6 +3009,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
|
||||
|
||||
|
|
@ -3361,8 +3398,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,
|
||||
):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -20,12 +21,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,
|
||||
|
|
@ -59,8 +61,11 @@ from litellm.proxy.vector_store_endpoints.endpoints import (
|
|||
from litellm.proxy.vector_store_endpoints.utils import (
|
||||
assert_user_can_access_vector_store_id,
|
||||
)
|
||||
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
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -179,6 +184,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",
|
||||
|
|
@ -660,6 +762,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,
|
||||
|
|
@ -684,15 +787,24 @@ async def rag_ingest(
|
|||
general_settings=general_settings,
|
||||
llm_router=llm_router,
|
||||
model="",
|
||||
route=get_request_route(request=request),
|
||||
)
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
from typing import Final, Optional
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
|
@ -34,6 +35,7 @@ from litellm.proxy.auth.auth_utils import (
|
|||
get_project_model_tpm_limit,
|
||||
get_request_route_template,
|
||||
is_request_body_safe,
|
||||
pre_db_read_auth_checks,
|
||||
)
|
||||
from litellm.router import Router
|
||||
from litellm.types.utils import oauth_token_exchange_litellm_params, server_owned_wif_litellm_params
|
||||
|
|
@ -4978,6 +4980,426 @@ 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",
|
||||
)
|
||||
|
||||
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",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -252,6 +252,9 @@ 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
|
||||
|
||||
|
|
@ -264,6 +267,17 @@ 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") -> list[object]:
|
||||
return []
|
||||
|
||||
|
||||
class FakeSettings:
|
||||
def is_set(self, name: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
class FakeStorageBackend:
|
||||
def __init__(self, contents: Mapping[str, bytes]) -> None:
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -6374,7 +6374,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",
|
||||
),
|
||||
],
|
||||
)
|
||||
|
|
@ -6400,7 +6404,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
|
||||
|
|
@ -6492,6 +6498,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(
|
||||
|
|
|
|||
|
|
@ -7674,7 +7674,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,
|
||||
)
|
||||
|
|
@ -7702,7 +7702,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,
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -7740,6 +7742,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"],
|
||||
|
|
@ -7761,7 +7824,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}
|
||||
|
||||
|
|
@ -7811,7 +7878,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,
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -7876,14 +7947,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"]}}}
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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. "
|
||||
|
|
@ -990,6 +992,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"}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue