From 4bfe9878f2adba441ea5d6c33035207cbe12be65 Mon Sep 17 00:00:00 2001 From: jesus Date: Sat, 12 Sep 2026 01:01:35 +0000 Subject: [PATCH 1/4] fix(auth): resolve model_group_alias before the zero-cost budget bypass check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 12 +++-- .../test_unmapped_model_budget_enforcement.py | 54 +++++++++++++++++++ 2 files changed, 63 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1efc9611fe6..d9020cc5daf 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -120,6 +120,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.common_utils import resolve_model_group_alias from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.utils import get_utc_datetime @@ -366,6 +367,10 @@ def _get_router_zero_cost_cache(llm_router: Router) -> dict[str, bool] | None: return cache if isinstance(cache, dict) else None +def _resolve_cost_model_group(model_name: str, llm_router: "Router") -> str: + return resolve_model_group_alias(llm_router.model_group_alias, model_name) or model_name + + def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None) -> bool: """ Check if a model has zero cost (no configured pricing). @@ -395,8 +400,9 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None return False continue try: + target_group: Final = _resolve_cost_model_group(model_name, llm_router) # Use router's get_model_group_info method directly for better reliability - model_group_info = llm_router.get_model_group_info(model_group=model_name) + model_group_info = llm_router.get_model_group_info(model_group=target_group) if model_group_info is None: # Model not found or no pricing info available @@ -436,7 +442,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None # not from defaulted sparse auto-registration entries. # See: https://github.com/BerriAI/litellm/issues/24770 safe_name = str(model_name).replace("\n", "").replace("\r", "") - if not _is_cost_explicitly_configured(model_name, llm_router): + if not _is_cost_explicitly_configured(target_group, llm_router): verbose_proxy_logger.debug( "Model %s has zero cost but no explicit cost " "configuration in model_cost entry — treating as unknown " @@ -447,7 +453,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None zero_cost_cache[model_name] = False return False - if _has_ptu_flat_cost(model_name, llm_router): + if _has_ptu_flat_cost(target_group, llm_router): verbose_proxy_logger.debug( "Model %s prices reserved PTU capacity as a flat cost, so its zero per-token " "rate is not a free model (enforce budget)", diff --git a/tests/test_litellm/proxy/auth/test_unmapped_model_budget_enforcement.py b/tests/test_litellm/proxy/auth/test_unmapped_model_budget_enforcement.py index bbe343bcede..4c17331984b 100644 --- a/tests/test_litellm/proxy/auth/test_unmapped_model_budget_enforcement.py +++ b/tests/test_litellm/proxy/auth/test_unmapped_model_budget_enforcement.py @@ -211,3 +211,57 @@ class TestUnmappedModelBudgetEnforcement: result = _is_model_cost_zero(model="paid-model", llm_router=mock_router) assert result is False + + +class TestModelGroupAliasZeroCost: + """model_group_alias names must resolve to their target group before the + zero-cost check, for both the string form and the hidden dict form.""" + + def setup_method(self): + """Snapshot litellm.model_cost before each test.""" + self._saved_model_cost = copy.deepcopy(litellm.model_cost) + + def teardown_method(self): + """Restore litellm.model_cost after each test.""" + litellm.model_cost = self._saved_model_cost + + def _router(self) -> Router: + return Router( + model_list=[ + { + "model_name": "free-model", + "litellm_params": { + "model": "ollama/llama2", + "api_base": "http://localhost:11434", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + }, + { + "model_name": "paid-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-fake"}, + }, + ], + model_group_alias={ + "free-alias": "free-model", + "free-hidden": {"model": "free-model", "hidden": True}, + "paid-alias": "paid-model", + "paid-hidden": {"model": "paid-model", "hidden": True}, + }, + ) + + def test_visible_alias_to_free_model_bypasses_budget(self): + router = self._router() + assert _is_model_cost_zero(model="free-alias", llm_router=router) is True + + def test_hidden_alias_to_free_model_bypasses_budget(self): + router = self._router() + assert _is_model_cost_zero(model="free-hidden", llm_router=router) is True + + def test_visible_alias_to_paid_model_enforces_budget(self): + router = self._router() + assert _is_model_cost_zero(model="paid-alias", llm_router=router) is False + + def test_hidden_alias_to_paid_model_enforces_budget(self): + router = self._router() + assert _is_model_cost_zero(model="paid-hidden", llm_router=router) is False From d9ef8746a92dfe4a7a250e65c7312dc3acb5247f Mon Sep 17 00:00:00 2001 From: jesus Date: Sat, 12 Sep 2026 19:15:12 +0000 Subject: [PATCH 2/4] fix(auth): drop Final annotation on loop-local target_group for basedpyright Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index d9020cc5daf..95c044fdab2 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -400,7 +400,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None return False continue try: - target_group: Final = _resolve_cost_model_group(model_name, llm_router) + target_group = _resolve_cost_model_group(model_name, llm_router) # Use router's get_model_group_info method directly for better reliability model_group_info = llm_router.get_model_group_info(model_group=target_group) From 05ad4c93f86cfbb42e22fd08ebaee0461e8fba80 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 28 Sep 2026 16:18:23 -0700 Subject: [PATCH 3/4] fix(proxy): judge the zero-cost budget bypass on the model key and team aliases dispatch to --- litellm/proxy/auth/auth_checks.py | 12 +++++ litellm/proxy/auth/user_api_key_auth.py | 24 +++++----- tests/unit/proxy/auth/test_auth_checks.py | 25 +++++++++++ .../unit/proxy/auth/test_user_api_key_auth.py | 44 +++++++++++++++++++ 4 files changed, 93 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 10f95cafff6..7c8cca2ff65 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -564,6 +564,18 @@ 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 is_dispatched_model_cost_zero( + model: str | list[str] | None, llm_router: Router | None, valid_token: UserAPIKeyAuth +) -> bool: + dispatched_model: Final = _dispatched_model_name(model, valid_token) if isinstance(model, str) else model + return _is_model_cost_zero(model=dispatched_model, llm_router=llm_router) + + _NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({}) _TEAM_GRANT_RELATIONS: Final[Mapping[str, object]] = MappingProxyType({"litellm_model_table": True}) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index e3ce9bcd850..fd4509dae66 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -48,7 +48,6 @@ from litellm.proxy.auth.auth_checks import ( _check_end_user_budget, _delete_cache_key_object, _get_user_role, - _is_model_cost_zero, _is_user_proxy_admin, _team_member_max_budget_alert_check, _virtual_key_max_budget_alert_check, @@ -65,6 +64,7 @@ from litellm.proxy.auth.auth_checks import ( get_team_membership, get_team_object, get_user_object, + is_dispatched_model_cost_zero, is_valid_fallback_model, jwt_key_mapping_cache_key, key_model_aliases_for_auth_check, @@ -1795,9 +1795,9 @@ async def _user_api_key_auth_builder( ) skip_budget_checks = False if model is not None and llm_router is not None: - from litellm.proxy.auth.auth_checks import _is_model_cost_zero - - skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router) + skip_budget_checks = is_dispatched_model_cost_zero( + model=model, llm_router=llm_router, valid_token=valid_token + ) if skip_budget_checks: verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model) @@ -2237,9 +2237,9 @@ async def _user_api_key_auth_builder( ) skip_budget_checks = False if model is not None and llm_router is not None: - from litellm.proxy.auth.auth_checks import _is_model_cost_zero - - skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router) + skip_budget_checks = is_dispatched_model_cost_zero( + model=model, llm_router=llm_router, valid_token=valid_token + ) if skip_budget_checks: verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model) @@ -2965,7 +2965,7 @@ async def _run_centralized_common_checks( route=route, request=request, llm_router=llm_router, - team_id=user_api_key_auth_obj.team_id, + valid_token=user_api_key_auth_obj, ) # Pin the metadata variable name (litellm_metadata vs metadata) before @@ -3136,17 +3136,17 @@ def _should_skip_budget_checks( route: str, request: Request | None, llm_router: Any | None, - team_id: str | None = None, + valid_token: UserAPIKeyAuth, ) -> bool: model: Final = _get_model_from_request_context( request_data=request_data, route=route, request=request, llm_router=llm_router, - team_id=team_id, + team_id=valid_token.team_id, ) if model is not None and llm_router is not None: - return _is_model_cost_zero(model=model, llm_router=llm_router) + return is_dispatched_model_cost_zero(model=model, llm_router=llm_router, valid_token=valid_token) return False @@ -3715,7 +3715,7 @@ 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_model_cost_zero(model=current_model, llm_router=llm_router) + is_dispatched_model_cost_zero(model=current_model, llm_router=llm_router, valid_token=valid_token) 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 7b92494a6dc..d63c4c88de4 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -25,6 +25,7 @@ from litellm.proxy.utils import PrismaClient from litellm.proxy.auth.auth_checks import ( can_team_access_model, _is_model_cost_zero, + is_dispatched_model_cost_zero, _virtual_key_soft_budget_check, _team_soft_budget_check, ) @@ -1574,3 +1575,27 @@ def test_zero_cost_check_follows_an_alias_repointed_at_runtime( assert _is_model_cost_zero(model=alias, llm_router=router) is not expected_after_repoint router.update_settings(model_group_alias=_aliases_to(repointed_target)) assert _is_model_cost_zero(model=alias, llm_router=router) is expected_after_repoint + + +@pytest.mark.parametrize( + "team_model_aliases, key_aliases, model, expected", + [ + (None, {"visible": "paid-model"}, "visible", False), + (None, {"free-model": "paid-model"}, "free-model", False), + ({"free-model": "paid-model"}, None, "free-model", False), + (None, {"my-free": "free-model"}, "my-free", True), + (None, {"my-free": "visible"}, "my-free", True), + ({"team-name": "key-name"}, {"key-name": "free-model"}, "team-name", True), + ({"team-name": "key-name"}, {"key-name": "paid-model"}, "team-name", False), + (None, {"other": "paid-model"}, "visible", True), + ], +) +def test_zero_cost_check_prices_the_model_the_key_and_team_aliases_dispatch_to( + team_model_aliases: dict[str, str] | None, + key_aliases: dict[str, str] | None, + model: str, + expected: bool, +) -> 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 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 9cdac341b1f..2feec498b56 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -1794,3 +1794,47 @@ def test_mapped_key_jwt_falls_through_to_the_shared_user_budget_attach(): "the mapped-key branch returns before the shared virtual-key checks, so the " "user's per-model budget is never attached and never enforced" ) + + +@pytest.mark.parametrize( + "key_aliases, expected", + [ + ({}, True), + ({"free-alias": "paid-model"}, False), + ], +) +def test_budget_skip_judges_the_model_a_key_alias_dispatches_to(key_aliases: dict[str, str], expected: bool) -> None: + from litellm.proxy.auth.user_api_key_auth import _should_skip_budget_checks + from litellm.router import Router + + router = Router( + model_list=[ + { + "model_name": "free-model", + "litellm_params": { + "model": "ollama/llama2", + "api_base": "http://localhost:11434", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + }, + { + "model_name": "paid-model", + "litellm_params": { + "model": "openai/paid-model", + "api_key": "sk-fake", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + }, + }, + ], + model_group_alias={"free-alias": "free-model"}, + ) + skipped = _should_skip_budget_checks( + request_data={"model": "free-alias"}, + route="/chat/completions", + request=None, + llm_router=router, + valid_token=UserAPIKeyAuth(aliases=key_aliases), + ) + assert skipped is expected From 5b84fc8df98b3022960529edf6a2e6f3c2b900dc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 28 Sep 2026 17:35:57 -0700 Subject: [PATCH 4/4] fix(proxy): price the model every alias hop dispatches to in the zero-cost budget check --- litellm/proxy/auth/auth_checks.py | 30 ++++++++--- litellm/proxy/auth/user_api_key_auth.py | 38 +++++++++++--- tests/unit/proxy/auth/test_auth_checks.py | 51 ++++++++++++++++++- .../unit/proxy/auth/test_user_api_key_auth.py | 43 +++++++++++++--- 4 files changed, 139 insertions(+), 23 deletions(-) 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