From bded600d6d9f9ebf2020f8eb7f22882505f5a9cc Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 29 Sep 2026 17:55:41 +0000 Subject: [PATCH] fix(auth): key the zero-cost cache by the resolved model group Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 24 +++++------ .../test_unmapped_model_budget_enforcement.py | 40 ++++++++++++++----- 2 files changed, 42 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index f43a4a92a57..b3433b569ea 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -472,15 +472,15 @@ 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 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: @@ -488,7 +488,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None # 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 +505,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 +514,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 +529,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 +539,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 +549,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 ebb97c6a29b..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): """ @@ -344,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