feat(rate-limit): enforce per-customer per-model tpm and rpm limits in the v3 limiter (#50857)

* feat(rate-limit): enforce per-customer per-model tpm and rpm limits in the v3 limiter

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

* refactor(rate-limit): build model scope settlement targets from one owner tuple

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

---------

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

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -2507,6 +2507,26 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
             },
         )
 
+    def _end_user_model_rate_limit_descriptor(
+        self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None
+    ) -> RateLimitDescriptor | None:
+        model: Final = self._rate_limited_model(requested_model)
+        if model is None or not user_api_key_dict.end_user_id:
+            return None
+        tpm_limit: Final = model.limit_in(user_api_key_dict.end_user_model_tpm_limit)
+        rpm_limit: Final = model.limit_in(user_api_key_dict.end_user_model_rpm_limit)
+        if tpm_limit is None and rpm_limit is None:
+            return None
+        return RateLimitDescriptor(
+            key="model_per_end_user",
+            value=f"{user_api_key_dict.end_user_id}:{model.group}",
+            rate_limit={
+                "requests_per_unit": rpm_limit,
+                "tokens_per_unit": tpm_limit,
+                "window_size": self.window_size,
+            },
+        )
+
     def _add_model_per_key_rate_limit_descriptor(
         self,
         user_api_key_dict: UserAPIKeyAuth,
@@ -2999,6 +3019,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
                 requested_model=candidate,
                 descriptors=descriptors,
             )
+            if (end_user_model := self._end_user_model_rate_limit_descriptor(user_api_key_dict, candidate)) is not None:
+                descriptors.append(end_user_model)
 
         # Per-request-tag rate limits scoped to this key
         self._add_tag_per_key_rate_limit_descriptor(
@@ -4525,26 +4547,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
             targets.append(("end_user", user_api_key_end_user_id))
         if user_api_key_organization_id:
             targets.append(("organization", user_api_key_organization_id))
+        model_scope_owners: Final = (
+            ("model_per_key", user_api_key),
+            ("model_per_team", user_api_key_team_id),
+            ("model_per_end_user", f"{user_api_key_end_user_id}" if user_api_key_end_user_id else None),
+            ("model_per_organization", user_api_key_organization_id),
+            ("model_per_project", user_api_key_project_id),
+        )
         if model_group:
             for group in dict.fromkeys(model_groups or (model_group,)):
-                if user_api_key:
-                    targets.append(("model_per_key", f"{user_api_key}:{group}"))
-                if user_api_key_team_id:
-                    targets.append(("model_per_team", f"{user_api_key_team_id}:{group}"))
-                if user_api_key_organization_id:
-                    targets.append(
-                        (
-                            "model_per_organization",
-                            f"{user_api_key_organization_id}:{group}",
-                        )
-                    )
-                if user_api_key_project_id:
-                    targets.append(
-                        (
-                            "model_per_project",
-                            f"{user_api_key_project_id}:{group}",
-                        )
-                    )
+                targets.extend((scope, f"{owner}:{group}") for scope, owner in model_scope_owners if owner)
         if agent_id:
             targets.append(("agent", agent_id))
             if session_id:

tests/integration/management/test_customer_model_limits.py

@@ -0,0 +1,55 @@
+import uuid
+from typing import Final
+
+import httpx
+import pytest
+
+from tests.integration._support.client import JSON_OBJECT, Gateway, Scenario, object_value, string_value
+
+pytestmark: Final = pytest.mark.shares_worker_state
+
+
+def _customer(scenario: Scenario, **fields: object) -> str:
+    customer: Final = f"integration-customer-{uuid.uuid4().hex}"
+    created: Final = scenario.gateway.post("/customer/new", {"user_id": customer, **fields})  # pyright: ignore[reportArgumentType]  # JsonValue kwargs
+    scenario.cleanups.callback(scenario.gateway.post, "/customer/delete", {"user_ids": [customer]})
+    budget_id: Final = created.get("budget_id")
+    if isinstance(budget_id, str):
+        scenario.cleanups.callback(scenario.delete_budget, budget_id)
+    return customer
+
+
+def _chat(gateway: Gateway, key: str, model: str, customer: str) -> httpx.Response:
+    return gateway.request(
+        "POST",
+        "/v1/chat/completions",
+        {
+            "model": model,
+            "messages": [{"role": "user", "content": f"customer quota probe {uuid.uuid4().hex}"}],
+            "user": customer,
+        },
+        key=key,
+    )
+
+
+def _assert_customer_model_limit_error(response: httpx.Response, customer: str, model: str) -> None:
+    assert response.status_code == 429, response.text
+    message: Final = string_value(object_value(JSON_OBJECT.validate_json(response.content)["error"])["message"])
+    assert message.startswith(
+        f"Rate limit exceeded for model_per_end_user: {customer}:{model}. Limit type: requests. Current limit: 1,"
+    ), message
+
+
+def test_customer_per_model_rpm_limit_is_shared_across_keys_and_isolated_per_customer(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario:
+        capped: Final = scenario.model()
+        open_model: Final = scenario.model()
+        acme: Final = _customer(scenario, model_rpm_limit={capped: 1})
+        globex: Final = _customer(scenario, model_rpm_limit={capped: 1})
+        key_one: Final = scenario.key()
+        key_two: Final = scenario.key()
+
+        assert _chat(gateway, key_one, capped, acme).status_code == 200
+        _assert_customer_model_limit_error(_chat(gateway, key_two, capped, acme), acme, capped)
+        assert _chat(gateway, key_one, capped, globex).status_code == 200
+        assert _chat(gateway, key_one, open_model, acme).status_code == 200

tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py

@@ -8247,6 +8247,11 @@ async def test_batch_increment_refunds_pipelined_groups_declared_after_the_one_t
             {},
             "model_per_organization",
         ),
+        (
+            {"end_user_id": "customer", "end_user_model_rpm_limit": {"test-model": 1}},
+            {},
+            "model_per_end_user",
+        ),
     ],
 )
 @pytest.mark.parametrize("request_kind", ["count", "generation"])
@@ -8695,8 +8700,9 @@ async def test_model_group_alias_shares_deployment_default_rpm_bucket_with_its_t
         ),
         ({"org_id": "o", "organization_metadata": {"model_rpm_limit": {"target": 1}}}, "model_per_organization"),
         ({"project_id": "p", "project_metadata": {"model_rpm_limit": {"target": 1}}}, "model_per_project"),
+        ({"end_user_id": "customer", "end_user_model_rpm_limit": {"target": 1}}, "model_per_end_user"),
     ],
-    ids=["key_metadata", "team_metadata", "organization_metadata", "project_metadata"],
+    ids=["key_metadata", "team_metadata", "organization_metadata", "project_metadata", "end_user_budget"],
 )
 async def test_model_group_alias_shares_metadata_model_rpm_bucket_with_its_target(
     limits: dict[str, object], counter_scope: str, first_name: str, second_name: str
@@ -9911,3 +9917,27 @@ def test_media_models_resolve_to_each_deployment_base_model():
     assert _resolve_media_models(None, "whisper", None) == ("whisper",)
     assert _resolve_media_models(router, "whisper", None) == ("azure/whisper-1",)
     assert _resolve_media_models(router, "unknown", None) == ("unknown",)
+
+
+@pytest.mark.asyncio
+async def test_end_user_model_rpm_limit_is_shared_across_keys_and_isolated_per_customer() -> None:
+    cache: Final = DualCache()
+    handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
+    acme_via_key_one: Final = UserAPIKeyAuth(
+        api_key=hash_token("sk-acme-1"), end_user_id="acme", end_user_model_rpm_limit={"target": 1}
+    )
+    acme_via_key_two: Final = UserAPIKeyAuth(
+        api_key=hash_token("sk-acme-2"), end_user_id="acme", end_user_model_rpm_limit={"target": 1}
+    )
+    globex: Final = UserAPIKeyAuth(
+        api_key=hash_token("sk-globex"), end_user_id="globex", end_user_model_rpm_limit={"target": 1}
+    )
+
+    await _rpm_request(handler, cache, acme_via_key_one, "target")
+    with pytest.raises(HTTPException) as exc:
+        await _rpm_request(handler, cache, acme_via_key_two, "target")
+    assert exc.value.status_code == 429
+    assert "model_per_end_user: acme:target" in str(exc.value.detail), exc.value.detail
+
+    await _rpm_request(handler, cache, globex, "target")
+    await _rpm_request(handler, cache, acme_via_key_one, "other-model")