mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(proxy): match provider-qualified group members against bare request models
This commit is contained in:
parent
c27cd0fce2
commit
35824fd3f9
2 changed files with 26 additions and 1 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue