feat(proxy): support shared model-group budgets for API keys

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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(

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

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]

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 */