mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-08 22:25:25 +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
This commit is contained in:
Ishaan Jaff
@@ -49,6 +49,14 @@ DEFAULT_REPLICATE_POLLING_DELAY_SECONDS = int(
|
||||
)
|
||||
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
|
||||
)
|
||||
|
||||
# Maximum number of base64 characters to keep in logging payloads.
|
||||
# Data URIs exceeding this are replaced with a size placeholder.
|
||||
# Set to 0 to disable truncation.
|
||||
|
||||
@@ -96,6 +96,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", {})
|
||||
@@ -161,6 +162,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) -> Iterator["ModelResponseStream"]:
|
||||
return self
|
||||
|
||||
@@ -1236,27 +1251,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"]
|
||||
@@ -1743,6 +1758,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()
|
||||
@@ -1755,7 +1771,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.decode('utf-8', errors='replace') if isinstance(chunk, bytes) else chunk}; custom_llm_provider: {self.custom_llm_provider}"
|
||||
@@ -1917,12 +1933,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
|
||||
|
||||
@@ -2004,7 +2021,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 = self.chunk_creator(chunk=chunk)
|
||||
if processed_chunk is None:
|
||||
|
||||
@@ -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,7 @@ 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 +57,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 +84,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:
|
||||
@@ -357,6 +371,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:
|
||||
@@ -365,6 +380,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
||||
self.finished = True
|
||||
raise StopAsyncIteration
|
||||
|
||||
self._check_max_streaming_duration()
|
||||
result = self._process_chunk(chunk)
|
||||
|
||||
if self.finished:
|
||||
@@ -460,6 +476,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
||||
|
||||
def __next__(self):
|
||||
try:
|
||||
self._check_max_streaming_duration()
|
||||
while True:
|
||||
# Get the next chunk from the stream
|
||||
try:
|
||||
@@ -468,6 +485,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()
|
||||
Reference in New Issue
Block a user