Add SSE-keepalive to not time out on long prefill on clients

Author: rltakashige

src/exo/api/keepalive.py

@@ -0,0 +1,34 @@
+from collections.abc import AsyncIterator
+from typing import Final
+
+import anyio
+
+_DONE: Final = object()
+
+
+async def with_sse_keepalive(
+    generator: AsyncIterator[str],
+    keepalive_message: str = ": keep-alive\n\n",
+    interval: float = 10.0,
+) -> AsyncIterator[str]:
+    yield keepalive_message
+    send, recv = anyio.create_memory_object_stream[str | object]()
+
+    async def _consume() -> None:
+        async for item in generator:
+            await send.send(item)
+        await send.send(_DONE)
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(_consume)
+        while True:
+            item: str | object | None = None
+            with anyio.move_on_after(interval):
+                item = await recv.receive()
+            if item is None:
+                yield keepalive_message
+            elif item is _DONE:
+                break
+            else:
+                assert isinstance(item, str)
+                yield item

src/exo/api/main.py

@@ -46,6 +46,7 @@
     generate_responses_stream,
     responses_request_to_text_generation,
 )
+from exo.api.keepalive import with_sse_keepalive
 from exo.api.types import (
     AddCustomModelParams,
     AdvancedImageParams,
@@ -780,9 +781,11 @@ async def chat_completions(
 
         if payload.stream:
             return StreamingResponse(
-                generate_chat_stream(
-                    command.command_id,
-                    self._token_chunk_stream(command.command_id),
+                with_sse_keepalive(
+                    generate_chat_stream(
+                        command.command_id,
+                        self._token_chunk_stream(command.command_id),
+                    ),
                 ),
                 media_type="text/event-stream",
                 headers={
@@ -1385,10 +1388,12 @@ async def claude_messages(
 
         if payload.stream:
             return StreamingResponse(
-                generate_claude_stream(
-                    command.command_id,
-                    payload.model,
-                    self._token_chunk_stream(command.command_id),
+                with_sse_keepalive(
+                    generate_claude_stream(
+                        command.command_id,
+                        payload.model,
+                        self._token_chunk_stream(command.command_id),
+                    ),
                 ),
                 media_type="text/event-stream",
                 headers={
@@ -1419,10 +1424,12 @@ async def openai_responses(
 
         if payload.stream:
             return StreamingResponse(
-                generate_responses_stream(
-                    command.command_id,
-                    payload.model,
-                    self._token_chunk_stream(command.command_id),
+                with_sse_keepalive(
+                    generate_responses_stream(
+                        command.command_id,
+                        payload.model,
+                        self._token_chunk_stream(command.command_id),
+                    ),
                 ),
                 media_type="text/event-stream",
                 headers={