scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

fix(proxy): match namespaced model ids in group budgets

Browse source
This commit is contained in:
mateo-berri 2026-07-23 07:55:46 +00:00
parent 9133bd95c3
commit 8f12dbc89f
No known key found for this signature in database
2 changed files with 53 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
litellm/proxy/hooks/model_max_budget_limiter.py
View file

@ -287,19 +287,30 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
def _get_matching_model_group_budget_configs(
self, model: str, internal_model_max_budget: GenericBudgetConfigType
) -> tuple[tuple[str, BudgetConfig], ...]:
model_without_provider = self._get_model_without_custom_llm_provider(model)
request_variants = self._model_name_variants(model)
return tuple(
(_group_name, _config)
for _group_name, _config in internal_model_max_budget.items()
if _config.models
and any(
_member == model
or _member == model_without_provider
or self._get_model_without_custom_llm_provider(_member) == model
for _member in _config.models
_member in request_variants or model in self._model_name_variants(_member) for _member in _config.models
)
)
def _model_name_variants(self, model: str) -> frozenset[str]:
"""
The name itself plus its provider-stripped forms: without the first
segment (`huggingface/meta-llama/Llama-3.1-8B` -> `meta-llama/Llama-3.1-8B`)
and without everything before the last slash (`openai/gpt-4` -> `gpt-4`).
Group matching compares one side's raw name against the other side's
variants, never stripped-vs-stripped, so a provider-prefixed member
still pins its group to that provider's route.
"""
if "/" not in model:
return frozenset({model})
return frozenset({model, model.split("/", 1)[1], self._get_model_without_custom_llm_provider(model)})
def _get_model_without_custom_llm_provider(self, model: str) -> str:
if "/" in model:
return model.split("/")[-1]

37
tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py
View file

@ -253,3 +253,40 @@ async def test_end_user_group_budget_within_budget_passes():
assert await limiter.is_end_user_within_model_budget(END_USER_ID, OPUS_GROUP_BUDGET, "anthropic-opus-4-7") is True
assert await limiter.is_end_user_within_model_budget(END_USER_ID, OPUS_GROUP_BUDGET, "anthropic-opus-4-8") is True
@pytest.mark.asyncio
async def test_group_budget_matches_namespaced_model_with_provider_prefix():
namespaced_group_budget = {
"llama-family": {
"models": ["meta-llama/Llama-3.1-8B", "meta-llama/Llama-3.1-70B"],
"budget_limit": 10.0,
"time_period": "30d",
}
}
limiter = _make_limiter()
key = _make_key(namespaced_group_budget)
await _log_spend(limiter, "huggingface/meta-llama/Llama-3.1-8B", 11.0, namespaced_group_budget)
with pytest.raises(litellm.BudgetExceededError, match="model group=llama-family"):
await limiter.is_key_within_model_budget(key, "huggingface/meta-llama/Llama-3.1-70B")
with pytest.raises(litellm.BudgetExceededError, match="model group=llama-family"):
await limiter.is_key_within_model_budget(key, "meta-llama/Llama-3.1-8B")
@pytest.mark.asyncio
async def test_group_member_with_provider_prefix_does_not_match_other_provider():
pinned_group_budget = {
"openai-gpt4": {
"models": ["openai/gpt-4"],
"budget_limit": 10.0,
"time_period": "30d",
}
}
limiter = _make_limiter()
key = _make_key(pinned_group_budget)
await _log_spend(limiter, "gpt-4", 11.0, pinned_group_budget)
assert await limiter.is_key_within_model_budget(key, "azure/gpt-4") is True