scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-03 02:22:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

merge: fold #27707 (provider/llm_provider on internal rate-limit errors) into #27687

Browse source 
Fold the provider-attribution work from #27707 onto #27687's unified
ProxyRateLimitError so a single 429 carries:

  - category (LITELLM_RATE_LIMIT / LITELLM_BATCH_RATE_LIMIT / vendor)
  - rate_limit_type (REQUESTS / TOKENS / CONCURRENT_REQUESTS / BUDGET / MAX_ITERATIONS)
  - model (resolved via litellm.get_llm_provider so the prefix is stripped)
  - llm_provider (with 'litellm_proxy' fallback when model is missing/unparseable)

Mechanics
- Drop ProxyHTTPRateLimitError from rate_limiter_utils (it was a parallel
  HTTPException+RateLimitError class introduced in #27707). Keep
  resolve_llm_provider_for_rate_limit() as the single helper for resolving
  (model, llm_provider) at the limiter raise site.
- At every internal raise site (parallel_request_limiter v1/v3,
  dynamic_rate_limiter v1/v3, batch_rate_limiter, max_budget_limiter,
  max_budget_per_session_limiter, max_iterations_limiter) call the resolver
  and pass model/llm_provider into ProxyRateLimitError alongside the
  category/rate_limit_type fields.
- _PROXY_MaxParallelRequestsHandler.raise_rate_limit_error now takes both
  rate_limit_type and requested_model (NoReturn return type preserved).
- _PROXY_BatchRateLimiter._raise_rate_limit_error takes requested_model.
- ProxyRateLimitError.__init__ normalizes None model/llm_provider so callers
  can pass through resolver output without a defensive fallback at every site.

Tests
- Update test_proxy_rate_limit_provider_field.py (added by #27707) to import
  ProxyRateLimitError from common_utils, drop status_code=429 from
  constructor calls, and replace the prometheus exception_class_name
  assertions with the back-compat shim from #27687 (returns 'HTTPException'
  to preserve existing dashboards).
- All 34 #27707 tests + all 142 #27687 tests + all 218 proxy/hooks tests pass.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
This commit is contained in:
Cursor Agent 2026-05-14 00:35:38 +00:00
commit e2370b26d9
No known key found for this signature in database
11 changed files with 1107 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/proxy/common_utils/proxy_rate_limit_error.py
View file

@ -140,8 +140,15 @@ class ProxyRateLimitError(HTTPException, RateLimitError): # type: ignore[misc]
] = RateLimitErrorCategory.LITELLM_RATE_LIMIT,
rate_limit_type: Optional[Union[str, RateLimitType]] = None,
model: Optional[str] = None,
llm_provider: str = "litellm_proxy",
llm_provider: Optional[str] = "litellm_proxy",
):
# Normalize None → safe defaults so callers (and the resolver helper
# in `rate_limiter_utils`) can pass `None` without producing an
# instance whose `.llm_provider` attribute is `None` — that would
# break Prometheus' `_get_exception_class_name` (it calls
# `.capitalize()` on the provider string).
model = model or ""
llm_provider = llm_provider or "litellm_proxy"
message = _coerce_message(detail)
stringified_headers: Optional[Dict[str, str]] = (
{k: str(v) for k, v in headers.items()} if headers else None

9
litellm/proxy/hooks/batch_rate_limiter.py
View file

@ -33,6 +33,7 @@ from litellm.exceptions import RateLimitErrorCategory, RateLimitType
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -106,6 +107,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
descriptors: List["RateLimitDescriptor"],
batch_usage: BatchFileUsage,
limit_type: str,
requested_model: Optional[str] = None,
) -> NoReturn:
"""Raise :class:`ProxyRateLimitError` (a 429) for batch rate limit exceeded."""
from datetime import datetime
@ -150,6 +152,9 @@ class _PROXY_BatchRateLimiter(CustomLogger):
f"Limit resets at: {reset_time_formatted}"
)
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(
requested_model
)
raise ProxyRateLimitError(
detail=detail,
headers={
@ -166,6 +171,8 @@ class _PROXY_BatchRateLimiter(CustomLogger):
if limit_type == "tokens"
else RateLimitType.REQUESTS
),
model=resolved_model,
llm_provider=llm_provider,
)
async def _check_and_increment_batch_counters(
@ -207,6 +214,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
)
if rate_limit_response["overall_code"] == "OVER_LIMIT":
requested_model = data.get("model") if data else None
for status in rate_limit_response["statuses"]:
if status["code"] == "OVER_LIMIT":
self._raise_rate_limit_error(
@ -214,6 +222,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
descriptors,
batch_usage,
status["rate_limit_type"],
requested_model=requested_model,
)
async def count_input_file_usage(

20
litellm/proxy/hooks/dynamic_rate_limiter.py
View file

@ -11,15 +11,17 @@ from litellm import ModelResponse, Router
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.exceptions import RateLimitType
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.hooks.rate_limiter_utils import (
convert_priority_to_percent,
resolve_llm_provider_for_rate_limit,
)
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import CallTypesLiteral
from litellm.utils import get_utc_datetime
from .rate_limiter_utils import convert_priority_to_percent
class DynamicRateLimiterCache:
"""
@ -218,6 +220,9 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
)
### CHECK TPM ###
if available_tpm is not None and available_tpm == 0:
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(
data.get("model")
)
raise ProxyRateLimitError(
detail={
"error": "Key={} over available TPM={}. Model TPM={}, Active keys={}".format(
@ -228,10 +233,14 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
)
},
rate_limit_type=RateLimitType.TOKENS,
model=data.get("model"),
model=resolved_model,
llm_provider=llm_provider,
)
### CHECK RPM ###
elif available_rpm is not None and available_rpm == 0:
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(
data.get("model")
)
raise ProxyRateLimitError(
detail={
"error": "Key={} over available RPM={}. Model RPM={}, Active keys={}".format(
@ -242,7 +251,8 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
)
},
rate_limit_type=RateLimitType.REQUESTS,
model=data.get("model"),
model=resolved_model,
llm_provider=llm_provider,
)
elif available_rpm is not None or available_tpm is not None:
## UPDATE CACHE WITH ACTIVE PROJECT

15
litellm/proxy/hooks/dynamic_rate_limiter_v3.py
View file

@ -23,7 +23,10 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
RateLimitDescriptorRateLimitObject,
_PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.proxy.hooks.rate_limiter_utils import convert_priority_to_percent
from litellm.proxy.hooks.rate_limiter_utils import (
convert_priority_to_percent,
resolve_llm_provider_for_rate_limit,
)
from litellm.proxy.utils import InternalUsageCache
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import CallTypesLiteral
@ -491,6 +494,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
)
if atomic_response["overall_code"] == "OVER_LIMIT":
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(model)
for status in atomic_response["statuses"]:
if status["code"] != "OVER_LIMIT":
continue
@ -513,7 +517,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
rate_limit_type=map_v3_rate_limit_type(
status["rate_limit_type"]
),
model=model,
model=resolved_model,
llm_provider=llm_provider,
)
if descriptor_key == "priority_model":
verbose_proxy_logger.debug(
@ -540,7 +545,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
rate_limit_type=map_v3_rate_limit_type(
status["rate_limit_type"]
),
model=model,
model=resolved_model,
llm_provider=llm_provider,
)
# Fail-closed guard: overall_code says OVER_LIMIT but no status
@ -570,11 +576,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
rate_limit_type=map_v3_rate_limit_type(
offending["rate_limit_type"] if offending else None
),
model=model,
headers={
"retry-after": str(self.v3_limiter.window_size),
"x-litellm-priority": priority or "default",
},
model=resolved_model,
llm_provider=llm_provider,
)
# If priority is NOT enforced (saturation below threshold) but

8
litellm/proxy/hooks/max_budget_limiter.py
View file

@ -4,9 +4,10 @@ from litellm import verbose_logger
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.exceptions import RateLimitType
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
class _PROXY_MaxBudgetLimiter(CustomLogger):
@ -65,9 +66,14 @@ class _PROXY_MaxBudgetLimiter(CustomLogger):
# CHECK IF REQUEST ALLOWED
if curr_spend >= max_budget:
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(
data.get("model") if data else None
)
raise ProxyRateLimitError(
detail="Max budget limit reached.",
rate_limit_type=RateLimitType.BUDGET,
model=resolved_model,
llm_provider=llm_provider,
)
except HTTPException as e:
raise e

8
litellm/proxy/hooks/max_budget_per_session_limiter.py
View file

@ -20,9 +20,10 @@ from typing import TYPE_CHECKING, Any, Optional, Union
from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.exceptions import RateLimitType
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
if TYPE_CHECKING:
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
@ -112,6 +113,9 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger):
)
if current_spend >= max_budget:
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(
data.get("model") if data else None
)
raise ProxyRateLimitError(
detail=(
f"Session budget exceeded for session {session_id}. "
@ -119,6 +123,8 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger):
f"max_budget_per_session: ${max_budget:.2f}."
),
rate_limit_type=RateLimitType.BUDGET,
model=resolved_model,
llm_provider=llm_provider,
)
return None

8
litellm/proxy/hooks/max_iterations_limiter.py
View file

@ -16,9 +16,10 @@ from typing import TYPE_CHECKING, Any, Optional, Union
from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.exceptions import RateLimitType
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
if TYPE_CHECKING:
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
@ -116,12 +117,17 @@ class _PROXY_MaxIterationsHandler(CustomLogger):
current_count = await self._increment_and_get(cache_key)
if current_count > max_iterations:
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(
data.get("model") if data else None
)
raise ProxyRateLimitError(
detail=(
f"Max iterations exceeded for session {session_id}. "
f"Current count: {current_count}, max_iterations: {max_iterations}."
),
rate_limit_type=RateLimitType.MAX_ITERATIONS,
model=resolved_model,
llm_provider=llm_provider,
)
verbose_proxy_logger.debug(

26
litellm/proxy/hooks/parallel_request_limiter.py
View file

@ -12,12 +12,13 @@ from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs
from litellm.proxy._types import CommonProxyErrors, CurrentItemRateLimit, UserAPIKeyAuth
from litellm.exceptions import RateLimitType
from litellm.proxy.auth.auth_utils import (
get_key_model_rpm_limit,
get_key_model_tpm_limit,
)
from litellm.exceptions import RateLimitType
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -87,6 +88,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
self.raise_rate_limit_error(
additional_details=f"{CommonProxyErrors.max_parallel_request_limit_reached.value}. Hit limit for {rate_limit_type}. Current limits: max_parallel_requests: {max_parallel_requests}, tpm_limit: {tpm_limit}, rpm_limit: {rpm_limit}",
rate_limit_type=triggered_type,
requested_model=data.get("model") if data else None,
)
new_val = {
"current_requests": 1,
@ -117,10 +119,16 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
triggered_type = RateLimitType.TOKENS
else:
triggered_type = RateLimitType.REQUESTS
requested_model = data.get("model") if data else None
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(
requested_model
)
raise ProxyRateLimitError(
detail=f"LiteLLM Rate Limit Handler for rate limit type = {rate_limit_type}. {CommonProxyErrors.max_parallel_request_limit_reached.value}. current rpm: {current['current_rpm']}, rpm limit: {rpm_limit}, current tpm: {current['current_tpm']}, tpm limit: {tpm_limit}, current max_parallel_requests: {current['current_requests']}, max_parallel_requests: {max_parallel_requests}",
headers={"retry-after": str(self.time_to_next_minute())},
rate_limit_type=triggered_type,
model=resolved_model,
llm_provider=llm_provider,
)
await self.internal_usage_cache.async_batch_set_cache(
@ -147,6 +155,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
self,
additional_details: Optional[str] = None,
rate_limit_type: Optional[RateLimitType] = None,
requested_model: Optional[str] = None,
) -> NoReturn:
"""
Raise a 429 with a retry-after header for litellm-proxy parallel-request limits.
@ -163,6 +172,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
(the global-limit branch in ``async_pre_call_hook`` and the
all-zeros base case in ``check_key_in_limits``). Callers that know
the dimension exactly should pass it explicitly.
``requested_model`` is resolved via :func:`get_llm_provider` so the
raised exception carries ``llm_provider`` (and a stripped ``model``)
for downstream loggers (Prometheus failure metric, observability
callbacks). Falls back to ``llm_provider="litellm_proxy"`` when the
model is missing or unparseable — see
:func:`resolve_llm_provider_for_rate_limit`.
"""
# additional_details is optional; build the detail with a None-guard
# so callers that pass nothing don't get the literal string "None"
@ -170,10 +186,15 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
error_message = "Max parallel request limit reached"
if additional_details is not None:
error_message = error_message + " " + additional_details
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(
requested_model
)
raise ProxyRateLimitError(
detail=error_message,
headers={"retry-after": str(self.time_to_next_minute())},
rate_limit_type=rate_limit_type or RateLimitType.CONCURRENT_REQUESTS,
model=resolved_model,
llm_provider=llm_provider,
)
async def get_all_cache_objects(
@ -265,7 +286,8 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
# if above -> raise error
if current_global_requests >= global_max_parallel_requests:
self.raise_rate_limit_error(
additional_details=f"Hit Global Limit: Limit={global_max_parallel_requests}, current: {current_global_requests}"
additional_details=f"Hit Global Limit: Limit={global_max_parallel_requests}, current: {current_global_requests}",
requested_model=data.get("model") if data else None,
)
# if below -> increment
else:

9
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -36,6 +36,7 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import (
ProxyRateLimitError,
map_v3_rate_limit_type,
)
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject
from litellm.types.utils import ModelResponse, Usage
@ -1839,6 +1840,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self,
response: RateLimitResponse,
descriptors: List[RateLimitDescriptor],
requested_model: Optional[str] = None,
) -> None:
"""Handle rate limit exceeded by raising :class:`ProxyRateLimitError` (a 429)."""
for status in response["statuses"]:
@ -1871,6 +1873,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
f"Limit resets at: {reset_time_formatted}"
)
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(
requested_model
)
raise ProxyRateLimitError(
detail=detail,
headers={
@ -1879,6 +1884,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
"reset_at": reset_time_formatted,
},
rate_limit_type=map_v3_rate_limit_type(status["rate_limit_type"]),
model=resolved_model,
llm_provider=llm_provider,
)
async def async_pre_call_hook(
@ -1979,6 +1986,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self._handle_rate_limit_error(
response=response,
descriptors=descriptors,
requested_model=requested_model,
)
else:
# add descriptors to request headers
@ -2024,6 +2032,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self._handle_rate_limit_error(
response=tpm_response,
descriptors=descriptors,
requested_model=requested_model,
)
else:
data["_litellm_rate_limit_descriptors"] = descriptors

47
litellm/proxy/hooks/rate_limiter_utils.py
View file

@ -2,11 +2,56 @@
Shared utility functions for rate limiter hooks.
"""
from typing import Optional, Union
from typing import Optional, Tuple, Union
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import PriorityReservationDict
PROXY_LLM_PROVIDER_FALLBACK = "litellm_proxy"
def resolve_llm_provider_for_rate_limit(
model: Optional[str],
) -> Tuple[str, str]:
"""
Resolve ``(model, llm_provider)`` for a request being rejected by an
internal proxy-side rate-limit hook.
These hooks fire from ``async_pre_call_hook`` — well before
:func:`litellm.get_llm_provider` is invoked anywhere else in the request
lifecycle — so the raised 429 would otherwise have an empty
``llm_provider`` field, making the resulting Prometheus
``litellm_proxy_failed_requests_metric`` show up with
``exception_class="RateLimitError"`` and no provider attribution.
Wrapped defensively: if ``model`` is missing, malformed, or
``get_llm_provider`` raises (unknown alias, router-only model, etc.) we
fall back to ``("", "litellm_proxy")`` so we never break the request path
by piling a second exception on top of the rate-limit one we're trying to
raise.
"""
if not model:
return "", PROXY_LLM_PROVIDER_FALLBACK
try:
resolved_model, custom_llm_provider, _, _ = litellm.get_llm_provider(
model=model,
)
return (
resolved_model or model,
custom_llm_provider or PROXY_LLM_PROVIDER_FALLBACK,
)
except Exception as e:
verbose_proxy_logger.debug(
"rate_limiter_utils.resolve_llm_provider_for_rate_limit: "
"could not resolve provider for model=%s, falling back to %s. err=%s",
model,
PROXY_LLM_PROVIDER_FALLBACK,
str(e),
)
return model, PROXY_LLM_PROVIDER_FALLBACK
def convert_priority_to_percent(
value: Union[float, PriorityReservationDict], model_info: Optional[ModelGroupInfo]

964
tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py Normal file
View file

@ -0,0 +1,964 @@
"""
Regression tests for the "provider field missing" bug on proxy-side
rate-limit errors.
Background
----------
The proxy's internal rate-limit hooks (parallel_request_limiter,
parallel_request_limiter_v3, dynamic_rate_limiter, dynamic_rate_limiter_v3,
batch_rate_limiter, max_budget_limiter, max_iterations_limiter,
max_budget_per_session_limiter) all fire from ``async_pre_call_hook`` —
*before* :func:`litellm.get_llm_provider` runs anywhere else in the request
lifecycle.
Until now, those hooks raised a bare ``HTTPException(429, ...)`` which carries
no ``llm_provider`` / ``model`` attribute. Downstream:
- The Prometheus ``litellm_proxy_failed_requests_metric`` reads
``exception.llm_provider`` via ``_get_exception_class_name`` — it came back
empty, so dashboards showed ``exception_class="HTTPException"`` with no
provider attribution.
- Observability callbacks that ``isinstance(e, RateLimitError)`` for
category routing missed these entirely.
The fix wraps every internal raise site in
:class:`ProxyRateLimitError` (an ``HTTPException`` *and* a
``litellm.RateLimitError``), and resolves ``model`` / ``llm_provider`` from
``data["model"]`` via :func:`get_llm_provider`. When the model is missing or
unparseable we fall back to ``llm_provider="litellm_proxy"`` so we never break
the request path with a second exception.
These tests pin both the happy path (provider correctly resolved) and the
fallback path (unknown model, missing model) for every limiter.
"""
import sys
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
import litellm
from litellm.caching.caching import DualCache
from litellm.exceptions import RateLimitError
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.batch_rate_limiter import (
BatchFileUsage,
_PROXY_BatchRateLimiter,
)
from litellm.proxy.hooks.dynamic_rate_limiter import _PROXY_DynamicRateLimitHandler
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
_PROXY_DynamicRateLimitHandlerV3,
)
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
from litellm.proxy.hooks.max_budget_per_session_limiter import (
_PROXY_MaxBudgetPerSessionHandler,
)
from litellm.proxy.hooks.max_iterations_limiter import _PROXY_MaxIterationsHandler
from litellm.proxy.hooks.parallel_request_limiter import (
_PROXY_MaxParallelRequestsHandler,
)
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
_PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.hooks.rate_limiter_utils import (
PROXY_LLM_PROVIDER_FALLBACK,
resolve_llm_provider_for_rate_limit,
)
from litellm.proxy.utils import InternalUsageCache
from litellm.types.agents import AgentResponse
# ---------------------------------------------------------------------------
# Helper class itself
# ---------------------------------------------------------------------------
class TestProxyRateLimitErrorClass:
"""Pin the dual ``HTTPException`` + ``RateLimitError`` shape."""
def test_is_both_http_exception_and_rate_limit_error(self):
e = ProxyRateLimitError(
detail="boom",
model="gpt-4o-mini",
llm_provider="openai",
)
# FastAPI handler keys off HTTPException to render the 429.
assert isinstance(e, HTTPException)
# Prometheus / observability key off RateLimitError + .llm_provider.
assert isinstance(e, RateLimitError)
assert e.status_code == 429
assert e.model == "gpt-4o-mini"
assert e.llm_provider == "openai"
# ProxyRateLimitError prefixes message via RateLimitError.__init__.
assert "boom" in e.message
assert e.detail == "boom"
def test_dict_detail_is_stringified_for_message(self):
# Some hooks pass a dict detail (e.g. dynamic_rate_limiter v1) — the
# `message` attr (read by RateLimitError.__str__ and observability
# callbacks) must still be a string.
e = ProxyRateLimitError(
detail={"error": "over rpm"},
model="claude-3-5-sonnet",
llm_provider="anthropic",
)
assert isinstance(e.message, str)
assert "over rpm" in e.message
def test_defaults_to_litellm_proxy_provider(self):
e = ProxyRateLimitError(detail="x")
assert e.llm_provider == PROXY_LLM_PROVIDER_FALLBACK
assert e.model == ""
def test_none_provider_normalized_to_fallback(self):
e = ProxyRateLimitError(
detail="x",
model=None,
llm_provider=None,
)
assert e.llm_provider == PROXY_LLM_PROVIDER_FALLBACK
assert e.model == ""
class TestResolveLLMProviderForRateLimit:
@pytest.mark.parametrize(
"model, expected_provider",
[
("gpt-4o-mini", "openai"),
("anthropic/claude-3-5-sonnet", "anthropic"),
("bedrock/meta.llama3-1-70b-instruct-v1:0", "bedrock"),
],
)
def test_known_models_resolve_provider(self, model, expected_provider):
resolved_model, provider = resolve_llm_provider_for_rate_limit(model)
assert provider == expected_provider
assert resolved_model # non-empty
@pytest.mark.parametrize("model", [None, "", "totally-not-a-real-model-name"])
def test_missing_or_unknown_model_falls_back(self, model):
# Must never raise — the resolver wraps `get_llm_provider` defensively
# because raising here would mask the rate-limit error we're trying
# to surface to the user.
resolved_model, provider = resolve_llm_provider_for_rate_limit(model)
assert provider == PROXY_LLM_PROVIDER_FALLBACK
# Resolver returns the input model verbatim on the unknown branch so
# the `.model` attribute is never silently swapped to a different one.
if not model:
assert resolved_model == ""
else:
assert resolved_model == model
def test_get_llm_provider_raising_is_swallowed(self):
# If get_llm_provider itself blows up (unexpected error), we still
# fall back rather than letting the secondary exception escape.
with patch.object(
litellm,
"get_llm_provider",
side_effect=RuntimeError("boom"),
):
resolved_model, provider = resolve_llm_provider_for_rate_limit("anything")
assert provider == PROXY_LLM_PROVIDER_FALLBACK
assert resolved_model == "anything"
# ---------------------------------------------------------------------------
# parallel_request_limiter v1
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_parallel_request_limiter_v1_populates_provider_when_at_rpm_limit():
"""
Trip the per-key RPM cap and assert the raised exception carries
``model`` / ``llm_provider`` resolved from ``data["model"]``.
"""
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache())
)
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-rl-test",
max_parallel_requests=10,
rpm_limit=1,
tpm_limit=10,
)
data = {"model": "gpt-4o-mini"}
# First request consumes the budget.
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data=data,
call_type="completion",
)
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data=data,
call_type="completion",
)
exc = exc_info.value
assert exc.status_code == 429
assert isinstance(exc, RateLimitError)
assert exc.llm_provider == "openai"
assert exc.model == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_parallel_request_limiter_v1_zero_limit_path_populates_provider():
"""
When tpm_limit / rpm_limit is 0 the limiter takes the
``raise_rate_limit_error`` path. That path receives ``requested_model``
via the call-site change and must pass it through.
"""
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache())
)
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-rl-zero",
max_parallel_requests=0,
rpm_limit=10,
tpm_limit=10,
)
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={"model": "anthropic/claude-3-5-sonnet"},
call_type="completion",
)
exc = exc_info.value
assert exc.status_code == 429
assert isinstance(exc, RateLimitError)
assert exc.llm_provider == "anthropic"
assert exc.model == "claude-3-5-sonnet"
@pytest.mark.asyncio
async def test_parallel_request_limiter_v1_global_limit_populates_provider():
"""global_max_parallel_requests path also threads the model through."""
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache())
)
user_api_key_dict = UserAPIKeyAuth(api_key="sk-global")
# Pre-fill the global counter so the next call exceeds it.
await handler.internal_usage_cache.async_set_cache(
key="global_max_parallel_requests",
value=5,
local_only=True,
litellm_parent_otel_span=None,
)
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={
"model": "bedrock/meta.llama3-1-70b-instruct-v1:0",
"metadata": {"global_max_parallel_requests": 1},
},
call_type="completion",
)
exc = exc_info.value
assert exc.status_code == 429
assert exc.llm_provider == "bedrock"
assert exc.model == "meta.llama3-1-70b-instruct-v1:0"
@pytest.mark.asyncio
async def test_parallel_request_limiter_v1_unknown_model_falls_back():
"""
When ``data["model"]`` is unparseable, the resolver falls back to
``litellm_proxy`` — and crucially does *not* leak a secondary exception.
"""
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache())
)
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-rl-unknown",
max_parallel_requests=10,
rpm_limit=1,
tpm_limit=10,
)
data = {"model": "totally-not-a-real-model"}
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data=data,
call_type="completion",
)
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data=data,
call_type="completion",
)
exc = exc_info.value
assert exc.status_code == 429
assert exc.llm_provider == PROXY_LLM_PROVIDER_FALLBACK
# Resolver returns the input verbatim so we don't silently relabel the
# model in the user-facing 429 detail.
assert exc.model == "totally-not-a-real-model"
@pytest.mark.asyncio
async def test_parallel_request_limiter_v1_missing_model_falls_back():
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache())
)
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-rl-no-model",
max_parallel_requests=10,
rpm_limit=1,
tpm_limit=10,
)
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
exc = exc_info.value
assert exc.llm_provider == PROXY_LLM_PROVIDER_FALLBACK
assert exc.model == ""
# ---------------------------------------------------------------------------
# parallel_request_limiter v3
# ---------------------------------------------------------------------------
def _v3_over_limit_response(rate_limit_type: str = "rpm") -> dict:
return {
"overall_code": "OVER_LIMIT",
"statuses": [
{
"code": "OVER_LIMIT",
"descriptor_key": "key",
"current_limit": 1,
"limit_remaining": -1,
"rate_limit_type": rate_limit_type,
}
],
}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model, expected_provider",
[
("gpt-4o-mini", "openai"),
("anthropic/claude-3-5-sonnet", "anthropic"),
],
)
async def test_parallel_request_limiter_v3_populates_provider(model, expected_provider):
handler = _PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=InternalUsageCache(DualCache())
)
descriptors = [{"key": "key", "value": "v", "rate_limit": {"requests_per_unit": 1}}]
over = _v3_over_limit_response()
with pytest.raises(HTTPException) as exc_info:
handler._handle_rate_limit_error(
response=over,
descriptors=descriptors,
requested_model=model,
)
exc = exc_info.value
assert exc.status_code == 429
assert isinstance(exc, RateLimitError)
assert exc.llm_provider == expected_provider
# v3 may strip the "anthropic/" prefix in the resolved model — accept
# either; we only care that the provider field is correct and the model
# is non-empty.
assert exc.model
@pytest.mark.asyncio
async def test_parallel_request_limiter_v3_unknown_model_falls_back():
handler = _PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=InternalUsageCache(DualCache())
)
descriptors = [{"key": "key", "value": "v", "rate_limit": {"requests_per_unit": 1}}]
with pytest.raises(HTTPException) as exc_info:
handler._handle_rate_limit_error(
response=_v3_over_limit_response(),
descriptors=descriptors,
requested_model="totally-bogus",
)
assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK
assert exc_info.value.model == "totally-bogus"
@pytest.mark.asyncio
async def test_parallel_request_limiter_v3_missing_model_falls_back():
handler = _PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=InternalUsageCache(DualCache())
)
descriptors = [{"key": "key", "value": "v", "rate_limit": {"requests_per_unit": 1}}]
with pytest.raises(HTTPException) as exc_info:
handler._handle_rate_limit_error(
response=_v3_over_limit_response(),
descriptors=descriptors,
requested_model=None,
)
assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK
assert exc_info.value.model == ""
# ---------------------------------------------------------------------------
# dynamic_rate_limiter v1
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_dynamic_rate_limiter_v1_tpm_zero_populates_provider():
handler = _PROXY_DynamicRateLimitHandler(internal_usage_cache=DualCache())
handler.check_available_usage = AsyncMock(return_value=(0, 5, 100, 5, 1))
user_api_key_dict = UserAPIKeyAuth(api_key="sk-dyn")
user_api_key_dict.metadata = {}
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={"model": "gpt-4o-mini"},
call_type="completion",
)
exc = exc_info.value
assert exc.status_code == 429
assert isinstance(exc, RateLimitError)
assert exc.llm_provider == "openai"
assert exc.model == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_dynamic_rate_limiter_v1_rpm_zero_populates_provider():
handler = _PROXY_DynamicRateLimitHandler(internal_usage_cache=DualCache())
handler.check_available_usage = AsyncMock(return_value=(5, 0, 5, 100, 1))
user_api_key_dict = UserAPIKeyAuth(api_key="sk-dyn")
user_api_key_dict.metadata = {}
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={"model": "anthropic/claude-3-5-sonnet"},
call_type="completion",
)
exc = exc_info.value
assert exc.llm_provider == "anthropic"
assert exc.model == "claude-3-5-sonnet"
@pytest.mark.asyncio
async def test_dynamic_rate_limiter_v1_unknown_model_falls_back():
handler = _PROXY_DynamicRateLimitHandler(internal_usage_cache=DualCache())
handler.check_available_usage = AsyncMock(return_value=(0, 5, 100, 5, 1))
user_api_key_dict = UserAPIKeyAuth(api_key="sk-dyn")
user_api_key_dict.metadata = {}
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={"model": "no-such-model"},
call_type="completion",
)
assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK
assert exc_info.value.model == "no-such-model"
# ---------------------------------------------------------------------------
# dynamic_rate_limiter v3 — exercise just the raise path via the helper, not
# the full Redis/Lua stack.
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_dynamic_rate_limiter_v3_model_capacity_path_populates_provider():
"""
The v3 dynamic limiter has three raise sites: model_saturation_check,
priority_model, and the fail-closed unknown-descriptor branch. We patch
the atomic increment to short-circuit straight into the model_saturation
path — that's the most common production trip — and confirm the
raised exception carries provider info.
"""
from litellm.types.router import ModelGroupInfo
handler = _PROXY_DynamicRateLimitHandlerV3(internal_usage_cache=DualCache())
handler.v3_limiter.atomic_check_and_increment_by_n = AsyncMock(
return_value={
"overall_code": "OVER_LIMIT",
"statuses": [
{
"code": "OVER_LIMIT",
"descriptor_key": "model_saturation_check",
"current_limit": 100,
"limit_remaining": 0,
"rate_limit_type": "rpm",
}
],
}
)
handler._create_priority_based_descriptors = MagicMock(return_value=[])
handler._create_model_tracking_descriptor = MagicMock(
return_value={
"key": "model_saturation_check",
"value": "gpt-4o-mini",
"rate_limit": {"requests_per_unit": 100},
}
)
user_api_key_dict = UserAPIKeyAuth(api_key="sk-dyn-v3")
user_api_key_dict.metadata = {}
model_info = ModelGroupInfo(model_group="gpt-4o-mini", providers=["openai"])
with pytest.raises(HTTPException) as exc_info:
await handler._check_rate_limits(
model="gpt-4o-mini",
model_group_info=model_info,
user_api_key_dict=user_api_key_dict,
priority="default",
saturation=1.0,
data={"model": "gpt-4o-mini"},
)
exc = exc_info.value
assert exc.status_code == 429
assert isinstance(exc, RateLimitError)
assert exc.llm_provider == "openai"
assert exc.model == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_dynamic_rate_limiter_v3_unknown_descriptor_path_populates_provider():
"""Fail-closed unknown-descriptor branch must still attribute provider."""
from litellm.types.router import ModelGroupInfo
handler = _PROXY_DynamicRateLimitHandlerV3(internal_usage_cache=DualCache())
handler.v3_limiter.atomic_check_and_increment_by_n = AsyncMock(
return_value={
"overall_code": "OVER_LIMIT",
"statuses": [
{
"code": "OVER_LIMIT",
"descriptor_key": "something_we_dont_handle",
"current_limit": 1,
"limit_remaining": 0,
"rate_limit_type": "rpm",
}
],
}
)
handler._create_priority_based_descriptors = MagicMock(return_value=[])
handler._create_model_tracking_descriptor = MagicMock(
return_value={
"key": "model_saturation_check",
"value": "gpt-4o-mini",
"rate_limit": {"requests_per_unit": 1},
}
)
user_api_key_dict = UserAPIKeyAuth(api_key="sk-dyn-v3-unknown")
user_api_key_dict.metadata = {}
model_info = ModelGroupInfo(model_group="gpt-4o-mini", providers=["openai"])
with pytest.raises(HTTPException) as exc_info:
await handler._check_rate_limits(
model="gpt-4o-mini",
model_group_info=model_info,
user_api_key_dict=user_api_key_dict,
priority="default",
saturation=1.0,
data={"model": "gpt-4o-mini"},
)
assert exc_info.value.llm_provider == "openai"
# ---------------------------------------------------------------------------
# batch_rate_limiter
# ---------------------------------------------------------------------------
def _batch_over_limit_response() -> dict:
return {
"overall_code": "OVER_LIMIT",
"statuses": [
{
"code": "OVER_LIMIT",
"descriptor_key": "key",
"current_limit": 10,
"limit_remaining": -5,
"rate_limit_type": "requests",
}
],
}
@pytest.mark.asyncio
async def test_batch_rate_limiter_populates_provider():
"""
batch_rate_limiter trips when the file's request/token count exceeds the
remaining window. The raise must thread `data["model"]` through the
helper.
"""
parallel_limiter = MagicMock()
parallel_limiter.window_size = 60
parallel_limiter._create_rate_limit_descriptors = MagicMock(
return_value=[
{"key": "key", "value": "v", "rate_limit": {"requests_per_unit": 10}}
]
)
parallel_limiter.atomic_check_and_increment_by_n = AsyncMock(
return_value=_batch_over_limit_response()
)
handler = _PROXY_BatchRateLimiter(
internal_usage_cache=InternalUsageCache(DualCache()),
parallel_request_limiter=parallel_limiter,
)
with pytest.raises(HTTPException) as exc_info:
await handler._check_and_increment_batch_counters(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-batch"),
data={"model": "gpt-4o-mini"},
batch_usage=BatchFileUsage(total_tokens=100, request_count=15),
)
exc = exc_info.value
assert exc.status_code == 429
assert isinstance(exc, RateLimitError)
assert exc.llm_provider == "openai"
assert exc.model == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_batch_rate_limiter_unknown_model_falls_back():
parallel_limiter = MagicMock()
parallel_limiter.window_size = 60
parallel_limiter._create_rate_limit_descriptors = MagicMock(
return_value=[
{"key": "key", "value": "v", "rate_limit": {"requests_per_unit": 10}}
]
)
parallel_limiter.atomic_check_and_increment_by_n = AsyncMock(
return_value=_batch_over_limit_response()
)
handler = _PROXY_BatchRateLimiter(
internal_usage_cache=InternalUsageCache(DualCache()),
parallel_request_limiter=parallel_limiter,
)
with pytest.raises(HTTPException) as exc_info:
await handler._check_and_increment_batch_counters(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-batch"),
data={"model": "fake-model-xyz"},
batch_usage=BatchFileUsage(total_tokens=100, request_count=15),
)
assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK
# ---------------------------------------------------------------------------
# max_budget_limiter
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_max_budget_limiter_populates_provider():
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-budget",
user_id="user-1",
user_max_budget=10.0,
)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=10.0),
):
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={"model": "gpt-4o-mini"},
call_type="completion",
)
exc = exc_info.value
assert exc.status_code == 429
assert isinstance(exc, RateLimitError)
assert exc.llm_provider == "openai"
assert exc.model == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_max_budget_limiter_no_model_falls_back():
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-budget",
user_id="user-1",
user_max_budget=10.0,
)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=10.0),
):
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK
assert exc_info.value.model == ""
# ---------------------------------------------------------------------------
# max_iterations_limiter
# ---------------------------------------------------------------------------
def _make_iter_agent(max_iterations: int) -> AgentResponse:
return AgentResponse(
agent_id="agent-iter",
agent_name="iter-agent",
litellm_params={"max_iterations": max_iterations},
agent_card_params={"name": "iter-agent", "version": "1.0.0"},
)
@pytest.mark.asyncio
async def test_max_iterations_limiter_populates_provider():
local_cache = DualCache()
handler = _PROXY_MaxIterationsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
user_api_key_dict = UserAPIKeyAuth(api_key="sk-iter", agent_id="agent-iter")
with patch(
"litellm.proxy.agent_endpoints.agent_registry.global_agent_registry"
) as mock_registry:
mock_registry.get_agent_by_id.return_value = _make_iter_agent(max_iterations=1)
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={
"model": "gpt-4o-mini",
"metadata": {"session_id": "session-iter-1"},
},
call_type="completion",
)
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={
"model": "gpt-4o-mini",
"metadata": {"session_id": "session-iter-1"},
},
call_type="completion",
)
exc = exc_info.value
assert exc.status_code == 429
assert isinstance(exc, RateLimitError)
assert exc.llm_provider == "openai"
assert exc.model == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_max_iterations_limiter_unknown_model_falls_back():
local_cache = DualCache()
handler = _PROXY_MaxIterationsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
user_api_key_dict = UserAPIKeyAuth(api_key="sk-iter", agent_id="agent-iter")
with patch(
"litellm.proxy.agent_endpoints.agent_registry.global_agent_registry"
) as mock_registry:
mock_registry.get_agent_by_id.return_value = _make_iter_agent(max_iterations=1)
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={
"model": "no-such-model",
"metadata": {"session_id": "session-iter-2"},
},
call_type="completion",
)
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={
"model": "no-such-model",
"metadata": {"session_id": "session-iter-2"},
},
call_type="completion",
)
assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK
# ---------------------------------------------------------------------------
# max_budget_per_session_limiter
# ---------------------------------------------------------------------------
def _make_session_budget_agent(max_budget: float) -> AgentResponse:
return AgentResponse(
agent_id="agent-session-budget",
agent_name="session-budget-agent",
litellm_params={"max_budget_per_session": max_budget},
agent_card_params={"name": "session-budget-agent", "version": "1.0.0"},
)
@pytest.mark.asyncio
async def test_max_budget_per_session_limiter_populates_provider():
handler = _PROXY_MaxBudgetPerSessionHandler(
internal_usage_cache=InternalUsageCache(DualCache())
)
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-session-budget", agent_id="agent-session-budget"
)
with patch(
"litellm.proxy.agent_endpoints.agent_registry.global_agent_registry"
) as mock_registry:
mock_registry.get_agent_by_id.return_value = _make_session_budget_agent(
max_budget=1.0
)
with patch.object(
handler, "_get_current_spend", new=AsyncMock(return_value=5.0)
):
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={
"model": "anthropic/claude-3-5-sonnet",
"metadata": {"session_id": "session-budget-1"},
},
call_type="completion",
)
exc = exc_info.value
assert exc.status_code == 429
assert isinstance(exc, RateLimitError)
assert exc.llm_provider == "anthropic"
@pytest.mark.asyncio
async def test_max_budget_per_session_limiter_unknown_model_falls_back():
handler = _PROXY_MaxBudgetPerSessionHandler(
internal_usage_cache=InternalUsageCache(DualCache())
)
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-session-budget", agent_id="agent-session-budget"
)
with patch(
"litellm.proxy.agent_endpoints.agent_registry.global_agent_registry"
) as mock_registry:
mock_registry.get_agent_by_id.return_value = _make_session_budget_agent(
max_budget=1.0
)
with patch.object(
handler, "_get_current_spend", new=AsyncMock(return_value=5.0)
):
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={
"model": "no-such-model",
"metadata": {"session_id": "session-budget-2"},
},
call_type="completion",
)
assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK
# ---------------------------------------------------------------------------
# Prometheus integration: failure metric reads exception.llm_provider
# via _get_exception_class_name. With the fix, this returns
# "Openai.RateLimitError" instead of plain "HTTPException" for proxy-side
# 429s on a known model. Pin that contract — that's what dashboards see.
# ---------------------------------------------------------------------------
def test_prometheus_exception_class_name_back_compat_for_proxy_rate_limit_error():
"""
`_get_exception_class_name` deliberately returns the literal string
``"HTTPException"`` for every ``ProxyRateLimitError`` instance so that
pre-existing dashboards / alerts (which key off the historical value)
keep working after the unified rate-limit error class landed in #27687.
Provider attribution is now surfaced separately via the
``rate_limit_category`` / ``rate_limit_type`` labels — this test pins
the back-compat shim itself.
"""
from litellm.integrations.prometheus import PrometheusLogger
exc = ProxyRateLimitError(
detail="over limit",
model="gpt-4o-mini",
llm_provider="openai",
)
assert PrometheusLogger._get_exception_class_name(exc) == "HTTPException"
# Same back-compat path even when the resolver fell back to litellm_proxy.
exc_no_model = ProxyRateLimitError(detail="over limit")
assert PrometheusLogger._get_exception_class_name(exc_no_model) == "HTTPException"
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-vv", "-x"]))