scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

fix(rate_limiter): attach v3 priority rate limit headers on /v1/messages

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-09 06:29:51 +00:00
parent 802e526cf9
commit 68bbbe35d5
5 changed files with 162 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -32,6 +32,10 @@ from litellm.proxy.hooks.rate_limiter_utils import (
resolve_llm_provider_for_rate_limit,
)
from litellm.proxy.utils import InternalUsageCache
from litellm.router_utils.add_retry_fallback_headers import (
ensure_response_additional_headers,
response_has_hidden_params,
)
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import CallTypesLiteral
@ -659,22 +663,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
data=data, user_api_key_dict=user_api_key_dict, response=response
)
# Add additional priority-specific headers
if isinstance(response, ModelResponse):
if response_has_hidden_params(response):
priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict)
# Get existing additional headers
additional_headers: Final = getattr(response, "_hidden_params", {}).get("additional_headers", {}) or {}
# Add priority information
additional_headers: Final = ensure_response_additional_headers(response)
additional_headers["x-litellm-priority"] = priority or "default"
additional_headers["x-litellm-rate-limiter-version"] = "v3"
# Update response
if not hasattr(response, "_hidden_params"):
response._hidden_params = {}
response._hidden_params["additional_headers"] = additional_headers
return response
except Exception as e:

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

@ -52,6 +52,10 @@ from litellm.proxy.hooks.batch_enqueued_tokens import (
canonical_provider_batch_id,
)
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
from litellm.router_utils.add_retry_fallback_headers import (
ensure_response_additional_headers,
response_has_hidden_params,
)
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage
from litellm.types.utils import (
@ -4677,34 +4681,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
Post-call hook to update rate limit headers in the response.
"""
try:
from pydantic import BaseModel
stash: Final = get_request_stash()
litellm_proxy_rate_limit_response: Final = stash.rate_limit_response if stash is not None else None
if litellm_proxy_rate_limit_response is not None:
# Update response headers
if hasattr(response, "_hidden_params"):
_hidden_params = getattr(response, "_hidden_params")
else:
_hidden_params = None
if _hidden_params is not None and (
isinstance(_hidden_params, BaseModel) or isinstance(_hidden_params, dict)
):
if isinstance(_hidden_params, BaseModel):
_hidden_params = _hidden_params.model_dump()
_additional_headers: Final = self._merge_ratelimit_statuses_into_additional_headers(
additional_headers=_hidden_params.get("additional_headers", {}) or {},
if litellm_proxy_rate_limit_response is not None and response_has_hidden_params(response):
additional_headers: Final = ensure_response_additional_headers(response)
additional_headers.update(
self._merge_ratelimit_statuses_into_additional_headers(
additional_headers={},
statuses=litellm_proxy_rate_limit_response["statuses"],
)
setattr(
response,
"_hidden_params",
{**_hidden_params, "additional_headers": _additional_headers},
)
)
except Exception as e:
verbose_proxy_logger.exception("Error in rate limit post-call hook: %s", e)

6
litellm/router_utils/add_retry_fallback_headers.py
View file

@ -50,6 +50,12 @@ def prepare_response_for_header_attachment(response: object) -> object | None:
return response
def response_has_hidden_params(response: object) -> bool:
if isinstance(response, dict):
return "_hidden_params" in response
return hasattr(response, "_hidden_params")
def ensure_response_additional_headers(response: object) -> dict[str, object]:
hidden_params: Final = get_hidden_params_dict(response, create=isinstance(response, dict))
_write_hidden_params(response, hidden_params)

66
tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py
View file

@ -1861,3 +1861,69 @@ async def test_tpm_only_model_enforces_priority_and_model_capacity(monkeypatch):
)
assert capacity_blocked.value.status_code == 429
assert "Model capacity reached" in capacity_blocked.value.detail["error"]
@pytest.mark.asyncio
async def test_post_call_success_hook_attaches_priority_headers_to_dict_response():
"""
Non-streaming /v1/messages hands the success hook a plain dict, not a ModelResponse.
The router has already stamped ``_hidden_params`` on it, so both the v3 descriptor
headers and the priority headers must land next to the router's headers there.
"""
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
RateLimitResponse,
RateLimitStatus,
get_or_create_request_stash,
)
handler = DynamicRateLimitHandler(internal_usage_cache=DualCache())
get_or_create_request_stash().rate_limit_response = RateLimitResponse(
overall_code="OK",
statuses=[
RateLimitStatus(
code="OK",
current_limit=75,
limit_remaining=74,
rate_limit_type="requests",
descriptor_key="priority_model",
)
],
)
response = {
"id": "msg_123",
"type": "message",
"role": "assistant",
"content": [],
"_hidden_params": {"additional_headers": {"x-litellm-attempted-retries": 0}},
}
await handler.async_post_call_success_hook(
data={"model": "anthropic-haiku"},
user_api_key_dict=UserAPIKeyAuth(metadata={"priority": "premium"}),
response=response,
)
additional_headers = response["_hidden_params"]["additional_headers"]
assert additional_headers["x-litellm-attempted-retries"] == 0
assert additional_headers["x-ratelimit-priority_model-limit-requests"] == 75
assert additional_headers["x-ratelimit-priority_model-remaining-requests"] == 74
assert additional_headers["x-litellm-priority"] == "premium"
assert additional_headers["x-litellm-rate-limiter-version"] == "v3"
@pytest.mark.asyncio
async def test_post_call_success_hook_leaves_raw_provider_dict_untouched():
"""
Passthrough routes hand the hook the raw provider JSON body and serialize whatever
comes back, so a dict without ``_hidden_params`` must not have one injected.
"""
handler = DynamicRateLimitHandler(internal_usage_cache=DualCache())
response = {"id": "msg_123", "type": "message", "role": "assistant", "content": []}
await handler.async_post_call_success_hook(
data={"model": "anthropic-haiku"},
user_api_key_dict=UserAPIKeyAuth(metadata={"priority": "premium"}),
response=response,
)
assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []}

74
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -6171,3 +6171,77 @@ async def test_success_hook_leaves_stash_untouched_for_non_batch_responses():
data={}, user_api_key_dict=user, response=ModelResponse(usage=Usage(total_tokens=5))
)
assert get_request_stash().batch_enqueued_reservation == reservation
@pytest.mark.asyncio
async def test_post_call_success_hook_attaches_ratelimit_headers_to_dict_response():
"""
Non-streaming /v1/messages responses are plain dicts carrying the router's
``_hidden_params`` as a key, not an attribute, so the descriptor headers used to
be dropped for that route.
"""
from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitResponse, RateLimitStatus
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
get_or_create_request_stash().rate_limit_response = RateLimitResponse(
overall_code="OK",
statuses=[
RateLimitStatus(
code="OK",
current_limit=100,
limit_remaining=99,
rate_limit_type="requests",
descriptor_key="model_saturation_check",
)
],
)
response = {
"id": "msg_123",
"type": "message",
"role": "assistant",
"content": [],
"_hidden_params": {"additional_headers": {"x-litellm-attempted-retries": 0}},
}
await handler.async_post_call_success_hook(
data={"model": "anthropic-haiku"},
user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-dict-response")),
response=response,
)
additional_headers = response["_hidden_params"]["additional_headers"]
assert additional_headers["x-litellm-attempted-retries"] == 0
assert additional_headers["x-ratelimit-model_saturation_check-limit-requests"] == 100
assert additional_headers["x-ratelimit-model_saturation_check-remaining-requests"] == 99
@pytest.mark.asyncio
async def test_post_call_success_hook_leaves_raw_provider_dict_untouched():
"""
Passthrough routes hand the hook the raw provider JSON body and serialize whatever
comes back, so a dict without ``_hidden_params`` must not have one injected.
"""
from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitResponse, RateLimitStatus
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
get_or_create_request_stash().rate_limit_response = RateLimitResponse(
overall_code="OK",
statuses=[
RateLimitStatus(
code="OK",
current_limit=100,
limit_remaining=99,
rate_limit_type="requests",
descriptor_key="model_saturation_check",
)
],
)
response = {"id": "msg_123", "type": "message", "role": "assistant", "content": []}
await handler.async_post_call_success_hook(
data={"model": "anthropic-haiku"},
user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-raw-dict")),
response=response,
)
assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []}