[None][feat] Anthropic Messages: keepalive pings while a stream is silent

tensorrt_llm/serve/anthropic_adapter.py

@@ -24,14 +24,15 @@
 chat templates, or the engine.
 """
 
+import asyncio
 import base64
 import hashlib
 import json
 import re
 import traceback
 import uuid
 from dataclasses import dataclass, field
-from typing import Any, AsyncIterator, Dict, List, Optional, Union
+from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Union
 
 from fastapi.responses import JSONResponse
 
@@ -52,6 +53,7 @@
     AnthropicMessagesResponse,
     AnthropicMessageStartEvent,
     AnthropicMessageStopEvent,
+    AnthropicPingEvent,
     AnthropicSignatureDelta,
     AnthropicStopReason,
     AnthropicTextBlock,
@@ -1073,10 +1075,12 @@ def error(self, message: str) -> List[str]:
 
 
 async def _iter_openai_sse_lines(
-    openai_sse: AsyncIterator[Union[str, bytes]],
-) -> AsyncIterator[str]:
+    openai_sse: AsyncIterator[Optional[Union[str, bytes]]],
+) -> AsyncIterator[Optional[str]]:
     """Yield complete lines from string frames or arbitrarily chunked bytes.
 
+    A None item, a tick from ``_with_ticks``, is passed through as None.
+
     The in-process server produces complete SSE strings, while the
     disaggregated frontend relays raw ``aiohttp`` byte chunks.  Network chunks
     need not align with either UTF-8 characters or SSE line boundaries, so
@@ -1088,42 +1092,139 @@ async def _iter_openai_sse_lines(
     buffer = ""
     input_type: Optional[type] = None
 
-    async for payload in openai_sse:
-        if not isinstance(payload, (str, bytes)):
-            raise TypeError(f"OpenAI SSE payload must be str or bytes, got {type(payload)!r}")
+    try:
+        async for payload in openai_sse:
+            if payload is None:
+                yield None
+                continue
+            if not isinstance(payload, (str, bytes)):
+                raise TypeError(f"OpenAI SSE payload must be str or bytes, got {type(payload)!r}")
 
-        current_type = type(payload)
-        if input_type is None:
-            input_type = current_type
-        elif current_type is not input_type:
-            raise TypeError("OpenAI SSE stream cannot mix str and bytes payloads")
+            current_type = type(payload)
+            if input_type is None:
+                input_type = current_type
+            elif current_type is not input_type:
+                raise TypeError("OpenAI SSE stream cannot mix str and bytes payloads")
 
-        if isinstance(payload, bytes):
-            buffer += decoder.decode(payload)
-        else:
-            buffer += payload
+            if isinstance(payload, bytes):
+                buffer += decoder.decode(payload)
+            else:
+                buffer += payload
 
-        while "\n" in buffer:
-            line, buffer = buffer.split("\n", 1)
-            yield line.rstrip("\r")
+            while "\n" in buffer:
+                line, buffer = buffer.split("\n", 1)
+                yield line.rstrip("\r")
+    finally:
+        # Leaving these lines early leaves the upstream too.
+        aclose = getattr(openai_sse, "aclose", None)
+        if aclose is not None:
+            await aclose()
 
     if input_type is bytes:
         buffer += decoder.decode(b"", final=True)
     if buffer:
         yield buffer.rstrip("\r")
 
 
+# Seconds a stream may stay silent before it sends a ping. Clients and proxies
+# reset their timers on any event; without one the stream is silent until the
+# first token (queueing and prefill), and for as long as the model writes a tool
+# call, which is sent only once complete.
+STREAM_PING_INTERVAL_S = 10.0
+
+
+async def _with_ticks(
+    source: AsyncIterator[Any], timeout: Callable[[], Optional[float]]
+) -> AsyncIterator[Any]:
+    """Yield the items of ``source``, and None each time ``timeout()`` seconds pass first.
+
+    ``timeout()`` is asked before every wait; None waits for the next item. One
+    task reads ``source`` to its end, so a tick never interrupts a read and the
+    source is never resumed from another task; that task is cancelled when the
+    consumer stops.
+    """
+    # One item at a time: the reader waits for the consumer, so a slow client
+    # holds back the source as reading it directly would.
+    items: asyncio.Queue = asyncio.Queue(maxsize=1)
+    end = object()
+    consumer_stopped = False
+
+    async def read_all() -> None:
+        try:
+            async for item in source:
+                await items.put(item)
+        except asyncio.CancelledError:
+            if not consumer_stopped:
+                # Raised inside the source, not by the consumer stopping: the
+                # consumer still waits, and the end marker ends its stream
+                # instead of leaving it to pings.
+                await items.put(end)
+            raise
+        except Exception:
+            await items.put(end)
+            raise
+        else:
+            await items.put(end)
+        finally:
+            # The reader closes the source in every case: stopped while handing
+            # over an item, it leaves the source suspended at its yield.
+            aclose = getattr(source, "aclose", None)
+            if aclose is not None:
+                await aclose()
+
+    reader = asyncio.create_task(read_all())
+    try:
+        while True:
+            try:
+                item = await asyncio.wait_for(items.get(), timeout())
+            except asyncio.TimeoutError:
+                yield None
+                continue
+            if item is end:
+                # Raises what reading the source raised.
+                await reader
+                return
+            yield item
+    finally:
+        consumer_stopped = True
+        reader.cancel()
+        # Return only once the reader has stopped and closed the source.
+        await asyncio.gather(reader, return_exceptions=True)
+
+
 async def reframe_openai_stream(
     openai_sse: AsyncIterator[Union[str, bytes]], model: str
 ) -> AsyncIterator[str]:
-    """Translate an OpenAI SSE stream into Anthropic SSE frames."""
+    """Translate an OpenAI SSE stream into Anthropic SSE frames.
+
+    Until [DONE], a ping goes out whenever no other event has for
+    ``STREAM_PING_INTERVAL_S`` seconds, before ``message_start`` too: the message
+    starts at the first upstream chunk, which carries the prompt's usage, so a
+    slow first token would otherwise leave the stream silent. Clients skip
+    pings wherever they appear.
+    """
     reframer = AnthropicStreamReframer(model=model)
     done = False
     # The frames that close the message, held from [DONE] until the upstream
     # has been read to its end.
     closing_frames: List[str] = []
+    loop = asyncio.get_running_loop()
+    last_event = loop.time()
+
+    def until_ping() -> Optional[float]:
+        if done:
+            return None
+        return max(0.0, STREAM_PING_INTERVAL_S - (loop.time() - last_event))
+
+    # The ticks wrap the upstream's payloads, not the lines split from them:
+    # one hand-off between the reading task and this one per payload.
+    lines = _iter_openai_sse_lines(_with_ticks(openai_sse, until_ping))
     try:
-        async for line in _iter_openai_sse_lines(openai_sse):
+        async for line in lines:
+            if line is None:
+                yield anthropic_sse(AnthropicPingEvent())
+                last_event = loop.time()
+                continue
             if done:
                 # The upstream is read to its end before the message is closed:
                 # the chat stream records the request's metrics after its last
@@ -1148,6 +1249,7 @@ async def reframe_openai_stream(
                 raise AnthropicResponseError("Malformed upstream stream chunk") from e
             for frame in reframer.process_chunk(chunk):
                 yield frame
+                last_event = loop.time()
         if done:
             for frame in closing_frames:
                 yield frame
@@ -1177,3 +1279,6 @@ async def reframe_openai_stream(
             return
         for frame in reframer.error("Internal server error"):
             yield frame
+    finally:
+        # Stops the task reading the upstream when this stream ends early.
+        await lines.aclose()

tensorrt_llm/serve/anthropic_protocol.py

@@ -501,6 +501,12 @@ class AnthropicMessageStopEvent(AnthropicBaseModel):
     type: Literal["message_stop"] = "message_stop"
 
 
+class AnthropicPingEvent(AnthropicBaseModel):
+    """A keepalive: the stream format allows any number of them after message_start."""
+
+    type: Literal["ping"] = "ping"
+
+
 class AnthropicErrorEvent(AnthropicBaseModel):
     type: Literal["error"] = "error"
     error: AnthropicError