Merge pull request #36024 from BerriAI/litellm_anthropic_sse_keepalive

fix(proxy): send keepalive pings on anthropic messages SSE streams during upstream silence

litellm/__init__.py

@@ -244,6 +244,7 @@ use_chat_completions_url_for_anthropic_messages: bool = bool(
 # Or via `litellm_settings.strip_anthropic_total_tokens: true` in
 # config.yaml.
 strip_anthropic_total_tokens: bool = False
+anthropic_sse_ping_interval_seconds: float = 15.0
 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

litellm/proxy/common_request_processing.py

@@ -46,6 +46,7 @@ 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.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
@@ -1980,7 +1981,10 @@ class ProxyBaseLLMRequestProcessing:
                             request=request,
                         )
                         return await create_response(
-                            generator=selected_data_generator,
+                            generator=wrap_sse_stream_with_keepalive_pings(
+                                stream=selected_data_generator,
+                                ping_interval_seconds=litellm.anthropic_sse_ping_interval_seconds,
+                            ),
                             media_type="text/event-stream",
                             headers=custom_headers,
                             request=request,

litellm/proxy/common_utils/sse_keepalive.py

@@ -0,0 +1,57 @@
+import asyncio
+import contextlib
+import math
+from collections.abc import AsyncGenerator
+from typing import Final
+
+import anyio
+
+ANTHROPIC_PING_SSE_CHUNK: Final = 'event: ping\ndata: {"type": "ping"}\n\n'
+
+
+def _coerce_interval(ping_interval_seconds: float | str | None) -> float | None:
+    if ping_interval_seconds is None:
+        return None
+    try:
+        interval: Final = float(ping_interval_seconds)
+    except (TypeError, ValueError):
+        return None
+    if not math.isfinite(interval) or interval <= 0:
+        return None
+    return interval
+
+
+def wrap_sse_stream_with_keepalive_pings(
+    stream: AsyncGenerator[str, None],
+    ping_interval_seconds: float | str | None,
+) -> 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)
+
+
+async def _keepalive_ping_stream(
+    stream: AsyncGenerator[str, None],
+    ping_interval_seconds: float,
+) -> AsyncGenerator[str, None]:
+    pending = asyncio.ensure_future(
+        stream.__anext__()
+    )  # rebind-ok: re-armed with the next __anext__ after each delivered chunk
+    try:
+        while True:
+            await asyncio.wait({pending}, timeout=ping_interval_seconds)
+            if not pending.done():
+                yield ANTHROPIC_PING_SSE_CHUNK
+                continue
+            try:
+                yield pending.result()
+            except StopAsyncIteration:
+                return
+            pending = asyncio.ensure_future(stream.__anext__())
+    finally:
+        pending.cancel()
+        with anyio.CancelScope(shield=True):
+            with contextlib.suppress(BaseException):
+                await pending
+            await stream.aclose()

tests/test_litellm/proxy/common_utils/test_sse_keepalive.py

@@ -0,0 +1,158 @@
+import asyncio
+from collections.abc import AsyncGenerator
+from typing import Final, cast
+
+import pytest
+from fastapi.responses import StreamingResponse
+
+from litellm.proxy.common_request_processing import create_response
+from litellm.proxy.common_utils.sse_keepalive import (
+    ANTHROPIC_PING_SSE_CHUNK,
+    wrap_sse_stream_with_keepalive_pings,
+)
+
+MESSAGE_START_CHUNK: Final = 'data: {"type": "message_start"}\n\n'
+TEXT_DELTA_CHUNK: Final = 'data: {"type": "content_block_delta"}\n\n'
+
+
+@pytest.mark.asyncio
+async def test_pings_fill_mid_stream_silence_and_preserve_chunk_order():
+    async def gappy_stream() -> AsyncGenerator[str, None]:
+        yield MESSAGE_START_CHUNK
+        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)
+    collected: Final = [chunk async for chunk in wrapped]
+
+    assert collected[0] == MESSAGE_START_CHUNK
+    assert collected[-1] == TEXT_DELTA_CHUNK
+    assert ANTHROPIC_PING_SSE_CHUNK in collected[1:-1]
+    assert [chunk for chunk in collected if chunk != ANTHROPIC_PING_SSE_CHUNK] == [
+        MESSAGE_START_CHUNK,
+        TEXT_DELTA_CHUNK,
+    ]
+
+
+@pytest.mark.asyncio
+async def test_ping_emitted_while_waiting_for_first_chunk():
+    async def slow_start_stream() -> AsyncGenerator[str, None]:
+        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)
+    collected: Final = [chunk async for chunk in wrapped]
+
+    assert collected[0] == ANTHROPIC_PING_SSE_CHUNK
+    assert collected[-1] == MESSAGE_START_CHUNK
+
+
+@pytest.mark.asyncio
+async def test_no_pings_when_chunks_arrive_faster_than_interval():
+    async def fast_stream() -> AsyncGenerator[str, None]:
+        yield MESSAGE_START_CHUNK
+        yield TEXT_DELTA_CHUNK
+        yield TEXT_DELTA_CHUNK
+
+    wrapped: Final = wrap_sse_stream_with_keepalive_pings(stream=fast_stream(), ping_interval_seconds=1.0)
+    collected: Final = [chunk async for chunk in wrapped]
+
+    assert collected == [MESSAGE_START_CHUNK, TEXT_DELTA_CHUNK, TEXT_DELTA_CHUNK]
+
+
+@pytest.mark.asyncio
+async def test_upstream_exception_propagates():
+    async def failing_stream() -> AsyncGenerator[str, None]:
+        yield MESSAGE_START_CHUNK
+        raise ValueError("upstream broke")
+
+    wrapped: Final = wrap_sse_stream_with_keepalive_pings(stream=failing_stream(), ping_interval_seconds=5.0)
+
+    assert await wrapped.__anext__() == MESSAGE_START_CHUNK
+    with pytest.raises(ValueError, match="upstream broke"):
+        await wrapped.__anext__()
+
+
+@pytest.mark.asyncio
+async def test_aclose_mid_silence_cancels_upstream_and_runs_its_cleanup():
+    upstream_cleaned_up: Final = asyncio.Event()
+
+    async def hung_stream() -> AsyncGenerator[str, None]:
+        try:
+            yield MESSAGE_START_CHUNK
+            await asyncio.Event().wait()
+            yield TEXT_DELTA_CHUNK
+        finally:
+            upstream_cleaned_up.set()
+
+    wrapped: Final = wrap_sse_stream_with_keepalive_pings(stream=hung_stream(), ping_interval_seconds=0.05)
+
+    assert await wrapped.__anext__() == MESSAGE_START_CHUNK
+    assert await wrapped.__anext__() == ANTHROPIC_PING_SSE_CHUNK
+    await wrapped.aclose()
+
+    assert upstream_cleaned_up.is_set()
+
+
+@pytest.mark.asyncio
+async def test_non_positive_interval_returns_stream_unwrapped():
+    async def any_stream() -> AsyncGenerator[str, None]:
+        yield MESSAGE_START_CHUNK
+
+    stream: Final = any_stream()
+    assert wrap_sse_stream_with_keepalive_pings(stream=stream, ping_interval_seconds=0) is stream
+    await stream.aclose()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "bad_interval",
+    [
+        None,
+        "abc",
+        "",
+        float("inf"),
+        float("nan"),
+        "-3",
+        cast("float | str | None", [15]),
+        cast("float | str | None", {"seconds": 15}),
+    ],
+)
+async def test_invalid_config_interval_returns_stream_unwrapped(bad_interval: float | str | None):
+    async def any_stream() -> AsyncGenerator[str, None]:
+        yield MESSAGE_START_CHUNK
+
+    stream: Final = any_stream()
+    assert wrap_sse_stream_with_keepalive_pings(stream=stream, ping_interval_seconds=bad_interval) is stream
+    await stream.aclose()
+
+
+@pytest.mark.asyncio
+async def test_numeric_string_interval_from_yaml_config_enables_pings():
+    async def slow_start_stream() -> AsyncGenerator[str, None]:
+        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")
+    collected: Final = [chunk async for chunk in wrapped]
+
+    assert collected[0] == ANTHROPIC_PING_SSE_CHUNK
+    assert collected[-1] == MESSAGE_START_CHUNK
+
+
+@pytest.mark.asyncio
+async def test_create_response_streams_ping_first_for_slow_upstream():
+    async def slow_start_stream() -> AsyncGenerator[str, None]:
+        await asyncio.sleep(0.2)
+        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),
+        media_type="text/event-stream",
+        headers={},
+    )
+
+    assert isinstance(response, StreamingResponse)
+    collected: Final = [chunk async for chunk in response.body_iterator]
+    assert collected[0] == ANTHROPIC_PING_SSE_CHUNK
+    assert collected[-1] == MESSAGE_START_CHUNK