mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
[Feat] Add control for setting upperbound on chunk processing time (#22209)
* add LITELLM_MAX_STREAMING_DURATION_SECONDS * add add LITELLM_MAX_STREAMING_DURATION_SECONDS * fix: address Greptile review - rename constant, add sync check, add tests - Rename MAX_STREAMING_CHUNK_DURATION_S → MAX_STREAMING_DURATION_S (misleading "CHUNK") - Add _check_max_streaming_duration to SyncResponsesAPIStreamingIterator.__next__ - Add 8 unit tests covering both CustomStreamWrapper and ResponsesAPI paths - Fix pre-existing pyright errors in streaming_handler.py Made-with: Cursor * add add LITELLM_MAX_STREAMING_DURATION_SECONDS Made-with: Cursor
This commit is contained in:
Ishaan Jaff
Ishaan Jaffer
parent
699911dfc5
commit
bbf2c4fd0c
5 changed files with 205 additions and 16 deletions
|
|
@ -46,6 +46,14 @@ DEFAULT_REPLICATE_POLLING_DELAY_SECONDS = int(
|
|||
os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)
|
||||
)
|
||||
DEFAULT_IMAGE_TOKEN_COUNT = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250))
|
||||
|
||||
# Maximum wall-clock seconds a streaming response is allowed to run.
|
||||
# Streams exceeding this duration are terminated with a Timeout error.
|
||||
# None (default) = no limit. Set env var to a number of seconds to enable globally.
|
||||
_max_stream_duration_env = os.getenv("LITELLM_MAX_STREAMING_DURATION_SECONDS", None)
|
||||
LITELLM_MAX_STREAMING_DURATION_SECONDS = (
|
||||
float(_max_stream_duration_env) if _max_stream_duration_env is not None else None
|
||||
)
|
||||
DEFAULT_IMAGE_WIDTH = int(os.getenv("DEFAULT_IMAGE_WIDTH", 300))
|
||||
DEFAULT_IMAGE_HEIGHT = int(os.getenv("DEFAULT_IMAGE_HEIGHT", 300))
|
||||
# Maximum size for image URL downloads in MB (default 50MB, set to 0 to disable limit)
|
||||
|
|
|
|||
|
|
@ -85,6 +85,7 @@ class CustomStreamWrapper:
|
|||
self.completion_stream = completion_stream
|
||||
self.sent_first_chunk = False
|
||||
self.sent_last_chunk = False
|
||||
self._stream_created_time: float = time.time()
|
||||
|
||||
litellm_params: GenericLiteLLMParams = GenericLiteLLMParams(
|
||||
**self.logging_obj.model_call_details.get("litellm_params", {})
|
||||
|
|
@ -150,6 +151,20 @@ class CustomStreamWrapper:
|
|||
self.is_function_call = self.check_is_function_call(logging_obj=logging_obj)
|
||||
self.created: Optional[int] = None
|
||||
|
||||
def _check_max_streaming_duration(self) -> None:
|
||||
"""Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS."""
|
||||
from litellm.constants import LITELLM_MAX_STREAMING_DURATION_SECONDS
|
||||
|
||||
if LITELLM_MAX_STREAMING_DURATION_SECONDS is None:
|
||||
return
|
||||
elapsed = time.time() - self._stream_created_time
|
||||
if elapsed > LITELLM_MAX_STREAMING_DURATION_SECONDS:
|
||||
raise litellm.Timeout(
|
||||
message=f"Stream exceeded max streaming duration of {LITELLM_MAX_STREAMING_DURATION_SECONDS}s (elapsed {elapsed:.1f}s)",
|
||||
model=self.model or "",
|
||||
llm_provider=self.custom_llm_provider or "",
|
||||
)
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
|
|
@ -1204,27 +1219,27 @@ class CustomStreamWrapper:
|
|||
else:
|
||||
completion_obj["content"] = str(chunk)
|
||||
elif self.custom_llm_provider == "petals":
|
||||
if len(self.completion_stream) == 0:
|
||||
if self.completion_stream is None or len(self.completion_stream) == 0:
|
||||
if self.received_finish_reason is not None:
|
||||
raise StopIteration
|
||||
else:
|
||||
self.received_finish_reason = "stop"
|
||||
chunk_size = 30
|
||||
new_chunk = self.completion_stream[:chunk_size]
|
||||
new_chunk = self.completion_stream[:chunk_size] # type: ignore[index]
|
||||
completion_obj["content"] = new_chunk
|
||||
self.completion_stream = self.completion_stream[chunk_size:]
|
||||
self.completion_stream = self.completion_stream[chunk_size:] # type: ignore[index]
|
||||
elif self.custom_llm_provider == "palm":
|
||||
# fake streaming
|
||||
response_obj = {}
|
||||
if len(self.completion_stream) == 0:
|
||||
if self.completion_stream is None or len(self.completion_stream) == 0:
|
||||
if self.received_finish_reason is not None:
|
||||
raise StopIteration
|
||||
else:
|
||||
self.received_finish_reason = "stop"
|
||||
chunk_size = 30
|
||||
new_chunk = self.completion_stream[:chunk_size]
|
||||
new_chunk = self.completion_stream[:chunk_size] # type: ignore[index]
|
||||
completion_obj["content"] = new_chunk
|
||||
self.completion_stream = self.completion_stream[chunk_size:]
|
||||
self.completion_stream = self.completion_stream[chunk_size:] # type: ignore[index]
|
||||
elif self.custom_llm_provider == "triton":
|
||||
response_obj = self.handle_triton_stream(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
|
|
@ -1711,6 +1726,7 @@ class CustomStreamWrapper:
|
|||
and self.custom_llm_provider == "cached_response"
|
||||
):
|
||||
cache_hit = True
|
||||
self._check_max_streaming_duration()
|
||||
try:
|
||||
if self.completion_stream is None:
|
||||
self.fetch_sync_stream()
|
||||
|
|
@ -1723,7 +1739,7 @@ class CustomStreamWrapper:
|
|||
):
|
||||
chunk = self.completion_stream
|
||||
else:
|
||||
chunk = next(self.completion_stream)
|
||||
chunk = next(self.completion_stream) # type: ignore[arg-type]
|
||||
if chunk is not None and chunk != b"":
|
||||
print_verbose(
|
||||
f"PROCESSED CHUNK PRE CHUNK CREATOR: {chunk}; custom_llm_provider: {self.custom_llm_provider}"
|
||||
|
|
@ -1885,12 +1901,13 @@ class CustomStreamWrapper:
|
|||
and self.custom_llm_provider == "cached_response"
|
||||
):
|
||||
cache_hit = True
|
||||
self._check_max_streaming_duration()
|
||||
try:
|
||||
if self.completion_stream is None:
|
||||
await self.fetch_stream()
|
||||
|
||||
if is_async_iterable(self.completion_stream):
|
||||
async for chunk in self.completion_stream:
|
||||
async for chunk in self.completion_stream: # type: ignore[union-attr]
|
||||
if chunk == "None" or chunk is None:
|
||||
continue # skip None chunks
|
||||
|
||||
|
|
@ -1965,7 +1982,7 @@ class CustomStreamWrapper:
|
|||
):
|
||||
chunk = self.completion_stream
|
||||
else:
|
||||
chunk = next(self.completion_stream)
|
||||
chunk = next(self.completion_stream) # type: ignore[arg-type]
|
||||
if chunk is not None and chunk != b"":
|
||||
processed_chunk: Optional[
|
||||
ModelResponseStream
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import json
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, Optional
|
||||
|
|
@ -7,7 +8,10 @@ from typing import Any, Dict, Optional
|
|||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.constants import STREAM_SSE_DONE_STRING
|
||||
from litellm.constants import (
|
||||
LITELLM_MAX_STREAMING_DURATION_SECONDS,
|
||||
STREAM_SSE_DONE_STRING,
|
||||
)
|
||||
from litellm.litellm_core_utils.asyncify import run_async_function
|
||||
from litellm.litellm_core_utils.core_helpers import process_response_headers
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
|
@ -56,6 +60,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self.completed_response: Optional[ResponsesAPIStreamingResponse] = None
|
||||
self.start_time = getattr(logging_obj, "start_time", datetime.now())
|
||||
self._failure_handled = False # Track if failure handler has been called
|
||||
self._stream_created_time: float = time.time()
|
||||
|
||||
# track request context for hooks
|
||||
self.litellm_metadata = litellm_metadata
|
||||
|
|
@ -82,6 +87,18 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self.response.headers or {}
|
||||
) # GUARANTEE OPENAI HEADERS IN RESPONSE
|
||||
|
||||
def _check_max_streaming_duration(self) -> None:
|
||||
"""Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS."""
|
||||
if LITELLM_MAX_STREAMING_DURATION_SECONDS is None:
|
||||
return
|
||||
elapsed = time.time() - self._stream_created_time
|
||||
if elapsed > LITELLM_MAX_STREAMING_DURATION_SECONDS:
|
||||
raise litellm.Timeout(
|
||||
message=f"Stream exceeded max streaming duration of {LITELLM_MAX_STREAMING_DURATION_SECONDS}s (elapsed {elapsed:.1f}s)",
|
||||
model=self.model or "",
|
||||
llm_provider=self.custom_llm_provider or "",
|
||||
)
|
||||
|
||||
def _process_chunk(self, chunk) -> Optional[ResponsesAPIStreamingResponse]:
|
||||
"""Process a single chunk of data from the stream"""
|
||||
if not chunk:
|
||||
|
|
@ -363,6 +380,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
|
||||
async def __anext__(self) -> ResponsesAPIStreamingResponse:
|
||||
try:
|
||||
self._check_max_streaming_duration()
|
||||
while True:
|
||||
# Get the next chunk from the stream
|
||||
try:
|
||||
|
|
@ -371,6 +389,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
self.finished = True
|
||||
raise StopAsyncIteration
|
||||
|
||||
self._check_max_streaming_duration()
|
||||
result = self._process_chunk(chunk)
|
||||
|
||||
if self.finished:
|
||||
|
|
@ -461,6 +480,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
|
||||
def __next__(self):
|
||||
try:
|
||||
self._check_max_streaming_duration()
|
||||
while True:
|
||||
# Get the next chunk from the stream
|
||||
try:
|
||||
|
|
@ -469,6 +489,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
self.finished = True
|
||||
raise StopIteration
|
||||
|
||||
self._check_max_streaming_duration()
|
||||
result = self._process_chunk(chunk)
|
||||
|
||||
if self.finished:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,126 @@
|
|||
"""
|
||||
Tests for LITELLM_MAX_STREAMING_DURATION_SECONDS — the global cap on streaming response wall-clock time.
|
||||
|
||||
Covers:
|
||||
- CustomStreamWrapper (chat/completions) sync + async
|
||||
- BaseResponsesAPIStreamingIterator (responses) sync + async
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _make_custom_stream_wrapper() -> CustomStreamWrapper:
|
||||
"""Build a minimal CustomStreamWrapper for testing."""
|
||||
return CustomStreamWrapper(
|
||||
completion_stream=None,
|
||||
model="test-model",
|
||||
logging_obj=MagicMock(),
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CustomStreamWrapper (chat/completions)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCustomStreamWrapperMaxDuration:
|
||||
def test_should_not_raise_when_duration_is_none(self):
|
||||
"""No limit configured → never raises."""
|
||||
wrapper = _make_custom_stream_wrapper()
|
||||
with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", None):
|
||||
wrapper._check_max_streaming_duration() # should not raise
|
||||
|
||||
def test_should_not_raise_when_under_limit(self):
|
||||
"""Stream is under the limit → no error."""
|
||||
wrapper = _make_custom_stream_wrapper()
|
||||
with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 60.0):
|
||||
wrapper._check_max_streaming_duration() # should not raise
|
||||
|
||||
def test_should_raise_timeout_when_exceeded(self):
|
||||
"""Stream exceeded the limit → litellm.Timeout."""
|
||||
wrapper = _make_custom_stream_wrapper()
|
||||
wrapper._stream_created_time = time.time() - 20 # simulate 20s elapsed
|
||||
with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0):
|
||||
with pytest.raises(litellm.Timeout, match="max streaming duration"):
|
||||
wrapper._check_max_streaming_duration()
|
||||
|
||||
def test_should_raise_on_sync_next_when_exceeded(self):
|
||||
"""__next__ should check the limit before iterating."""
|
||||
wrapper = _make_custom_stream_wrapper()
|
||||
wrapper._stream_created_time = time.time() - 20
|
||||
with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0):
|
||||
with pytest.raises(litellm.Timeout):
|
||||
wrapper.__next__()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_raise_on_async_anext_when_exceeded(self):
|
||||
"""__anext__ should check the limit before iterating."""
|
||||
wrapper = _make_custom_stream_wrapper()
|
||||
wrapper._stream_created_time = time.time() - 20
|
||||
with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0):
|
||||
with pytest.raises(litellm.Timeout):
|
||||
await wrapper.__anext__()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# BaseResponsesAPIStreamingIterator (responses)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestResponsesStreamingIteratorMaxDuration:
|
||||
def _make_base_iterator(self):
|
||||
"""Build a minimal BaseResponsesAPIStreamingIterator for testing."""
|
||||
from litellm.responses.streaming_iterator import (
|
||||
BaseResponsesAPIStreamingIterator,
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.headers = {}
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.model_call_details = {"litellm_params": {}}
|
||||
mock_logging_obj.start_time = time.time()
|
||||
|
||||
mock_provider_config = MagicMock()
|
||||
return BaseResponsesAPIStreamingIterator(
|
||||
response=mock_response,
|
||||
model="test-model",
|
||||
responses_api_provider_config=mock_provider_config,
|
||||
logging_obj=mock_logging_obj,
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
def test_should_not_raise_when_duration_is_none(self):
|
||||
it = self._make_base_iterator()
|
||||
with patch(
|
||||
"litellm.responses.streaming_iterator.LITELLM_MAX_STREAMING_DURATION_SECONDS", None
|
||||
):
|
||||
it._check_max_streaming_duration()
|
||||
|
||||
def test_should_not_raise_when_under_limit(self):
|
||||
it = self._make_base_iterator()
|
||||
with patch(
|
||||
"litellm.responses.streaming_iterator.LITELLM_MAX_STREAMING_DURATION_SECONDS", 60.0
|
||||
):
|
||||
it._check_max_streaming_duration()
|
||||
|
||||
def test_should_raise_timeout_when_exceeded(self):
|
||||
it = self._make_base_iterator()
|
||||
it._stream_created_time = time.time() - 20
|
||||
with patch(
|
||||
"litellm.responses.streaming_iterator.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0
|
||||
):
|
||||
with pytest.raises(litellm.Timeout, match="max streaming duration"):
|
||||
it._check_max_streaming_duration()
|
||||
|
|
@ -1,6 +1,10 @@
|
|||
{
|
||||
"compilerOptions": {
|
||||
"lib": ["dom", "dom.iterable", "esnext"],
|
||||
"lib": [
|
||||
"dom",
|
||||
"dom.iterable",
|
||||
"esnext"
|
||||
],
|
||||
"allowJs": true,
|
||||
"skipLibCheck": true,
|
||||
"strict": true,
|
||||
|
|
@ -10,7 +14,7 @@
|
|||
"moduleResolution": "bundler",
|
||||
"resolveJsonModule": true,
|
||||
"isolatedModules": true,
|
||||
"jsx": "preserve",
|
||||
"jsx": "react-jsx",
|
||||
"incremental": true,
|
||||
"plugins": [
|
||||
{
|
||||
|
|
@ -18,9 +22,22 @@
|
|||
}
|
||||
],
|
||||
"paths": {
|
||||
"@/*": ["./src/*"]
|
||||
}
|
||||
"@/*": [
|
||||
"./src/*"
|
||||
]
|
||||
},
|
||||
"target": "ES2017"
|
||||
},
|
||||
"include": ["next-env.d.ts", "**/*.ts", "**/*.tsx", ".next/types/**/*.ts"],
|
||||
"exclude": ["node_modules", "e2e_tests", "scripts"]
|
||||
"include": [
|
||||
"next-env.d.ts",
|
||||
"**/*.ts",
|
||||
"**/*.tsx",
|
||||
".next/types/**/*.ts",
|
||||
".next/dev/types/**/*.ts"
|
||||
],
|
||||
"exclude": [
|
||||
"node_modules",
|
||||
"e2e_tests",
|
||||
"scripts"
|
||||
]
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue