From a9db53e189f2c23586bfc4a6f472448afc4807ef Mon Sep 17 00:00:00 2001 From: oobabooga Date: Tue, 7 Jul 2026 19:50:40 -0300 Subject: [PATCH] Studio: stream reasoning tokens in the tool-loop generator (fixes DeepSeek thinking not streaming with a pill on) (#6947) --- .../core/inference/anthropic_compat.py | 28 +++ studio/backend/core/inference/llama_cpp.py | 54 +++- .../backend/tests/test_anthropic_messages.py | 43 ++++ .../backend/tests/test_llama_cpp_tool_loop.py | 231 +++++++++++++++++- 4 files changed, 338 insertions(+), 18 deletions(-) diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index 7b572a28ff..3c7a4cb182 100644 --- a/studio/backend/core/inference/anthropic_compat.py +++ b/studio/backend/core/inference/anthropic_compat.py @@ -258,6 +258,10 @@ class AnthropicStreamEmitter: self._open_tool_use_id: Optional[str] = None self._open_tool_args_sent: bool = False self._prev_text: str = "" + # Net minus in the text emitted to the client. Tracked + # from emitted deltas (not _prev_text, which a final bare shrink clobbers) + # so an unclosed reasoning-only block can be balanced before close. + self._open_think_tags: int = 0 self._usage: dict = {} def start( @@ -317,6 +321,7 @@ class AnthropicStreamEmitter: """Close any open block and emit message_delta + message_stop.""" events = [] if self._text_block_open or self._open_tool_call_id is not None: + events.extend(self._close_open_think()) events.append(self._close_block()) self._open_tool_call_id = None self._open_tool_use_id = None @@ -344,12 +349,33 @@ class AnthropicStreamEmitter: ) return events + def _close_open_think(self) -> list[str]: + """Emit a ```` delta when the streamed text left a ```` + open. This emitter diffs cumulative snapshots and drops the generator's + final bare shrink, so a reasoning-only reply would otherwise end on an + unclosed tag. Mirrors the chat route's reasoning extractor, which closes + the block on finish; balances the block before it is closed.""" + if not self._text_block_open or self._open_think_tags <= 0: + return [] + self._open_think_tags = 0 + return [ + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": {"type": "text_delta", "text": ""}, + }, + ) + ] + def _handle_content(self, event: dict) -> list[str]: cumulative = event.get("text", "") new_text = cumulative[len(self._prev_text) :] self._prev_text = cumulative if not new_text: return [] + self._open_think_tags += new_text.count("") - new_text.count("") if not self._text_block_open: events = self._open_text_block() else: @@ -374,6 +400,7 @@ class AnthropicStreamEmitter: events = [] if self._text_block_open: + events.extend(self._close_open_think()) events.append(self._close_block()) # Defensive: close a stale open tool_use block before starting another. elif self._open_tool_call_id is not None: @@ -452,6 +479,7 @@ class AnthropicStreamEmitter: events.extend(self._open_text_block()) # Reset text tracking for the next synthesis turn self._prev_text = "" + self._open_think_tags = 0 return events def _open_text_block(self) -> list[str]: diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 11a4ebb3ec..757467a008 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -8603,13 +8603,31 @@ class LlamaCppBackend: } def _flush_reasoning_and_buffer(): - """Append buffered reasoning (as a block) then the held + """Close a live-streamed block (or emit the buffered reasoning + as one block if it never streamed), then append the held content_buffer to the cumulative display text.""" - nonlocal cumulative_display - if reasoning_accum: + nonlocal cumulative_display, in_thinking + if in_thinking: + cumulative_display += "" + in_thinking = False + elif reasoning_accum: cumulative_display += "" + reasoning_accum + "" cumulative_display += content_buffer + def _close_streamed_think() -> bool: + """Close a live-streamed before a tool call drains, so + consumers without a reasoning extractor (Anthropic) get a balanced + block. Returns True when the caller should yield the result.""" + nonlocal cumulative_display, in_thinking, _last_emitted + if not in_thinking: + return False + cumulative_display += "" + in_thinking = False + if len(cumulative_display) > len(_last_emitted) and not _suppress_visible_output: + _last_emitted = cumulative_display + return True + return False + def _looks_like_enabled_bare_json(text: str, enabled_tool_names: set) -> bool: """True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False.""" probe = strip_llama3_leading_sentinels(text.lstrip()) @@ -8797,6 +8815,10 @@ class LlamaCppBackend: # the structured tool call. has_structured_tc = True detect_state = _S_DRAINING + # Close the reasoning prefix before the tool card + # (mirrors the is_match path). + if _close_streamed_think(): + yield {"type": "content", "text": cumulative_display} for tc_d in tc_deltas: idx = tc_d.get("index", 0) if idx not in tool_calls_acc: @@ -8882,17 +8904,17 @@ class LlamaCppBackend: continue # ── Reasoning tokens ── - # Yield only in STREAMING. In BUFFERING and - # DRAINING, accumulate silently so we don't - # corrupt the consumer's prev_text tracker - # (routes/inference.py never resets it - # between tool iterations). + # Stream live except while DRAINING: reasoning is + # orthogonal to tool detection (content_buffer + # only), and the route resets prev_text on + # tool_start, so the block stays a + # monotonic prefix like the no-tool path. reasoning = delta.get("reasoning_content", "") if reasoning: if _reasoning_started_at is None: _reasoning_started_at = time.monotonic() reasoning_accum += reasoning - if detect_state == _S_STREAMING: + if detect_state != _S_DRAINING: if not in_thinking: cumulative_display += "" in_thinking = True @@ -9020,9 +9042,15 @@ class LlamaCppBackend: _hold_buffer = True if _drain_silently: - # No visible prefix -- the buffered text IS - # the call; drain without yielding it. + # The buffered content IS the call; drain it + # without yielding. A live prefix is + # separate from it -- close that. detect_state = _S_DRAINING + if _close_streamed_think(): + yield { + "type": "content", + "text": cumulative_display, + } elif is_match: # Tool signal -- flush any visible # prefix before DRAINING so the @@ -9115,7 +9143,9 @@ class LlamaCppBackend: ), } elif reasoning_accum and not has_content_tokens: - # Reasoning-only reply: show it as plain text. + # Reasoning-only reply: show it as the main response, + # not a thinking block (mirrors the no-tool path; the + # route's extractor closes the streamed ). if _reasoning_started_at is not None and not _reasoning_summary_emitted: _reasoning_summary_emitted = True yield _reasoning_summary_event(_reasoning_started_at) diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 170b456eac..0c6550a3bb 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -52,6 +52,49 @@ from io import BytesIO as _BytesIO from types import SimpleNamespace +def _emitter_client_text(events: list[str]) -> str: + """Concatenate the text_delta payloads an SSE event list carries.""" + text = "" + for line in events: + for raw in line.split("\n"): + raw = raw.strip() + if not raw.startswith("data: "): + continue + data = json.loads(raw[len("data: ") :]) + delta = data.get("delta", {}) + if delta.get("type") == "text_delta": + text += delta.get("text", "") + return text + + +def test_anthropic_emitter_closes_reasoning_only_think_block(): + # A reasoning-only reply streams X live then shrinks to bare X at EOF. + # This emitter diffs cumulative snapshots and drops the shrink, so without a + # closing pass the client text would end on an unclosed . finish() + # must balance it. + emitter = AnthropicStreamEmitter() + events = emitter.start("msg_1", "m") + events += emitter.feed({"type": "content", "text": "The capital"}) + events += emitter.feed({"type": "content", "text": "The capital of France is Paris."}) + # The generator's final bare-text shrink (dropped by the cumulative diff). + events += emitter.feed({"type": "content", "text": "The capital of France is Paris."}) + events += emitter.finish() + + assert _emitter_client_text(events) == "The capital of France is Paris." + + +def test_anthropic_emitter_does_not_double_close_balanced_think(): + # A reasoning-then-answer reply already closes its own ; the balancer + # must not append a second one. + emitter = AnthropicStreamEmitter() + events = emitter.start("msg_1", "m") + events += emitter.feed({"type": "content", "text": "Thinking."}) + events += emitter.feed({"type": "content", "text": "Thinking.Answer."}) + events += emitter.finish() + + assert _emitter_client_text(events) == "Thinking.Answer." + + def test_streamed_anthropic_tool_use_records_api_monitor_reply(monkeypatch): import routes.inference as inf_mod diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index fb1b0e52b7..afac1f5249 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -221,7 +221,7 @@ def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch): assert assistant_messages[-1]["tool_calls"][0]["function"]["name"] == "render_html" -def test_buffered_reasoning_answer_emits_backend_summary(monkeypatch): +def test_streamed_reasoning_answer_emits_backend_summary(monkeypatch): stream = [ _sse({"reasoning_content": "I am thinking."}), _sse({"reasoning_content": " Still thinking."}), @@ -240,17 +240,236 @@ def test_buffered_reasoning_answer_emits_backend_summary(monkeypatch): ) ) + content_texts = [e["text"] for e in events if e["type"] == "content"] + # Reasoning streams live during BUFFERING instead of arriving as one block: + # each reasoning delta is emitted immediately, wrapped in . + assert content_texts[0] == "I am thinking." + assert content_texts[1] == "I am thinking. Still thinking." + # The final event closes the block and appends the answer. + assert content_texts[-1] == "I am thinking. Still thinking.Final answer." + summary_index = next( i for i, event in enumerate(events) if event["type"] == "reasoning_summary" ) - content_index = next(i for i, event in enumerate(events) if event["type"] == "content") - assert summary_index < content_index + final_content_index = max(i for i, event in enumerate(events) if event["type"] == "content") + assert summary_index < final_content_index assert events[summary_index]["duration_ms"] == 62000 - assert ( - events[content_index]["text"] - == "I am thinking. Still thinking.Final answer." + + +def test_reasoning_streams_incrementally_with_tools(monkeypatch): + # Regression (DeepSeek "thinking doesn't stream"): with a tool/pill active the + # tool-loop generator must stream reasoning token-by-token like the no-tool + # path, not accumulate it and dump one buffered block. + stream = [ + _sse({"reasoning_content": "Step one."}), + _sse({"reasoning_content": " Step two."}), + _sse({"reasoning_content": " Step three."}), + _sse({"content": "Done."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "think then answer"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) ) + reasoning_stage = [ + e["text"] + for e in events + if e["type"] == "content" + and e["text"].startswith("") + and "" not in e["text"] + ] + # One live emission per reasoning delta -- not a single dump. + assert reasoning_stage == [ + "Step one.", + "Step one. Step two.", + "Step one. Step two. Step three.", + ] + final = [e["text"] for e in events if e["type"] == "content"][-1] + assert final == "Step one. Step two. Step three.Done." + + +def test_reasoning_only_reply_matches_no_tool_path_with_tools(monkeypatch): + # A reasoning-only turn (whole answer in reasoning_content, no content, no + # tool) with a tool active streams the reasoning live, then resolves to the + # bare reasoning text -- identical to the no-tool generate_chat_completion + # path -- so the non-streaming drain still returns it as `content`, not an + # empty answer. + stream = [ + _sse({"reasoning_content": "The capital of France is Paris."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 5.0, 5.0]) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "just think"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e["text"] for e in events if e["type"] == "content"] + # Reasoning streamed live during BUFFERING (the fix). + assert content_texts[0] == "The capital of France is Paris." + # Resolves to bare reasoning, matching the no-tool sibling. + assert content_texts[-1] == "The capital of France is Paris." + + +def test_reasoning_before_structured_tool_closes_think_block(monkeypatch): + # Regression: reasoning streamed live during BUFFERING must be closed with + # before a structured tool_call drains, so consumers without a + # reasoning extractor (Anthropic /v1/messages) never receive an unclosed + # . Mirrors the is_match (XML tool signal) path. + tool_stream = [ + _sse({"reasoning_content": "Let me search."}), + *_structured_tool_call("web_search", {"query": "weather"}, "call_1"), + ] + final_stream = [ + _sse({"content": "It is sunny."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", lambda name, arguments, **_kwargs: "sunny" + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + tool_start_index = next(i for i, e in enumerate(events) if e["type"] == "tool_start") + content_before_tool = [e["text"] for e in events[:tool_start_index] if e["type"] == "content"] + # Reasoning streamed live, then closed before the tool -- balanced block. + assert content_before_tool[0] == "Let me search." + assert content_before_tool[-1] == "Let me search." + + +def _replay_route_reasoning_extractor(cumulatives: list[str]) -> tuple[str, str]: + """Replay the route's cumulative suffix-diff + reasoning extractor (the + shared core of routes/inference.py gguf_stream_chunks and the tool-loop + consumer) over content snapshots. Returns (visible, reasoning).""" + from routes.inference import _ResponsesReasoningExtractor + + extractor = _ResponsesReasoningExtractor(parse_think_markers = True) + prev_text = "" + visible: list[str] = [] + reasoning: list[str] = [] + for cumulative in cumulatives: + new_text = cumulative[len(prev_text) :] + prev_text = cumulative + if not new_text: + continue + reasoning_delta, visible_delta = extractor.feed(new_text) + if reasoning_delta: + reasoning.append(reasoning_delta) + if visible_delta: + visible.append(visible_delta) + final_reasoning, final_visible = extractor.finish() + if final_reasoning: + reasoning.append(final_reasoning) + if final_visible: + visible.append(final_visible) + return "".join(visible), "".join(reasoning) + + +def test_reasoning_only_route_output_matches_no_tool_path(monkeypatch): + # Parity contract: a reasoning-only reply must reach the client identically + # whether tools are on or off. Both generators stream live then + # resolve to the bare reasoning text; the route's suffix-diff + extractor + # must therefore produce the same (visible, reasoning) split for both. + stream = [ + _sse({"reasoning_content": "The capital"}), + _sse({"reasoning_content": " of France is Paris."}), + _done(), + ] + + tool_backend = _make_backend(monkeypatch, [list(stream)], []) + _patch_monotonic(monkeypatch, [1.0, 2.0, 2.0]) + tool_cumulatives = [ + e["text"] + for e in tool_backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "capital of France?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + if e.get("type") == "content" + ] + + no_tool_backend = _make_backend(monkeypatch, [list(stream)], []) + no_tool_cumulatives = [ + y + for y in no_tool_backend.generate_chat_completion( + messages = [{"role": "user", "content": "capital of France?"}], + ) + if isinstance(y, str) + ] + + # Both paths stream the reasoning live with the same leading shape. (Raw + # yield lists aren't compared verbatim: the tool path emits a pre-existing + # duplicate trailing event that the route's suffix-diff dedupes.) + assert tool_cumulatives[:3] == no_tool_cumulatives[:3] + # The contract that matters: identical route-level output. + tool_out = _replay_route_reasoning_extractor(tool_cumulatives) + no_tool_out = _replay_route_reasoning_extractor(no_tool_cumulatives) + assert tool_out == no_tool_out + # Pin the shared contract so a change to either path shows up here. + _visible, reasoning = tool_out + assert reasoning == "The capital of France is Paris." + + +def test_reasoning_before_bare_json_tool_closes_think_block(monkeypatch): + # _drain_silently sibling of the structured-tool close: a bare-JSON tool call + # with a live reasoning prefix must also close before draining, and + # must never leak the drained call text as content. + tool_stream = [ + _sse({"reasoning_content": "Searching now."}), + _sse({"content": '{"name":"web_search","arguments":{"query":"weather"}}'}), + _done(), + ] + final_stream = [ + _sse({"content": "It is sunny."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", lambda name, arguments, **_kwargs: "sunny" + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + tool_start_index = next(i for i, e in enumerate(events) if e["type"] == "tool_start") + content_before_tool = [e["text"] for e in events[:tool_start_index] if e["type"] == "content"] + assert content_before_tool[0] == "Searching now." + assert content_before_tool[-1] == "Searching now." + # The bare-JSON call text was drained, never surfaced as content. + assert not any('"name"' in t for t in content_before_tool) + def test_consumed_tool_final_pass_emits_latest_reasoning_summary(monkeypatch): tool_stream = [