fix(stream): robust reasoning/content split + flush un-tagged output
The incremental </think> search missed tags split across tokens, inverting reasoning/content classification. Detect the first </think> on the full buffer and track the content boundary as a byte offset. If the model never closes </think>, flush the buffered text as content so clients always receive the response. Chat UI now renders reasoning_content too. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
b636496497
commit
7cec411cba
@@ -1,6 +1,6 @@
|
|||||||
pub mod use_cases;
|
pub mod use_cases;
|
||||||
|
|
||||||
pub use use_cases::{
|
pub use use_cases::{
|
||||||
build_prompt, build_sampler, clean_text, parse_tool_calls, split_stream_chunk, validate_model,
|
build_prompt, build_sampler, clean_text, parse_tool_calls, split_stream_chunk, strip_markup,
|
||||||
SamplerParams,
|
validate_model, SamplerParams,
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -335,7 +335,7 @@ pub fn parse_tool_calls(text: &str) -> (String, Vec<ToolCall>) {
|
|||||||
/// Handles both fixed tags (`<|im_end|>`, `<think>`, `<tool_call>`, …) and the
|
/// Handles both fixed tags (`<|im_end|>`, `<think>`, `<tool_call>`, …) and the
|
||||||
/// attribute-bearing openers used by this model's tool format (`<function=…>`,
|
/// attribute-bearing openers used by this model's tool format (`<function=…>`,
|
||||||
/// `<parameter=…>`), even when a tag straddles a token boundary.
|
/// `<parameter=…>`), even when a tag straddles a token boundary.
|
||||||
fn strip_markup(text: &str) -> String {
|
pub fn strip_markup(text: &str) -> String {
|
||||||
let mut out = String::with_capacity(text.len());
|
let mut out = String::with_capacity(text.len());
|
||||||
let mut rest = text;
|
let mut rest = text;
|
||||||
|
|
||||||
@@ -383,22 +383,42 @@ fn strip_markup(text: &str) -> String {
|
|||||||
out
|
out
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Split an incremental streamed text fragment into `(reasoning, content)`
|
/// Compute `(reasoning, content)` deltas for an incremental streamed fragment.
|
||||||
/// deltas for SSE.
|
|
||||||
///
|
///
|
||||||
/// * `think_done` means the `</think>` boundary was already crossed **before**
|
/// * `new_text` — the not-yet-emitted fragment (`text_buf[sent_len..]`).
|
||||||
/// this fragment (i.e. it is not the chunk containing the first `</think>`).
|
/// * `sent_len` — byte offset in the full buffer where `new_text` begins.
|
||||||
/// * Whitespace **inside** a fragment is preserved — only the whitespace
|
/// * `content_start` — byte offset in the full buffer where the content phase
|
||||||
/// sitting immediately around the `</think>` boundary is trimmed, so
|
/// begins (immediately after the first `</think>`); `None` while still
|
||||||
/// reasoning does not end with a dangling newline and content does not begin
|
/// thinking.
|
||||||
/// with one. (Trimming every fragment corrupted inter-word spaces.)
|
///
|
||||||
/// * Before the first `</think>`, everything is emitted as `reasoning_content`;
|
/// The boundary is a position in the *full* buffer, not a string search in the
|
||||||
/// after it, as `content`. A stray second `</think>` is stripped, not split.
|
/// fragment — this stays correct even when `</think>` is split across tokens.
|
||||||
pub fn split_stream_chunk(new_text: &str, think_done: bool) -> (Option<String>, String) {
|
/// Whitespace inside a fragment is preserved; only the edges around the
|
||||||
if !think_done {
|
/// boundary are trimmed. A stray second `</think>` is stripped, not re-split.
|
||||||
if let Some(pos) = new_text.find("</think>") {
|
pub fn split_stream_chunk(
|
||||||
let mut before = strip_markup(&new_text[..pos]);
|
new_text: &str,
|
||||||
let mut after = strip_markup(&new_text[pos + 8..]);
|
sent_len: usize,
|
||||||
|
content_start: Option<usize>,
|
||||||
|
) -> (Option<String>, String) {
|
||||||
|
match content_start {
|
||||||
|
// Still thinking — everything is reasoning.
|
||||||
|
None => {
|
||||||
|
let cleaned = strip_markup(new_text);
|
||||||
|
let reasoning = if cleaned.is_empty() {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some(cleaned)
|
||||||
|
};
|
||||||
|
(reasoning, String::new())
|
||||||
|
}
|
||||||
|
// Boundary already emitted — everything is content.
|
||||||
|
Some(cs) if cs <= sent_len => (None, strip_markup(new_text)),
|
||||||
|
// Boundary falls inside this fragment (or beyond it, defensively).
|
||||||
|
Some(cs) => {
|
||||||
|
let rel = (cs - sent_len).min(new_text.len());
|
||||||
|
let (before, after) = new_text.split_at(rel);
|
||||||
|
let mut before = strip_markup(before);
|
||||||
|
let mut after = strip_markup(after);
|
||||||
|
|
||||||
while before.ends_with(['\n', ' ', '\t']) {
|
while before.ends_with(['\n', ' ', '\t']) {
|
||||||
before.pop();
|
before.pop();
|
||||||
@@ -413,19 +433,7 @@ pub fn split_stream_chunk(new_text: &str, think_done: bool) -> (Option<String>,
|
|||||||
Some(before)
|
Some(before)
|
||||||
};
|
};
|
||||||
(reasoning, after)
|
(reasoning, after)
|
||||||
} else {
|
|
||||||
// Still thinking — everything is reasoning.
|
|
||||||
let cleaned = strip_markup(new_text);
|
|
||||||
let reasoning = if cleaned.is_empty() {
|
|
||||||
None
|
|
||||||
} else {
|
|
||||||
Some(cleaned)
|
|
||||||
};
|
|
||||||
(reasoning, String::new())
|
|
||||||
}
|
}
|
||||||
} else {
|
|
||||||
// Think phase already ended — everything is content; strip stray markup.
|
|
||||||
(None, strip_markup(new_text))
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -471,7 +479,7 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn split_chunk_before_think_is_reasoning() {
|
fn split_chunk_before_think_is_reasoning() {
|
||||||
let (reasoning, content) = split_stream_chunk("Hello ", false);
|
let (reasoning, content) = split_stream_chunk("Hello ", 0, None);
|
||||||
assert_eq!(reasoning.as_deref(), Some("Hello "));
|
assert_eq!(reasoning.as_deref(), Some("Hello "));
|
||||||
assert_eq!(content, "");
|
assert_eq!(content, "");
|
||||||
}
|
}
|
||||||
@@ -479,21 +487,23 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
fn split_chunk_preserves_internal_spaces() {
|
fn split_chunk_preserves_internal_spaces() {
|
||||||
// Regression: trimming every fragment used to eat inter-word spaces.
|
// Regression: trimming every fragment used to eat inter-word spaces.
|
||||||
let (r1, _) = split_stream_chunk("Hello", false);
|
let (r1, _) = split_stream_chunk("Hello", 0, None);
|
||||||
let (r2, _) = split_stream_chunk(" world", false);
|
let (r2, _) = split_stream_chunk(" world", 5, None);
|
||||||
assert_eq!(format!("{}{}", r1.unwrap(), r2.unwrap()), "Hello world");
|
assert_eq!(format!("{}{}", r1.unwrap(), r2.unwrap()), "Hello world");
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn split_chunk_boundary_trims_only_edges() {
|
fn split_chunk_boundary_trims_only_edges() {
|
||||||
let (reasoning, content) = split_stream_chunk("question\n</think>\n\nAnswer ", false);
|
// text_buf = "question\n</think>\n\nAnswer "; boundary right after the
|
||||||
|
// tag at byte 17.
|
||||||
|
let (reasoning, content) = split_stream_chunk("question\n</think>\n\nAnswer ", 0, Some(17));
|
||||||
assert_eq!(reasoning.as_deref(), Some("question"));
|
assert_eq!(reasoning.as_deref(), Some("question"));
|
||||||
assert_eq!(content, "Answer ");
|
assert_eq!(content, "Answer ");
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn split_chunk_after_think_is_content() {
|
fn split_chunk_after_think_is_content() {
|
||||||
let (reasoning, content) = split_stream_chunk(" answer", true);
|
let (reasoning, content) = split_stream_chunk(" answer", 0, Some(0));
|
||||||
assert_eq!(reasoning, None);
|
assert_eq!(reasoning, None);
|
||||||
assert_eq!(content, " answer");
|
assert_eq!(content, " answer");
|
||||||
}
|
}
|
||||||
@@ -502,16 +512,29 @@ mod tests {
|
|||||||
fn split_chunk_stray_think_tag_is_stripped_not_split() {
|
fn split_chunk_stray_think_tag_is_stripped_not_split() {
|
||||||
// A second </think> (already past the boundary) must not restart
|
// A second </think> (already past the boundary) must not restart
|
||||||
// reasoning classification.
|
// reasoning classification.
|
||||||
let (reasoning, content) = split_stream_chunk("...</think>more", true);
|
let (reasoning, content) = split_stream_chunk("...</think>more", 0, Some(0));
|
||||||
assert_eq!(reasoning, None);
|
assert_eq!(reasoning, None);
|
||||||
assert_eq!(content, "...more");
|
assert_eq!(content, "...more");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn split_chunk_boundary_straddling_tokens() {
|
||||||
|
// `</think>` split as "</think" + ">" across two fragments: the boundary
|
||||||
|
// is detected on the full buffer, so the answer still becomes content.
|
||||||
|
let (r1, _) = split_stream_chunk("reasoning...</think", 0, None);
|
||||||
|
assert!(r1.is_some());
|
||||||
|
let (r2, c2) = split_stream_chunk("\n\n2 + 2 = 4.", 18, Some(18));
|
||||||
|
assert_eq!(r2, None);
|
||||||
|
assert_eq!(c2, "\n\n2 + 2 = 4.");
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn split_chunk_strips_special_and_markup() {
|
fn split_chunk_strips_special_and_markup() {
|
||||||
|
// boundary at byte 30 (right after "</think>").
|
||||||
let (reasoning, content) = split_stream_chunk(
|
let (reasoning, content) = split_stream_chunk(
|
||||||
"<|im_end|><think>Hello</think>\n<tool_call><function=get_weather>",
|
"<|im_end|><think>Hello</think>\n<tool_call><function=get_weather>",
|
||||||
false,
|
0,
|
||||||
|
Some(30),
|
||||||
);
|
);
|
||||||
assert_eq!(reasoning.as_deref(), Some("Hello"));
|
assert_eq!(reasoning.as_deref(), Some("Hello"));
|
||||||
assert_eq!(content, "");
|
assert_eq!(content, "");
|
||||||
|
|||||||
@@ -72,6 +72,13 @@
|
|||||||
display: flex; align-items: center; gap: 6px;
|
display: flex; align-items: center; gap: 6px;
|
||||||
}
|
}
|
||||||
.msg .tool-call::before { content: '\1F527'; }
|
.msg .tool-call::before { content: '\1F527'; }
|
||||||
|
.msg .reasoning {
|
||||||
|
font-size: 12px; font-style: italic;
|
||||||
|
color: var(--text2);
|
||||||
|
white-space: pre-wrap; word-break: break-word;
|
||||||
|
border-bottom: 1px solid var(--border);
|
||||||
|
padding-bottom: 8px; margin-bottom: 8px;
|
||||||
|
}
|
||||||
.msg.error {
|
.msg.error {
|
||||||
background: #2a1818; border-color: #4a2828; color: #f08080;
|
background: #2a1818; border-color: #4a2828; color: #f08080;
|
||||||
}
|
}
|
||||||
@@ -186,8 +193,33 @@ async function send() {
|
|||||||
const decoder = new TextDecoder();
|
const decoder = new TextDecoder();
|
||||||
let buffer = '';
|
let buffer = '';
|
||||||
let full = '';
|
let full = '';
|
||||||
|
let reasoning = '';
|
||||||
let toolCalls = null;
|
let toolCalls = null;
|
||||||
|
|
||||||
|
// Re-render the assistant bubble: reasoning (muted) above the answer.
|
||||||
|
function render() {
|
||||||
|
el.innerHTML = '';
|
||||||
|
if (reasoning) {
|
||||||
|
const r = document.createElement('div');
|
||||||
|
r.className = 'reasoning';
|
||||||
|
r.textContent = reasoning;
|
||||||
|
el.appendChild(r);
|
||||||
|
}
|
||||||
|
if (full) {
|
||||||
|
const c = document.createElement('div');
|
||||||
|
c.textContent = full;
|
||||||
|
el.appendChild(c);
|
||||||
|
}
|
||||||
|
if (toolCalls?.length) {
|
||||||
|
for (const tc of toolCalls) {
|
||||||
|
const t = document.createElement('div');
|
||||||
|
t.className = 'tool-call';
|
||||||
|
t.textContent = 'Calling tool: ' + (tc.function?.name || 'tool');
|
||||||
|
el.appendChild(t);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
while (true) {
|
while (true) {
|
||||||
const { done, value } = await reader.read();
|
const { done, value } = await reader.read();
|
||||||
if (done) break;
|
if (done) break;
|
||||||
@@ -206,18 +238,20 @@ async function send() {
|
|||||||
const delta = chunk.choices?.[0]?.delta;
|
const delta = chunk.choices?.[0]?.delta;
|
||||||
const finish = chunk.choices?.[0]?.finish_reason;
|
const finish = chunk.choices?.[0]?.finish_reason;
|
||||||
|
|
||||||
|
if (delta?.reasoning_content) {
|
||||||
|
reasoning += delta.reasoning_content;
|
||||||
|
render();
|
||||||
|
}
|
||||||
if (delta?.content) {
|
if (delta?.content) {
|
||||||
full += delta.content;
|
full += delta.content;
|
||||||
el.textContent = full;
|
render();
|
||||||
}
|
}
|
||||||
if (delta?.tool_calls) {
|
if (delta?.tool_calls) {
|
||||||
toolCalls = delta.tool_calls;
|
toolCalls = delta.tool_calls;
|
||||||
|
render();
|
||||||
}
|
}
|
||||||
if (finish === 'tool_calls' && toolCalls) {
|
if (finish === 'tool_calls') {
|
||||||
const t = document.createElement('div');
|
render();
|
||||||
t.className = 'tool-call';
|
|
||||||
t.textContent = 'Calling tool: ' + toolCalls.map(tc => tc.function?.name).join(', ');
|
|
||||||
el.appendChild(t);
|
|
||||||
}
|
}
|
||||||
} catch (e) { /* skip malformed chunk */ }
|
} catch (e) { /* skip malformed chunk */ }
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -176,7 +176,9 @@ async fn handle_streaming(
|
|||||||
let mut sampler = SendSampler(chat::build_sampler(¶ms));
|
let mut sampler = SendSampler(chat::build_sampler(¶ms));
|
||||||
let mut text_buf = String::new();
|
let mut text_buf = String::new();
|
||||||
let mut sent_len: usize = 0;
|
let mut sent_len: usize = 0;
|
||||||
let mut think_done = false;
|
// Byte offset in text_buf where the content phase begins (right after
|
||||||
|
// the first `</think>`); None while still thinking.
|
||||||
|
let mut content_start: Option<usize> = None;
|
||||||
|
|
||||||
let outcome = engine.generate(
|
let outcome = engine.generate(
|
||||||
&input_tokens,
|
&input_tokens,
|
||||||
@@ -187,11 +189,13 @@ async fn handle_streaming(
|
|||||||
&mut |_token, piece| {
|
&mut |_token, piece| {
|
||||||
text_buf.push_str(piece);
|
text_buf.push_str(piece);
|
||||||
|
|
||||||
// `was_thinking` is passed to split_stream_chunk so the chunk
|
// Robust boundary detection on the *full* buffer — a `</think>`
|
||||||
// containing the first </think> is treated as the boundary.
|
// tag may be split across tokens, which would defeat a search
|
||||||
let was_thinking = !think_done;
|
// over the incremental fragment only.
|
||||||
if !think_done && text_buf.contains("</think>") {
|
if content_start.is_none() {
|
||||||
think_done = true;
|
if let Some(pos) = text_buf.find("</think>") {
|
||||||
|
content_start = Some(pos + 8);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
let new_text = &text_buf[sent_len..];
|
let new_text = &text_buf[sent_len..];
|
||||||
@@ -199,7 +203,8 @@ async fn handle_streaming(
|
|||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
|
|
||||||
let (reasoning, content) = chat::split_stream_chunk(new_text, was_thinking);
|
let (reasoning, content) =
|
||||||
|
chat::split_stream_chunk(new_text, sent_len, content_start);
|
||||||
|
|
||||||
if let Some(reasoning) = reasoning {
|
if let Some(reasoning) = reasoning {
|
||||||
let chunk = SseChunk::delta(
|
let chunk = SseChunk::delta(
|
||||||
@@ -251,6 +256,28 @@ async fn handle_streaming(
|
|||||||
total_tokens: prompt_tokens + completion_tokens,
|
total_tokens: prompt_tokens + completion_tokens,
|
||||||
};
|
};
|
||||||
|
|
||||||
|
// If the model never emitted `</think>`, everything was streamed
|
||||||
|
// as reasoning_content. Flush it as content so the client always
|
||||||
|
// receives the response text.
|
||||||
|
if content_start.is_none() && !text_buf.is_empty() {
|
||||||
|
let cleaned = chat::strip_markup(&text_buf);
|
||||||
|
if !cleaned.is_empty() {
|
||||||
|
let chunk = SseChunk::delta(
|
||||||
|
chat_id.clone(),
|
||||||
|
created,
|
||||||
|
model_name.clone(),
|
||||||
|
SseDelta {
|
||||||
|
role: None,
|
||||||
|
content: Some(cleaned),
|
||||||
|
tool_calls: None,
|
||||||
|
reasoning_content: None,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
let event = serde_json::to_string(&chunk).unwrap();
|
||||||
|
let _ = tx.blocking_send(Ok(Event::default().data(event)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// Single-shot tool-calls delta (this model emits whole blocks).
|
// Single-shot tool-calls delta (this model emits whole blocks).
|
||||||
let mut sent_tool_calls = false;
|
let mut sent_tool_calls = false;
|
||||||
if outcome.finish == FinishReason::ToolCalls {
|
if outcome.finish == FinishReason::ToolCalls {
|
||||||
|
|||||||
Reference in New Issue
Block a user