scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(auth): key the zero-cost cache by the resolved model group

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
jesus 2026-09-29 17:55:41 +00:00
parent 68b0d08b5f
commit bded600d6d
2 changed files with 42 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
litellm/proxy/auth/auth_checks.py
View file

@ -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)

40
tests/test_litellm/proxy/auth/test_unmapped_model_budget_enforcement.py
View file

@ -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