From 4bfe9878f2adba441ea5d6c33035207cbe12be65 Mon Sep 17 00:00:00 2001 From: jesus Date: Sat, 12 Sep 2026 01:01:35 +0000 Subject: [PATCH] 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