scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

Merge f899d2f1ec into 74050e03c5

Browse source
This commit is contained in:
devin-ai-integration[bot] 2026-08-27 20:00:05 -05:00 • committed by GitHub
commit 3515a01ca3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 63 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -32,6 +32,9 @@ 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,
)
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import CallTypesLiteral
@ -659,21 +662,10 @@ 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):
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["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
priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict)
additional_headers: Final = ensure_response_additional_headers(response)
additional_headers["x-litellm-priority"] = priority or "default"
additional_headers["x-litellm-rate-limiter-version"] = "v3"
return response

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

@ -52,6 +52,9 @@ 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,
)
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage
from litellm.types.utils import (
@ -4648,36 +4651,24 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response):
"""
Post-call hook to update rate limit headers in the response.
Works for both object responses (``ModelResponse`` and friends) and
the dict responses returned by /v1/messages, which hold their
``_hidden_params`` under a key instead of an attribute.
"""
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 None:
return
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 {},
statuses=litellm_proxy_rate_limit_response["statuses"],
)
setattr(
response,
"_hidden_params",
{**_hidden_params, "additional_headers": _additional_headers},
)
additional_headers: Final = ensure_response_additional_headers(response)
additional_headers.update(
self._merge_ratelimit_statuses_into_additional_headers(
additional_headers=additional_headers,
statuses=litellm_proxy_rate_limit_response["statuses"],
)
)
except Exception as e:
verbose_proxy_logger.exception("Error in rate limit post-call hook: %s", e)

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

@ -1861,3 +1861,43 @@ 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_priority_headers_attached_to_dict_response():
"""
/v1/messages returns a plain dict rather than a ModelResponse, and used to be skipped
by the success hook, so priority rate limit headers never reached the client.
"""
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",
)
],
)
user_api_key_dict = UserAPIKeyAuth(metadata={"priority": "premium"})
response: dict = {"id": "msg_123", "type": "message", "role": "assistant", "content": []}
await handler.async_post_call_success_hook(
data={"model": "claude"}, user_api_key_dict=user_api_key_dict, response=response
)
additional_headers = response["_hidden_params"]["additional_headers"]
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"