diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 9fb026a64d9..ae9fa548cc9 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/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: diff --git a/tests/integration/management/test_customer_model_limits.py b/tests/integration/management/test_customer_model_limits.py new file mode 100644 index 00000000000..f1c207661d2 --- /dev/null +++ b/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 diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 741b6cb74b5..a39d02a2788 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/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")