diff --git a/litellm/proxy/common_utils/sse_keepalive.py b/litellm/proxy/common_utils/sse_keepalive.py index 3ea17b7b1e8..e6d5da03a71 100644 --- a/litellm/proxy/common_utils/sse_keepalive.py +++ b/litellm/proxy/common_utils/sse_keepalive.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import math from collections.abc import AsyncGenerator from typing import Final @@ -8,13 +9,26 @@ 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 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, + ping_interval_seconds: float | str | None, ) -> AsyncGenerator[str, None]: - if ping_interval_seconds <= 0: + interval: Final = _coerce_interval(ping_interval_seconds) + if interval is None: return stream - return _keepalive_ping_stream(stream=stream, ping_interval_seconds=ping_interval_seconds) + return _keepalive_ping_stream(stream=stream, ping_interval_seconds=interval) async def _keepalive_ping_stream( 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 bc6f4060163..6e98248dd18 100644 --- a/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py +++ b/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py @@ -104,6 +104,30 @@ async def test_non_positive_interval_returns_stream_unwrapped(): await stream.aclose() +@pytest.mark.asyncio +@pytest.mark.parametrize("bad_interval", [None, "abc", "", float("inf"), float("nan"), "-3"]) +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]: