Studio: stream reasoning tokens in the tool-loop generator (fixes DeepSeek thinking not streaming with a pill on) (#6947)
This commit is contained in:
oobabooga
GitHub
parent
304b8eca7a
commit
a9db53e189
4 changed files with 338 additions and 18 deletions
|
|
@ -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 <think> minus </think> 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 ``</think>`` delta when the streamed text left a ``<think>``
|
||||
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": "</think>"},
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
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("<think>") - new_text.count("</think>")
|
||||
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]:
|
||||
|
|
|
|||
|
|
@ -8603,13 +8603,31 @@ class LlamaCppBackend:
|
|||
}
|
||||
|
||||
def _flush_reasoning_and_buffer():
|
||||
"""Append buffered reasoning (as a <think> block) then the held
|
||||
"""Close a live-streamed <think> 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 += "</think>"
|
||||
in_thinking = False
|
||||
elif reasoning_accum:
|
||||
cumulative_display += "<think>" + reasoning_accum + "</think>"
|
||||
cumulative_display += content_buffer
|
||||
|
||||
def _close_streamed_think() -> bool:
|
||||
"""Close a live-streamed <think> 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 += "</think>"
|
||||
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 <think> 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 += "<think>"
|
||||
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 <think> 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 <think>).
|
||||
if _reasoning_started_at is not None and not _reasoning_summary_emitted:
|
||||
_reasoning_summary_emitted = True
|
||||
yield _reasoning_summary_event(_reasoning_started_at)
|
||||
|
|
|
|||
|
|
@ -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 <think>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 <think>. finish()
|
||||
# must balance it.
|
||||
emitter = AnthropicStreamEmitter()
|
||||
events = emitter.start("msg_1", "m")
|
||||
events += emitter.feed({"type": "content", "text": "<think>The capital"})
|
||||
events += emitter.feed({"type": "content", "text": "<think>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) == "<think>The capital of France is Paris.</think>"
|
||||
|
||||
|
||||
def test_anthropic_emitter_does_not_double_close_balanced_think():
|
||||
# A reasoning-then-answer reply already closes its own </think>; the balancer
|
||||
# must not append a second one.
|
||||
emitter = AnthropicStreamEmitter()
|
||||
events = emitter.start("msg_1", "m")
|
||||
events += emitter.feed({"type": "content", "text": "<think>Thinking."})
|
||||
events += emitter.feed({"type": "content", "text": "<think>Thinking.</think>Answer."})
|
||||
events += emitter.finish()
|
||||
|
||||
assert _emitter_client_text(events) == "<think>Thinking.</think>Answer."
|
||||
|
||||
|
||||
def test_streamed_anthropic_tool_use_records_api_monitor_reply(monkeypatch):
|
||||
import routes.inference as inf_mod
|
||||
|
||||
|
|
|
|||
|
|
@ -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 <think>.
|
||||
assert content_texts[0] == "<think>I am thinking."
|
||||
assert content_texts[1] == "<think>I am thinking. Still thinking."
|
||||
# The final event closes the block and appends the answer.
|
||||
assert content_texts[-1] == "<think>I am thinking. Still thinking.</think>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"]
|
||||
== "<think>I am thinking. Still thinking.</think>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 <think> 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("<think>")
|
||||
and "</think>" not in e["text"]
|
||||
]
|
||||
# One live emission per reasoning delta -- not a single dump.
|
||||
assert reasoning_stage == [
|
||||
"<think>Step one.",
|
||||
"<think>Step one. Step two.",
|
||||
"<think>Step one. Step two. Step three.",
|
||||
]
|
||||
final = [e["text"] for e in events if e["type"] == "content"][-1]
|
||||
assert final == "<think>Step one. Step two. Step three.</think>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] == "<think>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
|
||||
# </think> before a structured tool_call drains, so consumers without a
|
||||
# reasoning extractor (Anthropic /v1/messages) never receive an unclosed
|
||||
# <think>. 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] == "<think>Let me search."
|
||||
assert content_before_tool[-1] == "<think>Let me search.</think>"
|
||||
|
||||
|
||||
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 <think> 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 </think> 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] == "<think>Searching now."
|
||||
assert content_before_tool[-1] == "<think>Searching now.</think>"
|
||||
# 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 = [
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue