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

1"""SSE encoder: ``data: {json}\\n\\n`` frames plus ``[DONE]`` terminator.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import AsyncIterator 

6 

7from lilbee.server.chat_completions_api.models import CompletionsStreamChunk 

8from lilbee.server.handlers.sse import frames_with_keepalive 

9 

10# Cadence for SSE comment frames emitted when no chat token has arrived. Keeps 

11# clients (notably opencode) from tripping their idle-stream timeout during 

12# slow first-token latency on local models. 

13_KEEPALIVE_INTERVAL_S = 5.0 

14_KEEPALIVE_FRAME = b": keepalive\n\n" 

15 

16 

17async def encode_completions_sse( 

18 chunks: AsyncIterator[CompletionsStreamChunk], 

19) -> AsyncIterator[bytes]: 

20 """Frame each chunk as an SSE ``data:`` event and append the ``[DONE]`` sentinel.""" 

21 

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

23 async for chunk in chunks: 

24 yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n".encode() 

25 

26 async for frame in frames_with_keepalive( 

27 _frames(), keepalive=_KEEPALIVE_FRAME, interval_s=_KEEPALIVE_INTERVAL_S 

28 ): 

29 yield frame 

30 yield b"data: [DONE]\n\n"