mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
feat(proxy): add key-level disallowed_models column and enforce it at auth and listing
This commit is contained in:
Jesus Velarde
parent
6fa0477d84
commit
3376197267
14 changed files with 494 additions and 56 deletions
|
|
@ -0,0 +1,7 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "disallowed_models" TEXT[] DEFAULT ARRAY[]::TEXT[];
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "disallowed_models" TEXT[] DEFAULT ARRAY[]::TEXT[];
|
||||
|
|
@ -530,6 +530,7 @@ model LiteLLM_VerificationToken {
|
|||
key_type String?
|
||||
policies String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
disallowed_models String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
model_max_budget Json @default("{}")
|
||||
budget_fallbacks Json @default("{}")
|
||||
|
|
@ -632,6 +633,7 @@ model LiteLLM_DeletedVerificationToken {
|
|||
key_type String?
|
||||
policies String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
disallowed_models String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
model_max_budget Json @default("{}")
|
||||
budget_fallbacks Json @default("{}")
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
|
|||
object_permission_id: str | None = None
|
||||
object_permission: LiteLLM_ObjectPermissionTable | None = None
|
||||
access_group_ids: list[str] | None = None
|
||||
disallowed_models: tuple[str, ...] | None = None
|
||||
rotation_count: int | None = 0
|
||||
auto_rotate: bool | None = False
|
||||
rotation_interval: str | None = None
|
||||
|
|
|
|||
|
|
@ -100,6 +100,7 @@ from litellm.proxy.auth.budget_throttle import (
|
|||
budget_throttle_percentage,
|
||||
should_throttle_budget_exceeded,
|
||||
)
|
||||
from litellm.proxy.auth.key_disallowed_models import ensure_key_may_call_model
|
||||
from litellm.proxy.auth.model_access_denied import (
|
||||
customer_model_access_denied_client_message,
|
||||
model_access_denied_client_message,
|
||||
|
|
@ -5521,6 +5522,7 @@ async def can_key_call_model(
|
|||
Raises:
|
||||
- Exception: If token not allowed to call model
|
||||
"""
|
||||
ensure_key_may_call_model(model, valid_token, llm_router)
|
||||
key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token)
|
||||
try:
|
||||
return can_object_call_model(
|
||||
|
|
@ -5596,7 +5598,9 @@ async def can_key_call_resolved_model(
|
|||
skip_key_model_check: Final = valid_token.config or (
|
||||
isinstance(valid_token.models, list) and SpecialModelNames.all_team_models.value in valid_token.models
|
||||
)
|
||||
if not skip_key_model_check:
|
||||
if skip_key_model_check:
|
||||
ensure_key_may_call_model(model, valid_token, llm_router)
|
||||
else:
|
||||
await can_key_call_model(
|
||||
model=model,
|
||||
llm_model_list=llm_model_list,
|
||||
|
|
|
|||
80
litellm/proxy/auth/key_disallowed_models.py
Normal file
80
litellm/proxy/auth/key_disallowed_models.py
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
from collections.abc import Collection, Iterable
|
||||
from typing import Final
|
||||
|
||||
from fastapi import status
|
||||
from pydantic import BaseModel
|
||||
|
||||
import litellm
|
||||
from litellm.proxy._types import ModelAccessDeniedProxyException, ProxyErrorTypes, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.model_access_denied import model_access_denied_client_message
|
||||
from litellm.proxy.common_utils.model_listing_utils import alias_map
|
||||
from litellm.router_api import RegistryReader
|
||||
|
||||
|
||||
def key_disallowed_models(valid_token: UserAPIKeyAuth) -> frozenset[str]:
|
||||
return frozenset(valid_token.disallowed_models or ())
|
||||
|
||||
|
||||
def _names_model_resolves_through(
|
||||
model: str, valid_token: UserAPIKeyAuth, llm_router: RegistryReader | None
|
||||
) -> frozenset[str]:
|
||||
team_aliases: Final = alias_map(valid_token.team_model_aliases) if valid_token.team_model_aliases else {}
|
||||
key_aliases: Final = alias_map(valid_token.aliases) if valid_token.aliases else {}
|
||||
after_team_alias: Final = team_aliases.get(model, model)
|
||||
after_key_alias: Final = key_aliases.get(after_team_alias, after_team_alias)
|
||||
after_global_alias: Final = litellm.model_alias_map.get(after_key_alias, after_key_alias)
|
||||
dispatched: Final = key_aliases.get(after_global_alias, after_global_alias)
|
||||
router_alias_target: Final = (
|
||||
llm_router.registry().resolve(model, kind="alias")
|
||||
if llm_router is not None and model in llm_router.registry().aliases()
|
||||
else None
|
||||
)
|
||||
return frozenset(
|
||||
name
|
||||
for name in (model, after_team_alias, after_key_alias, after_global_alias, dispatched, router_alias_target)
|
||||
if name is not None
|
||||
)
|
||||
|
||||
|
||||
def first_disallowed_model(
|
||||
model: str | Iterable[str], valid_token: UserAPIKeyAuth, llm_router: RegistryReader | None
|
||||
) -> str | None:
|
||||
disallowed: Final = key_disallowed_models(valid_token)
|
||||
if not disallowed:
|
||||
return None
|
||||
requested: Final = (model,) if isinstance(model, str) else tuple(model)
|
||||
return next(
|
||||
(m for m in requested if _names_model_resolves_through(m, valid_token, llm_router) & disallowed),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def ensure_key_may_call_model(
|
||||
model: str | Iterable[str], valid_token: UserAPIKeyAuth, llm_router: RegistryReader | None
|
||||
) -> None:
|
||||
denied: Final = first_disallowed_model(model, valid_token, llm_router)
|
||||
if denied is None:
|
||||
return
|
||||
raise ModelAccessDeniedProxyException(
|
||||
message=model_access_denied_client_message(model=denied),
|
||||
internal_message=f"Key is not allowed to call model={denied}: it is listed in the key's disallowed_models.",
|
||||
type=ProxyErrorTypes.key_model_access_denied,
|
||||
param="model",
|
||||
code=status.HTTP_403_FORBIDDEN,
|
||||
)
|
||||
|
||||
|
||||
def without_key_disallowed_models(
|
||||
models: Iterable[str], valid_token: UserAPIKeyAuth, llm_router: RegistryReader | None
|
||||
) -> list[str]:
|
||||
disallowed: Final = key_disallowed_models(valid_token)
|
||||
if not disallowed:
|
||||
return list(models)
|
||||
return [m for m in models if not _names_model_resolves_through(m, valid_token, llm_router) & disallowed]
|
||||
|
||||
|
||||
def disallowed_models_removed_by(data: BaseModel | None, stored: Collection[str] | None) -> frozenset[str]:
|
||||
if data is None or "disallowed_models" not in data.model_fields_set:
|
||||
return frozenset()
|
||||
requested: Final[Collection[str]] = getattr(data, "disallowed_models", None) or ()
|
||||
return frozenset(stored or ()) - frozenset(requested)
|
||||
|
|
@ -123,6 +123,7 @@ from litellm.proxy.auth.auth_utils import (
|
|||
route_in_additonal_public_routes,
|
||||
)
|
||||
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
|
||||
from litellm.proxy.auth.key_disallowed_models import ensure_key_may_call_model, key_disallowed_models
|
||||
from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_context
|
||||
from litellm.proxy.auth.oauth2_check import Oauth2Handler
|
||||
from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request
|
||||
|
|
@ -4628,49 +4629,50 @@ async def enforce_key_and_fallback_model_access(
|
|||
"""
|
||||
await _normalize_claude_model(request_data, valid_token, request, route)
|
||||
await _resolve_router_settings_model_group_alias(request_data, valid_token, request, route)
|
||||
config: Final = valid_token.config
|
||||
skips_key_allowlist: Final = valid_token.config != {} or (
|
||||
isinstance(valid_token.models, list) and "all-team-models" in valid_token.models
|
||||
)
|
||||
if skips_key_allowlist and not key_disallowed_models(valid_token):
|
||||
return
|
||||
|
||||
if config != {}:
|
||||
model_list: Final = config.get("model_list", [])
|
||||
new_model_list: Final = model_list
|
||||
verbose_proxy_logger.debug("\n new llm router model list %s", new_model_list)
|
||||
elif isinstance(valid_token.models, list) and "all-team-models" in valid_token.models:
|
||||
pass
|
||||
else:
|
||||
model: Final = _get_model_from_request_context(
|
||||
request_data=request_data,
|
||||
route=route,
|
||||
request=request,
|
||||
model: Final = _get_model_from_request_context(
|
||||
request_data=request_data,
|
||||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
fallback_names: Final = tuple(
|
||||
name
|
||||
for target in iter_request_fallback_targets(request_data)
|
||||
if (name := fallback_target_model_name(target)) is not None
|
||||
)
|
||||
|
||||
if skips_key_allowlist:
|
||||
requested_models: Final = () if model is None else (model,) if isinstance(model, str) else tuple(model)
|
||||
ensure_key_may_call_model((*requested_models, *fallback_names), valid_token, llm_router)
|
||||
return
|
||||
|
||||
if model is not None:
|
||||
await can_key_call_model(
|
||||
model=model,
|
||||
llm_model_list=llm_model_list,
|
||||
valid_token=valid_token,
|
||||
llm_router=llm_router,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
|
||||
if model is not None:
|
||||
await can_key_call_model(
|
||||
model=model,
|
||||
llm_model_list=llm_model_list,
|
||||
valid_token=valid_token,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
fallback_names: Final = tuple(
|
||||
name
|
||||
for target in iter_request_fallback_targets(request_data)
|
||||
if (name := fallback_target_model_name(target)) is not None
|
||||
for _name in dict.fromkeys(fallback_names): # dedupe, preserve order
|
||||
await can_key_call_model(
|
||||
model=_name,
|
||||
llm_model_list=llm_model_list,
|
||||
valid_token=valid_token,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
await is_valid_fallback_model(
|
||||
model=_name,
|
||||
llm_router=llm_router,
|
||||
user_model=None,
|
||||
)
|
||||
|
||||
for _name in dict.fromkeys(fallback_names): # dedupe, preserve order
|
||||
await can_key_call_model(
|
||||
model=_name,
|
||||
llm_model_list=llm_model_list,
|
||||
valid_token=valid_token,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
await is_valid_fallback_model(
|
||||
model=_name,
|
||||
llm_router=llm_router,
|
||||
user_model=None,
|
||||
)
|
||||
|
||||
|
||||
_enforce_key_and_fallback_model_access: Final = enforce_key_and_fallback_model_access # pyright: ignore[reportUnknownVariableType] # existing tests use this alias
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@ from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypeVar, c
|
|||
import fastapi
|
||||
import yaml
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status
|
||||
from pydantic import TypeAdapter
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
|
|
@ -62,6 +62,7 @@ from litellm.proxy.auth.auth_utils import (
|
|||
enforce_batch_limits_are_admin_only,
|
||||
enforce_output_token_estimates_are_admin_only,
|
||||
)
|
||||
from litellm.proxy.auth.key_disallowed_models import disallowed_models_removed_by
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
|
||||
evict_and_broadcast,
|
||||
|
|
@ -3274,6 +3275,37 @@ async def _check_single_key_update_team_permissions(
|
|||
)
|
||||
|
||||
|
||||
async def check_disallowed_models_caller_permission(
|
||||
data: BaseModel | None,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
existing_key_row: LiteLLM_VerificationToken,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
) -> None:
|
||||
removed: Final = disallowed_models_removed_by(data, existing_key_row.disallowed_models)
|
||||
if not removed or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||||
return
|
||||
team_table: Final = (
|
||||
await get_team_object(
|
||||
team_id=existing_key_row.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
check_db_only=True,
|
||||
)
|
||||
if existing_key_row.team_id is not None and prisma_client is not None
|
||||
else None
|
||||
)
|
||||
if team_table is not None and is_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_table):
|
||||
return
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail={
|
||||
"error": "Only proxy admins and admins of the key's team can remove models from a key's "
|
||||
f"`disallowed_models`. Removed: {sorted(removed)}"
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def _process_single_key_update(
|
||||
update_key_request: UpdateKeyRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
@ -3348,6 +3380,9 @@ async def _process_single_key_update(
|
|||
user_api_key_dict,
|
||||
existing_metadata=existing_key_row.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict
|
||||
)
|
||||
await check_disallowed_models_caller_permission(
|
||||
update_key_request, user_api_key_dict, existing_key_row, prisma_client, user_api_key_cache
|
||||
)
|
||||
|
||||
# Custom key update hook
|
||||
if user_custom_key_update is not None:
|
||||
|
|
@ -3894,6 +3929,9 @@ async def _validate_update_key_data(
|
|||
user_api_key_dict,
|
||||
existing_metadata=existing_key_row.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict
|
||||
)
|
||||
await check_disallowed_models_caller_permission(
|
||||
data, user_api_key_dict, existing_key_row, checked_prisma_client, user_api_key_cache
|
||||
)
|
||||
|
||||
# Check team limits if key has a team_id (from request or existing key)
|
||||
team_obj: LiteLLM_TeamTableCachedObj | None = None
|
||||
|
|
@ -5298,6 +5336,7 @@ async def generate_key_helper_fn(
|
|||
rotation_interval: str | None = None,
|
||||
router_settings: dict[str, object] | None = None,
|
||||
access_group_ids: list[str] | None = None,
|
||||
disallowed_models: Sequence[str] | None = None,
|
||||
budget_limits: list | None = None, # multiple concurrent budget windows
|
||||
*,
|
||||
llm_router: Router | None = None,
|
||||
|
|
@ -5434,6 +5473,7 @@ async def generate_key_helper_fn(
|
|||
"object_permission_id": object_permission_id,
|
||||
"router_settings": router_settings_json,
|
||||
"access_group_ids": access_group_ids or [],
|
||||
"disallowed_models": list(disallowed_models or ()),
|
||||
}
|
||||
|
||||
# Add rotation fields if auto_rotate is enabled
|
||||
|
|
@ -6677,6 +6717,9 @@ async def regenerate_key_fn(
|
|||
user_api_key_dict,
|
||||
existing_metadata=_key_in_db.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict
|
||||
)
|
||||
await check_disallowed_models_caller_permission(
|
||||
data, user_api_key_dict, _key_in_db, prisma_client, user_api_key_cache
|
||||
)
|
||||
|
||||
if data is not None and (data.access_group_ids or data.object_permission is not None):
|
||||
regenerate_team_table: LiteLLM_TeamTableCachedObj | None = None
|
||||
|
|
|
|||
|
|
@ -452,6 +452,7 @@ from litellm.proxy.auth.auth_utils import (
|
|||
from litellm.proxy.auth.fallback_budget import RouterFallbackBudgetCheck, router_fallback_budget_check
|
||||
from litellm.proxy.auth.fallback_model_access import router_fallback_access_check
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
from litellm.proxy.auth.key_disallowed_models import key_disallowed_models, without_key_disallowed_models
|
||||
from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY, LicenseCheck
|
||||
from litellm.proxy.auth.login_throttle import (
|
||||
LoginThrottle,
|
||||
|
|
@ -18311,7 +18312,7 @@ async def _v1_model_info_lists_every_deployment(
|
|||
model_access_groups=model_access_groups,
|
||||
):
|
||||
return False
|
||||
if personal_key_team_model_access_applies(user_api_key_dict):
|
||||
if personal_key_team_model_access_applies(user_api_key_dict) or key_disallowed_models(user_api_key_dict):
|
||||
return False
|
||||
return not user_api_key_dict.agent_id or await resolve_agent_ceiling(user_api_key_dict.agent_id) is None
|
||||
|
||||
|
|
@ -18915,17 +18916,21 @@ async def model_group_info(
|
|||
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
|
||||
)
|
||||
all_models_str: Final = (
|
||||
await run_model_listing(
|
||||
partial(
|
||||
get_complete_model_list,
|
||||
key_models=(),
|
||||
team_models=(),
|
||||
proxy_model_list=llm_router.registry().names(),
|
||||
user_model=user_model,
|
||||
infer_model_from_keys=general_settings.get("infer_model_from_keys", False),
|
||||
return_wildcard_routes=False,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
without_key_disallowed_models(
|
||||
await run_model_listing(
|
||||
partial(
|
||||
get_complete_model_list,
|
||||
key_models=(),
|
||||
team_models=(),
|
||||
proxy_model_list=llm_router.registry().names(),
|
||||
user_model=user_model,
|
||||
infer_model_from_keys=general_settings.get("infer_model_from_keys", False),
|
||||
return_wildcard_routes=False,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
),
|
||||
user_api_key_dict,
|
||||
llm_router,
|
||||
)
|
||||
if is_proxy_admin
|
||||
else await get_available_models_for_user(
|
||||
|
|
|
|||
|
|
@ -530,6 +530,7 @@ model LiteLLM_VerificationToken {
|
|||
key_type String?
|
||||
policies String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
disallowed_models String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
model_max_budget Json @default("{}")
|
||||
budget_fallbacks Json @default("{}")
|
||||
|
|
@ -632,6 +633,7 @@ model LiteLLM_DeletedVerificationToken {
|
|||
key_type String?
|
||||
policies String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
disallowed_models String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
model_max_budget Json @default("{}")
|
||||
budget_fallbacks Json @default("{}")
|
||||
|
|
|
|||
|
|
@ -10080,6 +10080,7 @@ async def get_available_models_for_user(
|
|||
personal_key_team_model_access_applies,
|
||||
personal_key_team_visible_models,
|
||||
)
|
||||
from litellm.proxy.auth.key_disallowed_models import without_key_disallowed_models
|
||||
from litellm.proxy.auth.model_checks import (
|
||||
get_complete_model_list,
|
||||
get_key_models,
|
||||
|
|
@ -10204,10 +10205,8 @@ async def get_available_models_for_user(
|
|||
team_id=effective_team_id,
|
||||
resolve_agent_ceiling=resolve_agent_ceiling,
|
||||
)
|
||||
if agent_visible is None:
|
||||
return all_models
|
||||
capped: Final = [m for m in all_models if m in agent_visible]
|
||||
return capped
|
||||
visible_models: Final = all_models if agent_visible is None else [m for m in all_models if m in agent_visible]
|
||||
return without_key_disallowed_models(visible_models, user_api_key_dict, llm_router)
|
||||
|
||||
|
||||
def _safe_get_model_info(model: str, get_model_info: Callable[[str], ModelInfo]) -> ModelInfo | None:
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ class KeyRequestBase(GenerateRequestBase):
|
|||
allowed_routes: list[Any] | None = []
|
||||
allowed_passthrough_routes: list[object] | None = None
|
||||
denied_passthrough_routes: list[str] | None = None
|
||||
disallowed_models: list[str] | None = None
|
||||
allowed_vector_store_indexes: list[AllowedVectorStoreIndexItem] | None = None
|
||||
rpm_limit_type: Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] | None = (
|
||||
None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm
|
||||
|
|
|
|||
|
|
@ -530,6 +530,7 @@ model LiteLLM_VerificationToken {
|
|||
key_type String?
|
||||
policies String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
disallowed_models String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
model_max_budget Json @default("{}")
|
||||
budget_fallbacks Json @default("{}")
|
||||
|
|
@ -632,6 +633,7 @@ model LiteLLM_DeletedVerificationToken {
|
|||
key_type String?
|
||||
policies String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
disallowed_models String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
model_max_budget Json @default("{}")
|
||||
budget_fallbacks Json @default("{}")
|
||||
|
|
|
|||
172
tests/integration/authorization/test_key_disallowed_models.py
Normal file
172
tests/integration/authorization/test_key_disallowed_models.py
Normal file
|
|
@ -0,0 +1,172 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import JsonValue
|
||||
|
||||
from tests.integration._support.client import (
|
||||
JSON_OBJECT,
|
||||
Gateway,
|
||||
Scenario,
|
||||
eventually,
|
||||
list_value,
|
||||
object_value,
|
||||
string_value,
|
||||
)
|
||||
from tests.integration._support.database import read_rows
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _TeamGrant:
|
||||
team_id: str
|
||||
allowed: str
|
||||
disallowed: str
|
||||
|
||||
|
||||
def _team_grant(scenario: Scenario) -> _TeamGrant:
|
||||
group: Final = f"integration-group-{uuid.uuid4().hex[:10]}"
|
||||
allowed: Final = scenario.model(model_info={"access_groups": [group]})
|
||||
disallowed: Final = scenario.model(model_info={"access_groups": [group]})
|
||||
return _TeamGrant(team_id=scenario.team(models=[group]), allowed=allowed, disallowed=disallowed)
|
||||
|
||||
|
||||
def _observed_models(gateway: Gateway) -> tuple[str, ...]:
|
||||
response: Final = httpx.get(f"{gateway.upstream_url}/__observations", timeout=10)
|
||||
response.raise_for_status()
|
||||
requests: Final = list_value(JSON_OBJECT.validate_json(response.content)["requests"])
|
||||
return tuple(str(object_value(object_value(entry).get("body") or {}).get("model")) for entry in requests)
|
||||
|
||||
|
||||
def _chat(gateway: Gateway, key: str, body: dict[str, JsonValue]) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"messages": [{"role": "user", "content": f"disallowed {uuid.uuid4().hex}"}], **body},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def _assert_served(gateway: Gateway, key: str, model: str) -> None:
|
||||
response: Final = _chat(gateway, key, {"model": model})
|
||||
assert response.status_code == 200, f"{model}: {response.status_code} {response.text}"
|
||||
|
||||
|
||||
def _assert_denied(gateway: Gateway, key: str, body: dict[str, JsonValue], denied: str) -> None:
|
||||
before: Final = _observed_models(gateway)
|
||||
response: Final = _chat(gateway, key, body)
|
||||
assert response.status_code == 403, f"{response.status_code} {response.text}"
|
||||
error: Final = object_value(JSON_OBJECT.validate_json(response.content)["error"])
|
||||
assert error["type"] == "key_model_access_denied", response.text
|
||||
assert error["param"] == "model", response.text
|
||||
assert denied in string_value(error["message"]), response.text
|
||||
assert _observed_models(gateway) == before
|
||||
|
||||
|
||||
def _listed_models(gateway: Gateway, key: str) -> frozenset[str]:
|
||||
response: Final = gateway.request("GET", "/v1/models", key=key)
|
||||
assert response.status_code == 200, response.text
|
||||
entries: Final = list_value(JSON_OBJECT.validate_json(response.content)["data"])
|
||||
return frozenset(string_value(object_value(entry)["id"]) for entry in entries)
|
||||
|
||||
|
||||
def _stored_disallowed_models(key: str) -> list[dict[str, JsonValue]]:
|
||||
return read_rows(
|
||||
'SELECT disallowed_models FROM "LiteLLM_VerificationToken" WHERE token=%s',
|
||||
(sha256(key.encode()).hexdigest(),),
|
||||
)
|
||||
|
||||
|
||||
def test_a_key_inheriting_a_team_access_group_cannot_call_or_list_its_disallowed_models(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
grant: Final = _team_grant(scenario)
|
||||
key: Final = scenario.key(team_id=grant.team_id, disallowed_models=[grant.disallowed])
|
||||
assert _stored_disallowed_models(key) == [{"disallowed_models": [grant.disallowed]}]
|
||||
_assert_served(gateway, key, grant.allowed)
|
||||
_assert_denied(gateway, key, {"model": grant.disallowed}, grant.disallowed)
|
||||
listed: Final = eventually(
|
||||
lambda: _listed_models(gateway, key),
|
||||
lambda models: grant.allowed in models,
|
||||
seconds=30,
|
||||
)
|
||||
assert grant.allowed in listed
|
||||
assert grant.disallowed not in listed
|
||||
|
||||
|
||||
def test_an_all_team_models_key_cannot_call_its_disallowed_models(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
grant: Final = _team_grant(scenario)
|
||||
key: Final = scenario.key(
|
||||
team_id=grant.team_id, models=["all-team-models"], disallowed_models=[grant.disallowed]
|
||||
)
|
||||
_assert_served(gateway, key, grant.allowed)
|
||||
_assert_denied(gateway, key, {"model": grant.disallowed}, grant.disallowed)
|
||||
|
||||
|
||||
def test_disallowed_models_win_over_the_keys_own_allow_list(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
allowed: Final = scenario.model()
|
||||
disallowed: Final = scenario.model()
|
||||
key: Final = scenario.key(models=[allowed, disallowed], disallowed_models=[disallowed])
|
||||
_assert_served(gateway, key, allowed)
|
||||
_assert_denied(gateway, key, {"model": disallowed}, disallowed)
|
||||
assert disallowed not in _listed_models(gateway, key)
|
||||
|
||||
|
||||
def test_an_alias_or_a_fallback_cannot_reach_a_disallowed_model(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
grant: Final = _team_grant(scenario)
|
||||
alias: Final = f"alias-{uuid.uuid4().hex[:8]}"
|
||||
key: Final = scenario.key(
|
||||
team_id=grant.team_id, aliases={alias: grant.disallowed}, disallowed_models=[grant.disallowed]
|
||||
)
|
||||
_assert_denied(gateway, key, {"model": alias}, alias)
|
||||
_assert_denied(gateway, key, {"model": grant.allowed, "fallbacks": [grant.disallowed]}, grant.disallowed)
|
||||
|
||||
|
||||
def test_adding_and_clearing_disallowed_models_takes_effect_on_the_next_request(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
grant: Final = _team_grant(scenario)
|
||||
key: Final = scenario.key(team_id=grant.team_id)
|
||||
_assert_served(gateway, key, grant.disallowed)
|
||||
gateway.post("/key/update", {"key": key, "disallowed_models": [grant.disallowed]})
|
||||
_assert_denied(gateway, key, {"model": grant.disallowed}, grant.disallowed)
|
||||
info: Final = object_value(gateway.get("/key/info", {"key": key})["info"])
|
||||
assert info["disallowed_models"] == [grant.disallowed], info
|
||||
gateway.post("/key/update", {"key": key, "disallowed_models": []})
|
||||
_assert_served(gateway, key, grant.disallowed)
|
||||
|
||||
|
||||
def test_a_key_owner_cannot_remove_disallowed_models_but_can_add_them(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
grant: Final = _team_grant(scenario)
|
||||
owner: Final = scenario.member(grant.team_id)
|
||||
key: Final = scenario.key(team_id=grant.team_id, user_id=owner, disallowed_models=[grant.disallowed])
|
||||
cleared: Final = gateway.request("POST", "/key/update", {"key": key, "disallowed_models": []}, key=key)
|
||||
assert cleared.status_code == 403, f"{cleared.status_code} {cleared.text}"
|
||||
assert _stored_disallowed_models(key) == [{"disallowed_models": [grant.disallowed]}]
|
||||
_assert_denied(gateway, key, {"model": grant.disallowed}, grant.disallowed)
|
||||
widened: Final = gateway.request(
|
||||
"POST",
|
||||
"/key/update",
|
||||
{"key": key, "disallowed_models": [grant.disallowed, grant.allowed]},
|
||||
key=key,
|
||||
)
|
||||
assert widened.status_code == 200, f"{widened.status_code} {widened.text}"
|
||||
_assert_denied(gateway, key, {"model": grant.allowed}, grant.allowed)
|
||||
|
||||
|
||||
def test_an_admin_of_the_keys_team_can_remove_disallowed_models(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
grant: Final = _team_grant(scenario)
|
||||
admin: Final = scenario.member(grant.team_id, role="admin")
|
||||
admin_key: Final = scenario.key(team_id=grant.team_id, user_id=admin)
|
||||
key: Final = scenario.key(team_id=grant.team_id, disallowed_models=[grant.disallowed])
|
||||
_assert_denied(gateway, key, {"model": grant.disallowed}, grant.disallowed)
|
||||
cleared: Final = gateway.request("POST", "/key/update", {"key": key, "disallowed_models": []}, key=admin_key)
|
||||
assert cleared.status_code == 200, f"{cleared.status_code} {cleared.text}"
|
||||
assert _stored_disallowed_models(key) == [{"disallowed_models": []}]
|
||||
_assert_served(gateway, key, grant.disallowed)
|
||||
118
tests/unit/proxy/auth/test_key_disallowed_models.py
Normal file
118
tests/unit/proxy/auth/test_key_disallowed_models.py
Normal file
|
|
@ -0,0 +1,118 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import can_key_call_model
|
||||
from litellm.proxy.auth.key_disallowed_models import (
|
||||
disallowed_models_removed_by,
|
||||
ensure_key_may_call_model,
|
||||
without_key_disallowed_models,
|
||||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import enforce_key_and_fallback_model_access
|
||||
from litellm.types.proxy.management_endpoints.key_requests import UpdateKeyRequest
|
||||
|
||||
|
||||
def _token(**fields: object) -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(token="hashed-token", **fields)
|
||||
|
||||
|
||||
def _denied_model(error: pytest.ExceptionInfo[ProxyException]) -> str:
|
||||
assert error.value.type == ProxyErrorTypes.key_model_access_denied
|
||||
assert str(error.value.code) == "403"
|
||||
return error.value.message
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"fields,requested",
|
||||
[
|
||||
pytest.param({}, "gpt-denied", id="exact-name"),
|
||||
pytest.param({"aliases": {"short": "gpt-denied"}}, "short", id="key-alias"),
|
||||
pytest.param({"team_model_aliases": {"team-short": "gpt-denied"}}, "team-short", id="team-alias"),
|
||||
pytest.param({}, ["gpt-allowed", "gpt-denied"], id="denied-not-first-in-list"),
|
||||
],
|
||||
)
|
||||
def test_a_disallowed_model_is_refused_however_it_is_named(fields: dict[str, object], requested: object) -> None:
|
||||
token: Final = _token(disallowed_models=("gpt-denied",), **fields)
|
||||
with pytest.raises(ProxyException) as error:
|
||||
ensure_key_may_call_model(requested, token, None) # pyright: ignore[reportArgumentType] # parametrized str or list
|
||||
assert "is not available for this API key" in _denied_model(error)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("disallowed", [None, (), ("gpt-other",)])
|
||||
def test_a_model_outside_disallowed_models_is_not_refused(disallowed: tuple[str, ...] | None) -> None:
|
||||
ensure_key_may_call_model("gpt-allowed", _token(disallowed_models=disallowed), None)
|
||||
|
||||
|
||||
def test_listing_drops_every_disallowed_model_and_every_alias_pointing_at_one() -> None:
|
||||
token: Final = _token(disallowed_models=("b", "d"), aliases={"alias-b": "b"})
|
||||
assert without_key_disallowed_models(["a", "b", "alias-b", "c", "d"], token, None) == ["a", "c"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disallowed_models_win_over_the_keys_own_allow_list() -> None:
|
||||
token: Final = _token(models=["gpt-allowed", "gpt-denied"], disallowed_models=("gpt-denied",))
|
||||
assert await can_key_call_model("gpt-allowed", None, token, None) is True
|
||||
with pytest.raises(ProxyException) as error:
|
||||
await can_key_call_model("gpt-denied", None, token, None)
|
||||
_denied_model(error)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"fields",
|
||||
[
|
||||
pytest.param({"models": ["all-team-models"], "team_id": "team-1"}, id="all-team-models"),
|
||||
pytest.param({"config": {"model_list": []}}, id="config-key"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"request_data",
|
||||
[
|
||||
pytest.param({"model": "gpt-denied"}, id="model"),
|
||||
pytest.param({"model": "gpt-allowed", "fallbacks": ["gpt-denied"]}, id="fallback"),
|
||||
],
|
||||
)
|
||||
async def test_keys_that_skip_the_allow_list_still_honor_disallowed_models(
|
||||
fields: dict[str, object], request_data: dict[str, object]
|
||||
) -> None:
|
||||
token: Final = _token(disallowed_models=("gpt-denied",), **fields)
|
||||
with pytest.raises(ProxyException) as error:
|
||||
await enforce_key_and_fallback_model_access(
|
||||
valid_token=token,
|
||||
request_data=request_data,
|
||||
route="/chat/completions",
|
||||
request=None,
|
||||
llm_model_list=None,
|
||||
llm_router=None,
|
||||
)
|
||||
_denied_model(error)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_keys_that_skip_the_allow_list_are_served_models_outside_disallowed_models() -> None:
|
||||
await enforce_key_and_fallback_model_access(
|
||||
valid_token=_token(models=["all-team-models"], team_id="team-1", disallowed_models=("gpt-denied",)),
|
||||
request_data={"model": "gpt-allowed"},
|
||||
route="/chat/completions",
|
||||
request=None,
|
||||
llm_model_list=None,
|
||||
llm_router=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"update,stored,removed",
|
||||
[
|
||||
pytest.param({}, ["a", "b"], frozenset(), id="field-not-sent"),
|
||||
pytest.param({"disallowed_models": []}, ["a", "b"], frozenset({"a", "b"}), id="cleared"),
|
||||
pytest.param({"disallowed_models": None}, ["a"], frozenset({"a"}), id="null"),
|
||||
pytest.param({"disallowed_models": ["b"]}, ["a", "b"], frozenset({"a"}), id="narrowed"),
|
||||
pytest.param({"disallowed_models": ["a", "b", "c"]}, ["a", "b"], frozenset(), id="added"),
|
||||
pytest.param({"disallowed_models": []}, None, frozenset(), id="nothing-stored"),
|
||||
],
|
||||
)
|
||||
def test_removed_disallowed_models_are_the_stored_entries_the_update_drops(
|
||||
update: dict[str, object], stored: list[str] | None, removed: frozenset[str]
|
||||
) -> None:
|
||||
assert disallowed_models_removed_by(UpdateKeyRequest(key="sk-1", **update), stored) == removed
|
||||
Loading…
Add table
Reference in a new issue