diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a5e26a02828..51fe3c3c93d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2767,7 +2767,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): ) 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. Untagged models are unaffected, and proxy admins' personal access is exempt", + 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, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 0fea82c2feb..b897293ebce 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -94,12 +94,6 @@ from litellm.proxy.auth.model_access_denied import ( customer_model_access_denied_client_message, model_access_denied_client_message, ) -from litellm.proxy.auth.model_access_group_grants import ( - access_group_model_denied, - model_access_group_deny_by_default, - requested_model_names, - ungranted_access_group_models, -) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( broadcast_evictions, @@ -165,6 +159,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 @@ -1293,12 +1288,7 @@ async def common_checks( llm_router=llm_router, user_object=user_object, key_model_aliases=key_model_aliases_for_auth_check(valid_token), - teams_grant_access_group_models=personal_key_team_model_access_applies(valid_token), ) - if _model and team_object is None and user_object is None and valid_token is not None: - await _require_key_access_group_grant( - model=_model, valid_token=valid_token, llm_router=llm_router, prisma_client=prisma_client - ) if end_user_object is not None and end_user_object.models: with tracer.trace("litellm.proxy.auth.common_checks.can_customer_call_model"): @@ -5498,10 +5488,6 @@ async def can_key_call_model( Raises: - Exception: If token not allowed to call model """ - if valid_token.user_id is None: - await _require_key_access_group_grant( - model=model, valid_token=valid_token, llm_router=llm_router, prisma_client=prisma_client - ) key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token) try: return can_object_call_model( @@ -5534,39 +5520,6 @@ async def can_key_call_model( raise -async def _require_key_access_group_grant( - model: str | Sequence[str], - valid_token: UserAPIKeyAuth, - llm_router: litellm.Router | None, - prisma_client: DatabaseClient | None, -) -> None: - """A key with no team and no loaded owner owns its own grant, so deny-by-default reads the key's models""" - if ( - valid_token.team_id is not None - or valid_token.user_role == LitellmUserRoles.PROXY_ADMIN - or not _is_strict_grant_identity(valid_token) - or personal_key_team_model_access_applies(valid_token) - or not _model_access_group_deny_by_default_enabled() - ): - return - ungranted: Final = await _ungranted_after_access_group_ids( - ungranted=ungranted_access_group_models( - names=requested_model_names( - model, valid_token.team_model_aliases, key_model_aliases_for_auth_check(valid_token) - ), - granted_models=_resolve_key_models_for_auth_check(valid_token=valid_token), - llm_router=llm_router, - team_id=None, - ), - access_group_ids=valid_token.access_group_ids, - llm_router=llm_router, - team_id=None, - prisma_client=prisma_client, - ) - if ungranted: - raise access_group_model_denied(model=model, ungranted=ungranted, object_type="key") - - async def _personal_key_owner( valid_token: UserAPIKeyAuth, prisma_client: PrismaClient | None, @@ -5595,40 +5548,6 @@ async def _personal_key_owner( return None -async def _require_teamless_owner_access_group_grant( - model: str, - valid_token: UserAPIKeyAuth, - llm_router: litellm.Router | None, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> None: - if ( - not _is_strict_grant_identity(valid_token) - or personal_key_team_model_access_applies(valid_token) - or not _model_access_group_deny_by_default_enabled() - ): - return - owner: Final = await _personal_key_owner( - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - if owner is None: - await _require_key_access_group_grant( - model=model, valid_token=valid_token, llm_router=llm_router, prisma_client=prisma_client - ) - return - _require_user_access_group_grant( - model=model, - user_role=owner.user_role, - user_models=_user_models(owner), - llm_router=llm_router, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - - async def can_key_call_resolved_model( model: str, llm_model_list: Sequence[object] | None, @@ -5708,16 +5627,6 @@ async def can_key_call_resolved_model( key_model_aliases=key_model_aliases_for_auth_check(valid_token), ) - if team_object is None: - await _require_teamless_owner_access_group_grant( - model=model, - valid_token=valid_token, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - if personal_key_team_model_access_applies(valid_token) and valid_token.user_id is not None: owner: Final = await _personal_key_owner( valid_token=valid_token, @@ -5801,18 +5710,7 @@ async def can_team_access_model( 1. First checks native team-level model permissions (current implementation) 2. If not allowed natively, falls back to access_group_ids on the team - 3. Under general_settings.model_access_group_deny_by_default, a model in a model access group also needs - the team to name that group or model explicitly """ - if team_object is not None: - await _require_team_access_group_grant( - model=model, - team_object=team_object, - llm_router=llm_router, - team_model_aliases=team_model_aliases, - key_model_aliases=key_model_aliases, - prisma_client=prisma_client, - ) try: return can_object_call_model( model=model, @@ -5844,68 +5742,6 @@ async def can_team_access_model( raise -def _proxy_general_settings() -> Mapping[str, object]: - from litellm.proxy.proxy_server import general_settings - - return typed_general_settings(general_settings) - - -def _model_access_group_deny_by_default_enabled() -> bool: - return model_access_group_deny_by_default(_proxy_general_settings()) - - -_MODEL_NAMES: Final = TypeAdapter(tuple[str, ...]) - - -def _user_models(user_object: LiteLLM_UserTable) -> tuple[str, ...]: - return _MODEL_NAMES.validate_python(user_object.models) - - -async def _ungranted_after_access_group_ids( - ungranted: tuple[str, ...], - access_group_ids: Sequence[str] | None, - llm_router: Router | None, - team_id: str | None, - prisma_client: DatabaseClient | None, -) -> tuple[str, ...]: - if not ungranted or not access_group_ids: - return ungranted - return ungranted_access_group_models( - names=ungranted, - granted_models=await get_models_from_access_groups( - access_group_ids=access_group_ids, prisma_client=prisma_client - ), - llm_router=llm_router, - team_id=team_id, - ) - - -async def _require_team_access_group_grant( - model: str | Sequence[str], - team_object: LiteLLM_TeamTable, - llm_router: Router | None, - team_model_aliases: Mapping[str, str] | None, - key_model_aliases: Mapping[str, str] | None, - prisma_client: DatabaseClient | None, -) -> None: - if team_object.team_id == UI_TEAM_ID or not _model_access_group_deny_by_default_enabled(): - return - ungranted: Final = await _ungranted_after_access_group_ids( - ungranted=ungranted_access_group_models( - names=requested_model_names(model, team_model_aliases, key_model_aliases), - granted_models=team_object.models, - llm_router=llm_router, - team_id=team_object.team_id, - ), - access_group_ids=team_object.access_group_ids, - llm_router=llm_router, - team_id=team_object.team_id, - prisma_client=prisma_client, - ) - if ungranted: - raise access_group_model_denied(model=model, ungranted=ungranted, object_type="team") - - async def get_authorized_resources_from_key_access_groups( valid_token: UserAPIKeyAuth | None, team_object: LiteLLM_TeamTable | None, @@ -5992,14 +5828,9 @@ async def key_access_group_grants_model( key_model_aliases=key_model_aliases, object_type="key", ) + return True except ProxyException: return False - return not _model_access_group_deny_by_default_enabled() or not ungranted_access_group_models( - names=requested_model_names(model, valid_token.team_model_aliases if valid_token else None, key_model_aliases), - granted_models=authorized_models, - llm_router=llm_router, - team_id=valid_token.team_id if valid_token else None, - ) _key_access_group_grants_model: Final = key_access_group_grants_model @@ -6060,7 +5891,6 @@ async def can_user_call_model( llm_router: Router | None, user_object: LiteLLM_UserTable | None, key_model_aliases: Mapping[str, str] | None = None, - teams_grant_access_group_models: bool = False, ) -> Literal[True]: if user_object is None: return True @@ -6078,42 +5908,13 @@ async def can_user_call_model( code=status.HTTP_403_FORBIDDEN, ) - user_models: Final = _user_models(user_object) - can_object_call_model( + return can_object_call_model( model=model, llm_router=llm_router, - models=list(user_models), + models=user_object.models, key_model_aliases=key_model_aliases, object_type="user", ) - if not teams_grant_access_group_models: - _require_user_access_group_grant( - model=model, - user_role=user_object.user_role, - user_models=user_models, - llm_router=llm_router, - key_model_aliases=key_model_aliases, - ) - return True - - -def _require_user_access_group_grant( - model: str | Sequence[str], - user_role: str | None, - user_models: Sequence[str], - llm_router: Router | None, - key_model_aliases: Mapping[str, str] | None, -) -> None: - if user_role == LitellmUserRoles.PROXY_ADMIN or not _model_access_group_deny_by_default_enabled(): - return - ungranted: Final = ungranted_access_group_models( - names=requested_model_names(model, key_model_aliases=key_model_aliases), - granted_models=user_models, - llm_router=llm_router, - team_id=None, - ) - if ungranted: - raise access_group_model_denied(model=model, ungranted=ungranted, object_type="user") def personal_key_team_model_access_applies(valid_token: UserAPIKeyAuth | None) -> bool: @@ -6127,6 +5928,121 @@ def personal_key_team_model_access_applies(valid_token: UserAPIKeyAuth | None) - ) +_MODEL_NAMES: Final[TypeAdapter[tuple[str, ...]]] = TypeAdapter(tuple[str, ...]) + + +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.""" @@ -6758,7 +6674,7 @@ async def can_caller_call_search_tool( async def can_token_call_search_tool(search_tool_name: str, valid_token: UserAPIKeyAuth) -> Literal[True]: """`can_caller_call_search_tool` against the proxy's own settings, team cache and database""" - from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + from litellm.proxy.proxy_server import general_settings, prisma_client, proxy_logging_obj, user_api_key_cache async def _load_team_object() -> LiteLLM_TeamTable | None: if not valid_token.team_id: @@ -6774,7 +6690,7 @@ async def can_token_call_search_tool(search_tool_name: str, valid_token: UserAPI return await can_caller_call_search_tool( search_tool_name=search_tool_name, valid_token=valid_token, - general_settings=_proxy_general_settings(), + general_settings=typed_general_settings(general_settings), load_team_object=_load_team_object, ) diff --git a/litellm/proxy/auth/model_access_group_grants.py b/litellm/proxy/auth/model_access_group_grants.py deleted file mode 100644 index b028d46a532..00000000000 --- a/litellm/proxy/auth/model_access_group_grants.py +++ /dev/null @@ -1,141 +0,0 @@ -from collections.abc import Iterable, Iterator, Mapping, Sequence -from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal - -from fastapi import status -from pydantic import ConfigDict, ValidationError - -import litellm -from litellm.proxy._types import ( - ConfigGeneralSettings, - ModelAccessDeniedProxyException, - ProxyErrorTypes, -) -from litellm.proxy.auth.model_access_denied import model_access_denied_client_message -from litellm.types.llms.base_model import LiteLLMBaseModel - -if TYPE_CHECKING: - from litellm.router import Router - -_SETTING: Final = "model_access_group_deny_by_default" - - -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 - model_info: _DeploymentModelInfo | None = None - - @property - def groups(self) -> frozenset[str]: - return frozenset(self.model_info.access_groups or ()) if self.model_info is not None else frozenset() - - -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 it fails closed - """ - value: Final = general_settings.get(_SETTING, False) - if isinstance(value, bool): - return value - try: - return ConfigGeneralSettings.model_validate( - MappingProxyType({_SETTING: value}) - ).model_access_group_deny_by_default - except ValidationError: - return True - - -def _alias_chain( - model: str, team_model_aliases: Mapping[str, str] | None, key_model_aliases: Mapping[str, str] | None -) -> Iterator[str]: - after_team: Final = team_model_aliases.get(model, model) if team_model_aliases else model - after_key: Final = key_model_aliases.get(after_team, after_team) if key_model_aliases else after_team - yield from (model, after_team, after_key, litellm.model_alias_map.get(after_key, after_key)) - - -def _alias_chains( - model: str | Sequence[str], - team_model_aliases: Mapping[str, str] | None, - key_model_aliases: Mapping[str, str] | None, -) -> Iterator[str]: - for requested in (model,) if isinstance(model, str) else model: - yield from _alias_chain(requested, team_model_aliases, key_model_aliases) - - -def requested_model_names( - model: str | Sequence[str], - team_model_aliases: Mapping[str, str] | None = None, - key_model_aliases: Mapping[str, str] | None = None, -) -> tuple[str, ...]: - """The requested models plus every model their team, key and proxy aliases point at""" - return tuple(dict.fromkeys(_alias_chains(model, team_model_aliases, key_model_aliases))) - - -def _rows(rows: Iterable[object]) -> Iterator[_DeploymentRow]: - for row in rows: - yield _DeploymentRow.model_validate(row) - - -def _routing_group_member_rows(name: str, llm_router: "Router", team_id: str | None) -> Iterator[_DeploymentRow]: - group: Final = llm_router.get_routing_group(llm_router.get_model_from_alias(model=name) or name) - for member in group.models if group is not None else (): - yield from _rows(llm_router.get_model_list(model_name=member, team_id=team_id) or ()) - - -def _reachable_deployments(name: str, llm_router: "Router", team_id: str | None) -> Iterator[_DeploymentRow]: - """Every deployment a request for `name` can be served by: its model group (aliases and wildcards - included), a routing group's members, a deployment id, or a deployment's litellm_params.model""" - yield from _rows(llm_router.get_model_list(model_name=name, team_id=team_id) or ()) - yield from _routing_group_member_rows(name, llm_router, team_id) - deployment: Final = llm_router.get_deployment(model_id=name) if llm_router.has_model_id(name) else None - if deployment is not None: - yield _DeploymentRow.model_validate(deployment.model_dump(exclude_none=True)) - if name in llm_router.deployment_names: - yield from _rows(row for row in llm_router.get_model_list() or () if row["litellm_params"].get("model") == name) - - -def ungranted_access_group_models( - names: Sequence[str], - granted_models: Sequence[str], - llm_router: "Router | None", - team_id: str | None, -) -> tuple[str, ...]: - """ - The names that can reach a deployment in a model access group which `granted_models` does not name, - by group name, the deployment's model name, or the requested name. Names match exactly, so a blanket - grant (empty list, "*", all-proxy-models, all-team-models) or a wildcard pattern names nothing - """ - if llm_router is None: - return () - explicit: Final = frozenset(granted_models) - return tuple(name for name in names if not _is_explicitly_granted(name, explicit, llm_router, team_id)) - - -def _is_explicitly_granted(name: str, explicit: frozenset[str], llm_router: "Router", team_id: str | None) -> bool: - return name in explicit or all( - not row.groups or row.model_name in explicit or not explicit.isdisjoint(row.groups) - for row in _reachable_deployments(name, llm_router, team_id) - ) - - -def access_group_model_denied( - model: str | Sequence[str], ungranted: Sequence[str], object_type: Literal["key", "user", "team"] -) -> ModelAccessDeniedProxyException: - return ModelAccessDeniedProxyException( - message=model_access_denied_client_message(model=model if isinstance(model, str) else list(model)), - internal_message=( - f"{object_type} not allowed to access model. {', '.join(ungranted)} reaches a model access group and " - f"general_settings.{_SETTING} requires the {object_type} to list that group or model explicitly" - ), - type=ProxyErrorTypes.get_model_access_error_type_for_object(object_type=object_type), - param="model", - code=status.HTTP_403_FORBIDDEN, - ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 27fa4731baf..9062a8db2a2 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -84,6 +84,7 @@ 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_grants, resolve_and_validate_end_user_id, resolve_default_end_user_budget, route_skips_budget_checks, @@ -3792,6 +3793,19 @@ 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" in general_settings + and general_settings["model_access_group_deny_by_default"] is True + ), + 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, diff --git a/litellm/router.py b/litellm/router.py index d0147b79285..400f65058ea 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -296,6 +296,7 @@ from litellm.router_utils.kubernetes_pod_discovery import ( resolve_pods_after, resolve_pods_after_bound, ) +from litellm.router_utils.model_access_group_grants import drop_ungranted_deployments, require_granted_deployment from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( MID_STREAM_CONTINUATION_MARKER, ContinuationOutput, @@ -1327,7 +1328,7 @@ class Router: self.assistants_config = assistants_config self.search_tools = search_tools or [] self.guardrail_list = guardrail_list or [] - self.deployment_names: list[str] = [] # names of models under litellm_params. ex. azure/chatgpt-v-2 + self.deployment_names: list = [] # names of models under litellm_params. ex. azure/chatgpt-v-2 self.deployment_latency_map = {} ### CACHING ### cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = "local" # default to an in-memory cache @@ -14323,6 +14324,22 @@ 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 isinstance(deployments, Sequence): + return resolved_model, list(drop_ungranted_deployments(deployments, model, request_kwargs)) + return resolved_model, cast( # cast-ok: contract requires a plain dict for a single deployment + dict[str, object], require_granted_deployment(deployments, model, request_kwargs) + ) + + 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. diff --git a/litellm/router_utils/model_access_group_grants.py b/litellm/router_utils/model_access_group_grants.py new file mode 100644 index 00000000000..270eacd10df --- /dev/null +++ b/litellm/router_utils/model_access_group_grants.py @@ -0,0 +1,84 @@ +""" +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 _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 require_granted_deployment( + deployment: _Deployment, requested_model: str, request_kwargs: Mapping[str, object] | None +) -> _Deployment: + grants: Final = request_model_access_group_grants(request_kwargs) + if grants is None or _is_granted(deployment, requested_model, grants): + return deployment + raise _denied(requested_model) + + +def drop_ungranted_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) + granted: Final = tuple(deployment for deployment in deployments if _is_granted(deployment, requested_model, grants)) + if deployments and not granted: + raise _denied(requested_model) + return granted diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index c2f7c487b7b..2f0b0eedd44 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -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: Any | None = None # Expanded user object when expand=user is used created_by_user: Any | 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 diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index e1afa81d372..e40e375faa5 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -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 diff --git a/tests/unit/proxy/auth/test_handle_jwt.py b/tests/unit/proxy/auth/test_handle_jwt.py index d0c800627ad..a0ea6ce0403 100644 --- a/tests/unit/proxy/auth/test_handle_jwt.py +++ b/tests/unit/proxy/auth/test_handle_jwt.py @@ -1416,7 +1416,6 @@ async def test_find_team_with_model_access_model_group(monkeypatch): proxy_server_module = types.ModuleType("proxy_server") proxy_server_module.llm_router = router - proxy_server_module.general_settings = {} monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) team = LiteLLM_TeamTable(team_id="team-1", models=["test-group"]) @@ -1468,7 +1467,6 @@ async def test_find_team_with_model_access_resolves_team_model_alias(monkeypatch router = Router(model_list=[{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}]) proxy_server_module = types.ModuleType("proxy_server") proxy_server_module.llm_router = router - proxy_server_module.general_settings = {} monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) team = LiteLLM_TeamTable( @@ -1528,7 +1526,6 @@ async def test_find_team_with_model_access_v1_messages_default_routes(monkeypatc proxy_server_module = types.ModuleType("proxy_server") proxy_server_module.llm_router = router - proxy_server_module.general_settings = {} monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) team = LiteLLM_TeamTable(team_id="coding-team", models=["coding_only_models"]) @@ -1600,7 +1597,6 @@ async def test_find_team_with_model_access_reports_team_route_allowlist_denial() prisma_team_table.find_unique = AsyncMock() proxy_server_module: Final = ModuleType("proxy_server") proxy_server_module.llm_router = None - proxy_server_module.general_settings = {} with patch.dict(sys.modules, {"litellm.proxy.proxy_server": proxy_server_module}): with pytest.raises(HTTPException) as exc_info: @@ -3711,7 +3707,6 @@ async def test_find_team_with_model_access_unresolved_group_claim_returns_none( ) proxy_server_module = types.ModuleType("proxy_server") proxy_server_module.llm_router = router - proxy_server_module.general_settings = {} monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) async def raise_404(*_args, **_kwargs): @@ -3849,7 +3844,6 @@ async def test_find_team_with_model_access_resolved_team_without_model_still_rai ) proxy_server_module = types.ModuleType("proxy_server") proxy_server_module.llm_router = router - proxy_server_module.general_settings = {} monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) team = LiteLLM_TeamTable(team_id="real-team", models=["gpt-3.5-turbo"]) @@ -3931,7 +3925,6 @@ async def test_find_team_with_model_access_unresolved_group_claim_default_raises ) proxy_server_module = types.ModuleType("proxy_server") proxy_server_module.llm_router = router - proxy_server_module.general_settings = {} monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) async def raise_404(*_args, **_kwargs): diff --git a/tests/unit/proxy/auth/test_model_access_group_grants.py b/tests/unit/proxy/auth/test_model_access_group_grants.py deleted file mode 100644 index 2ccd694c880..00000000000 --- a/tests/unit/proxy/auth/test_model_access_group_grants.py +++ /dev/null @@ -1,345 +0,0 @@ -from collections.abc import Iterator, Sequence -from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest -from fastapi import Request - -import litellm -from litellm.proxy import proxy_server -from litellm.proxy._types import ( - LiteLLM_AccessGroupTable, - LiteLLM_TeamTable, - LiteLLM_TeamTableCachedObj, - LiteLLM_UserTable, - LitellmUserRoles, - ModelAccessDeniedProxyException, - ProxyErrorTypes, - UserAPIKeyAuth, -) -from litellm.proxy.auth import auth_checks -from litellm.proxy.auth.auth_checks import ( - can_key_call_model, - can_key_call_resolved_model, - can_user_call_model, - common_checks, -) -from litellm.proxy.auth.model_access_group_grants import ungranted_access_group_models - -_ROUTER: Final = litellm.Router( - model_list=[ - { - "model_name": "private-model", - "litellm_params": {"model": "openai/private-only-upstream", "api_key": "sk-fake"}, - "model_info": {"id": "private-deployment-id", "access_groups": ["group-x"]}, - }, - { - "model_name": "openai/private-ft", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-fake"}, - "model_info": {"access_groups": ["group-x"]}, - }, - {"model_name": "public-model", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-fake"}}, - ], - model_group_alias={"alias-to-private": "private-model"}, - routing_groups=[ - {"group_name": "mixed-group", "models": ["private-model", "public-model"], "routing_strategy": "simple-shuffle"} - ], -) - - -@pytest.fixture -def deny_by_default(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: - monkeypatch.setitem(proxy_server.general_settings, "model_access_group_deny_by_default", True) - yield - - -@pytest.fixture -def no_team_membership(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: - async def _no_membership(**_: object) -> None: - return None - - monkeypatch.setattr(auth_checks, "get_team_membership", _no_membership) - yield - - -async def _common_checks( - model: str, - token: UserAPIKeyAuth, - team: LiteLLM_TeamTable | None = None, - user: LiteLLM_UserTable | None = None, -) -> bool: - return await common_checks( - request_body={"model": model}, - team_object=team, - user_object=user, - end_user_object=None, - global_proxy_spend=None, - general_settings=proxy_server.general_settings, - route="/chat/completions", - llm_router=_ROUTER, - proxy_logging_obj=MagicMock(), - valid_token=token, - request=MagicMock(spec=Request), - ) - - -def _team(models: Sequence[str], access_group_ids: Sequence[str] = ()) -> LiteLLM_TeamTable: - return LiteLLM_TeamTable(team_id="team-1", models=list(models), access_group_ids=list(access_group_ids)) - - -_TEAM_KEY: Final = UserAPIKeyAuth(token="hashed-team-key", team_id="team-1") - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "team_models, model, allowed", - [ - ([], "private-model", False), - (["*"], "private-model", False), - (["all-proxy-models"], "private-model", False), - (["openai/*"], "openai/private-ft", False), - (["group-x"], "private-model", True), - (["private-model"], "private-model", True), - ([], "public-model", True), - ], -) -async def test_team_needs_an_explicit_grant_for_access_group_models( - deny_by_default: None, no_team_membership: None, team_models: list[str], model: str, allowed: bool -) -> None: - if allowed: - assert await _common_checks(model, _TEAM_KEY, team=_team(team_models)) is True - return - with pytest.raises(ModelAccessDeniedProxyException) as exc_info: - await _common_checks(model, _TEAM_KEY, team=_team(team_models)) - assert exc_info.value.type == ProxyErrorTypes.team_model_access_denied - assert str(exc_info.value.code) == "403" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("team_models, model", [([], "private-model"), (["openai/*"], "openai/private-ft")]) -async def test_blanket_team_grant_is_unchanged_when_setting_is_off( - monkeypatch: pytest.MonkeyPatch, no_team_membership: None, team_models: list[str], model: str -) -> None: - monkeypatch.delitem(proxy_server.general_settings, "model_access_group_deny_by_default", raising=False) - assert await _common_checks(model, _TEAM_KEY, team=_team(team_models)) is True - - -@pytest.mark.asyncio -@pytest.mark.parametrize("model", ["private-deployment-id", "openai/private-only-upstream", "mixed-group"]) -async def test_blanket_team_cannot_reach_an_access_group_deployment_by_another_name( - deny_by_default: None, no_team_membership: None, model: str -) -> None: - with pytest.raises(ModelAccessDeniedProxyException): - await _common_checks(model, _TEAM_KEY, team=_team([])) - - -@pytest.mark.parametrize( - "name, granted, ungranted", - [ - ("private-deployment-id", ["private-model"], ()), - ("private-deployment-id", ["private-deployment-id"], ()), - ("openai/private-only-upstream", ["group-x"], ()), - ("mixed-group", ["private-model"], ()), - ("mixed-group", ["public-model"], ("mixed-group",)), - ("alias-to-private", ["public-model"], ("alias-to-private",)), - ("alias-to-private", ["group-x"], ()), - ("mixed-group", ["*", "all-proxy-models", "openai/*"], ("mixed-group",)), - ], -) -def test_another_name_for_a_deployment_is_granted_by_its_model_or_group( - name: str, granted: list[str], ungranted: tuple[str, ...] -) -> None: - assert ungranted_access_group_models(names=[name], granted_models=granted, llm_router=_ROUTER, team_id=None) == ( - ungranted - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("team_models, allowed", [([], False), (["group-x"], True)]) -async def test_team_alias_cannot_reach_an_ungranted_access_group_model( - deny_by_default: None, no_team_membership: None, team_models: list[str], allowed: bool -) -> None: - token: Final = UserAPIKeyAuth( - token="hashed-team-key", team_id="team-1", team_model_aliases={"nice-name": "private-model"} - ) - if allowed: - assert await _common_checks("nice-name", token, team=_team(team_models)) is True - return - with pytest.raises(ModelAccessDeniedProxyException): - await _common_checks("nice-name", token, team=_team(team_models)) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("granted_names, allowed", [(["private-model"], True), (["public-model"], False)]) -async def test_team_unified_access_group_counts_as_an_explicit_grant( - deny_by_default: None, no_team_membership: None, granted_names: list[str], allowed: bool -) -> None: - group: Final = LiteLLM_AccessGroupTable( - access_group_id="ag-1", access_group_name="private", access_model_names=granted_names - ) - with patch.object(auth_checks, "get_access_object", new_callable=AsyncMock, return_value=group): - team: Final = _team([], access_group_ids=["ag-1"]) - if allowed: - assert await _common_checks("private-model", _TEAM_KEY, team=team) is True - return - with pytest.raises(ModelAccessDeniedProxyException): - await _common_checks("private-model", _TEAM_KEY, team=team) - - -@pytest.mark.asyncio -async def test_empty_key_list_on_a_team_key_defers_to_the_team_grant(deny_by_default: None) -> None: - token: Final = UserAPIKeyAuth(token="hashed-team-key", team_id="team-1", models=[]) - assert await can_key_call_model(model="private-model", llm_model_list=None, valid_token=token, llm_router=_ROUTER) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "user, allowed", - [ - (LiteLLM_UserTable(user_id="u1", models=[]), False), - (LiteLLM_UserTable(user_id="u1", models=["group-x"]), True), - (LiteLLM_UserTable(user_id="u1", models=[], user_role=LitellmUserRoles.PROXY_ADMIN.value), True), - ], -) -async def test_personal_key_needs_an_explicit_user_grant( - deny_by_default: None, user: LiteLLM_UserTable, allowed: bool -) -> None: - token: Final = UserAPIKeyAuth(token="hashed-personal-key", user_id="u1") - if allowed: - assert await _common_checks("private-model", token, user=user) is True - return - with pytest.raises(ModelAccessDeniedProxyException) as exc_info: - await _common_checks("private-model", token, user=user) - assert exc_info.value.type == ProxyErrorTypes.user_model_access_denied - - -@pytest.mark.asyncio -async def test_user_grant_is_skipped_when_owner_teams_decide(deny_by_default: None) -> None: - assert await can_user_call_model( - model="private-model", - llm_router=_ROUTER, - user_object=LiteLLM_UserTable(user_id="u1", models=[]), - teams_grant_access_group_models=True, - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("owner_models, allowed", [([], False), (["private-model"], True)]) -async def test_resolved_model_for_a_personal_key_needs_the_owner_grant( - monkeypatch: pytest.MonkeyPatch, deny_by_default: None, owner_models: list[str], allowed: bool -) -> None: - async def _get_user_object(**_: object) -> LiteLLM_UserTable: - return LiteLLM_UserTable(user_id="u1", models=owner_models) - - monkeypatch.setattr(auth_checks, "get_user_object", _get_user_object) - token: Final = UserAPIKeyAuth(token="hashed-personal-key", user_id="u1") - if allowed: - await can_key_call_resolved_model( - model="private-model", llm_model_list=None, valid_token=token, llm_router=_ROUTER - ) - return - with pytest.raises(ModelAccessDeniedProxyException): - await can_key_call_resolved_model( - model="private-model", llm_model_list=None, valid_token=token, llm_router=_ROUTER - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "token, allowed", - [ - (UserAPIKeyAuth(token="hashed-service-key", models=[]), False), - (UserAPIKeyAuth(token="hashed-service-key", models=["group-x"]), True), - (UserAPIKeyAuth(token="hashed-service-key", models=[], user_role=LitellmUserRoles.PROXY_ADMIN), True), - ], -) -async def test_key_with_no_team_or_user_owns_its_grant( - deny_by_default: None, token: UserAPIKeyAuth, allowed: bool -) -> None: - if allowed: - assert await can_key_call_model( - model="private-model", llm_model_list=None, valid_token=token, llm_router=_ROUTER - ) - return - with pytest.raises(ModelAccessDeniedProxyException) as exc_info: - await can_key_call_model(model="private-model", llm_model_list=None, valid_token=token, llm_router=_ROUTER) - assert exc_info.value.type == ProxyErrorTypes.key_model_access_denied - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "token, allowed", - [ - (UserAPIKeyAuth(token="hashed-service-key", models=["all-team-models"]), False), - (UserAPIKeyAuth(token="hashed-personal-key", user_id="missing-user", models=[]), False), - (UserAPIKeyAuth(token="hashed-personal-key", user_id="missing-user", models=["group-x"]), True), - ], -) -async def test_teamless_key_without_a_loaded_owner_needs_its_own_grant( - deny_by_default: None, token: UserAPIKeyAuth, allowed: bool -) -> None: - if allowed: - assert await _common_checks("private-model", token) is True - return - with pytest.raises(ModelAccessDeniedProxyException) as exc_info: - await _common_checks("private-model", token) - assert exc_info.value.type == ProxyErrorTypes.key_model_access_denied - - -@pytest.mark.asyncio -async def test_resolved_model_for_an_all_team_models_key_with_no_team_needs_its_own_grant( - deny_by_default: None, -) -> None: - token: Final = UserAPIKeyAuth(token="hashed-service-key", models=["all-team-models"]) - with pytest.raises(ModelAccessDeniedProxyException): - await can_key_call_resolved_model( - model="private-model", llm_model_list=None, valid_token=token, llm_router=_ROUTER - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("key_group_models, allowed", [(["*"], False), (["private-model"], True)]) -async def test_key_access_group_overrides_the_team_only_with_an_explicit_grant( - monkeypatch: pytest.MonkeyPatch, - deny_by_default: None, - no_team_membership: None, - key_group_models: list[str], - allowed: bool, -) -> None: - monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) - group: Final = LiteLLM_AccessGroupTable( - access_group_id="key-ag", - access_group_name="key-group", - access_model_names=key_group_models, - assigned_key_ids=["hashed-team-key"], - ) - token: Final = UserAPIKeyAuth(token="hashed-team-key", team_id="team-1", access_group_ids=["key-ag"]) - with patch.object(auth_checks, "get_access_object", new_callable=AsyncMock, return_value=group): - if allowed: - assert await _common_checks("private-model", token, team=_team([])) is True - return - with pytest.raises(ModelAccessDeniedProxyException): - await _common_checks("private-model", token, team=_team([])) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("owner_team_models, allowed", [(["group-x"], True), ([], False)]) -async def test_personal_key_reaches_an_access_group_model_through_the_owner_team_grant( - monkeypatch: pytest.MonkeyPatch, - deny_by_default: None, - no_team_membership: None, - owner_team_models: list[str], - allowed: bool, -) -> None: - async def _get_team_object(**_: object) -> LiteLLM_TeamTableCachedObj: - return LiteLLM_TeamTableCachedObj(team_id="team-1", models=owner_team_models) - - monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) - monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) - token: Final = UserAPIKeyAuth(token="hashed-personal-key", user_id="u1") - owner: Final = LiteLLM_UserTable(user_id="u1", models=[], teams=["team-1"]) - if allowed: - assert await _common_checks("private-model", token, user=owner) is True - return - with pytest.raises(ModelAccessDeniedProxyException): - await _common_checks("private-model", token, user=owner) diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 4cdd7fefd28..22a94d14077 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -5202,6 +5202,60 @@ 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"}, None), + ({}, None), + ], + ids=["deny-by-default", "setting-not-true", "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) + + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock + ), # test-quality-ok: stubs the sibling check so only the grant stamping is under test + patch( # test-quality-ok: stubs the budget reservation so only the grant stamping is under test + "litellm.proxy.auth.user_api_key_auth._reserve_budget_after_common_checks", new_callable=AsyncMock + ), + ): + 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 @@ -12480,11 +12534,14 @@ async def test_lens_oauth2_validates_identity_without_inference_user_field( patch("litellm.proxy.proxy_server.master_key", "sk-fixture-master"), patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), - patch.dict(os.environ, { - "OAUTH_TOKEN_INFO_ENDPOINT": "https://oauth.fixture.test/introspect", - "OAUTH_CLIENT_ID": "fixture-client", - "OAUTH_CLIENT_SECRET": "fixture-client-secret", - }), + patch.dict( + os.environ, + { + "OAUTH_TOKEN_INFO_ENDPOINT": "https://oauth.fixture.test/introspect", + "OAUTH_CLIENT_ID": "fixture-client", + "OAUTH_CLIENT_SECRET": "fixture-client-secret", + }, + ), patch("httpx.AsyncClient.send", side_effect=token_info) as network, ): if active: @@ -12611,6 +12668,4 @@ def test_should_skip_budget_checks_spares_only_zero_spend_mcp_discovery( ) -> None: from litellm.proxy.auth.user_api_key_auth import _should_skip_budget_checks - assert ( - _should_skip_budget_checks(request_data=request_data, route=route, request=None, llm_router=None) is expected - ) + assert _should_skip_budget_checks(request_data=request_data, route=route, request=None, llm_router=None) is expected diff --git a/tests/unit/router_utils/test_model_access_group_grants.py b/tests/unit/router_utils/test_model_access_group_grants.py new file mode 100644 index 00000000000..20738aa182b --- /dev/null +++ b/tests/unit/router_utils/test_model_access_group_grants.py @@ -0,0 +1,129 @@ +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 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 1664374e37d..fdb0e67a4a2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30150,7 +30150,7 @@ export interface components { 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. Untagged models are unaffected, and proxy admins' personal access is exempt + * @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;