fix(proxy): match provider-qualified group members against bare request models

This commit is contained in:
mateo-berri 2026-07-23 07:35:57 +00:00
parent c27cd0fce2
commit 35824fd3f9
No known key found for this signature in database
2 changed files with 26 additions and 1 deletions

View file

@ -258,7 +258,13 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
return tuple(
(_group_name, _config)
for _group_name, _config in internal_model_max_budget.items()
if _config.models and (model in _config.models or model_without_provider in _config.models)
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
)
)
def _get_model_without_custom_llm_provider(self, model: str) -> str:

View file

@ -155,3 +155,22 @@ async def test_spend_at_exactly_group_budget_passes():
await _log_spend(limiter, "anthropic-opus-4-7", 10.0, OPUS_GROUP_BUDGET)
assert await limiter.is_key_within_model_budget(key, "anthropic-opus-4-8") is True
@pytest.mark.asyncio
async def test_group_budget_matches_bare_request_when_member_is_provider_qualified():
qualified_group_budget = {
"gpt4-family": {
"models": ["openai/gpt-4", "openai/gpt-4o"],
"budget_limit": 10.0,
"time_period": "30d",
}
}
limiter = _make_limiter()
key = _make_key(qualified_group_budget)
await _log_spend(limiter, "gpt-4", 11.0, qualified_group_budget)
with pytest.raises(litellm.BudgetExceededError, match="model group=gpt4-family"):
await limiter.is_key_within_model_budget(key, "gpt-4o")
assert await limiter.is_key_within_model_budget(key, "gpt-5.5") is True