fix(stream): separate reasoning_content from content in SSE streaming
- Before </think>: tokens sent as delta.reasoning_content - After </think>: tokens sent as delta.content - Handles </think> split across token boundaries - Clean <|im_start|>/<|im_end|>/<think>/</think> from streamed text - Normal models (no thinking) send all as content as before
This commit is contained in:
@@ -168,6 +168,8 @@ async fn handle_streaming(
|
|||||||
|
|
||||||
let mut count = 0u32;
|
let mut count = 0u32;
|
||||||
let mut text_buf = String::new();
|
let mut text_buf = String::new();
|
||||||
|
let mut sent_len: usize = 0; // how many chars of text_buf have been sent
|
||||||
|
let mut think_done: bool = false; // true once </think> seen
|
||||||
let mut current = inner.sample(&mut sampler);
|
let mut current = inner.sample(&mut sampler);
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
@@ -255,9 +257,97 @@ async fn handle_streaming(
|
|||||||
}
|
}
|
||||||
|
|
||||||
let piece = state.engine.decode_token(current);
|
let piece = state.engine.decode_token(current);
|
||||||
let (_reasoning, content) = chat::clean_text(&piece);
|
|
||||||
|
|
||||||
if !content.is_empty() {
|
// Push into buffer
|
||||||
|
text_buf.push_str(&piece);
|
||||||
|
|
||||||
|
// Detect </think> transition
|
||||||
|
if !think_done && text_buf.contains("</think>") {
|
||||||
|
think_done = true;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Find new text since last send
|
||||||
|
let new_text = &text_buf[sent_len..]; // everything not yet streamed
|
||||||
|
if new_text.is_empty() {
|
||||||
|
// Nothing new to send; skip straight to decode
|
||||||
|
let pos = input_tokens.len() as i32 + count as i32;
|
||||||
|
if let Err(e) = inner.decode(current, pos) {
|
||||||
|
info!(" Decode error: {e}");
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
count += 1;
|
||||||
|
current = inner.sample(&mut sampler);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Strip special tokens from the NEW text chunk
|
||||||
|
// Split at </think> if present — before goes to reasoning, after to content
|
||||||
|
let (reasoning_part, content_part) = if let Some(pos) = new_text.find("</think>") {
|
||||||
|
let before = new_text[..pos]
|
||||||
|
.replace("<|im_end|>", "")
|
||||||
|
.replace("<|im_start|>", "")
|
||||||
|
.replace("<think>", "")
|
||||||
|
.trim()
|
||||||
|
.to_string();
|
||||||
|
let after = new_text[pos + 8..]
|
||||||
|
.replace("<|im_end|>", "")
|
||||||
|
.replace("<|im_start|>", "")
|
||||||
|
.replace("<think>", "")
|
||||||
|
.trim()
|
||||||
|
.to_string();
|
||||||
|
think_done = true;
|
||||||
|
(Some(before), after)
|
||||||
|
} else if think_done {
|
||||||
|
let cleaned = new_text
|
||||||
|
.replace("<|im_end|>", "")
|
||||||
|
.replace("<|im_start|>", "")
|
||||||
|
.replace("<think>", "")
|
||||||
|
.trim()
|
||||||
|
.to_string();
|
||||||
|
(None, cleaned)
|
||||||
|
} else {
|
||||||
|
let cleaned = new_text
|
||||||
|
.replace("<|im_end|>", "")
|
||||||
|
.replace("<|im_start|>", "")
|
||||||
|
.replace("<think>", "")
|
||||||
|
.trim()
|
||||||
|
.to_string();
|
||||||
|
(Some(cleaned), String::new())
|
||||||
|
};
|
||||||
|
|
||||||
|
// Send reasoning part (before </think>, or entire text if still thinking)
|
||||||
|
if let Some(ref r) = reasoning_part {
|
||||||
|
if !r.is_empty() {
|
||||||
|
let delta = SseDelta {
|
||||||
|
role: None,
|
||||||
|
content: None,
|
||||||
|
reasoning_content: Some(r.clone()),
|
||||||
|
tool_calls: None,
|
||||||
|
};
|
||||||
|
let chunk = serde_json::to_string(&SseChunk {
|
||||||
|
id: chat_id.clone(),
|
||||||
|
object: "chat.completion.chunk".into(),
|
||||||
|
created,
|
||||||
|
model: model_name.clone(),
|
||||||
|
choices: vec![SseChoice {
|
||||||
|
index: 0,
|
||||||
|
delta,
|
||||||
|
finish_reason: None,
|
||||||
|
}],
|
||||||
|
})
|
||||||
|
.unwrap();
|
||||||
|
let _ = tx.send(Ok(Event::default().data(chunk))).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Send content part (after </think>, or never if model doesn't think)
|
||||||
|
if !content_part.is_empty() {
|
||||||
|
let delta = SseDelta {
|
||||||
|
role: None,
|
||||||
|
content: Some(content_part),
|
||||||
|
reasoning_content: None,
|
||||||
|
tool_calls: None,
|
||||||
|
};
|
||||||
let chunk = serde_json::to_string(&SseChunk {
|
let chunk = serde_json::to_string(&SseChunk {
|
||||||
id: chat_id.clone(),
|
id: chat_id.clone(),
|
||||||
object: "chat.completion.chunk".into(),
|
object: "chat.completion.chunk".into(),
|
||||||
@@ -265,12 +355,7 @@ async fn handle_streaming(
|
|||||||
model: model_name.clone(),
|
model: model_name.clone(),
|
||||||
choices: vec![SseChoice {
|
choices: vec![SseChoice {
|
||||||
index: 0,
|
index: 0,
|
||||||
delta: SseDelta {
|
delta,
|
||||||
role: None,
|
|
||||||
content: Some(content.clone()),
|
|
||||||
tool_calls: None,
|
|
||||||
reasoning_content: None,
|
|
||||||
},
|
|
||||||
finish_reason: None,
|
finish_reason: None,
|
||||||
}],
|
}],
|
||||||
})
|
})
|
||||||
@@ -280,7 +365,7 @@ async fn handle_streaming(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
text_buf.push_str(&piece);
|
sent_len = text_buf.len();
|
||||||
|
|
||||||
// Check stop sequences
|
// Check stop sequences
|
||||||
let mut stop_now = false;
|
let mut stop_now = false;
|
||||||
|
|||||||
Reference in New Issue
Block a user