diff --git a/litellm/__init__.py b/litellm/__init__.py index e0c7d56361c..a29ee73013b 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -245,6 +245,12 @@ use_chat_completions_url_for_anthropic_messages: bool = bool( # config.yaml. strip_anthropic_total_tokens: bool = False anthropic_sse_ping_interval_seconds: float = 15.0 +# Emit an SSE comment (": ping") on OpenAI-shaped streaming routes (/chat/completions, +# /v1/responses, ...) whenever the upstream has sent nothing for this many seconds, so +# intermediaries with an idle read timeout (e.g. nginx `proxy-read-timeout`) don't drop +# long time-to-first-token streams. Disabled unless set, via +# `litellm_settings.sse_keepalive_ping_interval_seconds` in config.yaml. +sse_keepalive_ping_interval_seconds: float | None = None route_all_chat_openai_to_responses: bool = ( os.getenv("LITELLM_ROUTE_ALL_CHAT_OPENAI_TO_RESPONSES", "false").lower() == "true" ) # When True, routes all OpenAI /chat/completions requests through the Responses API bridge diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 159d7508f4e..49627a91cf6 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -46,7 +46,11 @@ from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, get_remaining_tokens_and_requests_from_request_data, ) -from litellm.proxy.common_utils.sse_keepalive import wrap_sse_stream_with_keepalive_pings +from litellm.proxy.common_utils.sse_keepalive import ( + ANTHROPIC_PING_SSE_CHUNK, + SSE_COMMENT_PING_CHUNK, + wrap_sse_stream_with_keepalive_pings, +) from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails @@ -1984,6 +1988,7 @@ class ProxyBaseLLMRequestProcessing: generator=wrap_sse_stream_with_keepalive_pings( stream=selected_data_generator, ping_interval_seconds=litellm.anthropic_sse_ping_interval_seconds, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, ), media_type="text/event-stream", headers=custom_headers, @@ -2015,7 +2020,11 @@ class ProxyBaseLLMRequestProcessing: ) ) return await create_response( - generator=selected_data_generator, + generator=wrap_sse_stream_with_keepalive_pings( + stream=selected_data_generator, + ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, + ping_chunk=SSE_COMMENT_PING_CHUNK, + ), media_type="text/event-stream", headers=custom_headers, request=request, diff --git a/litellm/proxy/common_utils/sse_keepalive.py b/litellm/proxy/common_utils/sse_keepalive.py index 6700700ff7c..d3ff4b0ca30 100644 --- a/litellm/proxy/common_utils/sse_keepalive.py +++ b/litellm/proxy/common_utils/sse_keepalive.py @@ -7,6 +7,7 @@ from typing import Final import anyio ANTHROPIC_PING_SSE_CHUNK: Final = 'event: ping\ndata: {"type": "ping"}\n\n' +SSE_COMMENT_PING_CHUNK: Final = ": ping\n\n" def _coerce_interval(ping_interval_seconds: float | str | None) -> float | None: @@ -24,16 +25,18 @@ def _coerce_interval(ping_interval_seconds: float | str | None) -> float | None: def wrap_sse_stream_with_keepalive_pings( stream: AsyncGenerator[str, None], ping_interval_seconds: float | str | None, + ping_chunk: str, ) -> AsyncGenerator[str, None]: interval: Final = _coerce_interval(ping_interval_seconds) if interval is None: return stream - return _keepalive_ping_stream(stream=stream, ping_interval_seconds=interval) + return _keepalive_ping_stream(stream=stream, ping_interval_seconds=interval, ping_chunk=ping_chunk) async def _keepalive_ping_stream( stream: AsyncGenerator[str, None], ping_interval_seconds: float, + ping_chunk: str, ) -> AsyncGenerator[str, None]: pending = asyncio.ensure_future( stream.__anext__() @@ -42,7 +45,7 @@ async def _keepalive_ping_stream( while True: await asyncio.wait({pending}, timeout=ping_interval_seconds) if not pending.done(): - yield ANTHROPIC_PING_SSE_CHUNK + yield ping_chunk continue try: yield pending.result() diff --git a/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py b/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py index 9cca9bbfe12..e343e46e187 100644 --- a/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py +++ b/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py @@ -22,7 +22,11 @@ async def test_pings_fill_mid_stream_silence_and_preserve_chunk_order(): await asyncio.sleep(0.3) yield TEXT_DELTA_CHUNK - wrapped: Final = wrap_sse_stream_with_keepalive_pings(stream=gappy_stream(), ping_interval_seconds=0.05) + wrapped: Final = wrap_sse_stream_with_keepalive_pings( + stream=gappy_stream(), + ping_interval_seconds=0.05, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ) collected: Final = [chunk async for chunk in wrapped] assert collected[0] == MESSAGE_START_CHUNK @@ -40,7 +44,11 @@ async def test_ping_emitted_while_waiting_for_first_chunk(): await asyncio.sleep(0.2) yield MESSAGE_START_CHUNK - wrapped: Final = wrap_sse_stream_with_keepalive_pings(stream=slow_start_stream(), ping_interval_seconds=0.05) + wrapped: Final = wrap_sse_stream_with_keepalive_pings( + stream=slow_start_stream(), + ping_interval_seconds=0.05, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ) collected: Final = [chunk async for chunk in wrapped] assert collected[0] == ANTHROPIC_PING_SSE_CHUNK @@ -54,7 +62,11 @@ async def test_no_pings_when_chunks_arrive_faster_than_interval(): yield TEXT_DELTA_CHUNK yield TEXT_DELTA_CHUNK - wrapped: Final = wrap_sse_stream_with_keepalive_pings(stream=fast_stream(), ping_interval_seconds=1.0) + wrapped: Final = wrap_sse_stream_with_keepalive_pings( + stream=fast_stream(), + ping_interval_seconds=1.0, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ) collected: Final = [chunk async for chunk in wrapped] assert collected == [MESSAGE_START_CHUNK, TEXT_DELTA_CHUNK, TEXT_DELTA_CHUNK] @@ -66,7 +78,11 @@ async def test_upstream_exception_propagates(): yield MESSAGE_START_CHUNK raise ValueError("upstream broke") - wrapped: Final = wrap_sse_stream_with_keepalive_pings(stream=failing_stream(), ping_interval_seconds=5.0) + wrapped: Final = wrap_sse_stream_with_keepalive_pings( + stream=failing_stream(), + ping_interval_seconds=5.0, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ) assert await wrapped.__anext__() == MESSAGE_START_CHUNK with pytest.raises(ValueError, match="upstream broke"): @@ -85,7 +101,11 @@ async def test_aclose_mid_silence_cancels_upstream_and_runs_its_cleanup(): finally: upstream_cleaned_up.set() - wrapped: Final = wrap_sse_stream_with_keepalive_pings(stream=hung_stream(), ping_interval_seconds=0.05) + wrapped: Final = wrap_sse_stream_with_keepalive_pings( + stream=hung_stream(), + ping_interval_seconds=0.05, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ) assert await wrapped.__anext__() == MESSAGE_START_CHUNK assert await wrapped.__anext__() == ANTHROPIC_PING_SSE_CHUNK @@ -100,7 +120,11 @@ async def test_non_positive_interval_returns_stream_unwrapped(): yield MESSAGE_START_CHUNK stream: Final = any_stream() - assert wrap_sse_stream_with_keepalive_pings(stream=stream, ping_interval_seconds=0) is stream + assert wrap_sse_stream_with_keepalive_pings( + stream=stream, + ping_interval_seconds=0, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ) is stream await stream.aclose() @@ -123,7 +147,11 @@ async def test_invalid_config_interval_returns_stream_unwrapped(bad_interval: fl yield MESSAGE_START_CHUNK stream: Final = any_stream() - assert wrap_sse_stream_with_keepalive_pings(stream=stream, ping_interval_seconds=bad_interval) is stream + assert wrap_sse_stream_with_keepalive_pings( + stream=stream, + ping_interval_seconds=bad_interval, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ) is stream await stream.aclose() @@ -133,7 +161,11 @@ async def test_numeric_string_interval_from_yaml_config_enables_pings(): await asyncio.sleep(0.2) yield MESSAGE_START_CHUNK - wrapped: Final = wrap_sse_stream_with_keepalive_pings(stream=slow_start_stream(), ping_interval_seconds="0.05") + wrapped: Final = wrap_sse_stream_with_keepalive_pings( + stream=slow_start_stream(), + ping_interval_seconds="0.05", + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ) collected: Final = [chunk async for chunk in wrapped] assert collected[0] == ANTHROPIC_PING_SSE_CHUNK @@ -147,7 +179,11 @@ async def test_create_response_streams_ping_first_for_slow_upstream(): yield MESSAGE_START_CHUNK response: Final = await create_response( - generator=wrap_sse_stream_with_keepalive_pings(stream=slow_start_stream(), ping_interval_seconds=0.05), + generator=wrap_sse_stream_with_keepalive_pings( + stream=slow_start_stream(), + ping_interval_seconds=0.05, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ), media_type="text/event-stream", headers={}, ) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index aacc7498ccb..dd19560c166 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -34,10 +34,12 @@ from litellm.proxy.common_request_processing import ( _UpstreamClosingStreamingResponse, create_response, ) +from litellm.proxy.common_utils.sse_keepalive import SSE_COMMENT_PING_CHUNK from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy._types import ProxyException from litellm.proxy._types import UserAPIKeyAuth as ProxyUserAPIKeyAuth from litellm.proxy.utils import ProxyLogging +from litellm.types.utils import ModelResponseStream class TestProxyBaseLLMRequestProcessing: @@ -5746,3 +5748,82 @@ class TestPerRequestModelGroupAlias: ) assert merged_for == ["group-b"] + + +class TestOpenAISseKeepalivePings: + """ + Regression for a streaming request dying at an ingress idle read timeout + (e.g. nginx `proxy-read-timeout`) when time-to-first-token exceeds it: the + OpenAI-shaped SSE routes must emit a keepalive comment while the upstream is + silent, once `litellm.sse_keepalive_ping_interval_seconds` is configured. + """ + + async def _run(self, monkeypatch, first_chunk_delay: float): + import litellm.proxy.common_request_processing as crp + from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth + + logging_obj = MagicMock() + logging_obj.litellm_call_id = "call-keepalive" + logging_obj.cost_breakdown = None + processing_obj = ProxyBaseLLMRequestProcessing( + data={"model": "gpt-4o", "stream": True, "litellm_logging_obj": logging_obj} + ) + + async def upstream(): + await asyncio.sleep(first_chunk_delay) + yield ModelResponseStream() + + async def fake_route_request(**kwargs): + async def _llm_call(): + return upstream() + + return _llm_call() + + monkeypatch.setattr(crp, "route_request", fake_route_request) + + def select_data_generator(response, user_api_key_dict, request_data, request): + async def _gen(): + async for _ in response: + yield 'data: {"choices": []}\n\n' + + return _gen() + + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + proxy_logging_obj.update_request_status = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_call_success_hook = AsyncMock() + + return await processing_obj.base_process_llm_request( + request=MagicMock(spec=Request, headers={}), + fastapi_response=Response(), + user_api_key_dict=RealUserAPIKeyAuth(api_key="sk-test"), + route_type="acompletion", + proxy_logging_obj=proxy_logging_obj, + general_settings={}, + proxy_config=MagicMock(spec=ProxyConfig), + select_data_generator=select_data_generator, + llm_router=None, + skip_pre_call_logic=True, + ) + + @pytest.mark.asyncio + async def test_ping_precedes_slow_first_chunk_on_chat_completions(self, monkeypatch): + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", 0.05) + + result = await self._run(monkeypatch, first_chunk_delay=0.3) + + assert isinstance(result, StreamingResponse) + streamed = [chunk async for chunk in result.body_iterator] + assert streamed[0] == SSE_COMMENT_PING_CHUNK + assert streamed[-1] == 'data: {"choices": []}\n\n' + + @pytest.mark.asyncio + async def test_no_pings_emitted_when_interval_unset(self, monkeypatch): + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", None) + + result = await self._run(monkeypatch, first_chunk_delay=0.3) + + assert isinstance(result, StreamingResponse) + streamed = [chunk async for chunk in result.body_iterator] + assert streamed == ['data: {"choices": []}\n\n']