diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 46a9f9910d7..c6981015417 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1,10 +1,11 @@ import asyncio import contextlib +import itertools import json import logging import math import time -from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from functools import lru_cache @@ -129,6 +130,11 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, relayed_provider_specific_fields, ) +from litellm.proxy.common_utils.proxy_rate_limit_error import ( + ATTEMPTED_FALLBACKS_HEADER, + rate_limit_fallback_from_request_data, + rate_limit_fallback_headers, +) from litellm.proxy.common_utils.responses_stream_errors import ( SSE_LINE_BREAK, ResponsesStreamErrorState, @@ -2384,12 +2390,12 @@ class ProxyBaseLLMRequestProcessing: response_cost: float | str | None = None, hidden_params: Mapping[str, object] | None = None, fastest_response_batch_completion: bool | None = None, - request_data: dict | None = {}, + request_data: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, litellm_logging_obj: LiteLLMLoggingObj | None = None, read_timing_from_logging_obj: bool = True, - **kwargs, - ) -> dict: + **kwargs: object, + ) -> dict[str, str]: exclude_values: Final = {"", None, "None"} hidden_params = hidden_params or {} resolved_call_id: Final = ( @@ -2453,10 +2459,14 @@ class ProxyBaseLLMRequestProcessing: **{k: str(v) for k, v in kwargs.items()}, } if request_data: - remaining_tokens_header: Final = get_remaining_tokens_and_requests_from_request_data(request_data) + remaining_tokens_header: Final = get_remaining_tokens_and_requests_from_request_data(dict(request_data)) headers.update(remaining_tokens_header) - logging_caching_headers: Final = get_logging_caching_headers(request_data) + quota_fallback: Final = rate_limit_fallback_from_request_data(request_data) + if quota_fallback is not None: + headers.update(rate_limit_fallback_headers(quota_fallback, headers.get(ATTEMPTED_FALLBACKS_HEADER))) + + logging_caching_headers: Final = get_logging_caching_headers(dict(request_data)) if logging_caching_headers: headers.update(logging_caching_headers) @@ -2864,7 +2874,9 @@ class ProxyBaseLLMRequestProcessing: from litellm.proxy.common_utils.proxy_rate_limit_error import ( PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS, ProxyRateLimitError, + RateLimitFallback, per_model_rate_limits_disable_fallbacks, + with_rate_limit_fallback, ) general_settings_view: Final = cast( # cast-ok: this method keeps its legacy bare-dict settings parameter @@ -2933,6 +2945,7 @@ class ProxyBaseLLMRequestProcessing: if isinstance(raw_router_settings, Mapping) else None ) + attempt_numbers: Final[Iterator[int]] = itertools.count(1) try: for fallback_model in fallback_models: if fallback_model == original_model: @@ -2955,7 +2968,10 @@ class ProxyBaseLLMRequestProcessing: user_api_key_dict.team_id, ): continue - return await self.common_processing_pre_call_logic( + attempt = next(attempt_numbers) + admitted: tuple[ + Mapping[str, object], LiteLLMLoggingObj + ] = await self.common_processing_pre_call_logic( request=request, general_settings=general_settings, proxy_logging_obj=proxy_logging_obj, @@ -2976,6 +2992,21 @@ class ProxyBaseLLMRequestProcessing: if is_hard_per_model_limit(fallback_exc): raise continue + else: + admitted_data, admitted_logging_obj = admitted + stamped: dict = dict( # mutable-ok: caller-defined request-payload shape + with_rate_limit_fallback( + admitted_data, + get_metadata_variable_name(request, admitted_data), + RateLimitFallback( + original_model_group=original_model, + descriptor_key=original_exc.descriptor_key or "unknown", + attempts=attempt, + ), + ) + ) + self.data = stamped + return stamped, admitted_logging_obj except BaseException: self.data = rate_limited_data raise diff --git a/litellm/proxy/common_utils/proxy_rate_limit_error.py b/litellm/proxy/common_utils/proxy_rate_limit_error.py index 28d28a5c9bd..6f3ef4e03d4 100644 --- a/litellm/proxy/common_utils/proxy_rate_limit_error.py +++ b/litellm/proxy/common_utils/proxy_rate_limit_error.py @@ -37,10 +37,12 @@ This module provides a single proxy-side error class that: import json from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType from typing import Final from fastapi import HTTPException -from pydantic import TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.exceptions import RateLimitError, RateLimitErrorCategory, RateLimitType @@ -73,6 +75,94 @@ def per_model_rate_limits_disable_fallbacks(general_settings: Mapping[str, objec return False +RATE_LIMIT_FALLBACK_METADATA_KEY: Final = "rate_limit_fallback" +RATE_LIMIT_FALLBACK_REASON_HEADER: Final = "x-litellm-fallback-reason" +ATTEMPTED_FALLBACKS_HEADER: Final = "x-litellm-attempted-fallbacks" + + +@dataclass(frozen=True, slots=True) +class RateLimitFallback: + original_model_group: str + descriptor_key: str + attempts: int = 1 + + def stamp(self) -> Mapping[str, object]: + return MappingProxyType( + { + "original_model_group": self.original_model_group, + "descriptor_key": self.descriptor_key, + "attempts": self.attempts, + } + ) + + @property + def reason(self) -> str: + return f"rate_limit:{self.descriptor_key}" + + +class _StampedBucket(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + rate_limit_fallback: RateLimitFallback | None = None + + +class _StampedRequest(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + metadata: _StampedBucket | None = None + litellm_metadata: _StampedBucket | None = None + + +_STAMPED_REQUEST: Final[TypeAdapter[_StampedRequest]] = TypeAdapter(_StampedRequest) +_ROUTER_HOPS: Final[TypeAdapter[int]] = TypeAdapter(int) + + +def rate_limit_fallback_from_request_data(request_data: Mapping[str, object]) -> RateLimitFallback | None: + try: + stamped: Final = _STAMPED_REQUEST.validate_python(request_data) + except ValidationError: + return None + buckets: Final = (stamped.metadata, stamped.litellm_metadata) + return next( + (bucket.rate_limit_fallback for bucket in buckets if bucket is not None and bucket.rate_limit_fallback), None + ) + + +_METADATA_BUCKET: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) + + +def _metadata_bucket(value: object) -> Mapping[str, object]: + try: + return _METADATA_BUCKET.validate_python(value) + except ValidationError: + return MappingProxyType({}) + + +def with_rate_limit_fallback( + request_data: Mapping[str, object], bucket_name: str, fallback: RateLimitFallback +) -> Mapping[str, object]: + bucket: Final = _metadata_bucket(request_data.get(bucket_name)) + return MappingProxyType( + {**request_data, bucket_name: {**bucket, RATE_LIMIT_FALLBACK_METADATA_KEY: dict(fallback.stamp())}} + ) + + +def rate_limit_fallback_headers(fallback: RateLimitFallback, attempted_fallbacks: object) -> Mapping[str, str]: + try: + router_hops: Final = _ROUTER_HOPS.validate_python(attempted_fallbacks) + except ValidationError: + return MappingProxyType( + { + ATTEMPTED_FALLBACKS_HEADER: str(max(fallback.attempts, 1)), + RATE_LIMIT_FALLBACK_REASON_HEADER: fallback.reason, + } + ) + return MappingProxyType( + { + ATTEMPTED_FALLBACKS_HEADER: str(max(router_hops, 0) + max(fallback.attempts, 1)), + RATE_LIMIT_FALLBACK_REASON_HEADER: fallback.reason, + } + ) + + def map_v3_rate_limit_type( v3_value: str | None, ) -> RateLimitType | None: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 0b2ba93af71..ff0883b6958 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -536,6 +536,7 @@ ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset( "attempted_fallbacks", "original_model_group", "previous_models", + "rate_limit_fallback", "request_retry_count", CLIENT_OUTPUT_CEILING_METADATA_KEY, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, diff --git a/tests/integration/routing/test_rate_limit_fallback_headers.py b/tests/integration/routing/test_rate_limit_fallback_headers.py new file mode 100644 index 00000000000..18a42cb4d1d --- /dev/null +++ b/tests/integration/routing/test_rate_limit_fallback_headers.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import uuid +from collections.abc import Iterator +from contextlib import ExitStack +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import OwnedState, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import wire_server +from integration.routing.test_per_model_rate_limit_fallbacks import ( + assert_served, + chat, + healthy_reply, + model_entry, + proxy_config, +) + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + primary: str + fallback: str + second_fallback: str + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + suffix: Final = uuid.uuid4().hex[:10] + primary: Final = f"quota-primary-{suffix}" + fallback: Final = f"quota-fallback-{suffix}" + second_fallback: Final = f"quota-second-fallback-{suffix}" + directory: Final = tmp_path_factory.mktemp("rate-limit-fallback-headers") + with ExitStack() as stack: + rig_gateway: Final = stack.enter_context(gateway_from_environment()) + primary_wire: Final = stack.enter_context(wire_server(healthy_reply)) + fallback_wire: Final = stack.enter_context(wire_server(healthy_reply)) + second_fallback_wire: Final = stack.enter_context(wire_server(healthy_reply)) + config: Final = proxy_config( + directory, + model_list=( + model_entry(primary, primary_wire), + model_entry(fallback, fallback_wire), + model_entry(second_fallback, second_fallback_wire), + ), + fallbacks=({primary: [fallback, second_fallback]},), + general_settings={}, + ) + owned: Final = stack.enter_context( + owned_proxy_process( + rig_gateway, directory, {}, config=config, workers=2, state=OwnedState.SHARED_WITH_WORKER + ) + ) + yield Rig(owned.gateway, primary, fallback, second_fallback) + + +def test_quota_fallback_response_names_the_hop_and_its_reason(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_tpm_limit={rig.primary: 1}) + + response: Final = chat(rig.gateway, key, rig.primary, f"pmrl-{uuid.uuid4().hex}") + assert_served(response, rig.fallback, "served pmrl-") + assert response.headers.get("x-litellm-attempted-fallbacks") == "1", dict(response.headers) + assert response.headers.get("x-litellm-fallback-reason") == "rate_limit:model_per_key", dict(response.headers) + + +def test_quota_fallback_response_counts_every_candidate_a_limit_rejected(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_tpm_limit={rig.primary: 1, rig.fallback: 1}) + + response: Final = chat(rig.gateway, key, rig.primary, f"pmrl-{uuid.uuid4().hex}") + assert_served(response, rig.second_fallback, "served pmrl-") + assert response.headers.get("x-litellm-attempted-fallbacks") == "2", dict(response.headers) + assert response.headers.get("x-litellm-fallback-reason") == "rate_limit:model_per_key", dict(response.headers) diff --git a/tests/unit/proxy/common_utils/test_proxy_rate_limit_error.py b/tests/unit/proxy/common_utils/test_proxy_rate_limit_error.py new file mode 100644 index 00000000000..a7628b72341 --- /dev/null +++ b/tests/unit/proxy/common_utils/test_proxy_rate_limit_error.py @@ -0,0 +1,58 @@ +from typing import Final + +import pytest + +from litellm.proxy.common_utils.proxy_rate_limit_error import ( + RateLimitFallback, + rate_limit_fallback_from_request_data, + rate_limit_fallback_headers, +) + +STAMP: Final = {"original_model_group": "opus", "descriptor_key": "model_per_end_user"} + + +@pytest.mark.parametrize("bucket", ("metadata", "litellm_metadata")) +def test_rate_limit_fallback_is_read_from_either_metadata_bucket(bucket: str) -> None: + fallback: Final = rate_limit_fallback_from_request_data({"model": "luna", bucket: {"rate_limit_fallback": STAMP}}) + + assert fallback is not None + assert fallback == RateLimitFallback(original_model_group="opus", descriptor_key="model_per_end_user") + assert fallback.reason == "rate_limit:model_per_end_user" + + +@pytest.mark.parametrize( + "request_data", + ( + {"model": "luna"}, + {"model": "luna", "metadata": {}}, + {"model": "luna", "metadata": {"rate_limit_fallback": "opus"}}, + {"model": "luna", "metadata": {"rate_limit_fallback": {"original_model_group": "opus"}}}, + ), +) +def test_a_request_without_a_complete_stamp_has_no_rate_limit_fallback(request_data: dict[str, object]) -> None: + assert rate_limit_fallback_from_request_data(request_data) is None + + +@pytest.mark.parametrize(("router_hops", "expected"), ((None, "1"), (0, "1"), ("0", "1"), (2, "3"), ("junk", "1"))) +def test_rate_limit_fallback_headers_count_the_hop_on_top_of_the_router(router_hops: object, expected: str) -> None: + headers: Final = rate_limit_fallback_headers( + RateLimitFallback(original_model_group="opus", descriptor_key="model_per_key"), router_hops + ) + + assert headers == { + "x-litellm-attempted-fallbacks": expected, + "x-litellm-fallback-reason": "rate_limit:model_per_key", + } + + +@pytest.mark.parametrize( + ("attempts", "router_hops", "expected"), ((2, None, "2"), (2, 0, "2"), (3, 1, "4"), (0, 0, "1")) +) +def test_rate_limit_fallback_headers_count_every_candidate_a_limit_rejected( + attempts: int, router_hops: object, expected: str +) -> None: + headers: Final = rate_limit_fallback_headers( + RateLimitFallback(original_model_group="opus", descriptor_key="model_per_key", attempts=attempts), router_hops + ) + + assert headers["x-litellm-attempted-fallbacks"] == expected diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 23023560a41..bfe381cb848 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -13742,3 +13742,38 @@ class TestTotalTimeout: assert not prepared assert not called assert elapsed < 120 + + +@pytest.mark.parametrize("router_attempted_fallbacks", (0, 1)) +def test_get_custom_headers_count_a_quota_fallback_and_name_its_reason(router_attempted_fallbacks: int) -> None: + auth: Final = ProxyUserAPIKeyAuth(api_key="sk-quota-fallback", tpm_limit=None, rpm_limit=None, max_budget=None) + request_data: Final = { + "model": "luna", + "metadata": {"rate_limit_fallback": {"original_model_group": "opus", "descriptor_key": "model_per_key"}}, + } + + custom_headers: Callable[..., dict[str, str]] = ProxyBaseLLMRequestProcessing.get_custom_headers + headers: Final = custom_headers( + user_api_key_dict=auth, + call_id="quota-fallback", + request_data=request_data, + **{"x-litellm-attempted-fallbacks": router_attempted_fallbacks}, + ) + + assert headers["x-litellm-attempted-fallbacks"] == str(router_attempted_fallbacks + 1) + assert headers["x-litellm-fallback-reason"] == "rate_limit:model_per_key" + + +def test_get_custom_headers_leave_a_plain_request_without_a_fallback_reason() -> None: + auth: Final = ProxyUserAPIKeyAuth(api_key="sk-no-fallback", tpm_limit=None, rpm_limit=None, max_budget=None) + + custom_headers: Callable[..., dict[str, str]] = ProxyBaseLLMRequestProcessing.get_custom_headers + headers: Final = custom_headers( + user_api_key_dict=auth, + call_id="no-fallback", + request_data={"model": "opus", "metadata": {}}, + **{"x-litellm-attempted-fallbacks": 0}, + ) + + assert headers["x-litellm-attempted-fallbacks"] == "0" + assert "x-litellm-fallback-reason" not in headers diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index f5757a0ff18..7891a87b1dc 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -9369,6 +9369,7 @@ def _reserved_stamp_key(key_metadata: dict | None = None) -> UserAPIKeyAuth: _PLANTED_STAMPS = { "attempted_fallbacks": 99, "original_model_group": "spoofed-group", + "rate_limit_fallback": {"original_model_group": "spoofed-group", "descriptor_key": "model_per_key"}, "request_retry_count": -100, "previous_models": [{"exception_type": "Timeout", "deployment_id": "spoofed", "attempted_retries": 101}], "_client_output_ceiling": {"api_base": "https://attacker.example"}, @@ -9404,6 +9405,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo assert "litellm_metadata" not in updated assert "attempted_fallbacks" not in updated["metadata"] assert "original_model_group" not in updated["metadata"] + assert "rate_limit_fallback" not in updated["metadata"] assert "_client_output_ceiling" not in updated["metadata"] assert "request_retry_count" not in updated["metadata"] assert "previous_models" not in updated["metadata"]