diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3ec430332ee..9d3dcfd14ce 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -140,6 +140,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.auth.auth_checks import UserNotFoundError from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.utils import get_utc_datetime @@ -473,22 +474,23 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None zero_cost_cache: Final = _get_router_zero_cost_cache(llm_router) for model_name in model_list: - if zero_cost_cache is not None: - cached = zero_cost_cache.get(model_name) - if cached is not None: - if cached is False: - return False - continue try: # 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) + target_group = resolve_model_group_alias(llm_router.model_group_alias, model_name) or model_name + if zero_cost_cache is not None: + cached = zero_cost_cache.get(target_group) + if cached is not None: + if cached is False: + return False + continue + 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 # Conservative approach: assume it has cost verbose_proxy_logger.debug("No model group info found for %s, assuming it has cost", model_name) if zero_cost_cache is not None: - zero_cost_cache[model_name] = False + zero_cost_cache[target_group] = False return False # Check costs for this model @@ -505,7 +507,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None output_cost, ) if zero_cost_cache is not None: - zero_cost_cache[model_name] = False + zero_cost_cache[target_group] = False return False # If either cost is non-zero, return False @@ -514,7 +516,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None "Model %s has non-zero cost (input: %s, output: %s)", model_name, input_cost, output_cost ) if zero_cost_cache is not None: - zero_cost_cache[model_name] = False + zero_cost_cache[target_group] = False return False # Costs are 0 — verify this is from explicit configuration, @@ -529,7 +531,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None safe_name, ) if zero_cost_cache is not None: - zero_cost_cache[model_name] = False + zero_cost_cache[target_group] = False return False if _has_ptu_flat_cost(model_name, llm_router): @@ -539,7 +541,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None safe_name, ) if zero_cost_cache is not None: - zero_cost_cache[model_name] = False + zero_cost_cache[target_group] = False return False verbose_proxy_logger.debug( @@ -549,7 +551,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None output_cost, ) if zero_cost_cache is not None: - zero_cost_cache[model_name] = True + zero_cost_cache[target_group] = True except Exception as e: # If we can't determine the cost, assume it has cost (conservative approach) 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 7665008a6a6..29db7c8884a 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 @@ -39,10 +39,7 @@ class TestUnmappedModelBudgetEnforcement: ] ) result = _is_model_cost_zero(model="custom-model", llm_router=router) - assert result is False, ( - "Unmapped model should enforce budget (return False), " - "not bypass it (return True)" - ) + assert result is False, "Unmapped model should enforce budget (return False), not bypass it (return True)" def test_explicitly_free_model_bypasses_budget(self): """A model with explicit cost=0 in model_info should bypass budget.""" @@ -65,9 +62,7 @@ class TestUnmappedModelBudgetEnforcement: ] ) result = _is_model_cost_zero(model="free-model", llm_router=router) - assert ( - result is True - ), "Explicitly free model should bypass budget (return True)" + assert result is True, "Explicitly free model should bypass budget (return True)" def test_known_paid_model_enforces_budget(self): """A model in the cost map with non-zero costs should enforce budget.""" @@ -101,9 +96,7 @@ class TestUnmappedModelBudgetEnforcement: ] ) result = _is_model_cost_zero(model="free-via-params", llm_router=router) - assert ( - result is True - ), "Model with explicit cost=0 in litellm_params should bypass budget" + assert result is True, "Model with explicit cost=0 in litellm_params should bypass budget" def test_cache_invalidates_on_in_place_pricing_update(self): """ @@ -285,9 +278,12 @@ class TestUnmappedModelBudgetEnforcement: "An aliased PTU group must not be read as free" ) - def test_hidden_model_group_alias_enforces_budget(self): - """A hidden alias keeps budget enforced: get_model_group_info() returns None for it, - so the cost is unknown before the configuration gate is reached.""" + def test_hidden_model_group_alias_to_free_model_bypasses_budget(self): + """A hidden alias to an explicitly free group bypasses budget, like the group itself. + + ``get_model_group_info`` returns None for hidden aliases, so the alias must be + resolved to its target group before the cost lookup. + """ router = Router( model_list=[ { @@ -304,7 +300,22 @@ class TestUnmappedModelBudgetEnforcement: model_group_alias={"hidden-alias": {"model": "free-model", "hidden": True}}, ) - assert _is_model_cost_zero(model="hidden-alias", llm_router=router) is False + assert _is_model_cost_zero(model="hidden-alias", llm_router=router) is True + + def test_hidden_model_group_alias_to_paid_model_enforces_budget(self): + """A hidden alias to a priced group keeps budget enforced.""" + router = Router( + model_list=[ + { + "model_name": "paid-model", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-fake"}, + "model_info": {"id": "paid-model-id"}, + }, + ], + model_group_alias={"hidden-paid-alias": {"model": "paid-model", "hidden": True}}, + ) + + assert _is_model_cost_zero(model="hidden-paid-alias", llm_router=router) is False def test_dangling_model_group_alias_enforces_budget(self): """An alias pointing at a group that does not exist keeps budget enforced.""" @@ -326,6 +337,33 @@ class TestUnmappedModelBudgetEnforcement: assert _is_model_cost_zero(model="dangling-alias", llm_router=router) is False + def test_repointed_hidden_alias_does_not_reuse_cached_free_result(self): + """Repointing a hidden alias from a free group to a paid group re-evaluates the cost.""" + 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_info": {"id": "free-model-id"}, + }, + { + "model_name": "paid-model", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-fake"}, + "model_info": {"id": "paid-model-id"}, + }, + ], + model_group_alias={"hidden-alias": {"model": "free-model", "hidden": True}}, + ) + + assert _is_model_cost_zero(model="hidden-alias", llm_router=router) is True + router.model_group_alias["hidden-alias"] = {"model": "paid-model", "hidden": True} + assert _is_model_cost_zero(model="hidden-alias", llm_router=router) is False + def test_handles_router_without_zero_cost_cache_attribute(self): """Tolerate router-like objects (e.g. ``MagicMock`` stand-ins) that do not expose ``_zero_cost_cache`` — the auth check must still