diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index e5fe5e875ba..7e5afd58d79 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -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() diff --git a/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py index fd552fea934..a7aadf1192f 100644 --- a/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py @@ -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