diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index ccd5c1dd8f5..a7c4a4d3491 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -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 MAX_STREAMING_CHUNK_DURATION_S.""" + from litellm.constants import MAX_STREAMING_CHUNK_DURATION_S + + if MAX_STREAMING_CHUNK_DURATION_S is None: + return + elapsed = time.time() - self._stream_created_time + if elapsed > MAX_STREAMING_CHUNK_DURATION_S: + raise litellm.Timeout( + message=f"Stream exceeded max streaming duration of {MAX_STREAMING_CHUNK_DURATION_S}s (elapsed {elapsed:.1f}s)", + model=self.model or "", + llm_provider=self.custom_llm_provider or "", + ) + def __iter__(self) -> Iterator["ModelResponseStream"]: return self @@ -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() @@ -1917,6 +1933,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: await self.fetch_stream() diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index edcbb0d11b8..034e1e91a50 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -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 MAX_STREAMING_CHUNK_DURATION_S, 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 MAX_STREAMING_CHUNK_DURATION_S.""" + if MAX_STREAMING_CHUNK_DURATION_S is None: + return + elapsed = time.time() - self._stream_created_time + if elapsed > MAX_STREAMING_CHUNK_DURATION_S: + raise litellm.Timeout( + message=f"Stream exceeded max streaming duration of {MAX_STREAMING_CHUNK_DURATION_S}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: