Spaces:
Sleeping
Sleeping
File size: 5,335 Bytes
c97e8a9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | """Incremental parser for combined <RESPONSE>/<WIDGET> output.
Unlike the previous parser which buffered the entire widget body before
emitting anything, this yields widget deltas as they arrive — giving the
frontend Claude-style, token-by-token widget streaming.
Event types yielded
-------------------
- ("response_delta", str) : raw response tokens (inside <RESPONSE>)
- ("response_closed", str) : full response text when </RESPONSE> is seen
- ("widget_start", "") : first token of <WIDGET> opens
- ("widget_delta", str) : raw widget tokens (inside <WIDGET>)
- ("complete", response_text, widget_raw) : after </WIDGET> or EOF
Callers decide whether to forward response_delta (strict-primitive mode may
prefer to buffer until the full response is known before re-streaming it).
"""
from __future__ import annotations
from typing import Generator, Iterable, Tuple
_RESPONSE_OPEN = "<RESPONSE>"
_RESPONSE_CLOSE = "</RESPONSE>"
_WIDGET_OPEN = "<WIDGET>"
_WIDGET_CLOSE = "</WIDGET>"
# Keep enough buffer around tag-ends to avoid splitting a tag across emits.
_TAG_SAFETY = max(len(_RESPONSE_CLOSE), len(_WIDGET_CLOSE))
def parse_combined_stream(
chunks: Iterable[str],
*,
emit_response_deltas: bool = True,
emit_widget_deltas: bool = True,
) -> Generator[Tuple[str, ...], None, None]:
"""Yield structured events from an iterable of raw LLM chunks.
The function is a pure state machine: preamble -> response -> widget_pending ->
widget -> complete. Each chunk is appended to a rolling buffer and emitted
up to a safe boundary so we never emit half of a closing tag.
"""
buf = ""
state = "preamble"
widget_started_emitted = False
response_text = ""
response_emitted_len = 0
widget_text = ""
widget_emitted_len = 0
def _find(needle: str, haystack_upper: str) -> int:
return haystack_upper.find(needle)
for chunk in chunks:
if not chunk:
continue
buf += chunk
buf_upper = buf.upper()
# --- Enter <RESPONSE> ---
if state == "preamble":
idx = _find(_RESPONSE_OPEN, buf_upper)
if idx != -1:
buf = buf[idx + len(_RESPONSE_OPEN):]
buf_upper = buf.upper()
state = "response"
response_emitted_len = 0
# --- Stream <RESPONSE> content until </RESPONSE> ---
if state == "response":
close_idx = _find(_RESPONSE_CLOSE, buf_upper)
if close_idx != -1:
to_send = buf[:close_idx][response_emitted_len:]
if to_send and emit_response_deltas:
yield ("response_delta", to_send)
response_text = buf[:close_idx].strip()
buf = buf[close_idx + len(_RESPONSE_CLOSE):]
buf_upper = buf.upper()
state = "widget_pending"
yield ("response_closed", response_text)
else:
safe_len = max(0, len(buf) - _TAG_SAFETY)
if safe_len > response_emitted_len:
delta = buf[response_emitted_len:safe_len]
if emit_response_deltas:
yield ("response_delta", delta)
response_emitted_len = safe_len
# --- Wait for <WIDGET> tag open ---
if state == "widget_pending":
idx = _find(_WIDGET_OPEN, buf_upper)
if idx != -1:
buf = buf[idx + len(_WIDGET_OPEN):]
buf_upper = buf.upper()
state = "widget"
widget_emitted_len = 0
if not widget_started_emitted:
widget_started_emitted = True
yield ("widget_start", "")
# --- Stream <WIDGET> content until </WIDGET> ---
if state == "widget":
close_idx = _find(_WIDGET_CLOSE, buf_upper)
if close_idx != -1:
to_send = buf[:close_idx][widget_emitted_len:]
if to_send and emit_widget_deltas:
yield ("widget_delta", to_send)
widget_text = buf[:close_idx].strip()
yield ("complete", response_text, widget_text)
return
safe_len = max(0, len(buf) - _TAG_SAFETY)
if safe_len > widget_emitted_len:
delta = buf[widget_emitted_len:safe_len]
if emit_widget_deltas:
yield ("widget_delta", delta)
widget_emitted_len = safe_len
# Stream ended without clean closure — salvage what we can.
if state == "response":
remaining = buf[response_emitted_len:]
if remaining and emit_response_deltas:
yield ("response_delta", remaining)
response_text = (response_text or buf).strip()
elif state == "widget":
remaining = buf[widget_emitted_len:]
if remaining and emit_widget_deltas:
yield ("widget_delta", remaining)
widget_text = (widget_text or buf).strip()
elif state == "preamble" and buf.strip():
# Model didn't use XML tags — treat the whole thing as the response.
response_text = buf.strip()
if emit_response_deltas:
yield ("response_delta", response_text)
yield ("complete", response_text, widget_text)
|