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

feat(proxy): enforce model group budgets for end users

Browse source
This commit is contained in:
mateo-berri 2026-07-23 07:43:42 +00:00
parent d722eb015d
commit 34ba9944bc
No known key found for this signature in database
2 changed files with 97 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -146,10 +146,13 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
)
if _current_model_budget_info is None:
verbose_proxy_logger.debug(f"Model {model} not found in end_user_model_max_budget")
return True
# check if current model is within budget
if _current_model_budget_info.max_budget and _current_model_budget_info.max_budget > 0:
if (
_current_model_budget_info is not None
and _current_model_budget_info.max_budget
and _current_model_budget_info.max_budget > 0
):
_current_spend = await self._get_end_user_spend_for_model(
end_user_id=end_user_id,
model=model,
@ -168,6 +171,25 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
entity_id=end_user_id,
)
for _group_name, _group_budget_info in self._get_matching_model_group_budget_configs(
model=model, internal_model_max_budget=internal_model_max_budget
):
if not _group_budget_info.max_budget or _group_budget_info.max_budget <= 0:
continue
_group_spend = await self._get_end_user_spend_for_model_group(
end_user_id=end_user_id,
model_group_name=_group_name,
key_budget_config=_group_budget_info,
)
if _group_spend is not None and _group_spend > _group_budget_info.max_budget:
raise litellm.BudgetExceededError(
message=f"LiteLLM End User: {end_user_id}, exceeded budget for model group={_group_name}, model={model}",
current_cost=_group_spend,
max_budget=_group_budget_info.max_budget,
entity_type=Litellm_EntityType.END_USER.value,
entity_id=end_user_id,
)
return True
async def _get_end_user_spend_for_model(
@ -235,6 +257,18 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
)
return await self.dual_cache.async_get_cache(key=model_group_spend_cache_key)
async def _get_end_user_spend_for_model_group(
self,
end_user_id: str,
model_group_name: str,
key_budget_config: BudgetConfig,
) -> float | None:
model_group_spend_cache_key = (
f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model_group_name}:"
f"{key_budget_config.budget_duration}"
)
return await self.dual_cache.async_get_cache(key=model_group_spend_cache_key)
def _get_request_model_budget_config(
self, model: str, internal_model_max_budget: GenericBudgetConfigType
) -> Optional[BudgetConfig]:
@ -388,6 +422,22 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
start_time_key=end_user_start_time_key,
response_cost=response_cost,
)
for _group_name, _group_budget_config in self._get_matching_model_group_budget_configs(
model=model, internal_model_max_budget=internal_model_max_budget
):
if _group_budget_config.budget_duration is None:
continue
end_user_group_spend_key = (
f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{_group_name}:"
f"{_group_budget_config.budget_duration}"
)
end_user_group_start_time_key = f"end_user_budget_start_time:{end_user_id}:{_group_name}"
await self._increment_spend_for_key(
budget_config=_group_budget_config,
spend_key=end_user_group_spend_key,
start_time_key=end_user_group_start_time_key,
response_cost=response_cost,
)
if self.dual_cache.redis_cache is not None:
await self._push_in_memory_increments_to_redis()

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

@ -208,3 +208,48 @@ async def test_group_without_time_period_does_not_track_or_block():
assert limiter.dual_cache.in_memory_cache.cache_dict == {}
assert await limiter.is_key_within_model_budget(key, "anthropic-opus-4-7") is True
END_USER_ID = "end-user-1"
async def _log_end_user_spend(
limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter,
model: str,
response_cost: float,
end_user_model_max_budget: dict,
) -> None:
kwargs = {
"standard_logging_object": {
"response_cost": response_cost,
"model": model,
"end_user": END_USER_ID,
"metadata": {"user_api_key_end_user_id": END_USER_ID},
},
"litellm_params": {"metadata": {"user_api_key_end_user_model_max_budget": end_user_model_max_budget}},
}
await limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None)
@pytest.mark.asyncio
async def test_end_user_group_budget_shared_across_models():
limiter = _make_limiter()
await _log_end_user_spend(limiter, "anthropic-opus-4-7", 11.0, OPUS_GROUP_BUDGET)
with pytest.raises(litellm.BudgetExceededError, match="model group=opus-family"):
await limiter.is_end_user_within_model_budget(END_USER_ID, OPUS_GROUP_BUDGET, "anthropic-opus-4-7")
with pytest.raises(litellm.BudgetExceededError, match="model group=opus-family"):
await limiter.is_end_user_within_model_budget(END_USER_ID, OPUS_GROUP_BUDGET, "anthropic-opus-4-8")
assert await limiter.is_end_user_within_model_budget(END_USER_ID, OPUS_GROUP_BUDGET, "anthropic-sonnet-5") is True
@pytest.mark.asyncio
async def test_end_user_group_budget_within_budget_passes():
limiter = _make_limiter()
await _log_end_user_spend(limiter, "anthropic-opus-4-7", 4.0, OPUS_GROUP_BUDGET)
await _log_end_user_spend(limiter, "anthropic-opus-4-8", 5.0, OPUS_GROUP_BUDGET)
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