Merge bded600d6d into 264b09ac8d

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)

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