[Frontend] Enrich chat_parsing streaming events (#58603)

Give serving adapters the delimiter spans and opener captures they need,
and keep streaming after a region that fails to parse.

Signed-off-by: Yoni Gozlan <yonigozlan@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Fix mypy index types in the event-stream test.
This commit is contained in:
Yoni Gozlan
2026-10-01 11:13:13 +08:00
committed by GitHub
parent a37c402f32
commit 2eaa3bc5ac
2 changed files with 157 additions and 18 deletions
+105
View File
@@ -1160,6 +1160,111 @@ class ResponseEventStreamTest(unittest.TestCase):
self.assertEqual(result, {"role": "assistant"})
self.assertEqual(final_events, [])
def test_region_events_expose_delimiter_offsets(self):
spec = {
"start_anchor": "<assistant>",
"fields": {
"content": {"content_args": {"strip": False}},
"tag": {"open_pattern": r"<tag id=\d+>", "close": "</tag>"},
},
}
parser = ResponseParser(spec, prefix="")
events = parser.feed("plain<tag id=7>body</tag>tail")
events += parser.finalize()[1]
boundaries = [
(event["type"], event["field"], parser.input_text[event["start"] : event["end"]])
for event in events
if event["type"] != "region_chunk"
]
self.assertEqual(
boundaries,
[
("region_open", "content", ""),
("region_close", "content", ""),
("region_open", "tag", "<tag id=7>"),
("region_close", "tag", "</tag>"),
("region_open", "content", ""),
("region_close", "content", ""),
],
)
def test_open_captures_and_prefix_boundary(self):
spec = {
"start_anchor": "[BEGIN]",
"fields": {
"tool_calls": {
"open_pattern": r"<call:(?P<name>\w+)>",
"close": "</call>",
"content": "json",
"transform": {"name": "{name}", "arguments": "{content}"},
}
},
}
parser = ResponseParser(spec, prefix="history[BEGIN]<call:get_")
prefix_end = len(parser.input_text)
events = parser.feed('weather>{"city":"Paris"}</call>')
opening = events[0]
self.assertEqual(prefix_end, len("<call:get_"))
self.assertEqual((opening["type"], opening["start"], opening["end"]), ("region_open", 0, len("<call:get_weather>")))
self.assertEqual(opening["captures"], {"name": "get_weather"})
def test_held_prefix_bytes_are_not_rechunked(self):
spec = {"start_anchor": "[BEGIN]", "fields": {"content": {}, "tag": {"open": "<tag>", "close": "</tag>"}}}
parser = ResponseParser(spec, prefix="[BEGIN]Existing <")
chunks = [event["text"] for event in parser.feed("b> tail") if event["type"] == "region_chunk"]
message, _ = parser.finalize()
self.assertEqual(chunks, ["b> tail"])
self.assertEqual(message["content"], "Existing <b> tail")
def test_malformed_region_events_recover_and_continue(self):
spec = {
"start_anchor": "[BEGIN]",
"fields": {
"tool_calls": {
"open_pattern": r"<call:(?P<name>\w+)>",
"close": "</call>",
"content": "json",
"repeats": True,
"transform": {
"type": "function",
"function": {"name": "{name}", "arguments": "{content}"},
},
}
},
}
parser = ResponseParser(spec, prefix="")
events = parser.feed('<call:bad>{"x":</call><call:good>{"x":1}</call>')
message, _ = parser.finalize()
opening, malformed = events[0], events[2]
self.assertEqual(malformed["type"], "region_malformed")
self.assertIsInstance(malformed["error"], ValueError)
self.assertEqual(parser.input_text[opening["end"] : malformed["start"]], '{"x":')
self.assertEqual(parser.input_text[malformed["start"] : malformed["end"]], "</call>")
self.assertEqual(
message["tool_calls"],
[{"type": "function", "function": {"name": "good", "arguments": {"x": 1}}}],
)
parser = ResponseParser(spec, prefix="")
parser.feed('<call:bad>{"x":')
_, events = parser.finalize()
self.assertEqual(events[-1]["type"], "region_malformed")
self.assertEqual(events[-1]["start"], len(parser.input_text))
def test_parse_response_raises_for_malformed_region(self):
spec = {
"start_anchor": "<|assistant|>",
"fields": {"x": {"open": "<x>", "close": "</x>", "content": "json"}},
}
with self.assertRaisesRegex(ValueError, "could not parse region as JSON"):
parse_response('<x>{"name":', spec, prefix="")
with self.assertRaisesRegex(ValueError, "could not parse region as JSON"):
parse_response("", spec, prefix='<|assistant|><x>{"name":</x>')
class PrefixAndTruncationTest(unittest.TestCase):
def test_prefix_lands_inside_explicit_region(self):
+52 -18
View File
@@ -85,8 +85,11 @@ def parse_response(
"""
response_template = load_response_template(response_template)
stream = ResponseParser(response_template, prefix=prefix, tools=tools)
stream.feed(text)
message, _ = stream.finalize()
events = stream.feed(text)
message, final_events = stream.finalize()
for event in stream.initial_events + events + final_events:
if event["type"] == "region_malformed":
raise event["error"]
return message
@@ -110,7 +113,12 @@ class ResponseParser:
Pass OpenAI-style `tools=` dictionaries to cast tool-call arguments
using the calling tool's JSON schema as each region closes.
Events can be either "region_open", "region_chunk", or "region_close".
Events can be "region_open", "region_chunk", "region_close", or "region_malformed".
Open, close and malformed events carry `start` and `end`, the span of `input_text` the
event consumed: its opening or closing delimiter, empty for implicit boundaries.
Explicit opens also carry the opener's named `captures`. A region whose value fails to
parse ends with "region_malformed" instead of "region_close", carrying the `error`, and
parsing continues after it.
ResponseParser requires the chat `prefix` (i.e. the chat history, the prefill before the current generation).
This is because chat templates or assistant prefills can sometimes write part of the message, and if we
@@ -155,10 +163,18 @@ class ResponseParser:
self._body: str = ""
self._opened: bool = False
self._finalized: bool = False
self._prefix_len: int = 0
self._malformed_fields: set[str] = set()
self.initial_events: list[dict] = []
if prefix:
self._consume_prefix(prefix)
@property
def input_text(self) -> str:
"""Raw parser input after start-anchor truncation. Before the first `feed()`, this is
exactly the truncated prefix, so its length marks where generated text begins."""
return self._buffer
def _consume_prefix(self, prefix: str) -> None:
"""Loads the prefix (the chat prefill sent to the model), right-truncates it to the start of the
assistant message (as determined by start_anchor) and then runs the remainder through the parser.
@@ -172,6 +188,7 @@ class ResponseParser:
return
self._buffer = truncated
self._process(self.initial_events, eos=False)
self._prefix_len = len(truncated)
def feed(self, text: str) -> list[dict]:
"""Feeds more text/tokens from the model output into the tokenizer, and returns any events that result
@@ -197,7 +214,11 @@ class ResponseParser:
raise RuntimeError("ResponseParser already finalized")
events: list[dict] = []
self._process(events, eos=True)
missing = [n for n, f in self._spec.fields.items() if not f.optional and n not in self._output]
missing = [
n
for n, f in self._spec.fields.items()
if not f.optional and n not in self._output and n not in self._malformed_fields
]
if missing:
raise ValueError(f"Required response_template fields missing from parsed output: {missing}")
defaults = self._spec.defaults
@@ -216,12 +237,12 @@ class ResponseParser:
self._accumulate(events, self._buffer[self._pos : m.start()])
self._pos = m.end()
if kind == "open":
self._close_current(events)
self._close_current(events, m.start(), m.start())
self._open_explicit(events, field, m)
else: # "close" (always the implicit region's close here,
# since explicit regions only expose their own close)
had_content = self._opened
self._close_current(events)
self._close_current(events, m.start(), m.end())
# Zero-width close on an already-empty region would just
# re-fire next iteration -- bail out to make progress.
if not had_content and m.start() == m.end():
@@ -233,7 +254,7 @@ class ResponseParser:
if self._pos < len(self._buffer):
self._accumulate(events, self._buffer[self._pos :])
self._pos = len(self._buffer)
self._close_current(events)
self._close_current(events, self._pos, self._pos)
break
# Stream everything up to the earliest still-pending delimiter. When
# nothing is pending `hold_start == len(self._buffer)`, so this flushes
@@ -341,20 +362,25 @@ class ResponseParser:
return
field = self._spec.fields[self._current]
if not self._opened:
events.append({"type": "region_open", "field": self._current})
events.append({"type": "region_open", "field": self._current, "start": self._pos, "end": self._pos})
self._opened = True
self._body += text
dirty = field.content not in STREAMABLE_PARSERS
events.append({"type": "region_chunk", "field": self._current, "text": text, "dirty": dirty})
# Prefix bytes held back as a possible delimiter belong to the prompt, not the chunk.
text = text[max(0, self._prefix_len - self._pos) :]
if text:
dirty = field.content not in STREAMABLE_PARSERS
events.append({"type": "region_chunk", "field": self._current, "text": text, "dirty": dirty})
def _open_explicit(self, events: list[dict], field: ResponseTemplateField, m: Any) -> None:
self._current = field.name
self._captures = {k: v for k, v in m.groupdict().items() if v is not None}
self._body = ""
self._opened = True
events.append({"type": "region_open", "field": field.name})
events.append(
{"type": "region_open", "field": field.name, "start": m.start(), "end": m.end(), "captures": self._captures}
)
def _close_current(self, events: list[dict]) -> None:
def _close_current(self, events: list[dict], close_start: int, close_end: int) -> None:
"""Close the current region and reset to the implicit/null region.
Skipped (aside from the reset) when the current region never opened --
avoids vacuous open/close pairs at every explicit boundary."""
@@ -362,22 +388,30 @@ class ResponseParser:
self._reset_to_implicit()
return
field = self._spec.fields[self._current]
value = process_field(self._body, field, self._captures)
if self._tool_params:
value = self._coerce_tool_calls(value)
if field.join is not None:
if not isinstance(value, str):
try:
value = process_field(self._body, field, self._captures)
if self._tool_params:
value = self._coerce_tool_calls(value)
if field.join is not None and not isinstance(value, str):
raise ValueError(
f"Field '{field.name}': 'join' requires each match to parse to a string, "
f"got {type(value).__name__}."
)
except (KeyError, TypeError, ValueError) as error:
self._malformed_fields.add(self._current)
events.append(
{"type": "region_malformed", "field": self._current, "start": close_start, "end": close_end, "error": error}
)
self._reset_to_implicit()
return
if field.join is not None:
previous = self._output.get(self._current)
self._output[self._current] = value if previous is None else previous + field.join + value
elif field.repeats:
self._output.setdefault(self._current, []).append(value)
else:
self._output[self._current] = value
events.append({"type": "region_close", "field": self._current, "value": value})
events.append({"type": "region_close", "field": self._current, "start": close_start, "end": close_end, "value": value})
self._reset_to_implicit()
def _coerce_tool_calls(self, value: Any) -> Any: