diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index ddb6f6bb184..608985910c5 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -2,7 +2,7 @@ import asyncio import hashlib import json from collections import OrderedDict -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass from datetime import datetime from types import MappingProxyType @@ -710,6 +710,22 @@ def _merge_router_settings_override( return _ROUTER_ONLY_SETTINGS.intersection(router_settings_from_override) +def _models_to_authorize(model: object) -> Iterator[str]: + """Each model the dispatch can reach, split the way auth splits the request model. + + A comma-separated value is a batch the router fans out, and a list is a set of + names, so every non-empty part is authorized on its own. + """ + items: Final[tuple[object, ...]] = tuple(model) if isinstance(model, list) else (model,) # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType] # a list body model holds unknown element types + for item in items: + if not isinstance(item, str): + continue + for part in item.split(","): + name = part.strip() + if name: + yield name + + async def route_request( data: dict, llm_router: LitellmRouter | None, @@ -725,15 +741,16 @@ async def route_request( # resolve it into data["model"] before dispatch, so this is the single place # that covers every route, including ones added later. dispatched_model: Final = data.get("model") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # route_request takes an untyped dict body - if user_api_key_dict is not None and llm_router is not None and isinstance(dispatched_model, str): + if user_api_key_dict is not None and llm_router is not None: from litellm.proxy.auth.auth_checks import can_key_call_resolved_model - await can_key_call_resolved_model( - model=dispatched_model, - llm_model_list=llm_router.model_list, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Router.model_list is an untyped list - valid_token=user_api_key_dict, - llm_router=llm_router, - ) + for model_name in _models_to_authorize(dispatched_model): # pyright: ignore[reportUnknownArgumentType] # dispatched_model comes from the untyped dict body + await can_key_call_resolved_model( + model=model_name, + llm_model_list=llm_router.model_list, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Router.model_list is an untyped list + valid_token=user_api_key_dict, + llm_router=llm_router, + ) router_only_settings_from_override: Final = _merge_router_settings_override(data) # pyright: ignore[reportUnknownArgumentType] # route_request takes an untyped dict body try: diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index f560e6543b6..ad411995600 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -3027,64 +3027,74 @@ async def test_user_config_routers_discard_the_least_recently_used_router_past_c routers.clear() -def _allowlist_router() -> "litellm.Router": +def _allowlist_router() -> litellm.Router: return litellm.Router( model_list=[ { - "model_name": "allowed-model", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake", "mock_response": "ok"}, - }, - { - "model_name": "denied-model", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake", "mock_response": "ok"}, - }, + "model_name": name, + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake", "mock_response": f"from {name}"}, + } + for name in ("model-a", "model-b", "model-c") ] ) -@pytest.mark.asyncio -async def test_route_request_denies_resolved_model_outside_key_allowlist(): - from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth +def _hi() -> list[dict[str, str]]: + return [{"role": "user", "content": "hi"}] - router = _allowlist_router() - key = UserAPIKeyAuth(api_key="sk-test", models=["allowed-model"]) - data = {"model": "denied-model", "messages": [{"role": "user", "content": "hi"}]} + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "requested", + ["model-c", "model-a,model-c", ["model-a", "model-c"]], + ids=["single", "comma_batch", "list"], +) +async def test_route_request_denies_any_routed_model_outside_key_allowlist(requested): + from litellm.proxy._types import ProxyErrorTypes, ProxyException + + key = UserAPIKeyAuth(api_key="sk-test", models=["model-a", "model-b"]) with pytest.raises(ProxyException) as exc: - await route_request(data, router, None, "acompletion", user_api_key_dict=key) + await route_request( + {"model": requested, "messages": _hi()}, _allowlist_router(), None, "acompletion", user_api_key_dict=key + ) assert exc.value.type == ProxyErrorTypes.key_model_access_denied.value @pytest.mark.asyncio -async def test_route_request_allows_resolved_model_inside_key_allowlist(): - from litellm.proxy._types import UserAPIKeyAuth +async def test_route_request_serves_a_model_inside_key_allowlist(): + key = UserAPIKeyAuth(api_key="sk-test", models=["model-a", "model-b"]) - router = _allowlist_router() - key = UserAPIKeyAuth(api_key="sk-test", models=["allowed-model"]) - data = {"model": "allowed-model", "messages": [{"role": "user", "content": "hi"}]} + llm_call = await route_request( + {"model": "model-a", "messages": _hi()}, _allowlist_router(), None, "acompletion", user_api_key_dict=key + ) + response = await llm_call - response = await route_request(data, router, None, "acompletion", user_api_key_dict=key) - assert await response if asyncio.iscoroutine(response) else response + assert response.choices[0].message.content == "from model-a" @pytest.mark.asyncio -async def test_route_request_without_a_key_skips_the_allowlist_check(): - router = _allowlist_router() - data = {"model": "denied-model", "messages": [{"role": "user", "content": "hi"}]} +async def test_route_request_serves_a_comma_batch_when_key_allows_every_model(): + key = UserAPIKeyAuth(api_key="sk-test", models=["model-a", "model-b"]) - response = await route_request(data, router, None, "acompletion") - assert await response if asyncio.iscoroutine(response) else response + llm_call = await route_request( + {"model": "model-a, model-b", "messages": _hi()}, + _allowlist_router(), + None, + "acompletion", + user_api_key_dict=key, + ) + responses = await llm_call + + assert sorted(response.choices[0].message.content for response in responses) == [ + "from model-a", + "from model-b", + ] @pytest.mark.asyncio -async def test_route_request_list_model_is_not_blocked_by_the_allowlist_check(): - from litellm.proxy._types import UserAPIKeyAuth +async def test_route_request_without_a_key_leaves_model_access_to_auth(): + llm_call = await route_request({"model": "model-c", "messages": _hi()}, _allowlist_router(), None, "acompletion") + response = await llm_call - router = _allowlist_router() - key = UserAPIKeyAuth(api_key="sk-test", models=["allowed-model"]) - data = {"model": ["allowed-model", "denied-model"], "messages": [{"role": "user", "content": "hi"}]} - - # A list model value is handled by the per-model auth path, not this backstop, - # so the check must not raise on it. - with contextlib.suppress(Exception): - await route_request(data, router, None, "acompletion", user_api_key_dict=key) + assert response.choices[0].message.content == "from model-c"