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

Merge pull request #24106 from BerriAI/Sameerlite/pre-ratelimit-bg

Browse source 
fix(polling): check rate limits before creating polling ID
This commit is contained in:
Sameer Kankute 2026-03-20 17:41:24 +05:30 • committed by GitHub
commit 8ad2068711
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 243 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
litellm/proxy/common_request_processing.py
View file

@ -900,6 +900,7 @@ class ProxyBaseLLMRequestProcessing:
version: Optional[str] = None,
is_streaming_request: Optional[bool] = False,
contents: Optional[list] = None, # Add contents parameter
skip_pre_call_logic: bool = False,
) -> Any:
"""
Common request processing logic for both chat completions and responses API endpoints
@ -909,22 +910,30 @@ class ProxyBaseLLMRequestProcessing:
)
self._debug_log_request_payload()
self.data, logging_obj = await self.common_processing_pre_call_logic(
request=request,
general_settings=general_settings,
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=user_api_key_dict,
version=version,
proxy_config=proxy_config,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
model=model,
route_type=route_type,
llm_router=llm_router,
)
if skip_pre_call_logic:
logging_obj = self.data.get("litellm_logging_obj")
if logging_obj is None:
raise ValueError(
"skip_pre_call_logic=True requires litellm_logging_obj to be set in data. "
"Ensure common_processing_pre_call_logic was called before using this parameter."
)
else:
self.data, logging_obj = await self.common_processing_pre_call_logic(
request=request,
general_settings=general_settings,
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=user_api_key_dict,
version=version,
proxy_config=proxy_config,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
model=model,
route_type=route_type,
llm_router=llm_router,
)
tasks = []
# Start the moderation check (during_call_hook) as early as possible

33
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -119,6 +119,35 @@ async def responses_api(
f"Starting background response with polling for model={data.get('model')}"
)
# Run pre-call checks (rate limits, guardrails, budget) BEFORE creating
# polling ID. This ensures rate-limited requests get a synchronous 429
# instead of a polling ID that immediately fails in the background task.
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
data, _logging_obj = await processor.common_processing_pre_call_logic(
request=request,
general_settings=general_settings,
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=user_api_key_dict,
version=version,
proxy_config=proxy_config,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
model=None,
route_type="aresponses",
llm_router=llm_router,
)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
version=version,
)
# Initialize polling handler with configured TTL (from global config)
polling_handler = ResponsePollingHandler(
redis_cache=redis_usage_cache,
@ -134,7 +163,9 @@ async def responses_api(
request_data=data,
)
# Start background task to stream and update cache
# Start background task to stream and update cache.
# Pass pre-processed data so the background task skips pre-call logic
# (rate limits, guardrails already checked above).
asyncio.create_task(
background_streaming_task(
polling_id=polling_id,

5
litellm/proxy/response_polling/background_streaming.py
View file

@ -65,7 +65,9 @@ async def background_streaming_task( # noqa: PLR0915
# Create processor
processor = ProxyBaseLLMRequestProcessing(data=data)
# Make streaming request
# Make streaming request.
# Pre-call checks (rate limits, guardrails, budget) were already run
# before polling ID creation, so skip them here to avoid double-counting.
response = await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
@ -83,6 +85,7 @@ async def background_streaming_task( # noqa: PLR0915
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
skip_pre_call_logic=True,
)
# Process streaming response following OpenAI events format

182
tests/proxy_unit_tests/test_response_polling_pre_call_checks.py Normal file
View file

@ -0,0 +1,182 @@
"""
Unit tests for pre-call checks running before polling ID creation.
Tests that rate limits, guardrails, and budget checks are enforced
BEFORE a polling ID is created, so rate-limited requests get a
synchronous error instead of a polling ID that immediately fails.
"""
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException, Request, Response
sys.path.insert(0, os.path.abspath("../.."))
import litellm
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
class TestSkipPreCallLogic:
"""Test that skip_pre_call_logic parameter works correctly"""
@pytest.mark.asyncio
async def test_skip_pre_call_logic_skips_common_processing(self):
"""When skip_pre_call_logic=True, common_processing_pre_call_logic should not be called"""
mock_logging_obj = MagicMock()
data = {
"model": "gpt-4",
"stream": True,
"litellm_logging_obj": mock_logging_obj,
}
processor = ProxyBaseLLMRequestProcessing(data=data)
mock_proxy_logging = AsyncMock()
mock_proxy_logging.during_call_hook = AsyncMock()
with (
patch.object(
processor, "common_processing_pre_call_logic", new_callable=AsyncMock
) as mock_pre_call,
patch(
"litellm.proxy.common_request_processing.route_request",
new_callable=AsyncMock,
return_value=MagicMock(),
),
):
try:
await processor.base_process_llm_request(
request=MagicMock(spec=Request),
fastapi_response=MagicMock(spec=Response),
user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
route_type="aresponses",
proxy_logging_obj=mock_proxy_logging,
llm_router=MagicMock(),
general_settings={},
proxy_config=MagicMock(),
skip_pre_call_logic=True,
)
except Exception:
pass # We only care that common_processing_pre_call_logic was not called
mock_pre_call.assert_not_called()
@pytest.mark.asyncio
async def test_without_skip_runs_common_processing(self):
"""When skip_pre_call_logic=False (default), common_processing_pre_call_logic should be called"""
data = {"model": "gpt-4"}
processor = ProxyBaseLLMRequestProcessing(data=data)
mock_logging_obj = MagicMock()
mock_proxy_logging = AsyncMock()
mock_proxy_logging.during_call_hook = AsyncMock()
with (
patch.object(
processor,
"common_processing_pre_call_logic",
new_callable=AsyncMock,
return_value=(data, mock_logging_obj),
) as mock_pre_call,
patch(
"litellm.proxy.common_request_processing.route_request",
new_callable=AsyncMock,
),
):
try:
await processor.base_process_llm_request(
request=MagicMock(spec=Request),
fastapi_response=MagicMock(spec=Response),
user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
route_type="aresponses",
proxy_logging_obj=mock_proxy_logging,
llm_router=MagicMock(),
general_settings={},
proxy_config=MagicMock(),
)
except Exception:
pass
mock_pre_call.assert_called_once()
class TestPollingEndpointPreCallGuard:
"""Test that the polling endpoint enforces pre-call checks before polling ID creation"""
@pytest.mark.asyncio
async def test_rate_limit_error_prevents_polling_id_creation(self):
"""responses_api() must raise 429 and never call generate_polling_id when rate-limited"""
from litellm.proxy.response_api_endpoints.endpoints import responses_api
from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler
rate_limit_exc = litellm.RateLimitError(
message="TPM limit exceeded",
llm_provider="",
model="gpt-4",
)
generate_polling_id_mock = MagicMock(return_value="litellm_poll_test")
proxy_server_patches = {
"litellm.proxy.proxy_server._read_request_body": AsyncMock(
return_value={"model": "gpt-4", "background": True}
),
"litellm.proxy.proxy_server.general_settings": {},
"litellm.proxy.proxy_server.llm_router": MagicMock(),
"litellm.proxy.proxy_server.native_background_mode": None,
"litellm.proxy.proxy_server.polling_cache_ttl": 3600,
"litellm.proxy.proxy_server.polling_via_cache_enabled": True,
"litellm.proxy.proxy_server.proxy_config": MagicMock(),
"litellm.proxy.proxy_server.proxy_logging_obj": AsyncMock(),
"litellm.proxy.proxy_server.redis_usage_cache": AsyncMock(),
"litellm.proxy.proxy_server.select_data_generator": None,
"litellm.proxy.proxy_server.user_api_base": None,
"litellm.proxy.proxy_server.user_max_tokens": None,
"litellm.proxy.proxy_server.user_model": None,
"litellm.proxy.proxy_server.user_request_timeout": None,
"litellm.proxy.proxy_server.user_temperature": None,
"litellm.proxy.proxy_server.version": "1.0.0",
}
with (
patch.multiple("litellm.proxy.proxy_server", **{
k.split(".")[-1]: v for k, v in proxy_server_patches.items()
}),
patch(
"litellm.proxy.response_polling.polling_handler.should_use_polling_for_request",
return_value=True,
),
patch.object(
ProxyBaseLLMRequestProcessing,
"common_processing_pre_call_logic",
new_callable=AsyncMock,
side_effect=rate_limit_exc,
),
patch.object(
ProxyBaseLLMRequestProcessing,
"_handle_llm_api_exception",
new_callable=AsyncMock,
return_value=HTTPException(status_code=429, detail="Rate limit exceeded"),
),
patch.object(ResponsePollingHandler, "generate_polling_id", generate_polling_id_mock),
# Prevent background task from running (avoids noise from incomplete mocks)
patch("asyncio.create_task"),
patch.object(
ResponsePollingHandler,
"create_initial_state",
new_callable=AsyncMock,
return_value=MagicMock(),
),
):
with pytest.raises(HTTPException) as exc_info:
await responses_api(
request=MagicMock(spec=Request),
fastapi_response=MagicMock(spec=Response),
user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
)
assert exc_info.value.status_code == 429
generate_polling_id_mock.assert_not_called()