diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 6fbeb681e48..d925b6c2298 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -472,11 +472,12 @@ class RateLimitStatus(TypedDict): limit_remaining: int rate_limit_type: Literal["requests", "tokens", "max_parallel_requests"] descriptor_key: str - # Only populated by the atomic_check_and_increment_by_n path. A caller - # matching a status back to its descriptor must key on (descriptor_key, - # descriptor_value) when this is present, not descriptor_key alone -- - # e.g. a batch charging several models' project ITPM/OTPM in one call - # produces multiple statuses sharing the same descriptor_key. + # Populated by the atomic_check_and_increment_by_n and windowed + # sliding-window paths. A caller matching a status back to its + # descriptor must key on (descriptor_key, descriptor_value) when this + # is present, not descriptor_key alone -- e.g. a batch charging several + # models' project ITPM/OTPM in one call, or a request carrying multiple + # rate-limited tags, produces statuses sharing the same descriptor_key. descriptor_value: NotRequired[ReadOnly[str]] @@ -506,6 +507,7 @@ class WindowKeyMetadata(TypedDict): tokens_limit: int | None window_size: int descriptor_key: str + descriptor_value: ReadOnly[str] class AtomicCounterMeta(TypedDict): @@ -1228,6 +1230,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): "limit_remaining": limit_remaining, "rate_limit_type": rate_limit_type, "descriptor_key": key_metadata[window_key]["descriptor_key"], + "descriptor_value": key_metadata[window_key]["descriptor_value"], } ) @@ -1532,6 +1535,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): "tokens_limit": int(tokens_limit) if tokens_limit is not None else None, "window_size": int(window_size), "descriptor_key": descriptor_key, + "descriptor_value": descriptor_value, } return keys_to_fetch, key_metadata, gauges @@ -3163,8 +3167,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for status in response["statuses"]: if status["code"] == "OVER_LIMIT": descriptor_key = status["descriptor_key"] + status_value: Final = status.get("descriptor_value") matching_descriptor = next( - (desc for desc in descriptors if desc["key"] == descriptor_key), + ( + desc + for desc in descriptors + if desc["key"] == descriptor_key + and (status_value is None or desc["value"] == status_value) + ), None, ) descriptor_value = matching_descriptor["value"] if matching_descriptor is not None else "unknown" diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 3c51297c5c3..2ac22d7ee03 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -5049,6 +5049,38 @@ async def test_tag_object_rpm_limit_enforced_v3(monkeypatch): assert calls == [("cell-1",), ("cell-1", "cell-2"), ("cell-1",), ("cell-2",), ("cell-2",), ("cell-2",)] +@pytest.mark.asyncio +async def test_tag_429_names_the_tag_that_is_over_its_limit(monkeypatch): + monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") + _request_stash.set(None) + resolver, _ = _static_tag_limits( + { + "cell-ok": TagRateLimit(rpm_limit=100, tpm_limit=None), + "cell-blocked": TagRateLimit(rpm_limit=1, tpm_limit=None), + } + ) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + tag_rate_limit_resolver=resolver, + ) + user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-tag-order")) + + async def call(tags: list[str]) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-3.5-turbo", "metadata": {"tags": tags}}, + call_type="", + ) + + await call(["cell-blocked"]) + with pytest.raises(HTTPException) as exc_info: + await call(["cell-ok", "cell-blocked"]) + assert exc_info.value.status_code == 429 + assert "cell-blocked" in str(exc_info.value.detail) + + @pytest.mark.asyncio async def test_tag_object_tpm_limit_enforced_v3(monkeypatch): monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") @@ -5096,9 +5128,16 @@ async def test_resolve_tag_rate_limits_from_db_reads_budget_row(monkeypatch): from litellm.models.tag import LiteLLM_TagTable from litellm.proxy import proxy_server from litellm.proxy.auth import auth_checks + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.hooks.parallel_request_limiter_v3 import resolve_tag_rate_limits_from_db + from litellm.proxy.utils import PrismaClient - async def fake_batch(tag_names, prisma_client, user_api_key_cache, **kwargs): + async def fake_batch( + tag_names: Sequence[str], + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + **kwargs: object, + ) -> dict[str, LiteLLM_TagTable]: return { "limited": LiteLLM_TagTable( tag_name="limited",