Coverage for src/lilbee/server/anthropic_api/streaming.py: 100%

1"""SSE encoder for the Anthropic Messages stream format.""" 

2 

3from __future__ import annotations 

4 

5import json 

6from collections.abc import AsyncIterator 

7from typing import Any 

8 

9from lilbee.server.anthropic_api.models import AnthropicEventType 

10from lilbee.server.handlers.sse import frames_with_keepalive 

11 

12# Cadence for ping events while no token has arrived (same rationale as the 

13# completions keepalive comment frames). 

14_KEEPALIVE_INTERVAL_S = 5.0 

15_PING_FRAME = b'event: ping\ndata: {"type": "ping"}\n\n' 

16 

17 

18def encode_anthropic_event(event_type: AnthropicEventType, payload: dict[str, Any]) -> bytes: 

19 """Frame one event as Anthropic SSE: ``event: <type>`` + ``data: <json>``.""" 

20 body = json.dumps(payload, separators=(",", ":")) 

21 return f"event: {event_type}\ndata: {body}\n\n".encode() 

22 

23 

24async def encode_anthropic_sse( 

25 events: AsyncIterator[tuple[AnthropicEventType, dict[str, Any]]], 

26) -> AsyncIterator[bytes]: 

27 """Frame each ``(type, payload)`` pair as SSE, pinging while upstream is slow. 

28 

29 There is no ``[DONE]`` sentinel in the Anthropic format; the translator's 

30 trailing ``message_stop`` terminates the stream. 

31 """ 

32 

33 async def _frames() -> AsyncIterator[bytes]: 

34 async for event_type, payload in events: 

35 yield encode_anthropic_event(event_type, payload) 

36 

37 async for frame in frames_with_keepalive( 

38 _frames(), keepalive=_PING_FRAME, interval_s=_KEEPALIVE_INTERVAL_S 

39 ): 

40 yield frame