mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
fix(proxy): keep auto routers from sending a key to a tier in its disallowed_models
The deny check only saw the requested route name, so a complexity or auto router could pick a tier the key disallows and dispatch it. The pre-routing step now asks the router host about the selected model, and the proxy host runs the key's deny list on it. Requests with no key in metadata are unchanged. Also adds unit tests for the deny list permission check, the null update and the batch row check
This commit is contained in:
parent
cb622a8542
commit
3d4b99f543
8 changed files with 181 additions and 0 deletions
|
|
@ -45,6 +45,18 @@ class ProxyRouterHost(RouterHost):
|
|||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
def authorize_routed_model(
|
||||
self,
|
||||
model: str,
|
||||
request_kwargs: Mapping[str, object],
|
||||
llm_router: "Router",
|
||||
) -> None:
|
||||
from litellm.proxy.auth import auto_router_checks
|
||||
|
||||
auto_router_checks.ensure_key_may_call_routed_model(
|
||||
model=model, request_kwargs=request_kwargs, llm_router=llm_router
|
||||
)
|
||||
|
||||
def compression_policy(
|
||||
self,
|
||||
llm_router: "Router | None",
|
||||
|
|
|
|||
|
|
@ -20,6 +20,16 @@ def _mapping(value: object) -> Mapping[str, object] | None:
|
|||
return None
|
||||
|
||||
|
||||
def ensure_key_may_call_routed_model(*, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> None:
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.key_disallowed_models import ensure_key_may_call_model
|
||||
|
||||
metadata: Final = _mapping(request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs)))
|
||||
actor: Final = metadata.get("user_api_key_auth") if metadata is not None else None
|
||||
if isinstance(actor, UserAPIKeyAuth):
|
||||
ensure_key_may_call_model(model, actor, llm_router)
|
||||
|
||||
|
||||
async def authorize_member_auto_router_inference(
|
||||
*,
|
||||
deployment: Mapping[str, object] | None,
|
||||
|
|
|
|||
|
|
@ -63,6 +63,14 @@ class RouterHost:
|
|||
response=httpx.Response(status_code=403, request=httpx.Request("POST", "https://litellm.invalid")),
|
||||
)
|
||||
|
||||
def authorize_routed_model(
|
||||
self,
|
||||
model: str,
|
||||
request_kwargs: Mapping[str, object],
|
||||
llm_router: "Router",
|
||||
) -> None:
|
||||
"""Raise when the caller may not be sent to ``model``, the one a routing strategy selected."""
|
||||
|
||||
def compression_policy(
|
||||
self,
|
||||
llm_router: "Router | None",
|
||||
|
|
|
|||
|
|
@ -921,6 +921,12 @@ class PreRouter:
|
|||
call_type=call_type,
|
||||
)
|
||||
hooked: Final[PreRoutingHookResponse | None] = routed
|
||||
if routed is not None:
|
||||
router_host().authorize_routed_model(
|
||||
model=routed.model,
|
||||
request_kwargs=request_kwargs,
|
||||
llm_router=public_router, # pyright: ignore[reportArgumentType] # the litellm.Router this core serves
|
||||
)
|
||||
traced: Final = (
|
||||
hooked.model_copy(update={"routing_decision": {**provenance, **hooked.routing_decision}})
|
||||
if hooked is not None and hooked.routing_decision is not None
|
||||
|
|
|
|||
|
|
@ -225,6 +225,27 @@ def test_a_key_alias_to_a_router_alias_cannot_reach_a_disallowed_model(gateway:
|
|||
_assert_denied(gateway, key, {"model": key_alias}, key_alias)
|
||||
|
||||
|
||||
def test_a_complexity_router_cannot_send_a_key_to_a_disallowed_tier(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
keyword: Final = f"simple-{uuid.uuid4().hex}"
|
||||
tier_upstream: Final = f"tier-{uuid.uuid4().hex[:8]}"
|
||||
disallowed: Final = scenario.model(model=f"openai/{tier_upstream}")
|
||||
other_tier: Final = scenario.model()
|
||||
router: Final = scenario.model(
|
||||
model="auto_router/complexity_router",
|
||||
complexity_router_config={
|
||||
"classifier_type": "heuristic",
|
||||
"keyword_tier_rules": [{"keywords": [keyword], "tier": "SIMPLE"}],
|
||||
"tiers": {"SIMPLE": disallowed, "MEDIUM": other_tier, "COMPLEX": other_tier, "REASONING": other_tier},
|
||||
},
|
||||
)
|
||||
body: Final[dict[str, JsonValue]] = {"model": router, "messages": [{"role": "user", "content": keyword}]}
|
||||
served: Final = _chat(gateway, scenario.key(), body)
|
||||
assert served.status_code == 200, f"{served.status_code} {served.text}"
|
||||
assert tier_upstream in _observed_models(gateway)
|
||||
_assert_denied(gateway, scenario.key(disallowed_models=[disallowed]), body, disallowed)
|
||||
|
||||
|
||||
def _deleted_disallowed_models(keys: tuple[str, ...]) -> list[dict[str, JsonValue]]:
|
||||
return read_rows(
|
||||
'SELECT disallowed_models FROM "LiteLLM_DeletedVerificationToken" WHERE token = ANY(%s) ORDER BY key_alias',
|
||||
|
|
|
|||
|
|
@ -980,6 +980,8 @@ def test_key_requires_batch_model_access_check_branches():
|
|||
check(UserAPIKeyAuth(api_key="sk", models=[], access_group_ids=["grp"])) is True
|
||||
)
|
||||
assert check(UserAPIKeyAuth(api_key="sk", models=[])) is False
|
||||
assert check(UserAPIKeyAuth(api_key="sk", models=[], disallowed_models=["gpt-denied"])) is True
|
||||
assert check(UserAPIKeyAuth(api_key="sk", models=["*"], disallowed_models=["gpt-denied"])) is True
|
||||
assert check(UserAPIKeyAuth(api_key="sk", models=["gpt-4o-mini"])) is True
|
||||
# Wildcard / all-proxy-models grant access to every model, so
|
||||
# can_key_call_model passes any model regardless of access groups (which
|
||||
|
|
|
|||
|
|
@ -2143,6 +2143,66 @@ async def test_update_key_keeps_a_due_budget_reset_when_the_duration_is_resent(
|
|||
assert ("budget_reset_at" in updated) is moves_reset
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_key_with_null_disallowed_models_clears_the_stored_deny_list() -> None:
|
||||
existing_key: Final = LiteLLM_VerificationToken(token="hashed", disallowed_models=("claude-*",))
|
||||
|
||||
updated: Final = await prepare_key_update_data(
|
||||
data=UpdateKeyRequest(key="sk-1", disallowed_models=None), existing_key_row=existing_key
|
||||
)
|
||||
|
||||
assert updated["disallowed_models"] == []
|
||||
|
||||
|
||||
_DENY_LIST_TEAM: Final = LiteLLM_TeamTableCachedObj(
|
||||
team_id="team-1",
|
||||
members_with_roles=[Member(user_id="team-admin", role="admin"), Member(user_id="member", role="user")],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("caller_id", "key_team_id", "requested", "allowed"),
|
||||
[
|
||||
("team-admin", "team-1", ["claude-*"], True),
|
||||
("member", "team-1", ["claude-*"], False),
|
||||
("member", None, ["claude-*"], False),
|
||||
("member", "team-1", ["claude-*", "gpt-denied", "another"], True),
|
||||
],
|
||||
ids=["team-admin-removes", "team-member-removes", "no-team-removes", "team-member-only-adds"],
|
||||
)
|
||||
async def test_only_admins_of_the_keys_team_may_remove_disallowed_models(
|
||||
monkeypatch: pytest.MonkeyPatch, caller_id: str, key_team_id: str | None, requested: list[str], allowed: bool
|
||||
) -> None:
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import check_disallowed_models_caller_permission
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
|
||||
AsyncMock(return_value=_DENY_LIST_TEAM),
|
||||
)
|
||||
|
||||
async def check() -> None:
|
||||
await check_disallowed_models_caller_permission(
|
||||
data=UpdateKeyRequest(key="sk-1", disallowed_models=requested),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-caller", user_id=caller_id
|
||||
),
|
||||
existing_key_row=LiteLLM_VerificationToken(
|
||||
token="hashed", team_id=key_team_id, disallowed_models=("claude-*", "gpt-denied")
|
||||
),
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=MagicMock(),
|
||||
)
|
||||
|
||||
if allowed:
|
||||
await check()
|
||||
return
|
||||
with pytest.raises(HTTPException) as refused:
|
||||
await check()
|
||||
assert refused.value.status_code == 403
|
||||
assert "gpt-denied" in str(refused.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_key_clears_end_user_budget_id_with_empty_string():
|
||||
"""Sending an empty end_user_budget_id detaches the key default without touching any budget row,
|
||||
|
|
|
|||
|
|
@ -75,3 +75,65 @@ def test_default_host_keeps_nothing_and_compresses_nothing():
|
|||
assert host.compression_policy(None, "auto", {}, ()) is None
|
||||
assert host.model_hop_compression_armed() is False
|
||||
assert asyncio.run(host.messages_for_routing(None, [{"role": "user", "content": "hi"}], {})) is None
|
||||
|
||||
|
||||
_TIER_KEYWORD = "zzsimpletier"
|
||||
|
||||
|
||||
def _complexity_router_reply(metadata: dict[str, object] | None) -> str:
|
||||
import asyncio
|
||||
|
||||
from litellm import Router
|
||||
|
||||
router = Router(
|
||||
model_list=[
|
||||
{"model_name": tier, "litellm_params": {"model": f"openai/{tier}", "api_key": "k", "mock_response": tier}}
|
||||
for tier in ("cheap", "smart")
|
||||
]
|
||||
+ [
|
||||
{
|
||||
"model_name": "auto",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {
|
||||
"classifier_type": "heuristic",
|
||||
"keyword_tier_rules": [{"keywords": [_TIER_KEYWORD], "tier": "SIMPLE"}],
|
||||
"tiers": {"SIMPLE": "cheap", "MEDIUM": "smart", "COMPLEX": "smart", "REASONING": "smart"},
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
extra = {} if metadata is None else {"metadata": metadata}
|
||||
response = asyncio.run(
|
||||
router.acompletion(model="auto", messages=[{"role": "user", "content": _TIER_KEYWORD}], **extra)
|
||||
)
|
||||
return response.choices[0].message.content # pyright: ignore[reportAttributeAccessIssue] # non-streaming reply
|
||||
|
||||
|
||||
@pytest.mark.parametrize("disallowed_models", [None, [], ["smart"]], ids=["no-key", "no-deny-list", "other-tier-denied"])
|
||||
def test_a_complexity_router_serves_its_selected_tier_when_the_key_does_not_disallow_it(
|
||||
disallowed_models: list[str] | None,
|
||||
):
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
metadata = (
|
||||
None
|
||||
if disallowed_models is None
|
||||
else {"user_api_key_auth": UserAPIKeyAuth(api_key="sk", disallowed_models=disallowed_models)}
|
||||
)
|
||||
|
||||
assert _complexity_router_reply(metadata) == "cheap"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("deny_entry", ["cheap", "che*"])
|
||||
def test_a_complexity_router_cannot_send_a_key_to_a_tier_it_disallows(deny_entry: str):
|
||||
from litellm.proxy._types import ModelAccessDeniedProxyException, ProxyErrorTypes, UserAPIKeyAuth
|
||||
|
||||
key = UserAPIKeyAuth(api_key="sk", disallowed_models=[deny_entry])
|
||||
|
||||
with pytest.raises(ModelAccessDeniedProxyException) as denied:
|
||||
_complexity_router_reply({"user_api_key_auth": key})
|
||||
|
||||
assert denied.value.type == ProxyErrorTypes.key_model_access_denied
|
||||
assert denied.value.code == "403"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue