diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index e8cf5fbc718..5aae3e25827 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -1,5 +1,5 @@ import json -from typing import List, Optional +from typing import List, Optional, Tuple import litellm from litellm._logging import verbose_proxy_logger @@ -53,18 +53,20 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ) # check if current model is in internal_model_max_budget - _current_model_budget_info = self._get_request_model_budget_config( + _budget_match = self._get_request_model_budget_key_and_config( model=model, internal_model_max_budget=internal_model_max_budget ) - if _current_model_budget_info is None: + if _budget_match is None: verbose_proxy_logger.debug(f"Model {model} not found in internal_model_max_budget") return True + _budget_key, _current_model_budget_info = _budget_match + # check if current model is within budget if _current_model_budget_info.max_budget and _current_model_budget_info.max_budget > 0: _current_spend = await self._get_virtual_key_spend_for_model( user_api_key_hash=user_api_key_dict.token, - model=model, + model=_budget_key, key_budget_config=_current_model_budget_info, ) if ( @@ -73,7 +75,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): and _current_spend > _current_model_budget_info.max_budget ): raise litellm.BudgetExceededError( - message=f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, exceeded budget for model={model}", + message=f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, exceeded budget for {self._describe_budget_scope(_budget_key, model)}", current_cost=_current_spend, max_budget=_current_model_budget_info.max_budget, entity_type=Litellm_EntityType.KEY.value, @@ -119,18 +121,20 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ) # check if current model is in internal_model_max_budget - _current_model_budget_info = self._get_request_model_budget_config( + _budget_match = self._get_request_model_budget_key_and_config( model=model, internal_model_max_budget=internal_model_max_budget ) - if _current_model_budget_info is None: + if _budget_match is None: verbose_proxy_logger.debug(f"Model {model} not found in end_user_model_max_budget") return True + _budget_key, _current_model_budget_info = _budget_match + # check if current model is within budget if _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, + model=_budget_key, key_budget_config=_current_model_budget_info, ) if ( @@ -139,7 +143,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): and _current_spend > _current_model_budget_info.max_budget ): raise litellm.BudgetExceededError( - message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}", + message=f"LiteLLM End User: {end_user_id}, exceeded budget for {self._describe_budget_scope(_budget_key, model)}", current_cost=_current_spend, max_budget=_current_model_budget_info.max_budget, entity_type=Litellm_EntityType.END_USER.value, @@ -209,10 +213,48 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): 1. Check if `model` is in `internal_model_max_budget` 2. If not, check if `model` without custom llm provider is in `internal_model_max_budget` + 3. If not, check if `model` belongs to a model-group budget (an entry whose + `models` list contains the request model) """ - return internal_model_max_budget.get(model, None) or internal_model_max_budget.get( - self._get_model_without_custom_llm_provider(model), None + _match = self._get_request_model_budget_key_and_config( + model=model, internal_model_max_budget=internal_model_max_budget ) + return _match[1] if _match is not None else None + + def _get_request_model_budget_key_and_config( + self, model: str, internal_model_max_budget: GenericBudgetConfigType + ) -> "Tuple[str, BudgetConfig] | None": + """ + Resolve the budget entry for the request model and the cache key its spend + is tracked under. + + Returns a ``(budget_key, config)`` tuple where ``budget_key`` is the model + name for per-model budgets and the group name for model-group budgets. Using + the group name as the cache key is what makes every model in the group draw + from one shared spend counter. + + Resolution order: + 1. Direct per-model match on `model` + 2. Per-model match on `model` without its custom llm provider prefix + 3. Model-group match: an entry whose `models` list contains the request + model (with or without the provider prefix) + """ + model_without_provider = self._get_model_without_custom_llm_provider(model) + for candidate in (model, model_without_provider): + config = internal_model_max_budget.get(candidate) + if config is not None and not config.models: + return candidate, 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): + return group_name, config + + return None + + def _describe_budget_scope(self, budget_key: str, model: str) -> str: + if budget_key == model or budget_key == self._get_model_without_custom_llm_provider(model): + return f"model={model}" + return f"model={model} (model_group={budget_key})" def _get_model_without_custom_llm_provider(self, model: str) -> str: if "/" in model: @@ -282,13 +324,12 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): internal_model_max_budget: GenericBudgetConfigType = {} for _model, _budget_info in user_api_key_model_max_budget.items(): internal_model_max_budget[_model] = BudgetConfig(**_budget_info) - key_budget_config = self._get_request_model_budget_config( + budget_match = self._get_request_model_budget_key_and_config( model=model, internal_model_max_budget=internal_model_max_budget ) - if key_budget_config is not None and key_budget_config.budget_duration: - virtual_spend_key = ( - f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{key_budget_config.budget_duration}" - ) + if budget_match is not None and budget_match[1].budget_duration: + budget_key, key_budget_config = budget_match + virtual_spend_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{budget_key}:{key_budget_config.budget_duration}" virtual_start_time_key = f"virtual_key_budget_start_time:{virtual_key}" await self._increment_spend_for_key( budget_config=key_budget_config, @@ -305,12 +346,13 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): internal_model_max_budget: GenericBudgetConfigType = {} for _model, _budget_info in user_api_key_end_user_model_max_budget.items(): internal_model_max_budget[_model] = BudgetConfig(**_budget_info) - key_budget_config = self._get_request_model_budget_config( + budget_match = self._get_request_model_budget_key_and_config( model=model, internal_model_max_budget=internal_model_max_budget ) - if key_budget_config is not None and key_budget_config.budget_duration: + if budget_match is not None and budget_match[1].budget_duration: + budget_key, key_budget_config = budget_match end_user_spend_key = ( - f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{key_budget_config.budget_duration}" + f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{budget_key}:{key_budget_config.budget_duration}" ) end_user_start_time_key = f"end_user_budget_start_time:{end_user_id}" await self._increment_spend_for_key( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 714ad372a5f..cad7af6db4a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3333,6 +3333,7 @@ class BudgetConfig(BaseModel): budget_duration: Optional[str] = None tpm_limit: Optional[int] = None rpm_limit: Optional[int] = None + models: Optional[List[str]] = None def __init__(self, **data: Any) -> None: # Map time_period to budget_duration if present diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index 55459721906..023768a52ca 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -558,3 +558,146 @@ async def test_async_log_success_event_skips_redis_push_without_redis(budget_lim kwargs, response_obj=None, start_time=None, end_time=None ) mock_push.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# Model-group (shared) budgets +# --------------------------------------------------------------------------- + + +def test_get_request_model_budget_key_and_config_group_match(budget_limiter): + """ + A budget entry that carries a `models` list defines a model-group budget. + Any request model in that list must resolve to the group name (the dict key), + so every model in the group shares one spend counter. + """ + internal_budget = { + "opus-family": GenericBudgetInfo( + budget_limit=50.0, + time_period="30d", + models=["claude-opus-4", "claude-opus-4-1"], + ), + "gpt-4": GenericBudgetInfo(budget_limit=100.0, time_period="1d"), + } + + # both group members resolve to the SAME budget key (the group name) + key_a, config_a = budget_limiter._get_request_model_budget_key_and_config( + model="claude-opus-4", internal_model_max_budget=internal_budget + ) + key_b, config_b = budget_limiter._get_request_model_budget_key_and_config( + model="claude-opus-4-1", internal_model_max_budget=internal_budget + ) + assert key_a == key_b == "opus-family" + assert config_a.max_budget == config_b.max_budget == 50.0 + + # provider-prefixed group member still resolves to the group + key_c, _ = budget_limiter._get_request_model_budget_key_and_config( + model="anthropic/claude-opus-4", internal_model_max_budget=internal_budget + ) + assert key_c == "opus-family" + + # a plain per-model entry resolves to the model name, not a group + key_d, config_d = budget_limiter._get_request_model_budget_key_and_config( + model="gpt-4", internal_model_max_budget=internal_budget + ) + assert key_d == "gpt-4" + assert config_d.max_budget == 100.0 + + # a model in no group and no per-model entry resolves to nothing + assert ( + budget_limiter._get_request_model_budget_key_and_config( + model="gemini-2.5-pro", internal_model_max_budget=internal_budget + ) + is None + ) + + +@pytest.mark.asyncio +async def test_is_key_within_model_budget_group_reads_group_counter(budget_limiter): + """ + Enforcement for any group member must read spend from the group counter + (budget_key == group name), so combined spend across the group is enforced. + """ + user_api_key = UserAPIKeyAuth( + token="test-key", + key_alias="test-alias", + model_max_budget={ + "opus-family": { + "budget_limit": 50.0, + "time_period": "30d", + "models": ["claude-opus-4", "claude-opus-4-1"], + } + }, + ) + + seen_models = [] + + async def _spend(user_api_key_hash, model, key_budget_config): + seen_models.append(model) + return 60.0 + + with patch.object( + budget_limiter, "_get_virtual_key_spend_for_model", side_effect=_spend + ): + # spend already over the shared 50.0 budget -> every member is blocked + for member in ("claude-opus-4", "claude-opus-4-1", "anthropic/claude-opus-4-1"): + with pytest.raises(litellm.BudgetExceededError): + await budget_limiter.is_key_within_model_budget(user_api_key, member) + + # all members were looked up under the single group counter + assert seen_models == ["opus-family", "opus-family", "opus-family"] + + +@pytest.mark.asyncio +async def test_async_log_success_event_group_members_share_one_counter(budget_limiter): + """ + Core regression for model-group budgets: spend from two DIFFERENT models in a + group must increment the SAME cache key (the group name) so the budget is + combined. Before this feature each model incremented its own key, letting a + developer exceed the intended total by spreading usage across the family. + """ + from litellm.proxy.hooks.model_max_budget_limiter import ( + VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX, + ) + + virtual_key = "test-key-hash" + budget_duration = "30d" + user_api_key_model_max_budget = { + "opus-family": { + "budget_limit": 50.0, + "time_period": budget_duration, + "models": ["claude-opus-4", "claude-opus-4-1"], + }, + } + + def _kwargs_for(model_group): + return { + "standard_logging_object": { + "response_cost": 0.10, + "model": model_group, + "model_group": model_group, + "metadata": {"user_api_key_hash": virtual_key}, + }, + "litellm_params": { + "metadata": { + "user_api_key_model_max_budget": user_api_key_model_max_budget, + }, + }, + } + + expected_key = ( + f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:opus-family:{budget_duration}" + ) + + with patch.object( + budget_limiter, "_increment_spend_for_key", new_callable=AsyncMock + ) as mock_increment: + await budget_limiter.async_log_success_event( + _kwargs_for("claude-opus-4"), None, None, None + ) + await budget_limiter.async_log_success_event( + _kwargs_for("claude-opus-4-1"), None, None, None + ) + + spend_keys = [c.kwargs["spend_key"] for c in mock_increment.call_args_list] + assert spend_keys == [expected_key, expected_key] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index e0223aa4af6..26a447f4c34 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -21395,6 +21395,8 @@ export interface components { budget_duration?: string | null; /** Max Budget */ max_budget?: number | null; + /** Models */ + models?: string[] | null; /** Rpm Limit */ rpm_limit?: number | null; /** Tpm Limit */