mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
Merge 07eb764096 into e16d79d635
This commit is contained in:
commit
8d62702bf8
12 changed files with 713 additions and 73 deletions
|
|
@ -2778,6 +2778,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
default=False,
|
||||
description="When True, a search tool must be explicitly listed in object_permission.search_tools: a virtual key needs its own grant plus its team's, a keyless team member needs the team's, and a user with neither needs their own. A missing permission record, an empty list, or an unresolved team grants nothing, and the unregistered search fallback is denied. The master key and dashboard sessions are exempt",
|
||||
)
|
||||
model_access_group_deny_by_default: bool = Field(
|
||||
default=False,
|
||||
description="When True, a model tagged with model_info.access_groups is only reachable by a caller whose grant names one of its groups or the model itself. The grant that counts is the team's for a team key, the user's for a personal key, and the key's own for a key with neither. An empty models list, '*', all-proxy-models and wildcard patterns no longer reach these models, by any name, alias, deployment id, routing group or fallback. A model name with both tagged and untagged deployments serves only the untagged ones to such callers. Untagged models are unaffected, and the master key and proxy admins are exempt",
|
||||
)
|
||||
missing_session_id: Literal["generate", "reject", "omit"] | None = Field(
|
||||
None,
|
||||
description="What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400; 'omit' leaves SpendLogs.session_id null, matching callbacks such as Langfuse that only record a client-established metadata.session_id. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id.",
|
||||
|
|
|
|||
|
|
@ -161,6 +161,7 @@ from litellm.repositories.table_repositories import (
|
|||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
from litellm.router import Router
|
||||
from litellm.router_utils.model_access_group_grants import explicit_model_grants
|
||||
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
|
||||
from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
|
||||
from litellm.types.router import DeploymentTypedDict
|
||||
|
|
@ -5941,6 +5942,140 @@ def personal_key_team_model_access_applies(valid_token: UserAPIKeyAuth | None) -
|
|||
)
|
||||
|
||||
|
||||
_MODEL_NAMES: Final[TypeAdapter[tuple[str, ...]]] = TypeAdapter(tuple[str, ...])
|
||||
|
||||
|
||||
def model_access_group_deny_by_default(general_settings: Mapping[str, object]) -> bool:
|
||||
"""
|
||||
Startup rejects a non-boolean value from the config file. A non-boolean value that reaches
|
||||
general_settings another way enables the policy, so only access-group models are denied.
|
||||
"""
|
||||
try:
|
||||
return ConfigGeneralSettings.model_validate(
|
||||
MappingProxyType(
|
||||
{
|
||||
"model_access_group_deny_by_default": general_settings.get(
|
||||
"model_access_group_deny_by_default", False
|
||||
)
|
||||
}
|
||||
)
|
||||
).model_access_group_deny_by_default
|
||||
except ValidationError:
|
||||
return True
|
||||
|
||||
|
||||
async def model_access_group_grants(
|
||||
deny_by_default: bool,
|
||||
valid_token: UserAPIKeyAuth,
|
||||
team_object: LiteLLM_TeamTable | None,
|
||||
user_object: LiteLLM_UserTable | None,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> frozenset[str] | None:
|
||||
"""
|
||||
The names that reach a model access group's deployments under general_settings.model_access_group_deny_by_default:
|
||||
the team's grant for a team key, the owner teams' for a personal key that takes its models from them, the user's
|
||||
for any other personal key, and the key's own otherwise, plus the key's unified access groups. None when the
|
||||
setting is off or the caller is the master key, a dashboard session or a proxy admin
|
||||
"""
|
||||
if (
|
||||
not deny_by_default
|
||||
or not _is_strict_grant_identity(valid_token)
|
||||
or valid_token.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
):
|
||||
return None
|
||||
owner_models: Final = await _model_access_group_owner_models(
|
||||
valid_token=valid_token,
|
||||
team_object=team_object,
|
||||
user_object=user_object,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
key_group_models: Final = await _access_group_model_names(
|
||||
access_group_ids=valid_token.access_group_ids,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
return explicit_model_grants((*owner_models, *key_group_models))
|
||||
|
||||
|
||||
async def _model_access_group_owner_models(
|
||||
valid_token: UserAPIKeyAuth,
|
||||
team_object: LiteLLM_TeamTable | None,
|
||||
user_object: LiteLLM_UserTable | None,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> tuple[str, ...]:
|
||||
if valid_token.team_id is not None:
|
||||
if team_object is None:
|
||||
return _MODEL_NAMES.validate_python(valid_token.team_models)
|
||||
return await _team_model_grants(team_object, prisma_client, user_api_key_cache, proxy_logging_obj)
|
||||
if personal_key_team_model_access_applies(valid_token):
|
||||
teams: Final = (
|
||||
await _load_personal_key_teams(
|
||||
user_object=user_object,
|
||||
valid_token=valid_token,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if user_object is not None
|
||||
else ()
|
||||
)
|
||||
team_grants: Final = await asyncio.gather(
|
||||
*(
|
||||
_team_model_grants(team, prisma_client, user_api_key_cache, proxy_logging_obj)
|
||||
for team in teams
|
||||
if team is not None
|
||||
)
|
||||
)
|
||||
return tuple(chain.from_iterable(team_grants))
|
||||
if valid_token.user_id is not None and user_object is not None:
|
||||
from litellm.proxy.auth.resolvers.grants import user_models
|
||||
|
||||
return user_models(user_object)
|
||||
return tuple(_resolve_key_models_for_auth_check(valid_token=valid_token))
|
||||
|
||||
|
||||
async def _team_model_grants(
|
||||
team_object: LiteLLM_TeamTable,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> tuple[str, ...]:
|
||||
return (
|
||||
*_MODEL_NAMES.validate_python(team_object.models),
|
||||
*await _access_group_model_names(
|
||||
access_group_ids=team_object.access_group_ids,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _access_group_model_names(
|
||||
access_group_ids: Sequence[str] | None,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> tuple[str, ...]:
|
||||
if not access_group_ids:
|
||||
return ()
|
||||
return tuple(
|
||||
await get_models_from_access_groups(
|
||||
access_group_ids=access_group_ids,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def personal_key_inherits_team_grants(valid_token: UserAPIKeyAuth | None) -> bool:
|
||||
"""MCP server, MCP tool and agent inheritance is for a virtual key only. A keyless auth with no
|
||||
team (one grant source of an admitted subject, a JWT user) resolves as the plain caller it is."""
|
||||
|
|
|
|||
|
|
@ -85,6 +85,8 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
|
|||
jwt_key_mapping_cache_key,
|
||||
key_access_group_grants_model,
|
||||
key_model_aliases_for_auth_check,
|
||||
model_access_group_deny_by_default,
|
||||
model_access_group_grants,
|
||||
resolve_and_validate_end_user_id,
|
||||
resolve_default_end_user_budget,
|
||||
route_skips_budget_checks,
|
||||
|
|
@ -3794,6 +3796,16 @@ async def run_centralized_common_checks(
|
|||
)
|
||||
|
||||
bind_admission_counter_keys(user_api_key_auth_obj, end_user_id=end_user_id)
|
||||
grants: Final = await model_access_group_grants(
|
||||
deny_by_default=model_access_group_deny_by_default(_GENERAL_SETTINGS_ADAPTER.validate_python(general_settings)),
|
||||
valid_token=user_api_key_auth_obj,
|
||||
team_object=typed_team_object,
|
||||
user_object=user_object if isinstance(user_object, LiteLLM_UserTable) else None,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
user_api_key_auth_obj.model_access_group_grants = grants # rebind-ok: carries server-only grants to the router
|
||||
try:
|
||||
_ = await common_checks(
|
||||
request=request,
|
||||
|
|
|
|||
|
|
@ -6959,6 +6959,7 @@ class ProxyConfig:
|
|||
for name in (
|
||||
"vector_store_deny_by_default",
|
||||
"search_tool_deny_by_default",
|
||||
"model_access_group_deny_by_default",
|
||||
"vertex_ai_pass_through_drop_anthropic_betas",
|
||||
"request_tag_mappings",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -301,6 +301,11 @@ from litellm.router_utils.kubernetes_pod_discovery import (
|
|||
resolve_pods_after,
|
||||
resolve_pods_after_bound,
|
||||
)
|
||||
from litellm.router_utils.model_access_group_grants import (
|
||||
deployment_is_granted,
|
||||
granted_deployments,
|
||||
model_access_denied,
|
||||
)
|
||||
from litellm.router_utils.pre_call_checks.continuation_prefill_check import (
|
||||
MID_STREAM_CONTINUATION_MARKER,
|
||||
ContinuationOutput,
|
||||
|
|
@ -14710,6 +14715,27 @@ class Router:
|
|||
specific_deployment: bool | None = False,
|
||||
request_kwargs: Mapping[str, object] | None = None,
|
||||
) -> tuple[str, list | dict]:
|
||||
"""The candidates `_resolve_available_deployments` picks, less any the caller's model access grants miss"""
|
||||
resolved_model, deployments = self._resolve_available_deployments(
|
||||
model=model, specific_deployment=specific_deployment, request_kwargs=request_kwargs
|
||||
)
|
||||
if not isinstance(deployments, Sequence):
|
||||
if not deployment_is_granted(deployments, model, request_kwargs):
|
||||
raise model_access_denied(model)
|
||||
return resolved_model, cast( # cast-ok: contract requires a plain dict for a single deployment
|
||||
dict[str, object], deployments
|
||||
)
|
||||
granted: Final = granted_deployments(deployments, model, request_kwargs)
|
||||
if deployments and not granted:
|
||||
raise model_access_denied(model)
|
||||
return resolved_model, list(granted)
|
||||
|
||||
def _resolve_available_deployments(
|
||||
self,
|
||||
model: str,
|
||||
specific_deployment: bool | None,
|
||||
request_kwargs: Mapping[str, object] | None,
|
||||
) -> tuple[str, Sequence[DeploymentTypedDict] | Mapping[str, object]]:
|
||||
"""
|
||||
Common checks for 'get_available_deployment' across sync + async call.
|
||||
|
||||
|
|
@ -15933,6 +15959,9 @@ class Router:
|
|||
model=registered_model_name, strategy_tags=selected_strategy.tags, request_kwargs=request_kwargs
|
||||
)
|
||||
definition: Final = selected_strategy.deployment
|
||||
strategy_source: Final = marker if marker is not None else definition.model_dump() if definition else None
|
||||
if strategy_source is not None and not deployment_is_granted(strategy_source, model, request_kwargs):
|
||||
raise model_access_denied(model)
|
||||
config_id: Final = definition.model_info.id if definition else None
|
||||
updated_at: Final = definition.model_info.updated_at if definition else None
|
||||
router_kind: Final = classify_strategy_router_model(definition.litellm_params.model) if definition else None
|
||||
|
|
|
|||
79
litellm/router_utils/model_access_group_grants.py
Normal file
79
litellm/router_utils/model_access_group_grants.py
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
"""
|
||||
Under general_settings.model_access_group_deny_by_default, auth stamps the caller's explicit model grants on its
|
||||
auth object, and the router drops every candidate deployment in a model access group those grants do not name.
|
||||
Filtering the deployments the router is about to pick, rather than predicting them at auth time, covers aliases,
|
||||
routing groups, deployment ids, default fallbacks and models loaded after auth alike
|
||||
"""
|
||||
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from typing import Final, TypeVar
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict
|
||||
|
||||
import litellm
|
||||
from litellm.router_utils.common_utils import request_user_api_key_auth
|
||||
from litellm.types.llms.base_model import LiteLLMBaseModel
|
||||
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
||||
|
||||
_BLANKET_GRANTS: Final = frozenset({"*", "all-proxy-models", "all-team-models"})
|
||||
|
||||
_Deployment: Final = TypeVar("_Deployment", bound=Mapping[str, object])
|
||||
|
||||
|
||||
class _DeploymentModelInfo(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
access_groups: tuple[str, ...] | None = None
|
||||
|
||||
|
||||
class _DeploymentRow(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
model_name: str | None = None
|
||||
model_info: _DeploymentModelInfo | None = None
|
||||
|
||||
|
||||
def explicit_model_grants(models: Iterable[str]) -> frozenset[str]:
|
||||
"""The names in `models` that grant one thing by name: blanket entries and wildcard patterns name nothing"""
|
||||
return frozenset(model for model in models if model not in _BLANKET_GRANTS and "*" not in model)
|
||||
|
||||
|
||||
def request_model_access_group_grants(request_kwargs: Mapping[str, object] | None) -> frozenset[str] | None:
|
||||
if request_kwargs is None:
|
||||
return None
|
||||
user_api_key_auth: Final = request_user_api_key_auth(request_kwargs)
|
||||
return user_api_key_auth.model_access_group_grants if isinstance(user_api_key_auth, UserAPIKeyAuth) else None
|
||||
|
||||
|
||||
def _is_granted(deployment: Mapping[str, object], requested_model: str, grants: frozenset[str]) -> bool:
|
||||
row: Final = _DeploymentRow.model_validate(deployment)
|
||||
groups: Final = (row.model_info.access_groups if row.model_info is not None else None) or ()
|
||||
return not groups or requested_model in grants or row.model_name in grants or not grants.isdisjoint(groups)
|
||||
|
||||
|
||||
def model_access_denied(model: str) -> litellm.PermissionDeniedError:
|
||||
from litellm.proxy.auth.model_access_denied import model_access_denied_client_message
|
||||
|
||||
return litellm.PermissionDeniedError(
|
||||
message=model_access_denied_client_message(model=model),
|
||||
llm_provider="",
|
||||
model=model,
|
||||
response=httpx.Response(status_code=403, request=httpx.Request(method="POST", url="https://litellm.ai")),
|
||||
)
|
||||
|
||||
|
||||
def deployment_is_granted(
|
||||
deployment: Mapping[str, object], requested_model: str, request_kwargs: Mapping[str, object] | None
|
||||
) -> bool:
|
||||
grants: Final = request_model_access_group_grants(request_kwargs)
|
||||
return grants is None or _is_granted(deployment, requested_model, grants)
|
||||
|
||||
|
||||
def granted_deployments(
|
||||
deployments: Sequence[_Deployment], requested_model: str, request_kwargs: Mapping[str, object] | None
|
||||
) -> tuple[_Deployment, ...]:
|
||||
grants: Final = request_model_access_group_grants(request_kwargs)
|
||||
if grants is None:
|
||||
return tuple(deployments)
|
||||
return tuple(deployment for deployment in deployments if _is_granted(deployment, requested_model, grants))
|
||||
|
|
@ -203,6 +203,15 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
|
|||
user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True)
|
||||
org_budget_snapshot: OrgBudgetSnapshot | None = Field(default=None, exclude=True)
|
||||
matched_model_access_groups: list[str] | None = Field(default=None, exclude=True)
|
||||
model_access_group_grants: frozenset[str] | None = Field(
|
||||
default=None,
|
||||
exclude=True,
|
||||
description=(
|
||||
"Server-only. Set by auth under general_settings.model_access_group_deny_by_default to the model and "
|
||||
"access group names the caller was explicitly granted; the router then drops tagged deployments these "
|
||||
"names miss. None means no such restriction"
|
||||
),
|
||||
)
|
||||
budget_throttle_pct: float | None = Field(default=None, exclude=True)
|
||||
user: object | None = None # Expanded user object when expand=user is used
|
||||
created_by_user: object | None = None # Expanded created_by user when expand=user is used
|
||||
|
|
@ -240,6 +249,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
|
|||
values.pop("invoked_agent_policy", None)
|
||||
values.pop("agent_invocation_cost", None)
|
||||
values.pop("billing_agent_policy", None)
|
||||
values.pop("model_access_group_grants", None)
|
||||
api_key: Final = values.get("api_key")
|
||||
if api_key is not None:
|
||||
values.update({"token": cls._safe_hash_litellm_api_key(api_key)}) # pyright: ignore[reportArgumentType] # non-str fails in the hasher
|
||||
|
|
|
|||
|
|
@ -31,10 +31,13 @@ from litellm.proxy._types import (
|
|||
LiteLLM_BudgetTable,
|
||||
LiteLLM_UserTable,
|
||||
LiteLLM_TeamTable,
|
||||
LiteLLM_TeamTableCachedObj,
|
||||
Litellm_EntityType,
|
||||
LitellmUserRoles,
|
||||
ProxyException,
|
||||
UI_TEAM_ID,
|
||||
)
|
||||
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
|
|
@ -160,12 +163,15 @@ async def test_team_require_trace_id_accepts_supported_trace_carriers(
|
|||
@pytest.mark.parametrize("route", ["/v1/traces", "/v1/logs", "/v1/traces/query"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_require_trace_id_accepts_trace_telemetry_routes(route: str) -> None:
|
||||
assert await _run_team_trace_id_check(
|
||||
route=route,
|
||||
method="POST",
|
||||
request_body={},
|
||||
team_metadata={"require_trace_id": True},
|
||||
) is True
|
||||
assert (
|
||||
await _run_team_trace_id_check(
|
||||
route=route,
|
||||
method="POST",
|
||||
request_body={},
|
||||
team_metadata={"require_trace_id": True},
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -197,11 +203,14 @@ async def test_team_require_trace_id_rejects_overridden_body_metadata_on_chat()
|
|||
async def test_team_require_trace_id_accepts_effective_metadata_on_responses_route(
|
||||
request_body: dict[str, object],
|
||||
) -> None:
|
||||
assert await _run_team_trace_id_check(
|
||||
route="/v1/responses",
|
||||
request_body=request_body,
|
||||
team_metadata={"require_trace_id": True},
|
||||
) is True
|
||||
assert (
|
||||
await _run_team_trace_id_check(
|
||||
route="/v1/responses",
|
||||
request_body=request_body,
|
||||
team_metadata={"require_trace_id": True},
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -240,12 +249,15 @@ async def test_team_require_trace_id_rejects_fallback_carriers_when_litellm_trac
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_require_trace_id_accepts_header_on_mcp_with_empty_body() -> None:
|
||||
assert await _run_team_trace_id_check(
|
||||
route="/mcp/",
|
||||
request_body={},
|
||||
headers={"x-litellm-trace-id": "trace-mcp-123"},
|
||||
team_metadata={"require_trace_id": True},
|
||||
) is True
|
||||
assert (
|
||||
await _run_team_trace_id_check(
|
||||
route="/mcp/",
|
||||
request_body={},
|
||||
headers={"x-litellm-trace-id": "trace-mcp-123"},
|
||||
team_metadata={"require_trace_id": True},
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -278,12 +290,15 @@ async def test_team_require_trace_id_rejects_nonheader_carriers_on_pass_through_
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_require_trace_id_accepts_header_on_pass_through_route() -> None:
|
||||
assert await _run_team_trace_id_check(
|
||||
route="/anthropic/v1/messages",
|
||||
request_body={},
|
||||
headers={"x-litellm-trace-id": "trace-provider-123"},
|
||||
team_metadata={"require_trace_id": True},
|
||||
) is True
|
||||
assert (
|
||||
await _run_team_trace_id_check(
|
||||
route="/anthropic/v1/messages",
|
||||
request_body={},
|
||||
headers={"x-litellm-trace-id": "trace-provider-123"},
|
||||
team_metadata={"require_trace_id": True},
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -329,7 +344,9 @@ async def test_team_require_trace_id_rejects_nonheader_carriers_on_registered_pa
|
|||
"jsonrpc": "2.0",
|
||||
"id": "1",
|
||||
"method": "message/send",
|
||||
"params": {"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}},
|
||||
"params": {
|
||||
"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}
|
||||
},
|
||||
},
|
||||
),
|
||||
],
|
||||
|
|
@ -370,7 +387,9 @@ async def test_team_require_trace_id_rejects_traceparent_on_mcp_and_agent_routes
|
|||
"jsonrpc": "2.0",
|
||||
"id": "1",
|
||||
"method": "message/send",
|
||||
"params": {"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}},
|
||||
"params": {
|
||||
"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}
|
||||
},
|
||||
"metadata": {"trace_id": "trace-body-123"},
|
||||
},
|
||||
),
|
||||
|
|
@ -463,7 +482,7 @@ def test_lens_management_does_not_require_an_inference_user_parameter(route: str
|
|||
async def test_get_end_user_object(customer_spend, customer_budget):
|
||||
"""
|
||||
Scenario 1: normal - get_end_user_object returns the cached user
|
||||
Scenario 2: user over budget - NOTE: budget enforcement now happens in
|
||||
Scenario 2: user over budget - NOTE: budget enforcement now happens in
|
||||
common_checks() via _check_end_user_budget(), not in get_end_user_object()
|
||||
|
||||
This test verifies that get_end_user_object correctly retrieves the end user
|
||||
|
|
@ -591,7 +610,7 @@ async def test_can_key_call_model(model, expect_to_work):
|
|||
if expect_to_work:
|
||||
await can_key_call_model(**args)
|
||||
else:
|
||||
with pytest.raises(Exception, match='is not available for this API key') as e:
|
||||
with pytest.raises(Exception, match="is not available for this API key") as e:
|
||||
await can_key_call_model(**args)
|
||||
|
||||
print(e)
|
||||
|
|
@ -863,16 +882,12 @@ async def test_is_valid_fallback_model():
|
|||
)
|
||||
|
||||
try:
|
||||
await is_valid_fallback_model(
|
||||
model="gpt-3.5-turbo", llm_router=router, user_model=None
|
||||
)
|
||||
await is_valid_fallback_model(model="gpt-3.5-turbo", llm_router=router, user_model=None)
|
||||
except Exception as e:
|
||||
pytest.fail(f"Expected is_valid_fallback_model to work, got exception: {e}")
|
||||
|
||||
with pytest.raises(Exception, match="Invalid") as exc_info:
|
||||
await is_valid_fallback_model(
|
||||
model="gpt-4o", llm_router=router, user_model=None
|
||||
)
|
||||
await is_valid_fallback_model(model="gpt-4o", llm_router=router, user_model=None)
|
||||
e = exc_info.value
|
||||
assert "Invalid" in str(e)
|
||||
|
||||
|
|
@ -886,9 +901,7 @@ async def test_is_valid_fallback_model():
|
|||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_virtual_key_max_budget_check(
|
||||
token_spend, max_budget, expect_budget_error
|
||||
):
|
||||
async def test_virtual_key_max_budget_check(token_spend, max_budget, expect_budget_error):
|
||||
"""
|
||||
Test if virtual key budget checks work as expected:
|
||||
1. Triggers budget alert for all cases
|
||||
|
|
@ -1003,9 +1016,7 @@ async def test_can_team_access_model(model, team_models, expect_to_work):
|
|||
team_model_aliases=None,
|
||||
)
|
||||
if not expect_to_work:
|
||||
pytest.fail(
|
||||
f"Expected model access check to fail for model={model}, team_models={team_models}"
|
||||
)
|
||||
pytest.fail(f"Expected model access check to fail for model={model}, team_models={team_models}")
|
||||
except Exception as e:
|
||||
if expect_to_work:
|
||||
pytest.fail(
|
||||
|
|
@ -1058,9 +1069,9 @@ async def test_virtual_key_soft_budget_check(spend, soft_budget, expect_alert):
|
|||
|
||||
await asyncio.sleep(0.1) # Allow time for the alert task to complete
|
||||
|
||||
assert (
|
||||
alert_triggered == expect_alert
|
||||
), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
|
||||
assert alert_triggered == expect_alert, (
|
||||
f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -1113,9 +1124,7 @@ async def test_virtual_key_soft_budget_check(spend, soft_budget, expect_alert):
|
|||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_soft_budget_check(
|
||||
spend, soft_budget, expect_alert, metadata, expected_alert_emails
|
||||
):
|
||||
async def test_team_soft_budget_check(spend, soft_budget, expect_alert, metadata, expected_alert_emails):
|
||||
"""
|
||||
Test cases for _team_soft_budget_check:
|
||||
1. Spend over soft budget, no alert_emails configured - should NOT trigger alert (alerts only sent when alert_emails configured)
|
||||
|
|
@ -1163,9 +1172,9 @@ async def test_team_soft_budget_check(
|
|||
|
||||
await asyncio.sleep(0.1) # Allow time for the alert task to complete
|
||||
|
||||
assert (
|
||||
alert_triggered == expect_alert
|
||||
), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
|
||||
assert alert_triggered == expect_alert, (
|
||||
f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
|
||||
)
|
||||
|
||||
if expect_alert:
|
||||
assert captured_call_info is not None
|
||||
|
|
@ -1177,10 +1186,7 @@ async def test_team_soft_budget_check(
|
|||
if expected_alert_emails is not None:
|
||||
assert captured_call_info.alert_emails == expected_alert_emails
|
||||
else:
|
||||
assert (
|
||||
captured_call_info.alert_emails is None
|
||||
or captured_call_info.alert_emails == []
|
||||
)
|
||||
assert captured_call_info.alert_emails is None or captured_call_info.alert_emails == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1311,9 +1317,7 @@ async def test_get_fuzzy_user_object():
|
|||
|
||||
# Test 5: Only email provided (no SSO ID)
|
||||
mock_prisma.db.litellm_usertable.find_first = AsyncMock(return_value=test_user)
|
||||
result = await _get_fuzzy_user_object(
|
||||
prisma_client=mock_prisma, user_email="test@example.com"
|
||||
)
|
||||
result = await _get_fuzzy_user_object(prisma_client=mock_prisma, user_email="test@example.com")
|
||||
assert result == test_user
|
||||
mock_prisma.db.litellm_usertable.find_first.assert_called_with(
|
||||
where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}},
|
||||
|
|
@ -1322,9 +1326,7 @@ async def test_get_fuzzy_user_object():
|
|||
|
||||
# Test 6: Only SSO ID provided (no email)
|
||||
mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=test_user)
|
||||
result = await _get_fuzzy_user_object(
|
||||
prisma_client=mock_prisma, sso_user_id="sso_123"
|
||||
)
|
||||
result = await _get_fuzzy_user_object(prisma_client=mock_prisma, sso_user_id="sso_123")
|
||||
assert result == test_user
|
||||
mock_prisma.db.litellm_usertable.find_unique.assert_called_with(
|
||||
where={"sso_user_id": "sso_123"}, include={"organization_memberships": True}
|
||||
|
|
@ -1371,7 +1373,7 @@ async def test_can_key_call_model_with_aliases(model, alias_map, expect_to_work)
|
|||
llm_router=router,
|
||||
)
|
||||
else:
|
||||
with pytest.raises(Exception, match='is not available for this API key') as e:
|
||||
with pytest.raises(Exception, match="is not available for this API key") as e:
|
||||
await can_key_call_model(
|
||||
model=model,
|
||||
llm_model_list=llm_model_list,
|
||||
|
|
@ -1457,9 +1459,7 @@ async def test_delete_cache_access_object():
|
|||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_resources_from_access_groups(
|
||||
resource_field, access_group_data, expected
|
||||
):
|
||||
async def test_get_resources_from_access_groups(resource_field, access_group_data, expected):
|
||||
"""Test _get_resources_from_access_groups returns correct resource list from access groups."""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -2049,7 +2049,9 @@ def test_is_model_cost_zero_judges_an_alias_chain_by_the_deployment_its_entry_ro
|
|||
model_group_alias={"chain-entry": "local-free", "local-free": "paid-gpt"},
|
||||
)
|
||||
expected: Final = {"chain-entry": True, "local-free": False, "paid-gpt": False}
|
||||
order: Final = ("chain-entry", "local-free", "paid-gpt") if entry_first else ("local-free", "paid-gpt", "chain-entry")
|
||||
order: Final = (
|
||||
("chain-entry", "local-free", "paid-gpt") if entry_first else ("local-free", "paid-gpt", "chain-entry")
|
||||
)
|
||||
|
||||
verdicts: Final = {name: is_model_cost_zero(model=name, llm_router=router) for name in order}
|
||||
|
||||
|
|
@ -2292,10 +2294,7 @@ async def test_bulk_eviction_clears_this_worker_before_waiting_on_a_wedged_redis
|
|||
monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5)
|
||||
client = _WedgedPublishRedisClient()
|
||||
cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000))
|
||||
hashed_tokens: Final = tuple(
|
||||
hashlib.sha256(f"token-{index}".encode()).hexdigest()
|
||||
for index in range(600)
|
||||
)
|
||||
hashed_tokens: Final = tuple(hashlib.sha256(f"token-{index}".encode()).hexdigest() for index in range(600))
|
||||
for hashed_token in hashed_tokens:
|
||||
cache.in_memory_cache_for(hashed_token).set_cache(key=hashed_token, value={"token": hashed_token})
|
||||
seeded: Final = tuple(
|
||||
|
|
@ -2355,10 +2354,7 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_
|
|||
monkeypatch.setattr(pubsub_module, "_pending_publishes", set())
|
||||
monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5)
|
||||
client = _WedgedPublishRedisClient()
|
||||
hashed_tokens: Final = tuple(
|
||||
hashlib.sha256(f"token-{index}".encode()).hexdigest()
|
||||
for index in range(600)
|
||||
)
|
||||
hashed_tokens: Final = tuple(hashlib.sha256(f"token-{index}".encode()).hexdigest() for index in range(600))
|
||||
redis = _AuthRedisCache({token: {"token": token} for token in hashed_tokens})
|
||||
cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000))
|
||||
cache.attach_redis_cache(redis) # pyright: ignore[reportArgumentType] # a dict-backed fake of the Redis layer
|
||||
|
|
@ -2380,3 +2376,136 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_
|
|||
|
||||
assert refilled == []
|
||||
assert redis.entries == {}
|
||||
|
||||
|
||||
_GRANT_ACCESS_GROUPS: Final = {"ag-team": ["team-group-model"], "ag-key": ["key-group-model"]}
|
||||
|
||||
|
||||
async def _model_access_group_grants(
|
||||
token: UserAPIKeyAuth,
|
||||
*,
|
||||
team: LiteLLM_TeamTable | None = None,
|
||||
user: LiteLLM_UserTable | None = None,
|
||||
deny_by_default: bool = True,
|
||||
cached_teams: Sequence[LiteLLM_TeamTableCachedObj] = (),
|
||||
) -> frozenset[str] | None:
|
||||
from litellm.proxy._types import LiteLLM_AccessGroupTable
|
||||
from litellm.proxy.auth.auth_checks import cache_access_object, model_access_group_grants
|
||||
|
||||
cache: Final = UserApiKeyCache()
|
||||
for group_id, models in _GRANT_ACCESS_GROUPS.items():
|
||||
await cache_access_object(
|
||||
access_group_id=group_id,
|
||||
access_group_table=LiteLLM_AccessGroupTable(
|
||||
access_group_id=group_id, access_group_name=group_id, access_model_names=models
|
||||
),
|
||||
user_api_key_cache=cache,
|
||||
)
|
||||
for cached_team in cached_teams:
|
||||
await cache.async_set_cache(
|
||||
key=f"team_id:{cached_team.team_id}", value=cached_team, model_type=LiteLLM_TeamTableCachedObj
|
||||
)
|
||||
return await model_access_group_grants(
|
||||
deny_by_default=deny_by_default,
|
||||
valid_token=token,
|
||||
team_object=team,
|
||||
user_object=user,
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=cache,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"token, deny_by_default",
|
||||
[
|
||||
(UserAPIKeyAuth(token="k1", models=["group-x"]), False),
|
||||
(UserAPIKeyAuth(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, models=["group-x"]), True),
|
||||
(UserAPIKeyAuth(token="k1", team_id=UI_TEAM_ID, team_models=["group-x"]), True),
|
||||
(UserAPIKeyAuth(token="k1", models=["group-x"], user_role=LitellmUserRoles.PROXY_ADMIN), True),
|
||||
],
|
||||
ids=["setting-off", "master-key", "dashboard-session", "proxy-admin"],
|
||||
)
|
||||
async def test_model_access_group_grants_is_unrestricted_when_off_or_for_admin_identities(
|
||||
token: UserAPIKeyAuth, deny_by_default: bool
|
||||
) -> None:
|
||||
assert await _model_access_group_grants(token, deny_by_default=deny_by_default) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_key_grants_are_the_team_models_and_team_access_group_models_without_blankets() -> None:
|
||||
token: Final = UserAPIKeyAuth(token="k1", team_id="t1", models=["key-model"])
|
||||
team: Final = LiteLLM_TeamTable(
|
||||
team_id="t1", models=["*", "openai/*", "all-proxy-models", "group-x"], access_group_ids=["ag-team"]
|
||||
)
|
||||
|
||||
assert await _model_access_group_grants(token, team=team) == frozenset({"group-x", "team-group-model"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_key_without_a_loaded_team_falls_back_to_the_token_team_models() -> None:
|
||||
token: Final = UserAPIKeyAuth(token="k1", team_id="t1", models=["key-model"], team_models=["group-x", "*"])
|
||||
|
||||
assert await _model_access_group_grants(token) == frozenset({"group-x"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_access_groups_add_their_models_to_the_owner_grants() -> None:
|
||||
token: Final = UserAPIKeyAuth(token="k1", team_id="t1", access_group_ids=["ag-key"])
|
||||
team: Final = LiteLLM_TeamTable(team_id="t1", models=["group-x"])
|
||||
|
||||
assert await _model_access_group_grants(token, team=team) == frozenset({"group-x", "key-group-model"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_personal_key_grants_are_the_owner_user_models_not_the_key_models() -> None:
|
||||
token: Final = UserAPIKeyAuth(token="k1", user_id="u1", models=["key-model"])
|
||||
user: Final = LiteLLM_UserTable(user_id="u1", models=["user-model", "all-proxy-models"])
|
||||
|
||||
assert await _model_access_group_grants(token, user=user) == frozenset({"user-model"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_personal_key_taking_models_from_teams_gets_the_union_of_its_owner_teams(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True)
|
||||
token: Final = UserAPIKeyAuth(token="k1", user_id="u1", models=["key-model"])
|
||||
user: Final = LiteLLM_UserTable(user_id="u1", models=["user-model"], teams=["team-a", "team-b"])
|
||||
owner_teams: Final = (
|
||||
LiteLLM_TeamTableCachedObj(team_id="team-a", models=["a-model", "*"]),
|
||||
LiteLLM_TeamTableCachedObj(team_id="team-b", models=["b-model"], access_group_ids=["ag-team"]),
|
||||
)
|
||||
|
||||
assert await _model_access_group_grants(token, user=user, cached_teams=owner_teams) == frozenset(
|
||||
{"a-model", "b-model", "team-group-model"}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"key_models, expected",
|
||||
[
|
||||
(["key-model", "all-proxy-models", "gpt-*"], frozenset({"key-model"})),
|
||||
(["all-team-models"], frozenset()),
|
||||
],
|
||||
)
|
||||
async def test_teamless_key_without_a_loaded_user_is_granted_its_own_models(
|
||||
key_models: list[str], expected: frozenset[str]
|
||||
) -> None:
|
||||
token: Final = UserAPIKeyAuth(token="k1", user_id="u1", models=key_models, team_models=["team-model"])
|
||||
|
||||
assert await _model_access_group_grants(token) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"build",
|
||||
[
|
||||
lambda: UserAPIKeyAuth(token="k1", model_access_group_grants=frozenset({"group-x"})),
|
||||
lambda: UserAPIKeyAuth.model_validate({"token": "k1", "model_access_group_grants": ["group-x"]}),
|
||||
],
|
||||
ids=["constructor", "model_validate"],
|
||||
)
|
||||
def test_caller_input_cannot_set_model_access_group_grants(build: Callable[[], UserAPIKeyAuth]) -> None:
|
||||
assert build().model_access_group_grants is None
|
||||
|
|
|
|||
|
|
@ -5202,6 +5202,54 @@ async def test_centralized_common_checks_enforces_team_model_max_budget_from_the
|
|||
assert limiter.calls == [("t1", team_caps, key_caps, "gpt-4o")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"general_settings, expected_grants",
|
||||
[
|
||||
({"model_access_group_deny_by_default": True}, frozenset({"group-x"})),
|
||||
({"model_access_group_deny_by_default": "true"}, frozenset({"group-x"})),
|
||||
({"model_access_group_deny_by_default": "not-a-bool"}, frozenset({"group-x"})),
|
||||
({"model_access_group_deny_by_default": False}, None),
|
||||
({}, None),
|
||||
],
|
||||
ids=["deny-by-default", "quoted-true", "invalid-value-fails-closed", "setting-false", "setting-absent"],
|
||||
)
|
||||
async def test_centralized_common_checks_stamps_the_team_model_access_group_grants_on_the_token(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
general_settings: Mapping[str, object],
|
||||
expected_grants: frozenset[str] | None,
|
||||
) -> None:
|
||||
from starlette.datastructures import URL
|
||||
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
token: Final = UserAPIKeyAuth(token="k1", team_id="t1", models=["key-model"])
|
||||
request: Final = Request(scope={"type": "http"})
|
||||
request._url = URL(url="/chat/completions")
|
||||
user_api_key_cache: Final = UserApiKeyCache()
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key="team_id:t1",
|
||||
value=LiteLLM_TeamTableCachedObj(team_id="t1", models=["group-x", "*"]),
|
||||
model_type=LiteLLM_TeamTableCachedObj,
|
||||
)
|
||||
for name, value in {
|
||||
**_proxy_attrs_for_centralized_checks(),
|
||||
"prisma_client": MagicMock(),
|
||||
"user_api_key_cache": user_api_key_cache,
|
||||
"general_settings": general_settings,
|
||||
}.items():
|
||||
monkeypatch.setattr(proxy_server, name, value)
|
||||
|
||||
await run_centralized_common_checks(
|
||||
user_api_key_auth_obj=token,
|
||||
request=request,
|
||||
request_data={"model": "group-x"},
|
||||
route="/chat/completions",
|
||||
)
|
||||
|
||||
assert token.model_access_group_grants == expected_grants
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_centralized_common_checks_loads_member_budget_and_exempts_discovery_routes():
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
|
|
|||
|
|
@ -2546,7 +2546,9 @@ async def test_load_config_logs_disabled_budget_reservation_once(tmp_path, monke
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("setting", ["vector_store_deny_by_default", "search_tool_deny_by_default"])
|
||||
@pytest.mark.parametrize(
|
||||
"setting", ["vector_store_deny_by_default", "search_tool_deny_by_default", "model_access_group_deny_by_default"]
|
||||
)
|
||||
@pytest.mark.parametrize(("yaml_value", "expected"), [("true", True), ("false", False)])
|
||||
async def test_load_config_yaml_deny_by_default_is_boolean(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, setting: str, yaml_value: str, expected: bool
|
||||
|
|
@ -2564,7 +2566,9 @@ async def test_load_config_yaml_deny_by_default_is_boolean(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("setting", ["vector_store_deny_by_default", "search_tool_deny_by_default"])
|
||||
@pytest.mark.parametrize(
|
||||
"setting", ["vector_store_deny_by_default", "search_tool_deny_by_default", "model_access_group_deny_by_default"]
|
||||
)
|
||||
@pytest.mark.parametrize("yaml_value", ["", "enabled"], ids=["null", "string"])
|
||||
async def test_load_config_rejects_non_boolean_deny_by_default(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, setting: str, yaml_value: str
|
||||
|
|
|
|||
183
tests/unit/router_utils/test_model_access_group_grants.py
Normal file
183
tests/unit/router_utils/test_model_access_group_grants.py
Normal file
|
|
@ -0,0 +1,183 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.router_utils.model_access_group_grants import explicit_model_grants
|
||||
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
||||
|
||||
|
||||
def _deployment(model_name: str, reply: str, model_id: str, access_groups: tuple[str, ...] = ()) -> dict[str, object]:
|
||||
return {
|
||||
"model_name": model_name,
|
||||
"litellm_params": {"model": f"openai/{model_id}-upstream", "api_key": "sk-fake", "mock_response": reply},
|
||||
"model_info": {"id": model_id, "access_groups": list(access_groups)},
|
||||
}
|
||||
|
||||
|
||||
def _router() -> litellm.Router:
|
||||
return litellm.Router(
|
||||
model_list=[
|
||||
_deployment("private-model", "private", "private-id", ("group-x",)),
|
||||
_deployment("public-model", "public", "public-id"),
|
||||
_deployment("shared-model", "shared-x", "shared-x-id", ("group-x",)),
|
||||
_deployment("shared-model", "shared-y", "shared-y-id", ("group-y",)),
|
||||
],
|
||||
model_group_alias={"alias-to-private": "private-model"},
|
||||
routing_groups=[
|
||||
{
|
||||
"group_name": "mixed-group",
|
||||
"models": ["private-model", "public-model"],
|
||||
"routing_strategy": "simple-shuffle",
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _metadata(grants: frozenset[str] | None) -> dict[str, object]:
|
||||
caller: Final = UserAPIKeyAuth(api_key="sk-caller")
|
||||
caller.model_access_group_grants = grants
|
||||
return {"user_api_key_auth": caller}
|
||||
|
||||
|
||||
async def _reply(router: litellm.Router, model: str, grants: frozenset[str] | None) -> str | None:
|
||||
response: Final = await router.acompletion(
|
||||
model=model, messages=[{"role": "user", "content": "hi"}], metadata=_metadata(grants)
|
||||
)
|
||||
assert isinstance(response, litellm.ModelResponse)
|
||||
return response.choices[0].message.content # pyright: ignore[reportAttributeAccessIssue] # non-streaming choices
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "grants", "reply"),
|
||||
[
|
||||
pytest.param("private-model", frozenset({"group-x"}), "private", id="group-grant"),
|
||||
pytest.param("private-model", frozenset({"private-model"}), "private", id="model-name-grant"),
|
||||
pytest.param("alias-to-private", frozenset({"private-model"}), "private", id="alias-of-granted-model"),
|
||||
pytest.param("private-model", None, "private", id="setting-off-or-exempt"),
|
||||
pytest.param("public-model", frozenset(), "public", id="untagged-model"),
|
||||
],
|
||||
)
|
||||
async def test_granted_callers_reach_the_deployment(model: str, grants: frozenset[str] | None, reply: str) -> None:
|
||||
assert await _reply(_router(), model, grants) == reply
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
pytest.param("private-model", id="model-name"),
|
||||
pytest.param("alias-to-private", id="alias"),
|
||||
pytest.param("private-id", id="deployment-id"),
|
||||
pytest.param("openai/private-id-upstream", id="provider-model-name"),
|
||||
],
|
||||
)
|
||||
async def test_ungranted_callers_get_403_on_every_name_for_a_tagged_deployment(model: str) -> None:
|
||||
with pytest.raises(litellm.PermissionDeniedError) as denied:
|
||||
await _reply(_router(), model, explicit_model_grants(["*", "openai/*", "all-proxy-models", "group-y"]))
|
||||
assert denied.value.status_code == 403
|
||||
assert f"The requested model '{model}' is not available for this API key" in denied.value.message
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("grants", "replies"),
|
||||
[
|
||||
pytest.param(frozenset(), {"public"}, id="ungranted-caller-gets-only-untagged-members"),
|
||||
pytest.param(frozenset({"mixed-group"}), {"private", "public"}, id="group-named-explicitly-reaches-all"),
|
||||
],
|
||||
)
|
||||
async def test_a_routing_group_serves_the_members_the_caller_is_granted(
|
||||
grants: frozenset[str], replies: set[str]
|
||||
) -> None:
|
||||
router: Final = _router()
|
||||
assert {await _reply(router, "mixed-group", grants) for _ in range(40)} == replies
|
||||
|
||||
|
||||
async def test_a_model_split_across_groups_serves_the_copy_in_the_callers_group() -> None:
|
||||
router: Final = _router()
|
||||
replies: Final = {await _reply(router, "shared-model", frozenset({"group-x"})) for _ in range(20)}
|
||||
assert replies == {"shared-x"}
|
||||
|
||||
|
||||
async def test_a_catch_all_deployment_in_a_group_is_not_granted_by_a_catch_all_grant() -> None:
|
||||
router: Final = litellm.Router(model_list=[_deployment("*", "catch-all", "catch-all-id", ("group-x",))])
|
||||
with pytest.raises(litellm.PermissionDeniedError):
|
||||
await _reply(router, "*", explicit_model_grants(["*"]))
|
||||
with pytest.raises(litellm.PermissionDeniedError):
|
||||
await _reply(router, "openai/gpt-4o", explicit_model_grants(["*"]))
|
||||
assert await _reply(router, "openai/gpt-4o", explicit_model_grants(["group-x"])) == "catch-all"
|
||||
|
||||
|
||||
async def test_the_default_fallback_does_not_reach_a_tagged_deployment() -> None:
|
||||
router: Final = litellm.Router(
|
||||
model_list=[_deployment("private-model", "private", "private-id", ("group-x",))],
|
||||
default_fallbacks=["private-model"],
|
||||
)
|
||||
with pytest.raises(litellm.PermissionDeniedError):
|
||||
await _reply(router, "unknown-model", frozenset())
|
||||
assert await _reply(router, "unknown-model", frozenset({"group-x"})) == "private"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("models", "grants"),
|
||||
[
|
||||
pytest.param(["*", "all-proxy-models", "all-team-models"], frozenset(), id="blanket"),
|
||||
pytest.param(["openai/*", "group-x", "private-model"], frozenset({"group-x", "private-model"}), id="mixed"),
|
||||
pytest.param([], frozenset(), id="empty"),
|
||||
],
|
||||
)
|
||||
def test_only_names_that_grant_one_thing_count_as_explicit(models: list[str], grants: frozenset[str]) -> None:
|
||||
assert explicit_model_grants(models) == grants
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "resolved_name", "deployment_ids"),
|
||||
[
|
||||
pytest.param("alias-to-private", "private-model", ("private-id",), id="alias"),
|
||||
pytest.param("mixed-group", "mixed-group", ("private-id", "public-id"), id="routing-group"),
|
||||
],
|
||||
)
|
||||
def test_resolution_finds_tagged_deployments_for_any_caller(
|
||||
model: str, resolved_name: str, deployment_ids: tuple[str, ...]
|
||||
) -> None:
|
||||
name, deployments = _router()._resolve_available_deployments(
|
||||
model=model, specific_deployment=False, request_kwargs={"metadata": _metadata(frozenset())}
|
||||
)
|
||||
assert isinstance(deployments, list)
|
||||
assert (name, tuple(d["model_info"]["id"] for d in deployments)) == (resolved_name, deployment_ids)
|
||||
|
||||
|
||||
def test_resolution_finds_a_tagged_deployment_by_id_for_any_caller() -> None:
|
||||
name, deployment = _router()._resolve_available_deployments(
|
||||
model="private-id", specific_deployment=False, request_kwargs={"metadata": _metadata(frozenset())}
|
||||
)
|
||||
assert isinstance(deployment, dict)
|
||||
assert (name, deployment["model_info"]["id"]) == ("openai/private-id-upstream", "private-id")
|
||||
|
||||
|
||||
def _router_with_tagged_complexity_router() -> litellm.Router:
|
||||
return litellm.Router(
|
||||
model_list=[
|
||||
_deployment("public-model", "public", "public-id"),
|
||||
{
|
||||
"model_name": "private-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {
|
||||
"tiers": {"SIMPLE": "public-model", "MEDIUM": "public-model", "COMPLEX": "public-model"}
|
||||
},
|
||||
"complexity_router_default_model": "public-model",
|
||||
},
|
||||
"model_info": {"id": "private-router-id", "access_groups": ["group-x"]},
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
async def test_a_tagged_auto_router_refuses_ungranted_callers_even_when_it_picks_a_public_target() -> None:
|
||||
with pytest.raises(litellm.PermissionDeniedError) as denied:
|
||||
await _reply(_router_with_tagged_complexity_router(), "private-router", frozenset({"public-model"}))
|
||||
assert denied.value.status_code == 403
|
||||
|
||||
|
||||
async def test_a_tagged_auto_router_serves_callers_granted_its_group() -> None:
|
||||
assert await _reply(_router_with_tagged_complexity_router(), "private-router", frozenset({"group-x"})) == "public"
|
||||
6
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
6
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -30225,6 +30225,12 @@ export interface components {
|
|||
* @description What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400; 'omit' leaves SpendLogs.session_id null, matching callbacks such as Langfuse that only record a client-established metadata.session_id. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id.
|
||||
*/
|
||||
missing_session_id?: ("generate" | "reject" | "omit") | null;
|
||||
/**
|
||||
* Model Access Group Deny By Default
|
||||
* @description When True, a model tagged with model_info.access_groups is only reachable by a caller whose grant names one of its groups or the model itself. The grant that counts is the team's for a team key, the user's for a personal key, and the key's own for a key with neither. An empty models list, '*', all-proxy-models and wildcard patterns no longer reach these models, by any name, alias, deployment id, routing group or fallback. A model name with both tagged and untagged deployments serves only the untagged ones to such callers. Untagged models are unaffected, and the master key and proxy admins are exempt
|
||||
* @default false
|
||||
*/
|
||||
model_access_group_deny_by_default: boolean;
|
||||
/**
|
||||
* Model List Healthy Only
|
||||
* @description When true, `/models`, `/v1/models/{id}` and `/model/info` hide models whose backing deployments are all unhealthy, for every caller, without needing `healthy_only=true` per request. Requires `background_health_checks: true`, and keeps deployment health state cached without turning on `enable_health_check_routing`, so routing is unaffected. With no health state nothing is hidden. Hiding is presentation-only, a hidden model can still be called.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue