diff --git a/tests/test_cancel.py b/tests/test_cancel.py index 5755a43b..0fc8ffae 100644 --- a/tests/test_cancel.py +++ b/tests/test_cancel.py @@ -418,8 +418,9 @@ class TestStreamFlushBeforeToolCalls: arguments_delta: str = "" def stream_content_then_tool(): - # Content long enough to leave chars in pending buffer - # (_MAX_TAG_LEN = 13, so _drain_pending retains last 13 chars) + # Content long enough to leave chars in the tag-scan carry + # buffer (ThinkTagSplitter retains the last MAX_TAG_LEN = 12 + # chars until a flush) yield FakeChunk(content_delta="Hello world, this is a test message") yield FakeChunk( tool_call_deltas=[FakeToolDelta(index=0, id="tc_1", name="bash")], diff --git a/tests/test_midstream_retry.py b/tests/test_midstream_retry.py index cc16d2da..5bed52a4 100644 --- a/tests/test_midstream_retry.py +++ b/tests/test_midstream_retry.py @@ -25,6 +25,7 @@ from turnstone.core.memory import load_last_error from turnstone.core.providers import StreamChunk, UsageInfo from turnstone.core.providers._protocol import IncompleteStreamError from turnstone.core.session import ChatSession +from turnstone.core.streaming_text import ThinkTagSplitter from turnstone.core.trajectory import dicts_from_turns @@ -175,10 +176,10 @@ class TestMidStreamRetry: assert len(assistant) == 1 assert assistant[0]["content"] == "Hello world" # The dead attempt's tokens streamed live (minus the trailing - # _MAX_TAG_LEN chars the tag-scan buffer held back — never + # MAX_TAG_LEN chars the tag-scan buffer held back — never # displayed, so nothing needs to finalize them); the retried text # is a fresh bubble, not an append onto the dead attempt's. - flushed = first_text[: len(first_text) - ChatSession._MAX_TAG_LEN] + flushed = first_text[: len(first_text) - ThinkTagSplitter.MAX_TAG_LEN] assert ui.of("content") == [flushed, "Hello world"] # The dead attempt is finalized EVERYWHERE before re-streaming: # stream_end (client finalize) -> turn_committed (server buffer diff --git a/tests/test_think_tag_split.py b/tests/test_think_tag_split.py new file mode 100644 index 00000000..144c69ed --- /dev/null +++ b/tests/test_think_tag_split.py @@ -0,0 +1,196 @@ +"""Behavior pins for the interactive think-tag splitting layer. + +``ChatSession._stream_attempt`` splits streamed content into content vs +reasoning around ````/```` tags, buffering potential +partial tags across chunk boundaries. These tables pin the CURRENT +emission behavior — exact UI token sequence and final message content — +so the logic can move into a standalone ``ThinkTagSplitter`` class with +byte-identical output. Every case drives the real chunk consumer end to +end; none reaches into the implementation, so the same rows must stay +green across the extraction. + +Pinned rules: + +- partial-tag buffering: a chunk ending in a possible tag prefix emits + nothing until the tag resolves or the safe-flush margin clears it; +- safe-flush: with no tag in sight, everything but the trailing + MAX-tag-length chars flushes immediately (live streaming), the tail + only at stream end; +- open/close tag selection by EARLIEST index among the tag variants; +- ``in_think`` transitions, including the ``reasoning_delta`` (path-1) + interplay and the raw pending flush when tool calls begin (pending + text can no longer be a partial tag). +""" + +import pytest + +from tests._session_helpers import make_session +from turnstone.core.providers import StreamChunk, ToolCallDelta +from turnstone.core.streaming_text import ThinkTagSplitter + + +class _TokenRecorderUI: + """Records dispatched token events; stubs the rest of the surface + ``_stream_attempt`` touches.""" + + def __init__(self): + self.tokens = [] + + def on_content_token(self, text): + self.tokens.append(("content", text)) + + def on_reasoning_token(self, text): + self.tokens.append(("reasoning", text)) + + def on_thinking_stop(self): + pass + + def on_stream_end(self): + pass + + def on_info(self, message): + pass + + def on_error(self, message): + pass + + +def _drive(chunks, *, show_reasoning=True): + session = make_session() + session.show_reasoning = show_reasoning + ui = _TokenRecorderUI() + session.ui = ui + msg = session._stream_attempt(iter(chunks)) + return msg, ui.tokens + + +def _c(text): + return StreamChunk(content_delta=text) + + +_FINISH = StreamChunk(finish_reason="stop") + +# (case id, chunks, expected ordered token events, expected final content) +CASES = [ + ( + "plain_short_held_until_end", + [_c("hello"), _FINISH], + [("content", "hello")], + "hello", + ), + ( + "think_block_single_chunk", + [_c("deepanswer"), _FINISH], + [("reasoning", "deep"), ("content", "answer")], + "answer", + ), + ( + "tag_split_across_chunk_boundary", + [_c("beforerafter"), _FINISH], + [("content", "before"), ("reasoning", "r"), ("content", "after")], + "beforeafter", + ), + ( + "reasoning_tag_variant", + [_c("xdone"), _FINISH], + [("reasoning", "x"), ("content", "done")], + "done", + ), + ( + "earliest_tag_index_wins", + [_c("abcde"), _FINISH], + [ + ("content", "a"), + ("reasoning", "b"), + ("content", "c"), + ("reasoning", "d"), + ("content", "e"), + ], + "ace", + ), + ( + "safe_flush_streams_all_but_max_tag_len", + [_c("x" * 30), _FINISH], + [("content", "x" * 18), ("content", "x" * 12)], + "x" * 30, + ), + ( + "path1_reasoning_then_content", + [StreamChunk(reasoning_delta="rr"), _c("cc"), _FINISH], + [("reasoning", "rr"), ("content", "cc")], + "cc", + ), + ( + "unterminated_think_tail_flushes_as_reasoning", + [_c("abc"), _FINISH], + [("reasoning", "abc")], + "", + ), + ( + "content_surrounds_think_block", + [_c("beforemidafter"), _FINISH], + [("content", "before"), ("reasoning", "mid"), ("content", "after")], + "beforeafter", + ), + ( + "safe_flush_inside_think_block", + [_c("" + "y" * 20), _c("ok"), _FINISH], + [("reasoning", "y" * 8), ("reasoning", "y" * 12), ("content", "ok")], + "ok", + ), +] + + +@pytest.mark.parametrize( + ("chunks", "expected_events", "expected_content"), + [c[1:] for c in CASES], + ids=[c[0] for c in CASES], +) +def test_tag_splitting_emissions(chunks, expected_events, expected_content): + msg, tokens = _drive(chunks) + assert tokens == expected_events + assert msg["content"] == expected_content + + +def test_show_reasoning_off_suppresses_reasoning_dispatch_only(): + msg, tokens = _drive([_c("deepanswer"), _FINISH], show_reasoning=False) + assert tokens == [("content", "answer")] + assert msg["content"] == "answer" + + +def test_splitter_standalone_contract(): + """The extracted class is consumable without a session: feed() emits + tag-resolved spans through the callback, flush_pending() drains the + carry buffer raw, and in_think is externally writable for + out-of-band transitions.""" + events = [] + splitter = ThinkTagSplitter(lambda text, is_reasoning: events.append((text, is_reasoning))) + splitter.feed("abc") + assert events == [("a", False), ("b", True)] + assert splitter.pending == "c" + splitter.flush_pending() + assert events == [("a", False), ("b", True), ("c", False)] + assert splitter.pending == "" + splitter.in_think = True + splitter.feed("tail") + splitter.flush_pending() + assert events[-1] == ("tail", True) + + +def test_tool_calls_flush_pending_raw_at_current_state(): + # Once tool calls begin, buffered text cannot be a partial tag: it + # flushes RAW (no tag scan) at the current in_think state. + chunks = [ + _c("part", "") - _THINK_CLOSE_TAGS = ("", "") - _MAX_TAG_LEN = max(len(t) for t in _THINK_OPEN_TAGS + _THINK_CLOSE_TAGS) - def _stream_attempt( self, stream: Iterator[StreamChunk], my_generation: int = 0 ) -> dict[str, Any]: @@ -7789,9 +7784,7 @@ class ChatSession: provider_blocks: list[dict[str, Any]] = [] usage_acc: UsageInfo | None = None first_token = True - in_think = False # inside a ... block path1_reasoning = False # last reasoning came via reasoning_delta field - pending = "" # buffer for partial tag detection def _flush_text(text: str, is_reasoning: bool) -> None: """Dispatch text to the appropriate UI callback.""" @@ -7805,53 +7798,9 @@ class ChatSession: content_parts.append(text) self.ui.on_content_token(text) - def _drain_pending() -> None: - """Process the pending buffer, flushing content and detecting tags.""" - nonlocal pending, in_think - - while pending: - if in_think: - # Look for any close tag - best_idx, best_tag = None, None - for tag in self._THINK_CLOSE_TAGS: - idx = pending.find(tag) - if idx != -1 and (best_idx is None or idx < best_idx): - best_idx, best_tag = idx, tag - - if best_idx is not None: - assert best_tag is not None - _flush_text(pending[:best_idx], True) - pending = pending[best_idx + len(best_tag) :] - in_think = False - continue - - # No close tag found — check if tail could be a partial tag - safe = len(pending) - self._MAX_TAG_LEN - if safe > 0: - _flush_text(pending[:safe], True) - pending = pending[safe:] - break - else: - # Look for any open tag - best_idx, best_tag = None, None - for tag in self._THINK_OPEN_TAGS: - idx = pending.find(tag) - if idx != -1 and (best_idx is None or idx < best_idx): - best_idx, best_tag = idx, tag - - if best_idx is not None: - assert best_tag is not None - _flush_text(pending[:best_idx], False) - pending = pending[best_idx + len(best_tag) :] - in_think = True - continue - - # No open tag found — flush all but potential partial tag - safe = len(pending) - self._MAX_TAG_LEN - if safe > 0: - _flush_text(pending[:safe], False) - pending = pending[safe:] - break + # Owns the partial-tag carry buffer and the in-think state; + # dispatch stays here in _flush_text. + splitter = ThinkTagSplitter(_flush_text) def _stop_spinner_once() -> None: """Stop the spinner on first real content. Call is idempotent.""" @@ -7878,8 +7827,7 @@ class ChatSession: appends to ``content``, hiding the partial-output signal from the model. """ - if pending: - _flush_text(pending, in_think) + splitter.flush_pending() self.ui.on_stream_end() partial: dict[str, Any] = {"role": "assistant"} partial["content"] = "".join(content_parts) or "" @@ -7930,7 +7878,7 @@ class ChatSession: if chunk.reasoning_delta: _stop_spinner_once() reasoning_parts.append(chunk.reasoning_delta) - in_think = True + splitter.in_think = True path1_reasoning = True if self.show_reasoning: self.ui.on_reasoning_token(chunk.reasoning_delta) @@ -7941,21 +7889,18 @@ class ChatSession: # Close reasoning if transitioning from Path 1 reasoning if path1_reasoning: path1_reasoning = False - in_think = False - pending += chunk.content_delta - _drain_pending() + splitter.in_think = False + splitter.feed(chunk.content_delta) # Handle tool call deltas if chunk.tool_call_deltas: _stop_spinner_once() # Flush any buffered content — model has moved to tool calls, # so pending text cannot be a partial tag. - if pending: - _flush_text(pending, in_think) - pending = "" + splitter.flush_pending() # Close reasoning if transitioning from reasoning - if in_think: - in_think = False + if splitter.in_think: + splitter.in_think = False for tcd in chunk.tool_call_deltas: # THE tool-call merge rule, shared with drain_stream # and the Google raw-fidelity capture — the chat @@ -7985,8 +7930,7 @@ class ChatSession: raise # Flush any remaining buffered text - if pending: - _flush_text(pending, in_think) + splitter.flush_pending() # Warn on non-standard finish reasons if finish_reason == "length": diff --git a/turnstone/core/streaming_text.py b/turnstone/core/streaming_text.py new file mode 100644 index 00000000..14c11fce --- /dev/null +++ b/turnstone/core/streaming_text.py @@ -0,0 +1,82 @@ +"""Streaming think-tag splitting for interactive content streams.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable + + +class ThinkTagSplitter: + """Split streamed content into content vs reasoning around think tags. + + Local-model servers without a reasoning parser emit reasoning inline + as ````/```` blocks inside the content stream, and + a tag can arrive split across chunk boundaries. This class owns the + carry buffer and the in-tag state: :meth:`feed` buffers a chunk and + emits every span that provably cannot be part of an unresolved tag; + the trailing :data:`MAX_TAG_LEN` chars are held until a later chunk + (or a final :meth:`flush_pending`) resolves them. + + Emission goes through the *emit* callback ``(text, is_reasoning)`` — + spans arrive in stream order, never empty, exactly once. The + consumer owns dispatch (UI callbacks, accumulation) and may + read/write :attr:`in_think` directly to mirror out-of-band + transitions (the provider-normalized ``reasoning_delta`` path, + tool-call starts). + + Tag selection: at each step the EARLIEST occurrence wins among the + tag variants for the current state (open tags outside a block, close + tags inside). + """ + + OPEN_TAGS: tuple[str, ...] = ("", "") + CLOSE_TAGS: tuple[str, ...] = ("", "") + MAX_TAG_LEN = max(len(t) for t in OPEN_TAGS + CLOSE_TAGS) + + def __init__(self, emit: Callable[[str, bool], None]) -> None: + self._emit = emit + self.pending = "" + self.in_think = False + + def feed(self, text: str) -> None: + """Buffer *text* and emit every tag-resolved span.""" + self.pending += text + self._drain() + + def flush_pending(self) -> None: + """Emit the raw carry buffer at the current state, with no tag scan. + + For stream boundaries where a partial tag is no longer possible: + end of stream, tool calls beginning, cancellation. + """ + if self.pending: + self._emit(self.pending, self.in_think) + self.pending = "" + + def _drain(self) -> None: + while self.pending: + tags = self.CLOSE_TAGS if self.in_think else self.OPEN_TAGS + best_idx, best_tag = None, None + for tag in tags: + idx = self.pending.find(tag) + if idx != -1 and (best_idx is None or idx < best_idx): + best_idx, best_tag = idx, tag + + if best_idx is not None: + assert best_tag is not None + if best_idx: + self._emit(self.pending[:best_idx], self.in_think) + self.pending = self.pending[best_idx + len(best_tag) :] + self.in_think = not self.in_think + continue + + # No tag in sight — everything but a possible partial-tag tail + # is provably safe to emit now (live streaming beats holding + # the whole buffer for a tag that may never come). + safe = len(self.pending) - self.MAX_TAG_LEN + if safe > 0: + self._emit(self.pending[:safe], self.in_think) + self.pending = self.pending[safe:] + break