Spaces:
Sleeping
Sleeping
| """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) | |