Problem 116 · LLM Systems

SSE Stream Parser that Reassembles Tool-Call Arguments

HardAI engineering

01Problem

A streaming LLM API sends Server-Sent Events. Each event's `data:` is JSON: `{"type":"text","delta":...}`, `{"type":"tool_call_start","index","id","name"}`, `{"type":"tool_call_delta","index","args":<partial JSON string>}`, `{"type":"tool_call_end","index"}`, and finally `data: [DONE]`. Bytes arrive in arbitrary chunks that can split lines, CRLFs and multi-byte UTF-8 characters. Write `StreamAssembler.feed(chunk: bytes)` and `result()` returning the full text and the parsed arguments of every tool call; `result()` must raise if the stream ended early or a tool call never closed.

02Key insight

There are three layers of framing and each can be split: bytes into UTF-8 characters, characters into lines, lines into events. Handle each with its own buffer (an incremental UTF-8 decoder, a line buffer, an event accumulator dispatched on a blank line) and the parser becomes chunk-size-invariant, which the test proves by feeding every chunk size from 1 to 7. Tool-call arguments arrive as fragments of a JSON string, so you concatenate per call index and parse only on the end event; parsing early fails, and parallel tool calls interleave by index. Treat a missing `[DONE]` or an unclosed call as an error, because a truncated stream that looks successful executes a tool with half its arguments.

03Solution

solution.pypython73 lines

import codecs
import json
from dataclasses import dataclass


class SSEParser:
    def __init__(self) -> None:
        self._dec = codecs.getincrementaldecoder("utf-8")()
        self._buf = ""
        self._event: str | None = None
        self._data: list[str] = []

    def feed(self, chunk: bytes) -> list[tuple[str, str]]:
        self._buf += self._dec.decode(chunk)
        events: list[tuple[str, str]] = []
        while (idx := self._buf.find("\n")) >= 0:
            line = self._buf[:idx].rstrip("\r")
            self._buf = self._buf[idx + 1:]
            if line == "":
                if self._data:
                    events.append((self._event or "message", "\n".join(self._data)))
                self._event, self._data = None, []
            elif line.startswith(":"):
                continue
            else:
                field, _, value = line.partition(":")
                value = value[1:] if value.startswith(" ") else value
                if field == "data":
                    self._data.append(value)
                elif field == "event":
                    self._event = value
        return events


@dataclass
class ToolCall:
    id: str
    name: str
    args_text: str = ""
    args: dict | None = None


class StreamAssembler:
    def __init__(self) -> None:
        self.parser = SSEParser()
        self.text: list[str] = []
        self.calls: dict[int, ToolCall] = {}
        self.done = False

    def feed(self, chunk: bytes) -> None:
        for _, data in self.parser.feed(chunk):
            if data == "[DONE]":
                self.done = True
                continue
            ev = json.loads(data)
            kind = ev["type"]
            if kind == "text":
                self.text.append(ev["delta"])
            elif kind == "tool_call_start":
                self.calls[ev["index"]] = ToolCall(ev["id"], ev["name"])
            elif kind == "tool_call_delta":
                self.calls[ev["index"]].args_text += ev["args"]
            elif kind == "tool_call_end":
                call = self.calls[ev["index"]]
                call.args = json.loads(call.args_text) if call.args_text.strip() else {}

    def result(self) -> tuple[str, list[ToolCall]]:
        if not self.done:
            raise ValueError("stream ended before [DONE]")
        open_calls = [c.id for c in self.calls.values() if c.args is None]
        if open_calls:
            raise ValueError(f"unterminated tool calls: {open_calls}")
        return "".join(self.text), [self.calls[i] for i in sorted(self.calls)]

Try it first.

TimeO(total bytes) amortised (line search only scans new data in a production version)

SpaceO(longest event + accumulated text and arguments)

tests.py · 5 passing

tests.pypython34 lines

def ev(obj):
    return "data: " + json.dumps(obj, ensure_ascii=False) + "\r\n\r\n"
stream = (": keep-alive\n\n"
          + ev({"type": "text", "delta": "Café ☕ "})
          + ev({"type": "tool_call_start", "index": 0, "id": "c1", "name": "search"})
          + ev({"type": "tool_call_start", "index": 1, "id": "c2", "name": "weather"})
          + ev({"type": "tool_call_delta", "index": 0, "args": "{\"q\": \"crè"})
          + ev({"type": "tool_call_delta", "index": 1, "args": "{\"city\":"})
          + ev({"type": "tool_call_delta", "index": 0, "args": "me brûlée\", \"k\": 3}"})
          + ev({"type": "tool_call_end", "index": 0})
          + ev({"type": "tool_call_delta", "index": 1, "args": " \"Paris\"}"})
          + ev({"type": "tool_call_end", "index": 1})
          + "event: done\ndata: [DONE]\n\n").encode("utf-8")
expected = None
for size in range(1, 8):
    a = StreamAssembler()
    for i in range(0, len(stream), size):
        a.feed(stream[i:i + size])
    text, calls = a.result()
    got = (text, [(c.id, c.name, c.args) for c in calls])
    expected = expected or got
    assert got == expected
assert expected[0] == "Café ☕ "
assert expected[1] == [("c1", "search", {"q": "crème brûlée", "k": 3}), ("c2", "weather", {"city": "Paris"})]
cut = stream[: stream.index(b"tool_call_end")]
a = StreamAssembler()
a.feed(cut)
try:
    a.result()
    assert False
except ValueError:
    pass
p = SSEParser()
assert p.feed(b"data: a\ndata: b\n\n") == [("message", "a\nb")]

04Staff-level follow-up

For UI responsiveness you often want to show arguments while they stream; that needs a tolerant partial-JSON parser that closes open strings and brackets, used only for display, never for execution. Also handle provider error events mid-stream, reconnection with Last-Event-ID, and backpressure: if the consumer is slow, buffer in bytes with a cap rather than growing unbounded.