diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 7c8cca2ff65..7673c3983b4 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -89,7 +89,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, _safe_get_request_query_params, ) -from litellm.proxy.common_utils.model_listing_utils import alias_map +from litellm.proxy.common_utils.model_listing_utils import alias_map, alias_target, caller_alias_maps from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import ( END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, @@ -564,15 +564,33 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None return True -def _dispatched_model_name(model_name: str, valid_token: UserAPIKeyAuth) -> str: - after_team_alias: Final = alias_map(valid_token.team_model_aliases).get(model_name, model_name) - return alias_map(valid_token.aliases).get(after_team_alias, after_team_alias) +def _dispatched_model_name( + model_name: str, valid_token: UserAPIKeyAuth, router_settings_aliases: object, router_settings_rewrite_pending: bool +) -> str: + requested: Final = ( + resolve_model_group_alias(router_settings_aliases, model_name) or model_name + if router_settings_rewrite_pending + else model_name + ) + caller_aliases: Final = caller_alias_maps( + valid_token.aliases, valid_token.team_model_aliases, valid_token.team_id, None + ) + aliased: Final = alias_target(requested, caller_aliases) or requested + return resolve_model_group_alias(router_settings_aliases, aliased) or aliased def is_dispatched_model_cost_zero( - model: str | list[str] | None, llm_router: Router | None, valid_token: UserAPIKeyAuth + model: str | list[str] | None, + llm_router: Router | None, + valid_token: UserAPIKeyAuth, + router_settings_aliases: object, + router_settings_rewrite_pending: bool, ) -> bool: - dispatched_model: Final = _dispatched_model_name(model, valid_token) if isinstance(model, str) else model + dispatched_model: Final = ( + _dispatched_model_name(model, valid_token, router_settings_aliases, router_settings_rewrite_pending) + if isinstance(model, str) + else model + ) return _is_model_cost_zero(model=dispatched_model, llm_router=llm_router) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index fd4509dae66..8ced8eed481 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1795,8 +1795,8 @@ async def _user_api_key_auth_builder( ) skip_budget_checks = False if model is not None and llm_router is not None: - skip_budget_checks = is_dispatched_model_cost_zero( - model=model, llm_router=llm_router, valid_token=valid_token + skip_budget_checks = await _is_dispatched_model_cost_zero( + model=model, llm_router=llm_router, valid_token=valid_token, request=request ) if skip_budget_checks: verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model) @@ -2237,8 +2237,8 @@ async def _user_api_key_auth_builder( ) skip_budget_checks = False if model is not None and llm_router is not None: - skip_budget_checks = is_dispatched_model_cost_zero( - model=model, llm_router=llm_router, valid_token=valid_token + skip_budget_checks = await _is_dispatched_model_cost_zero( + model=model, llm_router=llm_router, valid_token=valid_token, request=request ) if skip_budget_checks: verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model) @@ -2960,7 +2960,7 @@ async def _run_centralized_common_checks( keep_token_limits=user_custom_auth is not None, ) - skip_budget_checks: Final = _should_skip_budget_checks( + skip_budget_checks: Final = await _should_skip_budget_checks( request_data=request_data, route=route, request=request, @@ -3131,7 +3131,7 @@ async def _reserve_budget_after_common_checks( request.state.budget_reservation = reservation # rebind-ok: read by the release middleware -def _should_skip_budget_checks( +async def _should_skip_budget_checks( request_data: dict, route: str, request: Request | None, @@ -3146,10 +3146,30 @@ def _should_skip_budget_checks( team_id=valid_token.team_id, ) if model is not None and llm_router is not None: - return is_dispatched_model_cost_zero(model=model, llm_router=llm_router, valid_token=valid_token) + return await _is_dispatched_model_cost_zero( + model=model, llm_router=llm_router, valid_token=valid_token, request=request + ) return False +async def _is_dispatched_model_cost_zero( + model: str | list[str], llm_router: litellm.Router, valid_token: UserAPIKeyAuth, request: Request | None +) -> bool: + from litellm.proxy.proxy_server import prisma_client, proxy_config, proxy_logging_obj + + settings: Final = await proxy_config.get_hierarchical_router_settings( + user_api_key_dict=valid_token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj + ) + return is_dispatched_model_cost_zero( + model=model, + llm_router=llm_router, + valid_token=valid_token, + router_settings_aliases=settings.get("model_group_alias") if isinstance(settings, Mapping) else None, + router_settings_rewrite_pending=request is not None + and request.scope.get(MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY) is not True, + ) + + def _resolve_request_principal(request: Request, valid_token: UserAPIKeyAuth) -> Principal: """Project the resolved identity into one per-request Principal, off the key object the builder already fetched, and stamp the request network context @@ -3715,7 +3735,9 @@ async def _run_post_custom_auth_checks( # every budget check for these; this path did not, so the same request could # be refused under custom auth and served under the other two. skip_budget_checks: Final = ( - is_dispatched_model_cost_zero(model=current_model, llm_router=llm_router, valid_token=valid_token) + await _is_dispatched_model_cost_zero( + model=current_model, llm_router=llm_router, valid_token=valid_token, request=request + ) if current_model is not None and llm_router is not None else False ) diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index d63c4c88de4..3f1af71c06a 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -1598,4 +1598,53 @@ def test_zero_cost_check_prices_the_model_the_key_and_team_aliases_dispatch_to( ) -> None: router: Final = _alias_router(_aliases_to("free-model")) token: Final = UserAPIKeyAuth(aliases=key_aliases or {}, team_model_aliases=team_model_aliases) - assert is_dispatched_model_cost_zero(model=model, llm_router=router, valid_token=token) is expected + assert ( + is_dispatched_model_cost_zero( + model=model, + llm_router=router, + valid_token=token, + router_settings_aliases=None, + router_settings_rewrite_pending=False, + ) + is expected + ) + + +@pytest.mark.parametrize( + "key_aliases, router_settings_aliases, model_alias_map, rewrite_pending, model, expected", + [ + ({"my-free": "free-model"}, {"free-model": "paid-model"}, {}, True, "my-free", False), + ({"my-free": "free-model"}, {}, {"free-model": "paid-model"}, True, "my-free", False), + ({}, {"free-model": "paid-model"}, {}, True, "free-model", False), + ({}, {"router-free": "free-model"}, {}, True, "router-free", True), + ({}, {"router-free": {"model": "visible", "hidden": True}}, {}, True, "router-free", True), + ({"global-free": "free-model"}, {}, {"my-free": "global-free"}, True, "my-free", True), + ({"global-free": "paid-model"}, {}, {"my-free": "global-free"}, True, "my-free", False), + ({}, {}, {"my-free": "visible"}, True, "my-free", True), + ({"free-model": "paid-model"}, {"router-free": "free-model"}, {}, True, "router-free", False), + ({}, {"router-free": "free-model", "free-model": "paid-model"}, {}, True, "router-free", False), + ({}, {"router-free": "free-model", "free-model": "paid-model"}, {}, False, "router-free", True), + ], +) +def test_zero_cost_check_prices_the_model_after_every_alias_hop_dispatch_applies( + monkeypatch: pytest.MonkeyPatch, + key_aliases: dict[str, str], + router_settings_aliases: dict[str, str | dict[str, str | bool]], + model_alias_map: dict[str, str], + rewrite_pending: bool, + model: str, + expected: bool, +) -> None: + monkeypatch.setattr(litellm, "model_alias_map", model_alias_map) + router: Final = _alias_router(_aliases_to("free-model")) + token: Final = UserAPIKeyAuth(aliases=key_aliases) + assert ( + is_dispatched_model_cost_zero( + model=model, + llm_router=router, + valid_token=token, + router_settings_aliases=router_settings_aliases, + router_settings_rewrite_pending=rewrite_pending, + ) + is expected + ) diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 2feec498b56..aa09f6ae45f 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -1796,14 +1796,28 @@ def test_mapped_key_jwt_falls_through_to_the_shared_user_budget_attach(): ) +@pytest.mark.asyncio @pytest.mark.parametrize( - "key_aliases, expected", + "model, key_aliases, key_router_settings, router_settings_rewritten, expected", [ - ({}, True), - ({"free-alias": "paid-model"}, False), + ("free-alias", {}, None, None, True), + ("free-alias", {"free-alias": "paid-model"}, None, None, False), + ("my-alias", {"my-alias": "free-alias"}, None, None, True), + ("my-alias", {"my-alias": "free-alias"}, {"model_group_alias": {"free-alias": "paid-model"}}, None, False), + ("rs-alias", {"free-model": "paid-model"}, {"model_group_alias": {"rs-alias": "free-model"}}, False, False), + ("rs-alias", {"free-model": "paid-model"}, {"model_group_alias": {"rs-alias": "free-model"}}, True, True), ], ) -def test_budget_skip_judges_the_model_a_key_alias_dispatches_to(key_aliases: dict[str, str], expected: bool) -> None: +async def test_budget_skip_judges_the_model_the_key_aliases_dispatch_to( + model: str, + key_aliases: dict[str, str], + key_router_settings: dict[str, dict[str, str]] | None, + router_settings_rewritten: bool | None, + expected: bool, +) -> None: + from starlette.requests import Request + + from litellm.constants import MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY from litellm.proxy.auth.user_api_key_auth import _should_skip_budget_checks from litellm.router import Router @@ -1830,11 +1844,24 @@ def test_budget_skip_judges_the_model_a_key_alias_dispatches_to(key_aliases: dic ], model_group_alias={"free-alias": "free-model"}, ) - skipped = _should_skip_budget_checks( - request_data={"model": "free-alias"}, + skipped = await _should_skip_budget_checks( + request_data={"model": model}, route="/chat/completions", - request=None, + request=( + None + if router_settings_rewritten is None + else Request( + { + "type": "http", + "method": "POST", + "path": "/chat/completions", + "headers": [], + "query_string": b"", + MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY: router_settings_rewritten, + } + ) + ), llm_router=router, - valid_token=UserAPIKeyAuth(aliases=key_aliases), + valid_token=UserAPIKeyAuth(aliases=key_aliases, router_settings=key_router_settings), ) assert skipped is expected