From 21a22110404d98dcab5f51cd90e9acb0acc15e47 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Fri, 4 Sep 2026 04:05:05 +1000 Subject: [PATCH] [Rust Frontend] Use token-attributed text in reasoning and unified parsers (#54884) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Bugen Zhao --- rust/src/chat/src/output/default/unified.rs | 23 +- rust/src/parser/benches/utils/adapter.rs | 7 +- rust/src/parser/src/reasoning/cohere_cmd.rs | 4 +- rust/src/parser/src/reasoning/deepseek_r1.rs | 4 +- rust/src/parser/src/reasoning/delimited.rs | 58 ++-- rust/src/parser/src/reasoning/hy.rs | 4 +- rust/src/parser/src/reasoning/kimi.rs | 4 +- rust/src/parser/src/reasoning/minimax_m3.rs | 35 +- rust/src/parser/src/reasoning/mod.rs | 34 +- rust/src/parser/src/reasoning/qwen3.rs | 4 +- rust/src/parser/src/reasoning/seed_oss.rs | 43 +-- rust/src/parser/src/reasoning/step3p5.rs | 134 ++++---- rust/src/parser/src/reasoning/tests.rs | 303 +++++++++++++++--- rust/src/parser/src/unified/combined.rs | 85 ++++- rust/src/parser/src/unified/gemma4.rs | 124 +++++-- rust/src/parser/src/unified/hy/hy_v3.rs | 10 +- rust/src/parser/src/unified/hy/hy_v4.rs | 10 +- rust/src/parser/src/unified/inkling.rs | 131 +++++--- rust/src/parser/src/unified/kimi_k3.rs | 59 ++-- rust/src/parser/src/unified/mod.rs | 91 +++++- .../tokenizer/src/incremental/attribution.rs | 56 +++- .../src/incremental/attribution/tests.rs | 139 ++++++++ 22 files changed, 1012 insertions(+), 350 deletions(-) diff --git a/rust/src/chat/src/output/default/unified.rs b/rust/src/chat/src/output/default/unified.rs index b511bb2424..33160fa992 100644 --- a/rust/src/chat/src/output/default/unified.rs +++ b/rust/src/chat/src/output/default/unified.rs @@ -13,7 +13,7 @@ use futures::{StreamExt as _, pin_mut}; use thiserror_ext::AsReport; use tracing::warn; use vllm_parser::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput}; -use vllm_text::output::{DecodedTextEvent, SampledDelta}; +use vllm_text::output::{DecodedText, DecodedTextEvent, SampledDelta}; use crate::Result; use crate::error::Error; @@ -63,14 +63,17 @@ impl UnifiedParserState { } /// Convert one decoded text delta into zero or more parsed assistant events. - fn process_delta(&mut self, delta: String) -> Result> { + fn process_delta(&mut self, delta: DecodedText) -> Result> { if self.parser_failed { self.open_call_index = None; - return Ok(text_event(AssistantBlockKind::Text, delta).into_iter().collect()); + return Ok(text_event(AssistantBlockKind::Text, delta.text).into_iter().collect()); } let mut output = UnifiedParserOutput::default(); - match self.parser.parse_into(&delta, &mut output) { + // The parser consumes the delta; keep its text for the fallback path + // that re-emits it as plain text. + let fallback_text = delta.text.clone(); + match self.parser.parse_into(delta, &mut output) { Ok(()) => { let mut events = Vec::new(); self.process_parser_output(output, &mut events)?; @@ -89,7 +92,7 @@ impl UnifiedParserState { let recovered = self.parser.reset(); if recovered.is_empty() && events.is_empty() { - push_text_delta(&mut events, AssistantBlockKind::Text, delta); + push_text_delta(&mut events, AssistantBlockKind::Text, fallback_text); } else { push_text_delta(&mut events, AssistantBlockKind::Text, recovered); } @@ -137,7 +140,7 @@ impl UnifiedParserState { } UnifiedParserEvent::Reasoning(delta) => { self.open_call_index = None; - push_text_delta(events, AssistantBlockKind::Reasoning, delta); + push_text_delta(events, AssistantBlockKind::Reasoning, delta.text); } UnifiedParserEvent::ToolCall(item) => { self.process_tool_item(item, events)?; @@ -245,7 +248,7 @@ pub(crate) async fn unified_event_stream( }, finished, } => { - for next in state.process_delta(decoded.text)? { + for next in state.process_delta(decoded)? { y.yield_ok(next).await; } if logprobs.is_some() || !token_ids.is_empty() { @@ -284,7 +287,7 @@ mod tests { use vllm_parser::unified::{Gemma4UnifiedParser, UnifiedParserError, UnifiedParserOutput}; use vllm_tokenizer::test_utils::TestTokenizer; - use super::unified_event_stream; + use super::{DecodedText, unified_event_stream}; use crate::event::AssistantBlockKind; use crate::output::AssistantEvent; @@ -332,7 +335,7 @@ mod tests { fn parse_into( &mut self, - _delta: &str, + _delta: DecodedText, output: &mut UnifiedParserOutput, ) -> vllm_parser::unified::Result<()> { match self.steps.pop_front().expect("unexpected parser call") { @@ -418,7 +421,7 @@ mod tests { fn reasoning(delta: &str) -> UnifiedParserOutput { let mut output = UnifiedParserOutput::default(); - output.push_reasoning(delta.to_string()); + output.push_reasoning(DecodedText::unattributed(delta)); output } diff --git a/rust/src/parser/benches/utils/adapter.rs b/rust/src/parser/benches/utils/adapter.rs index 9cba325116..a49410f0bf 100644 --- a/rust/src/parser/benches/utils/adapter.rs +++ b/rust/src/parser/benches/utils/adapter.rs @@ -9,7 +9,7 @@ use vllm_parser::tool::{ use vllm_parser::unified::{ UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput, }; -use vllm_tokenizer::Tokenizer; +use vllm_tokenizer::{DecodedText, Tokenizer}; /// Tokenizer stub used by unified-parser benchmarks. struct BenchTokenizer; @@ -99,7 +99,10 @@ impl ToolParser for UnifiedToolParserAdapter { fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { let mut unified_output = UnifiedParserOutput::default(); - let result = self.inner.parse_into(chunk, &mut unified_output).map_err(map_unified_error); + let result = self + .inner + .parse_into(DecodedText::unattributed(chunk), &mut unified_output) + .map_err(map_unified_error); append_unified_output(unified_output, output)?; result } diff --git a/rust/src/parser/src/reasoning/cohere_cmd.rs b/rust/src/parser/src/reasoning/cohere_cmd.rs index c269c99823..e1003fae2d 100644 --- a/rust/src/parser/src/reasoning/cohere_cmd.rs +++ b/rust/src/parser/src/reasoning/cohere_cmd.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use vllm_tokenizer::DynTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer}; use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result}; @@ -38,7 +38,7 @@ impl ReasoningParser for CohereCmdReasoningParser { Ok(()) } - fn push(&mut self, delta: &str) -> Result { + fn push(&mut self, delta: DecodedText) -> Result { Ok(self.inner.push(delta)) } diff --git a/rust/src/parser/src/reasoning/deepseek_r1.rs b/rust/src/parser/src/reasoning/deepseek_r1.rs index 90a071c39c..f5de178a39 100644 --- a/rust/src/parser/src/reasoning/deepseek_r1.rs +++ b/rust/src/parser/src/reasoning/deepseek_r1.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use vllm_tokenizer::DynTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer}; use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result}; @@ -37,7 +37,7 @@ impl ReasoningParser for DeepSeekR1ReasoningParser { Ok(()) } - fn push(&mut self, delta: &str) -> Result { + fn push(&mut self, delta: DecodedText) -> Result { Ok(self.inner.push(delta)) } diff --git a/rust/src/parser/src/reasoning/delimited.rs b/rust/src/parser/src/reasoning/delimited.rs index 78213d291a..a6a573c847 100644 --- a/rust/src/parser/src/reasoning/delimited.rs +++ b/rust/src/parser/src/reasoning/delimited.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use vllm_tokenizer::{DynTokenizer, Tokenizer}; +use vllm_tokenizer::{DecodedText, DynTokenizer, Tokenizer}; use super::{ReasoningDelta, ReasoningError, Result}; @@ -19,7 +19,7 @@ use super::{ReasoningDelta, ReasoningError, Result}; pub(crate) struct DelimitedReasoningParser { tokenizer: DynTokenizer, current_in_reasoning: bool, - buffer: String, + buffer: DecodedText, start_token: String, end_token: String, start_token_id: u32, @@ -55,7 +55,7 @@ impl DelimitedReasoningParser { Ok(Self { tokenizer, current_in_reasoning: default_in_reasoning, - buffer: String::new(), + buffer: DecodedText::default(), start_token, end_token, start_token_id, @@ -81,44 +81,58 @@ impl DelimitedReasoningParser { } /// Parse one decoded text delta and return its reasoning/content split. - pub(crate) fn push(&mut self, delta: &str) -> ReasoningDelta { - self.buffer.push_str(delta); + pub(crate) fn push(&mut self, delta: DecodedText) -> ReasoningDelta { + self.buffer.append(delta); - let partial_suffix_len = self.partial_suffix_len(&self.buffer); - let stable_len = self.buffer.len() - partial_suffix_len; - let pending_suffix = self.buffer.split_off(stable_len); - let stable_text = std::mem::replace(&mut self.buffer, pending_suffix); + let partial_suffix_len = self.partial_suffix_len(&self.buffer.text); + let stable_len = self.buffer.text.len() - partial_suffix_len; + let stable = self.buffer.drain_prefix(stable_len); - self.parse_stable_text(&stable_text) + self.parse_stable_text(stable) } /// Flush any buffered partial delimiter suffix at end of stream. pub(crate) fn finish(&mut self) -> ReasoningDelta { - let stable_text = std::mem::take(&mut self.buffer); - self.parse_stable_text(&stable_text) + // `drain_prefix(text.len())` takes trailing zero-width tokens too. + let stable = self.buffer.drain_prefix(self.buffer.text.len()); + self.parse_stable_text(stable) } /// Parse text that is known not to end with a partial delimiter suffix. - fn parse_stable_text(&mut self, mut stable: &str) -> ReasoningDelta { + /// + /// Reasoning and content pieces keep the attributions of the tokens that + /// produced them; delimiter marker spans are drained and dropped, keeping + /// marker tokens out of any count. + fn parse_stable_text(&mut self, mut stable: DecodedText) -> ReasoningDelta { let mut delta = ReasoningDelta::default(); - while !stable.is_empty() { + while !stable.text.is_empty() { if self.current_in_reasoning { - if let Some(end_idx) = stable.find(&self.end_token) { - delta.push_reasoning(&stable[..end_idx]); - stable = &stable[end_idx + self.end_token.len()..]; + if let Some(end_idx) = stable.text.find(&self.end_token) { + delta.push_reasoning(stable.drain_prefix(end_idx)); + let _ = stable.drain_prefix(self.end_token.len()); self.current_in_reasoning = false; } else { delta.push_reasoning(stable); - break; + return delta; } - } else if let Some(start_idx) = stable.find(&self.start_token) { - delta.push_content(&stable[..start_idx]); - stable = &stable[start_idx + self.start_token.len()..]; + } else if let Some(start_idx) = stable.text.find(&self.start_token) { + delta.push_content(stable.drain_prefix(start_idx)); + let _ = stable.drain_prefix(self.start_token.len()); self.current_in_reasoning = true; } else { delta.push_content(stable); - break; + return delta; + } + } + + // A remainder with empty text may still carry zero-width tokens; + // attribute them to the current state. + if !stable.attributions.is_empty() { + if self.current_in_reasoning { + delta.push_reasoning(stable); + } else { + delta.push_content(stable); } } diff --git a/rust/src/parser/src/reasoning/hy.rs b/rust/src/parser/src/reasoning/hy.rs index b61d987e82..6e20a3bdc9 100644 --- a/rust/src/parser/src/reasoning/hy.rs +++ b/rust/src/parser/src/reasoning/hy.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use vllm_tokenizer::DynTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer}; use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningError, ReasoningParser, Result}; @@ -41,7 +41,7 @@ impl ReasoningParser for HyReasoningParser { Ok(()) } - fn push(&mut self, delta: &str) -> Result { + fn push(&mut self, delta: DecodedText) -> Result { Ok(self.inner.push(delta)) } diff --git a/rust/src/parser/src/reasoning/kimi.rs b/rust/src/parser/src/reasoning/kimi.rs index 324db3fb1c..9332389fa8 100644 --- a/rust/src/parser/src/reasoning/kimi.rs +++ b/rust/src/parser/src/reasoning/kimi.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use vllm_tokenizer::DynTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer}; use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result}; @@ -32,7 +32,7 @@ impl ReasoningParser for KimiReasoningParser { Ok(()) } - fn push(&mut self, delta: &str) -> Result { + fn push(&mut self, delta: DecodedText) -> Result { Ok(self.inner.push(delta)) } diff --git a/rust/src/parser/src/reasoning/minimax_m3.rs b/rust/src/parser/src/reasoning/minimax_m3.rs index 34b6fb1aab..c250184cf4 100644 --- a/rust/src/parser/src/reasoning/minimax_m3.rs +++ b/rust/src/parser/src/reasoning/minimax_m3.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use vllm_tokenizer::DynTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer}; use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result}; @@ -20,7 +20,7 @@ pub struct MiniMaxM3ReasoningParser { at_response_start: bool, /// Holds an initial suffix like `` while preserving later unmatched /// closers as ordinary content. - fn push_inner(&mut self, delta: &str) -> ReasoningDelta { + fn push_inner(&mut self, delta: DecodedText) -> ReasoningDelta { if self.at_response_start && !self.inner.in_reasoning() { - self.leading_end_buffer.push_str(delta); - let buffered = std::mem::take(&mut self.leading_end_buffer); + self.leading_end_buffer.append(delta); + let mut buffered = self.leading_end_buffer.take(); - if buffered.is_empty() { - return ReasoningDelta::default(); - } - if let Some(rest) = buffered.strip_prefix(M3_THINK_END) { + if buffered.text.starts_with(M3_THINK_END) { self.at_response_start = false; - return self.inner.push(rest); + // The dropped marker span keeps its tokens out of any count. + let _ = buffered.drain_prefix(M3_THINK_END.len()); + return self.inner.push(buffered); } - if M3_THINK_END.starts_with(buffered.as_str()) { + if M3_THINK_END.starts_with(buffered.text.as_str()) { self.leading_end_buffer = buffered; return ReasoningDelta::default(); } self.at_response_start = false; - return self.inner.push(&buffered); + return self.inner.push(buffered); } self.inner.push(delta) @@ -62,10 +61,10 @@ impl MiniMaxM3ReasoningParser { fn append_delta(target: &mut ReasoningDelta, delta: ReasoningDelta) { if let Some(reasoning) = delta.reasoning { - target.push_reasoning(&reasoning); + target.push_reasoning(reasoning); } if let Some(content) = delta.content { - target.push_content(&content); + target.push_content(content); } } @@ -84,16 +83,16 @@ impl ReasoningParser for MiniMaxM3ReasoningParser { Ok(()) } - fn push(&mut self, delta: &str) -> Result { + fn push(&mut self, delta: DecodedText) -> Result { Ok(self.push_inner(delta)) } fn finish(&mut self) -> Result { let mut delta = ReasoningDelta::default(); if !self.leading_end_buffer.is_empty() { - let pending = std::mem::take(&mut self.leading_end_buffer); + let pending = self.leading_end_buffer.take(); self.at_response_start = false; - append_delta(&mut delta, self.inner.push(&pending)); + append_delta(&mut delta, self.inner.push(pending)); } append_delta(&mut delta, self.inner.finish()); Ok(delta) diff --git a/rust/src/parser/src/reasoning/mod.rs b/rust/src/parser/src/reasoning/mod.rs index ed355f1caf..e5ff9081f7 100644 --- a/rust/src/parser/src/reasoning/mod.rs +++ b/rust/src/parser/src/reasoning/mod.rs @@ -28,7 +28,7 @@ mod seed_oss; mod step3p5; use thiserror::Error; -use vllm_tokenizer::DynTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer}; pub use self::cohere_cmd::CohereCmdReasoningParser; pub use self::deepseek_r1::DeepSeekR1ReasoningParser; @@ -61,10 +61,14 @@ pub type Step3ReasoningParser = Qwen3ReasoningParser; pub type Result = std::result::Result; /// One parsed streaming delta split into reasoning and visible content. +/// +/// Each portion carries the attributions of the generated tokens that produced +/// it, so downstream consumers can count reasoning tokens exactly. Marker +/// spans are dropped by the parsers, keeping marker tokens out of any count. #[derive(Debug, Default, Clone, PartialEq, Eq)] pub struct ReasoningDelta { - pub reasoning: Option, - pub content: Option, + pub reasoning: Option, + pub content: Option, } impl ReasoningDelta { @@ -74,24 +78,30 @@ impl ReasoningDelta { } /// Append text to the reasoning portion, creating it on first use. - pub(crate) fn push_reasoning(&mut self, text: &str) { - if text.is_empty() { + /// + /// A piece with empty text but non-empty attributions (zero-width tokens + /// only) is still kept: the tokens are attributed to the current state. + pub(crate) fn push_reasoning(&mut self, piece: DecodedText) { + if piece.is_empty() { return; } match &mut self.reasoning { - Some(existing) => existing.push_str(text), - None => self.reasoning = Some(text.to_string()), + Some(existing) => existing.append(piece), + None => self.reasoning = Some(piece), } } /// Append text to the visible content portion, creating it on first use. - pub(crate) fn push_content(&mut self, text: &str) { - if text.is_empty() { + /// + /// A piece with empty text but non-empty attributions (zero-width tokens + /// only) is still kept: the tokens are attributed to the current state. + pub(crate) fn push_content(&mut self, piece: DecodedText) { + if piece.is_empty() { return; } match &mut self.content { - Some(existing) => existing.push_str(text), - None => self.content = Some(text.to_string()), + Some(existing) => existing.append(piece), + None => self.content = Some(piece), } } } @@ -119,7 +129,7 @@ pub trait ReasoningParser: Send { } /// Feed one decoded text delta into the parser. - fn push(&mut self, delta: &str) -> Result; + fn push(&mut self, delta: DecodedText) -> Result; /// Flush any buffered partial delimiter state at end of stream. fn finish(&mut self) -> Result { diff --git a/rust/src/parser/src/reasoning/qwen3.rs b/rust/src/parser/src/reasoning/qwen3.rs index 1289b47f55..477c6468aa 100644 --- a/rust/src/parser/src/reasoning/qwen3.rs +++ b/rust/src/parser/src/reasoning/qwen3.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use vllm_tokenizer::DynTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer}; use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result}; @@ -36,7 +36,7 @@ impl ReasoningParser for Qwen3ReasoningParser { Ok(()) } - fn push(&mut self, delta: &str) -> Result { + fn push(&mut self, delta: DecodedText) -> Result { Ok(self.inner.push(delta)) } diff --git a/rust/src/parser/src/reasoning/seed_oss.rs b/rust/src/parser/src/reasoning/seed_oss.rs index 2de7a614c0..b0054779b1 100644 --- a/rust/src/parser/src/reasoning/seed_oss.rs +++ b/rust/src/parser/src/reasoning/seed_oss.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use vllm_tokenizer::DynTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer}; use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result}; @@ -38,7 +38,7 @@ impl ReasoningParser for SeedOssReasoningParser { Ok(()) } - fn push(&mut self, delta: &str) -> Result { + fn push(&mut self, delta: DecodedText) -> Result { Ok(self.inner.push(delta)) } @@ -53,17 +53,20 @@ mod tests { use super::SeedOssReasoningParser; use crate::reasoning::ReasoningParser; - use crate::reasoning::tests::{SEED_THINK_END_ID, SEED_THINK_START_ID, fake_tokenizer}; + use crate::reasoning::tests::{ + SEED_THINK_END_ID, SEED_THINK_START_ID, content_str, fake_tokenizer, push_str, + reasoning_str, + }; #[test] fn without_prompt_markers_expects_start_token() { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap(); - let delta = parser.push("implicit reasoninganswer").unwrap(); + let delta = push_str(&mut parser, "implicit reasoninganswer"); assert_eq!(delta.reasoning, None); assert_eq!( - delta.content.as_deref(), + content_str(&delta), Some("implicit reasoninganswer") ); } @@ -75,9 +78,9 @@ mod tests { // Prompt prefills ``, opening reasoning before the stream. parser.initialize(&[SEED_THINK_START_ID]).unwrap(); - let delta = parser.push("reasonanswer").unwrap(); - assert_eq!(delta.reasoning.as_deref(), Some("reason")); - assert_eq!(delta.content.as_deref(), Some("answer")); + let delta = push_str(&mut parser, "reasonanswer"); + assert_eq!(reasoning_str(&delta), Some("reason")); + assert_eq!(content_str(&delta), Some("answer")); } #[test] @@ -87,9 +90,9 @@ mod tests { // Prompt already closed reasoning with ``. parser.initialize(&[SEED_THINK_END_ID]).unwrap(); - let delta = parser.push("answer").unwrap(); + let delta = push_str(&mut parser, "answer"); assert_eq!(delta.reasoning, None); - assert_eq!(delta.content.as_deref(), Some("answer")); + assert_eq!(content_str(&delta), Some("answer")); } #[test] @@ -98,9 +101,9 @@ mod tests { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap(); - let delta = parser.push("reasonanswer").unwrap(); - assert_eq!(delta.reasoning.as_deref(), Some("reason")); - assert_eq!(delta.content.as_deref(), Some("answer")); + let delta = push_str(&mut parser, "reasonanswer"); + assert_eq!(reasoning_str(&delta), Some("reason")); + assert_eq!(content_str(&delta), Some("answer")); } #[test] @@ -121,12 +124,12 @@ mod tests { "Final ", "answer", ] { - let delta = parser.push(delta_str).unwrap(); + let delta = push_str(&mut parser, delta_str); if let Some(r) = delta.reasoning { - reasoning.push_str(&r); + reasoning.push_str(&r.text); } if let Some(c) = delta.content { - content.push_str(&c); + content.push_str(&c.text); } } assert_eq!(reasoning, "Some reasoning content"); @@ -140,12 +143,12 @@ mod tests { parser.initialize(&[SEED_THINK_START_ID]).unwrap(); // Closing delimiter `` arrives in two halves. - let first = parser.push("reasonanswer").unwrap(); + let second = push_str(&mut parser, "think>answer"); assert_eq!(second.reasoning, None); - assert_eq!(second.content.as_deref(), Some("answer")); + assert_eq!(content_str(&second), Some("answer")); } } diff --git a/rust/src/parser/src/reasoning/step3p5.rs b/rust/src/parser/src/reasoning/step3p5.rs index 8a826a74f9..17953f20d0 100644 --- a/rust/src/parser/src/reasoning/step3p5.rs +++ b/rust/src/parser/src/reasoning/step3p5.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use vllm_tokenizer::DynTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer}; use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result}; @@ -16,7 +16,7 @@ use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result}; pub struct Step3p5ReasoningParser { inner: DelimitedReasoningParser, /// `\n` at end of last reasoning delta, held in case `` follows. - pending_reasoning_newline: bool, + pending_reasoning_newline: Option, /// Last push ended on `` without emitting content; the next /// content delta's leading `\n` should be dropped. just_ended_reasoning: bool, @@ -27,7 +27,7 @@ impl Step3p5ReasoningParser { pub fn new(tokenizer: DynTokenizer) -> Result { Ok(Self { inner: DelimitedReasoningParser::new(tokenizer, "", "", false)?, - pending_reasoning_newline: false, + pending_reasoning_newline: None, just_ended_reasoning: false, }) } @@ -45,40 +45,42 @@ impl Step3p5ReasoningParser { !now_in_reasoning && (was_in_reasoning || inner_delta.reasoning.is_some()); // Replay or drop a previously-held trailing reasoning newline. - if self.pending_reasoning_newline { + if self.pending_reasoning_newline.is_some() { if let Some(reasoning) = inner_delta.reasoning.as_mut() { - reasoning.insert(0, '\n'); - self.pending_reasoning_newline = false; + let mut held = self.pending_reasoning_newline.take().unwrap(); + held.append(std::mem::take(reasoning)); + *reasoning = held; } else if transitioned { // The held `\n` was the one right before ``: drop it. - self.pending_reasoning_newline = false; + self.pending_reasoning_newline = None; } } // Hold back a trailing reasoning `\n` until we know if `` follows. if let Some(reasoning) = inner_delta.reasoning.as_mut() - && reasoning.ends_with('\n') + && reasoning.text.ends_with('\n') { - reasoning.pop(); + let kept = reasoning.drain_prefix(reasoning.text.len() - 1); + let newline = std::mem::replace(reasoning, kept); if !transitioned { - self.pending_reasoning_newline = true; + self.pending_reasoning_newline = Some(newline); } } // Drop a leading `\n` of content emitted right after ``. if let Some(content) = inner_delta.content.as_mut() && (transitioned || self.just_ended_reasoning) - && content.starts_with('\n') + && content.text.starts_with('\n') { - content.remove(0); + let _ = content.drain_prefix(1); } self.just_ended_reasoning = transitioned && inner_delta.content.is_none(); - if inner_delta.reasoning.as_deref() == Some("") { + if inner_delta.reasoning.as_ref().is_some_and(DecodedText::is_empty) { inner_delta.reasoning = None; } - if inner_delta.content.as_deref() == Some("") { + if inner_delta.content.as_ref().is_some_and(DecodedText::is_empty) { inner_delta.content = None; } @@ -99,7 +101,7 @@ impl ReasoningParser for Step3p5ReasoningParser { Ok(()) } - fn push(&mut self, delta: &str) -> Result { + fn push(&mut self, delta: DecodedText) -> Result { let was = self.inner.in_reasoning(); let inner_delta = self.inner.push(delta); let now = self.inner.in_reasoning(); @@ -113,12 +115,11 @@ impl ReasoningParser for Step3p5ReasoningParser { let mut delta = self.process(inner_delta, was, now); // Emit a still-held newline rather than silently dropping it. - if self.pending_reasoning_newline { + if let Some(held) = self.pending_reasoning_newline.take() { match delta.reasoning.as_mut() { - Some(existing) => existing.push('\n'), - None => delta.reasoning = Some("\n".to_string()), + Some(existing) => existing.append(held), + None => delta.reasoning = Some(held), } - self.pending_reasoning_newline = false; } Ok(delta) @@ -131,7 +132,9 @@ mod tests { use super::Step3p5ReasoningParser; use crate::reasoning::ReasoningParser; - use crate::reasoning::tests::{THINK_START_ID, fake_tokenizer}; + use crate::reasoning::tests::{ + THINK_START_ID, content_str, fake_tokenizer, push_str, reasoning_str, + }; #[test] fn picks_up_prompt_start_boundary() { @@ -140,12 +143,12 @@ mod tests { // Prompt prefills ``, opening reasoning before the stream. parser.initialize(&[THINK_START_ID]).unwrap(); - let delta = parser.push("This is a reasoning sectionThis is the rest").unwrap(); - assert_eq!( - delta.reasoning.as_deref(), - Some("This is a reasoning section") + let delta = push_str( + &mut parser, + "This is a reasoning sectionThis is the rest", ); - assert_eq!(delta.content.as_deref(), Some("This is the rest")); + assert_eq!(reasoning_str(&delta), Some("This is a reasoning section")); + assert_eq!(content_str(&delta), Some("This is the rest")); } #[test] @@ -153,8 +156,8 @@ mod tests { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap(); - let pushed = parser.push("reason without end").unwrap(); - assert_eq!(pushed.reasoning.as_deref(), Some("reason without end")); + let pushed = push_str(&mut parser, "reason without end"); + assert_eq!(reasoning_str(&pushed), Some("reason without end")); assert_eq!(pushed.content, None); let flushed = parser.finish().unwrap(); @@ -166,7 +169,7 @@ mod tests { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap(); - let pushed = parser.push("").unwrap(); + let pushed = push_str(&mut parser, ""); assert!(pushed.is_empty()); let flushed = parser.finish().unwrap(); assert!(flushed.is_empty()); @@ -180,14 +183,15 @@ mod tests { let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap(); parser.initialize(&[THINK_START_ID]).unwrap(); - let delta = parser - .push("\n This is a \n reasoning section\n\n\n\n\nThis is the rest") - .unwrap(); + let delta = push_str( + &mut parser, + "\n This is a \n reasoning section\n\n\n\n\nThis is the rest", + ); assert_eq!( - delta.reasoning.as_deref(), + reasoning_str(&delta), Some("\n This is a \n reasoning section\n\n") ); - assert_eq!(delta.content.as_deref(), Some("\nThis is the rest")); + assert_eq!(content_str(&delta), Some("\nThis is the rest")); } #[test] @@ -195,9 +199,9 @@ mod tests { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap(); - let delta = parser.push("reason\n\nanswer").unwrap(); - assert_eq!(delta.reasoning.as_deref(), Some("reason")); - assert_eq!(delta.content.as_deref(), Some("answer")); + let delta = push_str(&mut parser, "reason\n\nanswer"); + assert_eq!(reasoning_str(&delta), Some("reason")); + assert_eq!(content_str(&delta), Some("answer")); } #[test] @@ -207,18 +211,18 @@ mod tests { // The trailing `\n` from the first push is held until we know whether // `` follows. - let first = parser.push("reason\n").unwrap(); - assert_eq!(first.reasoning.as_deref(), Some("reason")); + let first = push_str(&mut parser, "reason\n"); + assert_eq!(reasoning_str(&first), Some("reason")); assert_eq!(first.content, None); // `` arrives standalone; the held newline should be dropped. - let second = parser.push("").unwrap(); + let second = push_str(&mut parser, ""); assert!(second.is_empty()); // The leading newline of the first content delta is dropped. - let third = parser.push("\nanswer").unwrap(); + let third = push_str(&mut parser, "\nanswer"); assert_eq!(third.reasoning, None); - assert_eq!(third.content.as_deref(), Some("answer")); + assert_eq!(content_str(&third), Some("answer")); } #[test] @@ -226,11 +230,11 @@ mod tests { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap(); - let first = parser.push("reason\n").unwrap(); - assert_eq!(first.reasoning.as_deref(), Some("reason")); + let first = push_str(&mut parser, "reason\n"); + assert_eq!(reasoning_str(&first), Some("reason")); - let second = parser.push("more reason").unwrap(); - assert_eq!(second.reasoning.as_deref(), Some("\nmore reason")); + let second = push_str(&mut parser, "more reason"); + assert_eq!(reasoning_str(&second), Some("\nmore reason")); assert_eq!(second.content, None); } @@ -239,11 +243,11 @@ mod tests { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap(); - let first = parser.push("reason\n").unwrap(); - assert_eq!(first.reasoning.as_deref(), Some("reason")); + let first = push_str(&mut parser, "reason\n"); + assert_eq!(reasoning_str(&first), Some("reason")); let flushed = parser.finish().unwrap(); - assert_eq!(flushed.reasoning.as_deref(), Some("\n")); + assert_eq!(reasoning_str(&flushed), Some("\n")); assert_eq!(flushed.content, None); } @@ -252,9 +256,9 @@ mod tests { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap(); - let delta = parser.push("line1\nline2tail").unwrap(); - assert_eq!(delta.reasoning.as_deref(), Some("line1\nline2")); - assert_eq!(delta.content.as_deref(), Some("tail")); + let delta = push_str(&mut parser, "line1\nline2tail"); + assert_eq!(reasoning_str(&delta), Some("line1\nline2")); + assert_eq!(content_str(&delta), Some("tail")); } #[test] @@ -264,9 +268,9 @@ mod tests { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap(); - let delta = parser.push("reason\n\nanswer").unwrap(); - assert_eq!(delta.reasoning.as_deref(), Some("reason\n")); - assert_eq!(delta.content.as_deref(), Some("answer")); + let delta = push_str(&mut parser, "reason\n\nanswer"); + assert_eq!(reasoning_str(&delta), Some("reason\n")); + assert_eq!(content_str(&delta), Some("answer")); } #[test] @@ -276,18 +280,18 @@ mod tests { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap(); - let first = parser.push("reason").unwrap(); - assert_eq!(first.reasoning.as_deref(), Some("reason")); + let first = push_str(&mut parser, "reason"); + assert_eq!(reasoning_str(&first), Some("reason")); assert_eq!(first.content, None); - let second = parser.push("\nfirst").unwrap(); + let second = push_str(&mut parser, "\nfirst"); assert_eq!(second.reasoning, None); - assert_eq!(second.content.as_deref(), Some("first")); + assert_eq!(content_str(&second), Some("first")); // A `\n` arriving in a later content delta must NOT be dropped. - let third = parser.push("\nsecond").unwrap(); + let third = push_str(&mut parser, "\nsecond"); assert_eq!(third.reasoning, None); - assert_eq!(third.content.as_deref(), Some("\nsecond")); + assert_eq!(content_str(&third), Some("\nsecond")); } #[test] @@ -295,9 +299,9 @@ mod tests { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap(); - let delta = parser.push("reasontail").unwrap(); - assert_eq!(delta.reasoning.as_deref(), Some("reason")); - assert_eq!(delta.content.as_deref(), Some("tail")); + let delta = push_str(&mut parser, "reasontail"); + assert_eq!(reasoning_str(&delta), Some("reason")); + assert_eq!(content_str(&delta), Some("tail")); } #[test] @@ -305,8 +309,8 @@ mod tests { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap(); - let delta = parser.push("answer").unwrap(); + let delta = push_str(&mut parser, "answer"); assert_eq!(delta.reasoning, None); - assert_eq!(delta.content.as_deref(), Some("answer")); + assert_eq!(content_str(&delta), Some("answer")); } } diff --git a/rust/src/parser/src/reasoning/tests.rs b/rust/src/parser/src/reasoning/tests.rs index f7c8e8e3d6..0c41a53b8d 100644 --- a/rust/src/parser/src/reasoning/tests.rs +++ b/rust/src/parser/src/reasoning/tests.rs @@ -4,10 +4,12 @@ use std::sync::Arc; use vllm_tokenizer::test_utils::TestTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer, TokenAnchor, TokenAttribution}; use super::{ - DeepSeekR1ReasoningParser, DelimitedReasoningParser, MiniMaxM3ReasoningParser, - Qwen3ReasoningParser, ReasoningParser, + CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DelimitedReasoningParser, + KimiReasoningParser, MiniMaxM3ReasoningParser, Qwen3ReasoningParser, ReasoningDelta, + ReasoningParser, Result, SeedOssReasoningParser, Step3p5ReasoningParser, }; pub(crate) const THINK_START_ID: u32 = 256; @@ -37,16 +39,29 @@ pub(crate) fn fake_tokenizer() -> TestTokenizer { .with_regular_token("", SEED_THINK_END_ID) } +/// Feed an unattributed text delta into a reasoning parser. +pub(crate) fn push_str(parser: &mut P, delta: &str) -> ReasoningDelta { + parser.push(DecodedText::unattributed(delta)).unwrap() +} + +/// Return the reasoning text of a delta, if any. +pub(crate) fn reasoning_str(delta: &ReasoningDelta) -> Option<&str> { + delta.reasoning.as_ref().map(|piece| piece.text.as_str()) +} + +/// Return the content text of a delta, if any. +pub(crate) fn content_str(delta: &ReasoningDelta) -> Option<&str> { + delta.content.as_ref().map(|piece| piece.text.as_str()) +} + #[test] fn delimited_content_only_stream() { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = DelimitedReasoningParser::new(tokenizer, "", "", false).unwrap(); - assert_eq!( - parser.push("plain content").content.as_deref(), - Some("plain content") - ); + let delta = parser.push(DecodedText::unattributed("plain content")); + assert_eq!(content_str(&delta), Some("plain content")); } #[test] @@ -55,9 +70,9 @@ fn delimited_single_chunk_with_reasoning_and_content() { let mut parser = DelimitedReasoningParser::new(tokenizer, "", "", false).unwrap(); - let delta = parser.push("reasonanswer"); - assert_eq!(delta.reasoning.as_deref(), Some("reason")); - assert_eq!(delta.content.as_deref(), Some("answer")); + let delta = parser.push(DecodedText::unattributed("reasonanswer")); + assert_eq!(reasoning_str(&delta), Some("reason")); + assert_eq!(content_str(&delta), Some("answer")); } #[test] @@ -66,10 +81,10 @@ fn delimited_partial_tokens_across_chunks() { let mut parser = DelimitedReasoningParser::new(tokenizer, "", "", false).unwrap(); - assert!(parser.push("reasonanswer"); - assert_eq!(delta.reasoning.as_deref(), Some("reason")); - assert_eq!(delta.content.as_deref(), Some("answer")); + assert!(parser.push(DecodedText::unattributed("reasonanswer")); + assert_eq!(reasoning_str(&delta), Some("reason")); + assert_eq!(content_str(&delta), Some("answer")); } #[test] @@ -79,10 +94,10 @@ fn delimited_finish_flushes_buffer() { DelimitedReasoningParser::new(tokenizer, "", "", false).unwrap(); parser.initialize(&[THINK_START_ID]); - let delta = parser.push("unfinishedanswer").unwrap(); + let delta = push_str(&mut parser, "reasonanswer"); assert_eq!(delta.reasoning, None); - assert_eq!(delta.content.as_deref(), Some("reasonanswer")); + assert_eq!(content_str(&delta), Some("reasonanswer")); } #[test] @@ -101,9 +116,9 @@ fn qwen3_prompt_end_marker_starts_in_content() { let mut parser = Qwen3ReasoningParser::new(tokenizer).unwrap(); parser.initialize(&[THINK_END_ID]).unwrap(); - let delta = parser.push("answer").unwrap(); + let delta = push_str(&mut parser, "answer"); assert_eq!(delta.reasoning, None); - assert_eq!(delta.content.as_deref(), Some("answer")); + assert_eq!(content_str(&delta), Some("answer")); } #[test] @@ -111,15 +126,15 @@ fn qwen3_tolerates_old_and_new_formats() { let tokenizer = Arc::new(fake_tokenizer()); let mut old_parser = Qwen3ReasoningParser::new(tokenizer.clone()).unwrap(); - let old = old_parser.push("reasonanswer").unwrap(); - assert_eq!(old.reasoning.as_deref(), Some("reason")); - assert_eq!(old.content.as_deref(), Some("answer")); + let old = push_str(&mut old_parser, "reasonanswer"); + assert_eq!(reasoning_str(&old), Some("reason")); + assert_eq!(content_str(&old), Some("answer")); let mut new_parser = Qwen3ReasoningParser::new(tokenizer).unwrap(); new_parser.initialize(&[THINK_START_ID]).unwrap(); - let new = new_parser.push("reasonanswer").unwrap(); - assert_eq!(new.reasoning.as_deref(), Some("reason")); - assert_eq!(new.content.as_deref(), Some("answer")); + let new = push_str(&mut new_parser, "reasonanswer"); + assert_eq!(reasoning_str(&new), Some("reason")); + assert_eq!(content_str(&new), Some("answer")); } #[test] @@ -129,9 +144,9 @@ fn qwen3_stops_scanning_at_last_special_token() { parser.initialize(&[THINK_START_ID, SPECIAL_BOUNDARY_ID]).unwrap(); - let delta = parser.push("answer").unwrap(); + let delta = push_str(&mut parser, "answer"); assert_eq!(delta.reasoning, None); - assert_eq!(delta.content.as_deref(), Some("answer")); + assert_eq!(content_str(&delta), Some("answer")); } #[test] @@ -139,9 +154,9 @@ fn deepseek_r1_defaults_to_reasoning_without_prompt_boundary() { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = DeepSeekR1ReasoningParser::new(tokenizer).unwrap(); - let delta = parser.push("reasonanswer").unwrap(); - assert_eq!(delta.reasoning.as_deref(), Some("reason")); - assert_eq!(delta.content.as_deref(), Some("answer")); + let delta = push_str(&mut parser, "reasonanswer"); + assert_eq!(reasoning_str(&delta), Some("reason")); + assert_eq!(content_str(&delta), Some("answer")); } #[test] @@ -151,9 +166,9 @@ fn deepseek_r1_stops_scanning_at_last_special_token() { parser.initialize(&[THINK_END_ID, SPECIAL_BOUNDARY_ID]).unwrap(); - let delta = parser.push("reasonanswer").unwrap(); - assert_eq!(delta.reasoning.as_deref(), Some("reason")); - assert_eq!(delta.content.as_deref(), Some("answer")); + let delta = push_str(&mut parser, "reasonanswer"); + assert_eq!(reasoning_str(&delta), Some("reason")); + assert_eq!(content_str(&delta), Some("answer")); } #[test] @@ -161,9 +176,9 @@ fn minimax_m3_handles_explicit_think_delimiters() { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap(); - let delta = parser.push("reasonanswer").unwrap(); - assert_eq!(delta.reasoning.as_deref(), Some("reason")); - assert_eq!(delta.content.as_deref(), Some("answer")); + let delta = push_str(&mut parser, "reasonanswer"); + assert_eq!(reasoning_str(&delta), Some("reason")); + assert_eq!(content_str(&delta), Some("answer")); } #[test] @@ -171,9 +186,9 @@ fn minimax_m3_drops_leading_end_marker() { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap(); - let delta = parser.push("answer").unwrap(); + let delta = push_str(&mut parser, "answer"); assert_eq!(delta.reasoning, None); - assert_eq!(delta.content.as_deref(), Some("answer")); + assert_eq!(content_str(&delta), Some("answer")); } #[test] @@ -181,9 +196,9 @@ fn minimax_m3_preserves_non_leading_end_marker() { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap(); - let delta = parser.push("XXXYYY").unwrap(); + let delta = push_str(&mut parser, "XXXYYY"); assert_eq!(delta.reasoning, None); - assert_eq!(delta.content.as_deref(), Some("XXXYYY")); + assert_eq!(content_str(&delta), Some("XXXYYY")); } #[test] @@ -191,10 +206,10 @@ fn minimax_m3_drops_split_leading_end_marker() { let tokenizer = Arc::new(fake_tokenizer()); let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap(); - assert!(parser.push("answer").unwrap(); + assert!(push_str(&mut parser, "answer"); assert_eq!(delta.reasoning, None); - assert_eq!(delta.content.as_deref(), Some("answer")); + assert_eq!(content_str(&delta), Some("answer")); } #[test] @@ -203,9 +218,9 @@ fn minimax_m3_uses_prompt_prefilled_start_marker() { let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap(); parser.initialize(&[MM_THINK_START_ID]).unwrap(); - let delta = parser.push("reasonanswer").unwrap(); - assert_eq!(delta.reasoning.as_deref(), Some("reason")); - assert_eq!(delta.content.as_deref(), Some("answer")); + let delta = push_str(&mut parser, "reasonanswer"); + assert_eq!(reasoning_str(&delta), Some("reason")); + assert_eq!(content_str(&delta), Some("answer")); } #[test] @@ -214,7 +229,197 @@ fn minimax_m3_uses_prompt_prefilled_end_marker() { let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap(); parser.initialize(&[MM_THINK_END_ID]).unwrap(); - let delta = parser.push("answer").unwrap(); + let delta = push_str(&mut parser, "answer"); assert_eq!(delta.reasoning, None); - assert_eq!(delta.content.as_deref(), Some("answer")); + assert_eq!(content_str(&delta), Some("answer")); +} + +#[test] +fn delimited_zero_width_only_piece_is_attributed_to_current_state() { + let tokenizer = Arc::new(fake_tokenizer()); + let mut parser = + DelimitedReasoningParser::new(tokenizer, "", "", false).unwrap(); + parser.initialize(&[THINK_START_ID]); + + // A filtered special token produces a zero-width attribution with no text; + // it must survive as a reasoning piece rather than being dropped. + let delta = parser.push(DecodedText { + text: String::new(), + attributions: [TokenAttribution { + token_id: 42, + anchor: TokenAnchor::ZeroWidth { byte_offset: 0 }, + }] + .into_iter() + .collect(), + }); + assert_eq!(reasoning_str(&delta), Some("")); + assert_eq!( + delta + .reasoning + .as_ref() + .unwrap() + .attributions + .iter() + .map(|attr| attr.token_id) + .collect::>(), + [42] + ); +} + +#[test] +fn delimited_marker_tokens_are_dropped_from_attributions() { + let tokenizer = Arc::new(fake_tokenizer()); + let mut parser = + DelimitedReasoningParser::new(tokenizer, "", "", false).unwrap(); + + let mut collected = CollectedAttributions::default(); + for chunk in attributed_chunks(&[ + Some(""), + Some("reason"), + Some(""), + Some("answer"), + ]) { + collected.record(parser.push(chunk)); + } + collected.record(parser.finish()); + + assert_eq!(collected.reasoning_text, "reason"); + assert_eq!(collected.content_text, "answer"); + // The marker tokens (1 and 3) are dropped with their spans; only the + // reasoning and content tokens remain attributed. + assert_eq!(collected.reasoning_ids, [2]); + assert_eq!(collected.content_ids, [4]); +} + +#[test] +fn reasoning_parsers_conserve_token_attributions() { + type ParserCase = ( + fn(DynTokenizer) -> Result>, + u32, + &'static str, + ); + + let parsers: Vec = vec![ + (Qwen3ReasoningParser::create, THINK_START_ID, "qwen3"), + ( + DeepSeekR1ReasoningParser::create, + THINK_START_ID, + "deepseek_r1", + ), + (KimiReasoningParser::create, MINIMAX_THINK_START_ID, "kimi"), + ( + CohereCmdReasoningParser::create, + START_THINKING_ID, + "cohere_cmd", + ), + ( + SeedOssReasoningParser::create, + SEED_THINK_START_ID, + "seed_oss", + ), + ( + MiniMaxM3ReasoningParser::create, + MM_THINK_START_ID, + "minimax_m3", + ), + (Step3p5ReasoningParser::create, THINK_START_ID, "step3p5"), + ]; + + for (create, think_start_id, name) in parsers { + let mut parser = create(Arc::new(fake_tokenizer())).unwrap(); + // Prompt prefills the start marker, opening reasoning before the stream. + parser.initialize(&[think_start_id]).unwrap(); + + // Script: reasoning text (1), a zero-width token inside reasoning (2), + // the end marker (3), then visible content (4). The zero-width token + // is attributed to reasoning, matching Python's current-state rule; + // the marker token is dropped with its span. + let end_marker = match name { + "kimi" => "◁/think▷", + "cohere_cmd" => "<|END_THINKING|>", + "seed_oss" => "", + "minimax_m3" => "", + _ => "", + }; + let chunks = attributed_chunks(&[Some("reason"), None, Some(end_marker), Some("answer")]); + + let collected = collect_attributed(parser.as_mut(), chunks); + + assert_eq!(collected.reasoning_text, "reason", "{name}"); + assert_eq!(collected.content_text, "answer", "{name}"); + assert_eq!(collected.reasoning_ids, [1, 2], "{name}"); + assert_eq!(collected.content_ids, [4], "{name}"); + } +} + +/// Build one chunk per piece, each carrying a single token attributed to it. +/// +/// A `Some(text)` piece produces one token visibly anchored at the piece start; +/// a `None` piece produces one zero-width token. Token IDs are 1-based +/// positions in the script. +pub(crate) fn attributed_chunks(pieces: &[Option<&str>]) -> Vec { + pieces + .iter() + .enumerate() + .map(|(index, piece)| { + let token_id = index as u32 + 1; + match piece { + Some(text) => DecodedText { + text: (*text).to_string(), + attributions: [TokenAttribution { + token_id, + anchor: TokenAnchor::Visible { byte_offset: 0 }, + }] + .into_iter() + .collect(), + }, + None => DecodedText { + text: String::new(), + attributions: [TokenAttribution { + token_id, + anchor: TokenAnchor::ZeroWidth { byte_offset: 0 }, + }] + .into_iter() + .collect(), + }, + } + }) + .collect() +} + +/// Token attributions and text collected from a stream of reasoning deltas. +#[derive(Default)] +pub(crate) struct CollectedAttributions { + pub reasoning_ids: Vec, + pub content_ids: Vec, + pub reasoning_text: String, + pub content_text: String, +} + +impl CollectedAttributions { + pub fn record(&mut self, delta: ReasoningDelta) { + if let Some(reasoning) = delta.reasoning { + self.reasoning_ids + .extend(reasoning.attributions.iter().map(|attr| attr.token_id)); + self.reasoning_text.push_str(&reasoning.text); + } + if let Some(content) = delta.content { + self.content_ids.extend(content.attributions.iter().map(|attr| attr.token_id)); + self.content_text.push_str(&content.text); + } + } +} + +/// Push attributed chunks through a reasoning parser and collect what is +/// attributed to reasoning and content pieces. +pub(crate) fn collect_attributed( + parser: &mut dyn ReasoningParser, + chunks: Vec, +) -> CollectedAttributions { + let mut collected = CollectedAttributions::default(); + for chunk in chunks { + collected.record(parser.push(chunk).unwrap()); + } + collected.record(parser.finish().unwrap()); + collected } diff --git a/rust/src/parser/src/unified/combined.rs b/rust/src/parser/src/unified/combined.rs index afecc31665..a8963517b3 100644 --- a/rust/src/parser/src/unified/combined.rs +++ b/rust/src/parser/src/unified/combined.rs @@ -3,7 +3,7 @@ //! Adapter that combines reasoning and tool parsers. -use vllm_tokenizer::DynTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer}; use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput}; use crate::reasoning::ReasoningParser; @@ -87,9 +87,9 @@ impl UnifiedParser for CombinedParser { self.tool.as_ref().and_then(|parser| parser.tool_call_id(tool_index)) } - fn parse_into(&mut self, delta: &str, output: &mut UnifiedParserOutput) -> Result<()> { + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> { let Some(reasoning) = self.reasoning.as_mut() else { - return self.parse_tool(delta, output); + return self.parse_tool(&delta.text, output); }; let reasoning_delta = reasoning.push(delta)?; @@ -97,7 +97,11 @@ impl UnifiedParser for CombinedParser { output.push_reasoning(reasoning); } if let Some(content) = reasoning_delta.content { - self.parse_tool(&content, output)?; + // Content attributions stop at this boundary: the tool parser trait + // consumes plain text. + if !content.text.is_empty() { + self.parse_tool(&content.text, output)?; + } } Ok(()) } @@ -109,8 +113,10 @@ impl UnifiedParser for CombinedParser { if let Some(reasoning) = reasoning_delta.reasoning { output.push_reasoning(reasoning); } - if let Some(content) = reasoning_delta.content { - self.parse_tool(&content, &mut output)?; + if let Some(content) = reasoning_delta.content + && !content.text.is_empty() + { + self.parse_tool(&content.text, &mut output)?; } } output.append(self.flush_tool()?); @@ -127,6 +133,7 @@ mod tests { use std::sync::Arc; use vllm_tokenizer::test_utils::TestTokenizer; + use vllm_tokenizer::{DecodedText, TokenAnchor, TokenAttribution}; use super::CombinedParser; use crate::reasoning::{Qwen3ReasoningParser, ReasoningDelta, ReasoningParser}; @@ -156,7 +163,7 @@ mod tests { fn collect(parser: &mut dyn UnifiedParser, chunks: &[&str]) -> UnifiedParserOutput { let mut output = UnifiedParserOutput::default(); for chunk in chunks { - parser.parse_into(chunk, &mut output).unwrap(); + parser.parse_into(DecodedText::unattributed(*chunk), &mut output).unwrap(); } output.append(parser.finish().unwrap()); output @@ -178,10 +185,10 @@ mod tests { true } - fn push(&mut self, delta: &str) -> crate::reasoning::Result { + fn push(&mut self, delta: DecodedText) -> crate::reasoning::Result { Ok(ReasoningDelta { reasoning: None, - content: Some(delta.to_string()), + content: Some(delta), }) } } @@ -259,7 +266,63 @@ mod tests { assert_eq!( output.events, vec![ - UnifiedParserEvent::Reasoning("work".to_string()), + UnifiedParserEvent::Reasoning(DecodedText::unattributed("work")), + UnifiedParserEvent::Text("answer".to_string()), + ] + ); + } + + #[test] + fn combined_parser_reasoning_events_carry_token_attributions() { + let tokenizer = Arc::new(tokenizer()); + let reasoning = Qwen3ReasoningParser::create(tokenizer).unwrap(); + let mut parser = CombinedParser::new(Some(reasoning), None); + + let chunk = |token_id: u32, text: &str| DecodedText { + text: text.to_string(), + attributions: [TokenAttribution { + token_id, + anchor: TokenAnchor::Visible { byte_offset: 0 }, + }] + .into_iter() + .collect(), + }; + + let mut output = UnifiedParserOutput::default(); + for (token_id, text) in [ + (1, ""), + (2, "reason"), + (3, ""), + (4, "answer"), + ] { + parser.parse_into(chunk(token_id, text), &mut output).unwrap(); + } + output.append(parser.finish().unwrap()); + + // The reasoning tokens keep their attributions through the combined + // parser; marker tokens (1 and 3) are dropped with their spans. + let reasoning_ids: Vec = output + .events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Reasoning(piece) => Some(piece), + _ => None, + }) + .flat_map(|piece| piece.attributions.iter().map(|attr| attr.token_id)) + .collect(); + assert_eq!(reasoning_ids, [2]); + assert_eq!( + output.events, + vec![ + UnifiedParserEvent::Reasoning(DecodedText { + text: "reason".to_string(), + attributions: [TokenAttribution { + token_id: 2, + anchor: TokenAnchor::Visible { byte_offset: 0 }, + }] + .into_iter() + .collect(), + }), UnifiedParserEvent::Text("answer".to_string()), ] ); @@ -300,7 +363,7 @@ mod tests { let mut parser = CombinedParser::new(None, Some(Box::new(PartialThenErrorToolParser))); let mut output = UnifiedParserOutput::default(); - let error = parser.parse_into("bad", &mut output).unwrap_err(); + let error = parser.parse_into(DecodedText::unattributed("bad"), &mut output).unwrap_err(); assert!(matches!(error, crate::unified::UnifiedParserError::Tool(_))); assert_eq!( diff --git a/rust/src/parser/src/unified/gemma4.rs b/rust/src/parser/src/unified/gemma4.rs index f6ec6b58ba..1adda930bf 100644 --- a/rust/src/parser/src/unified/gemma4.rs +++ b/rust/src/parser/src/unified/gemma4.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use serde_json::{Map, Number, Value}; -use vllm_tokenizer::DynTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer}; use winnow::ascii::multispace0 as ws0; use winnow::combinator::{alt, delimited, eof, opt, separated, seq, terminated}; use winnow::error::{ContextError, ErrMode, ModalResult}; @@ -29,8 +29,8 @@ type Gemma4Input<'i> = Partial<&'i str>; #[derive(Debug, Clone, PartialEq)] enum Gemma4Event { - Text { len: usize }, - Reasoning { len: usize }, + Text, + Reasoning, ReasoningStart, ReasoningEnd, ToolCallStart, @@ -69,7 +69,7 @@ enum Gemma4Mode { /// /// Arguments are emitted only after a full Gemma4 tool call is parsed. pub struct Gemma4UnifiedParser { - buffer: String, + buffer: DecodedText, mode: Gemma4Mode, emitted_tool_count: usize, tokenizer: DynTokenizer, @@ -84,7 +84,7 @@ impl Gemma4UnifiedParser { let channel_end_token_id = token_id(tokenizer.as_ref(), CHANNEL_END)?; Ok(Self { - buffer: String::new(), + buffer: DecodedText::default(), mode: Gemma4Mode::default(), emitted_tool_count: 0, channel_start_token_id, @@ -93,14 +93,16 @@ impl Gemma4UnifiedParser { }) } - fn apply_event(&mut self, event: Gemma4Event, output: &mut UnifiedParserOutput) -> Result<()> { + fn apply_event( + &mut self, + event: Gemma4Event, + piece: DecodedText, + output: &mut UnifiedParserOutput, + ) -> Result<()> { match event { - Gemma4Event::Text { len: consumed_len } => { - output.push_text(self.buffer[..consumed_len].to_string()); - } - Gemma4Event::Reasoning { len: consumed_len } => { - output.push_reasoning(self.buffer[..consumed_len].to_string()); - } + Gemma4Event::Text => output.push_text(piece.text), + Gemma4Event::Reasoning => output.push_reasoning(piece), + // Marker and tool-syntax spans are drained and dropped with their tokens. Gemma4Event::ReasoningStart => self.mode = Gemma4Mode::Reasoning, Gemma4Event::ReasoningEnd => self.mode = Gemma4Mode::Text, Gemma4Event::ToolCallStart => self.mode = Gemma4Mode::Header, @@ -144,22 +146,17 @@ impl Gemma4UnifiedParser { } fn reset(&mut self) -> String { + let buffer = self.buffer.take().text; let raw = match std::mem::replace(&mut self.mode, Gemma4Mode::Text) { - Gemma4Mode::Text => std::mem::take(&mut self.buffer), + Gemma4Mode::Text => buffer, Gemma4Mode::Reasoning => { - format!("{}{}", REASONING_START, std::mem::take(&mut self.buffer)) + format!("{REASONING_START}{buffer}") } Gemma4Mode::Header => { - format!("{}{}", TOOL_CALL_START, std::mem::take(&mut self.buffer)) + format!("{TOOL_CALL_START}{buffer}") } Gemma4Mode::ToolCall { name, .. } => { - format!( - "{}{}{}{{{}", - TOOL_CALL_START, - CALL_PREFIX, - name, - std::mem::take(&mut self.buffer) - ) + format!("{TOOL_CALL_START}{CALL_PREFIX}{name}{{{buffer}") } }; self.mode = Gemma4Mode::Text; @@ -187,16 +184,16 @@ impl UnifiedParser for Gemma4UnifiedParser { true } - fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> { - self.buffer.push_str(chunk); + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> { + self.buffer.append(delta); while let Some((event, consumed_len)) = { - parse_buffered_event(&self.buffer, |input| { + parse_buffered_event(&self.buffer.text, |input| { parse_next_gemma4_event(input, &mut self.mode) })? } { - self.apply_event(event, output)?; - self.buffer.drain(..consumed_len); + let piece = self.buffer.drain_prefix(consumed_len); + self.apply_event(event, piece, output)?; } Ok(()) @@ -206,8 +203,8 @@ impl UnifiedParser for Gemma4UnifiedParser { let mut output = UnifiedParserOutput::default(); match &self.mode { - Gemma4Mode::Text => output.push_text(std::mem::take(&mut self.buffer)), - Gemma4Mode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)), + Gemma4Mode::Text => output.push_text(self.buffer.take().text), + Gemma4Mode::Reasoning => output.push_reasoning(self.buffer.take()), Gemma4Mode::Header | Gemma4Mode::ToolCall { .. } => { return Err(parsing_failed!("incomplete Gemma4 tool call")); } @@ -306,14 +303,12 @@ fn gemma4_tool_name(input: &mut Gemma4Input<'_>) -> ModalResult { /// Parse a safe text run before the next Gemma4 marker. fn safe_text_event(input: &mut Gemma4Input<'_>) -> ModalResult { - safe_text_len_mul(input, &[REASONING_START, TOOL_CALL_START]) - .map(|len| Gemma4Event::Text { len }) + safe_text_len_mul(input, &[REASONING_START, TOOL_CALL_START]).map(|_| Gemma4Event::Text) } /// Parse a safe reasoning run before the next Gemma4 marker. fn safe_reasoning_event(input: &mut Gemma4Input<'_>) -> ModalResult { - safe_text_len_mul(input, &[CHANNEL_END, TOOL_CALL_START]) - .map(|len| Gemma4Event::Reasoning { len }) + safe_text_len_mul(input, &[CHANNEL_END, TOOL_CALL_START]).map(|_| Gemma4Event::Reasoning) } /// Parse raw Gemma4 arguments through the first end marker outside a Gemma string. @@ -514,6 +509,7 @@ mod tests { use serde_json::{Value, json}; use thiserror_ext::AsReport; use vllm_tokenizer::test_utils::TestTokenizer; + use vllm_tokenizer::{DecodedText, TokenAnchor, TokenAttribution}; use winnow::combinator::{eof, terminated}; use winnow::error::ErrMode; use winnow::prelude::*; @@ -545,7 +541,7 @@ mod tests { impl UnifiedParserTestExt for Gemma4UnifiedParser { fn parse_chunk(&mut self, chunk: &str) -> super::Result { let mut output = UnifiedParserOutput::default(); - self.parse_into(chunk, &mut output)?; + self.parse_into(DecodedText::unattributed(chunk), &mut output)?; Ok(output) } @@ -578,7 +574,7 @@ mod tests { self.events .iter() .filter_map(|event| match event { - UnifiedParserEvent::Reasoning(text) => Some(text.as_str()), + UnifiedParserEvent::Reasoning(text) => Some(text.text.as_str()), UnifiedParserEvent::Text(_) | UnifiedParserEvent::ToolCall(_) => None, }) .collect() @@ -1116,4 +1112,62 @@ mod tests { assert_eq!(raw, input); } + + #[test] + fn gemma4_reasoning_events_conserve_token_attributions() { + let mut parser = test_parser(); + + let chunk = |token_id: u32, text: &str| DecodedText { + text: text.to_string(), + attributions: [TokenAttribution { + token_id, + anchor: TokenAnchor::Visible { byte_offset: 0 }, + }] + .into_iter() + .collect(), + }; + // A zero-width token inside the reasoning channel counts as reasoning, + // matching Python's current-state rule. + let zero_width = |token_id: u32| DecodedText { + text: String::new(), + attributions: [TokenAttribution { + token_id, + anchor: TokenAnchor::ZeroWidth { byte_offset: 0 }, + }] + .into_iter() + .collect(), + }; + + let mut output = UnifiedParserOutput::default(); + for delta in [ + chunk(1, "<|channel>thought\n"), + zero_width(2), + chunk(3, "reason"), + chunk(4, ""), + chunk(5, "answer"), + chunk(6, "<|tool_call>"), + chunk(7, "call:get_weather{location:<|\"|>Paris<|\"|>}"), + chunk(8, ""), + ] { + parser.parse_into(delta, &mut output).unwrap(); + } + output.append(parser.finish().unwrap()); + + // The reasoning tokens keep their attributions; marker and tool-syntax + // tokens (1, 4, 6, 7, 8) are dropped with their spans, and text/tool + // pieces carry no attributions. + let reasoning_ids: Vec = output + .events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Reasoning(piece) => Some(piece), + _ => None, + }) + .flat_map(|piece| piece.attributions.iter().map(|attr| attr.token_id)) + .collect(); + assert_eq!(reasoning_ids, [2, 3]); + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "answer"); + assert_eq!(first_call(&output).name.as_deref(), Some("get_weather")); + } } diff --git a/rust/src/parser/src/unified/hy/hy_v3.rs b/rust/src/parser/src/unified/hy/hy_v3.rs index 069eb5de80..6aebf7f904 100644 --- a/rust/src/parser/src/unified/hy/hy_v3.rs +++ b/rust/src/parser/src/unified/hy/hy_v3.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use vllm_tokenizer::DynTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer}; use super::detect_hy_token_suffix; use crate::reasoning::HyReasoningParser; @@ -54,7 +54,7 @@ impl UnifiedParser for HyV3UnifiedParser { self.inner.tool_call_id(tool_index) } - fn parse_into(&mut self, delta: &str, output: &mut UnifiedParserOutput) -> Result<()> { + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> { self.inner.parse_into(delta, output) } @@ -72,7 +72,7 @@ mod tests { use std::sync::Arc; use serde_json::json; - use vllm_tokenizer::{Tokenizer, test_utils::TestTokenizer}; + use vllm_tokenizer::{DecodedText, Tokenizer, test_utils::TestTokenizer}; use xgrammar_structural_tag::builders::StructuralTagOptions; use xgrammar_structural_tag::{ FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag, @@ -131,14 +131,14 @@ mod tests { ]; let mut output = UnifiedParserOutput::default(); for chunk in chunks { - parser.parse_into(chunk, &mut output).unwrap(); + parser.parse_into(DecodedText::unattributed(chunk), &mut output).unwrap(); } output.append(parser.finish().unwrap()); assert_eq!( output.events, vec![ - UnifiedParserEvent::Reasoning("reasoning".to_string()), + UnifiedParserEvent::Reasoning(DecodedText::unattributed("reasoning")), UnifiedParserEvent::Text("answer".to_string()), UnifiedParserEvent::ToolCall(crate::tool::ToolCallDelta { tool_index: 0, diff --git a/rust/src/parser/src/unified/hy/hy_v4.rs b/rust/src/parser/src/unified/hy/hy_v4.rs index d3ad175bf9..58ee042f3c 100644 --- a/rust/src/parser/src/unified/hy/hy_v4.rs +++ b/rust/src/parser/src/unified/hy/hy_v4.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use vllm_tokenizer::DynTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer}; use super::detect_hy_token_suffix; use crate::reasoning::HyReasoningParser; @@ -54,7 +54,7 @@ impl UnifiedParser for HyV4UnifiedParser { self.inner.tool_call_id(tool_index) } - fn parse_into(&mut self, delta: &str, output: &mut UnifiedParserOutput) -> Result<()> { + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> { self.inner.parse_into(delta, output) } @@ -72,7 +72,7 @@ mod tests { use std::sync::Arc; use serde_json::json; - use vllm_tokenizer::{Tokenizer, test_utils::TestTokenizer}; + use vllm_tokenizer::{DecodedText, Tokenizer, test_utils::TestTokenizer}; use xgrammar_structural_tag::builders::StructuralTagOptions; use xgrammar_structural_tag::{ FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag, @@ -139,14 +139,14 @@ mod tests { ]; let mut output = UnifiedParserOutput::default(); for chunk in chunks { - parser.parse_into(chunk, &mut output).unwrap(); + parser.parse_into(DecodedText::unattributed(chunk), &mut output).unwrap(); } output.append(parser.finish().unwrap()); assert_eq!( output.events, vec![ - UnifiedParserEvent::Reasoning("reasoning".to_string()), + UnifiedParserEvent::Reasoning(DecodedText::unattributed("reasoning")), UnifiedParserEvent::Text("answer".to_string()), UnifiedParserEvent::ToolCall(crate::tool::ToolCallDelta { tool_index: 0, diff --git a/rust/src/parser/src/unified/inkling.rs b/rust/src/parser/src/unified/inkling.rs index d35fb48d9c..d1ed783671 100644 --- a/rust/src/parser/src/unified/inkling.rs +++ b/rust/src/parser/src/unified/inkling.rs @@ -8,7 +8,7 @@ use winnow::prelude::*; use winnow::stream::Partial; use winnow::token::literal; -use vllm_tokenizer::DynTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer}; use super::{Result, UnifiedParser, UnifiedParserOutput, token_id}; use crate::tool::json::{ @@ -56,15 +56,15 @@ type InklingInput<'i> = Partial<&'i str>; #[derive(Debug, Clone, PartialEq, Eq)] enum InklingEvent { - Text { len: usize }, - Reasoning { len: usize }, + Text, + Reasoning, TextStart, ReasoningStart, MessageStart, - Header { len: usize }, + Header, ToolJsonStart, ToolJsonHeader { name: String }, - ToolJsonArgs { len: usize, complete: bool }, + ToolJsonArgs { complete: bool }, BlockEnd, } @@ -73,7 +73,7 @@ enum InklingMode { #[default] Idle, MessageHeader { - pending: String, + pending: DecodedText, }, Text, Reasoning, @@ -86,7 +86,7 @@ enum InklingMode { /// Unified parser for Inkling typed content blocks. pub struct InklingUnifiedParser { - buffer: String, + buffer: DecodedText, mode: InklingMode, emitted_tool_count: usize, active_tool_index: Option, @@ -104,7 +104,7 @@ impl InklingUnifiedParser { let content_thinking_token_id = token_id(tokenizer.as_ref(), CONTENT_THINKING)?; Ok(Self { - buffer: String::new(), + buffer: DecodedText::default(), mode: InklingMode::Idle, emitted_tool_count: 0, active_tool_index: None, @@ -120,7 +120,7 @@ impl InklingUnifiedParser { for token_id in prompt_token_ids.iter().rev().copied() { if token_id == self.message_model_token_id { self.mode = InklingMode::MessageHeader { - pending: String::new(), + pending: DecodedText::default(), }; return; } @@ -138,24 +138,27 @@ impl InklingUnifiedParser { } } - fn apply_event(&mut self, event: InklingEvent, output: &mut UnifiedParserOutput) -> Result<()> { + fn apply_event( + &mut self, + event: InklingEvent, + piece: DecodedText, + output: &mut UnifiedParserOutput, + ) -> Result<()> { match event { - InklingEvent::Text { len } => output.push_text(self.buffer[..len].to_string()), - InklingEvent::Reasoning { len } => { - output.push_reasoning(self.buffer[..len].to_string()); - } + InklingEvent::Text => output.push_text(piece.text), + InklingEvent::Reasoning => output.push_reasoning(piece), InklingEvent::MessageStart => { self.mode = InklingMode::MessageHeader { - pending: String::new(), + pending: DecodedText::default(), }; } - InklingEvent::Header { len } => { + InklingEvent::Header => { let InklingMode::MessageHeader { pending } = &mut self.mode else { return Err(parsing_failed!( "Inkling header text outside a message header" )); }; - pending.push_str(&self.buffer[..len]); + pending.append(piece); } InklingEvent::TextStart => self.mode = InklingMode::Text, InklingEvent::ReasoningStart => self.mode = InklingMode::Reasoning, @@ -173,7 +176,7 @@ impl InklingUnifiedParser { arguments: String::new(), }); } - InklingEvent::ToolJsonArgs { len, complete } => { + InklingEvent::ToolJsonArgs { complete } => { let Some(tool_index) = self.active_tool_index else { return Err(parsing_failed!( "Inkling arguments without an active tool call" @@ -182,7 +185,7 @@ impl InklingUnifiedParser { output.push_call(ToolCallDelta { tool_index, name: None, - arguments: self.buffer[..len].to_string(), + arguments: piece.text, }); if complete { self.mode = InklingMode::ToolJsonClose; @@ -191,7 +194,7 @@ impl InklingUnifiedParser { InklingEvent::BlockEnd => { let mode = std::mem::take(&mut self.mode); if let InklingMode::MessageHeader { pending } = mode { - output.push_text(pending); + output.push_text(pending.text); } self.active_tool_index = None; } @@ -201,12 +204,12 @@ impl InklingUnifiedParser { fn reset(&mut self) -> String { let mut uncommitted = match std::mem::take(&mut self.mode) { - InklingMode::MessageHeader { pending } => pending, + InklingMode::MessageHeader { pending } => pending.text, _ => String::new(), }; self.active_tool_index = None; self.emitted_tool_count = 0; - uncommitted.push_str(&std::mem::take(&mut self.buffer)); + uncommitted.push_str(&self.buffer.take().text); uncommitted } } @@ -231,14 +234,14 @@ impl UnifiedParser for InklingUnifiedParser { true } - fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> { - self.buffer.push_str(chunk); + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> { + self.buffer.append(delta); - while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| { + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer.text, |input| { parse_next_inkling_event(input, &mut self.mode) })? { - self.apply_event(event, output)?; - self.buffer.drain(..consumed_len); + let piece = self.buffer.drain_prefix(consumed_len); + self.apply_event(event, piece, output)?; } Ok(()) @@ -248,15 +251,12 @@ impl UnifiedParser for InklingUnifiedParser { let mut output = UnifiedParserOutput::default(); match &mut self.mode { - InklingMode::Idle | InklingMode::Text => { - output.push_text(std::mem::take(&mut self.buffer)) - } + InklingMode::Idle | InklingMode::Text => output.push_text(self.buffer.take().text), InklingMode::MessageHeader { pending } => { - pending.push_str(&self.buffer); - self.buffer.clear(); - output.push_text(std::mem::take(pending)); + pending.append(self.buffer.take()); + output.push_text(std::mem::take(pending).text); } - InklingMode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)), + InklingMode::Reasoning => output.push_reasoning(self.buffer.take()), InklingMode::ToolJsonHeader | InklingMode::ToolJsonArgs { .. } | InklingMode::ToolJsonClose => { @@ -367,22 +367,22 @@ fn block_end_event(input: &mut InklingInput<'_>) -> ModalResult { /// Parse safe text while waiting for the next Inkling marker. fn safe_idle_text_event(input: &mut InklingInput<'_>) -> ModalResult { - safe_text_len_mul(input, IDLE_MARKERS).map(|len| InklingEvent::Text { len }) + safe_text_len_mul(input, IDLE_MARKERS).map(|_| InklingEvent::Text) } /// Parse safe header text before the next Inkling marker. fn safe_header_event(input: &mut InklingInput<'_>) -> ModalResult { - safe_text_len_mul(input, IDLE_MARKERS).map(|len| InklingEvent::Header { len }) + safe_text_len_mul(input, IDLE_MARKERS).map(|_| InklingEvent::Header) } /// Parse safe text before the end of a Inkling text block. fn safe_text_event(input: &mut InklingInput<'_>) -> ModalResult { - safe_text_len_mul(input, BLOCK_END_MARKERS).map(|len| InklingEvent::Text { len }) + safe_text_len_mul(input, BLOCK_END_MARKERS).map(|_| InklingEvent::Text) } /// Parse safe reasoning before the end of a Inkling reasoning block. fn safe_reasoning_event(input: &mut InklingInput<'_>) -> ModalResult { - safe_text_len_mul(input, BLOCK_END_MARKERS).map(|len| InklingEvent::Reasoning { len }) + safe_text_len_mul(input, BLOCK_END_MARKERS).map(|_| InklingEvent::Reasoning) } /// Parse a Inkling JSON tool-call header. @@ -400,9 +400,8 @@ fn parse_tool_json_args_event( input: &mut JsonToolInput<'_>, json_scan: &mut JsonObjectScanState, ) -> ModalResult { - let len = take_json_object(input, json_scan)?; + take_json_object(input, json_scan)?; Ok(InklingEvent::ToolJsonArgs { - len, complete: json_scan.complete(), }) } @@ -423,11 +422,11 @@ fn parse_tool_json_close_event(input: &mut InklingInput<'_>) -> ModalResult UnifiedParserTestExt for T { fn parse_chunk(&mut self, chunk: &str) -> super::Result { let mut output = UnifiedParserOutput::default(); - self.parse_into(chunk, &mut output)?; + self.parse_into(DecodedText::unattributed(chunk), &mut output)?; Ok(output) } @@ -529,7 +528,7 @@ mod tests { self.events .iter() .filter_map(|event| match event { - UnifiedParserEvent::Reasoning(text) => Some(text.as_str()), + UnifiedParserEvent::Reasoning(text) => Some(text.text.as_str()), _ => None, }) .collect() @@ -587,6 +586,50 @@ mod tests { assert!(output.calls().is_empty()); } + #[test] + fn inkling_reasoning_events_conserve_token_attributions() { + let chunk = |token_id: u32, text: &str| DecodedText { + text: text.to_string(), + attributions: [TokenAttribution { + token_id, + anchor: TokenAnchor::Visible { byte_offset: 0 }, + }] + .into_iter() + .collect(), + }; + + let mut parser = test_parser(); + let mut output = UnifiedParserOutput::default(); + for (token_id, text) in [ + (1, CONTENT_THINKING), + (2, "reason"), + (3, END_MESSAGE), + (4, MESSAGE_MODEL), + (5, CONTENT_TEXT), + (6, "answer"), + (7, END_MESSAGE), + ] { + parser.parse_into(chunk(token_id, text), &mut output).unwrap(); + } + output.append(parser.finish().unwrap()); + + // The reasoning tokens keep their attributions; marker tokens + // (1, 3, 4, 5, 7) are dropped with their spans, and text pieces + // carry no attributions. + let reasoning_ids: Vec = output + .events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Reasoning(piece) => Some(piece), + _ => None, + }) + .flat_map(|piece| piece.attributions.iter().map(|attr| attr.token_id)) + .collect(); + assert_eq!(reasoning_ids, [2]); + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "answer"); + } + #[test] fn inkling_streaming_holds_split_markers() { let output = collect_stream(&[ diff --git a/rust/src/parser/src/unified/kimi_k3.rs b/rust/src/parser/src/unified/kimi_k3.rs index 932bab656a..2f8278dc62 100644 --- a/rust/src/parser/src/unified/kimi_k3.rs +++ b/rust/src/parser/src/unified/kimi_k3.rs @@ -41,7 +41,7 @@ mod structural_tag; pub use structural_tag::KimiK3StructuralTagBuilder; use serde_json::{Map, Value}; -use vllm_tokenizer::DynTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer}; use winnow::ascii::{multispace0 as ws0, multispace1 as ws1}; use winnow::combinator::{alt, delimited, eof, preceded, repeat, seq, terminated}; use winnow::error::{ContextError, ErrMode, ModalResult, StrContext}; @@ -95,12 +95,8 @@ type KimiK3Input<'i> = Partial<&'i str>; #[derive(Debug, Clone, PartialEq, Eq)] enum KimiK3Event { - Text { - len: usize, - }, - Reasoning { - len: usize, - }, + Text, + Reasoning, /// Structural noise consumed without emitting anything. Skip, ThinkOpen, @@ -147,7 +143,7 @@ enum KimiK3Mode { /// Unified parser for Kimi K3 XTML think / response / tools channels. pub struct KimiK3UnifiedParser { - buffer: String, + buffer: DecodedText, mode: KimiK3Mode, /// Number of calls emitted in the current response. emitted_call_count: usize, @@ -163,7 +159,7 @@ impl KimiK3UnifiedParser { let sep_token_id = token_id(tokenizer.as_ref(), SEP)?; Ok(Self { - buffer: String::new(), + buffer: DecodedText::default(), mode: KimiK3Mode::default(), emitted_call_count: 0, tokenizer, @@ -203,12 +199,16 @@ impl KimiK3UnifiedParser { }; } - fn apply_event(&mut self, event: KimiK3Event, output: &mut UnifiedParserOutput) -> Result<()> { + fn apply_event( + &mut self, + event: KimiK3Event, + piece: DecodedText, + output: &mut UnifiedParserOutput, + ) -> Result<()> { match event { - KimiK3Event::Text { len } => output.push_text(self.buffer[..len].to_string()), - KimiK3Event::Reasoning { len } => { - output.push_reasoning(self.buffer[..len].to_string()); - } + KimiK3Event::Text => output.push_text(piece.text), + KimiK3Event::Reasoning => output.push_reasoning(piece), + // Marker and noise spans are drained and dropped with their tokens. KimiK3Event::Skip => {} KimiK3Event::ThinkOpen => self.mode = KimiK3Mode::Reasoning, KimiK3Event::ThinkClose => self.mode = KimiK3Mode::Idle, @@ -252,7 +252,7 @@ impl KimiK3UnifiedParser { fn reset_state(&mut self) -> String { self.mode = KimiK3Mode::Idle; self.emitted_call_count = 0; - std::mem::take(&mut self.buffer) + self.buffer.take().text } } @@ -279,14 +279,14 @@ impl UnifiedParser for KimiK3UnifiedParser { Some(&KIMI_K3_STRUCTURAL_TAG_BUILDER) } - fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> { - self.buffer.push_str(chunk); + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> { + self.buffer.append(delta); - while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| { + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer.text, |input| { parse_next_kimi_k3_event(input, &mut self.mode) })? { - self.apply_event(event, output)?; - self.buffer.drain(..consumed_len); + let piece = self.buffer.drain_prefix(consumed_len); + self.apply_event(event, piece, output)?; } Ok(()) @@ -297,13 +297,15 @@ impl UnifiedParser for KimiK3UnifiedParser { match &self.mode { KimiK3Mode::Idle | KimiK3Mode::Response => { - output.push_text(std::mem::take(&mut self.buffer)); + output.push_text(self.buffer.take().text); + } + KimiK3Mode::Reasoning => { + output.push_reasoning(self.buffer.take()); } - KimiK3Mode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)), KimiK3Mode::Epilogue | KimiK3Mode::Done => self.buffer.clear(), // A tools channel truncated between complete calls loses only its // closing markers; keep the calls already emitted. - KimiK3Mode::Tools if self.buffer.is_empty() => {} + KimiK3Mode::Tools if self.buffer.text.is_empty() => {} KimiK3Mode::Tools | KimiK3Mode::Call { .. } => { return Err(parsing_failed!("incomplete Kimi K3 tool call")); } @@ -406,17 +408,17 @@ fn parse_done_event(input: &mut KimiK3Input<'_>) -> ModalResult { /// Parse safe text while waiting for the next channel marker. fn safe_idle_text_event(input: &mut KimiK3Input<'_>) -> ModalResult { - safe_text_len_mul(input, IDLE_MARKERS).map(|len| KimiK3Event::Text { len }) + safe_text_len_mul(input, IDLE_MARKERS).map(|_| KimiK3Event::Text) } /// Parse safe reasoning before the think close marker. fn safe_reasoning_event(input: &mut KimiK3Input<'_>) -> ModalResult { - safe_text_len_mul(input, REASONING_MARKERS).map(|len| KimiK3Event::Reasoning { len }) + safe_text_len_mul(input, REASONING_MARKERS).map(|_| KimiK3Event::Reasoning) } /// Parse safe response text before the next channel marker. fn safe_response_text_event(input: &mut KimiK3Input<'_>) -> ModalResult { - safe_text_len_mul(input, RESPONSE_MARKERS).map(|len| KimiK3Event::Text { len }) + safe_text_len_mul(input, RESPONSE_MARKERS).map(|_| KimiK3Event::Text) } /// Skip non-content noise after the response channel closed. @@ -571,6 +573,7 @@ mod tests { use serde_json::{Value, json}; use thiserror_ext::AsReport; + use vllm_tokenizer::DecodedText; use vllm_tokenizer::Tokenizer as _; use vllm_tokenizer::test_utils::TestTokenizer; @@ -604,7 +607,7 @@ mod tests { impl UnifiedParserTestExt for KimiK3UnifiedParser { fn parse_chunk(&mut self, chunk: &str) -> super::Result { let mut output = UnifiedParserOutput::default(); - self.parse_into(chunk, &mut output)?; + self.parse_into(DecodedText::unattributed(chunk), &mut output)?; Ok(output) } @@ -636,7 +639,7 @@ mod tests { self.events .iter() .filter_map(|event| match event { - UnifiedParserEvent::Reasoning(text) => Some(text.as_str()), + UnifiedParserEvent::Reasoning(text) => Some(text.text.as_str()), _ => None, }) .collect() diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index 9f5547bcf2..17c3dd1b24 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -16,7 +16,7 @@ pub use inkling::InklingUnifiedParser; pub use kimi_k3::{KimiK3StructuralTagBuilder, KimiK3UnifiedParser}; use thiserror::Error; use thiserror_ext::Macro; -use vllm_tokenizer::DynTokenizer; +use vllm_tokenizer::{DecodedText, DynTokenizer}; use crate::reasoning::ReasoningError; use crate::tool::{ @@ -32,7 +32,11 @@ pub enum UnifiedParserEvent { /// Normal assistant-visible text. Text(String), /// Reasoning text hidden from the normal content stream. - Reasoning(String), + /// + /// Carries the attributions of the generated tokens that produced it, so + /// downstream consumers can count reasoning tokens exactly. Marker spans + /// are dropped by the parsers, keeping marker tokens out of any count. + Reasoning(DecodedText), /// A tool-call update extracted from visible assistant text. ToolCall(ToolCallDelta), } @@ -58,15 +62,18 @@ impl UnifiedParserOutput { } /// Append one reasoning text event if `delta` is non-empty. - pub fn push_reasoning(&mut self, delta: impl AsRef + Into) { - if delta.as_ref().is_empty() { + /// + /// A piece with empty text but non-empty attributions (zero-width tokens + /// only) is still kept: the tokens count as reasoning. + pub fn push_reasoning(&mut self, delta: DecodedText) { + if delta.is_empty() { return; } if let Some(UnifiedParserEvent::Reasoning(last_text)) = self.events.last_mut() { - last_text.push_str(delta.as_ref()); + last_text.append(delta); return; } - self.events.push(UnifiedParserEvent::Reasoning(delta.into())); + self.events.push(UnifiedParserEvent::Reasoning(delta)); } /// Append one tool-call event. @@ -98,6 +105,8 @@ impl UnifiedParserOutput { #[cfg(test)] mod tests { + use vllm_tokenizer::{DecodedText, TokenAnchor, TokenAttribution}; + use super::{UnifiedParserEvent, UnifiedParserOutput}; use crate::tool::ToolCallDelta; @@ -107,8 +116,8 @@ mod tests { output.push_text("hello"); output.push_text(" "); output.push_text("world"); - output.push_reasoning("think"); - output.push_reasoning("ing"); + output.push_reasoning(DecodedText::unattributed("think")); + output.push_reasoning(DecodedText::unattributed("ing")); output.push_call(ToolCallDelta { tool_index: 0, name: Some("lookup".to_string()), @@ -120,7 +129,7 @@ mod tests { output.events, vec![ UnifiedParserEvent::Text("hello world".to_string()), - UnifiedParserEvent::Reasoning("thinking".to_string()), + UnifiedParserEvent::Reasoning(DecodedText::unattributed("thinking")), UnifiedParserEvent::ToolCall(ToolCallDelta { tool_index: 0, name: Some("lookup".to_string()), @@ -139,11 +148,11 @@ mod tests { let mut other = UnifiedParserOutput::default(); other.push_text(" "); other.push_text("world"); - other.push_reasoning("think"); + other.push_reasoning(DecodedText::unattributed("think")); output.append(other); let mut after_reasoning = UnifiedParserOutput::default(); - after_reasoning.push_reasoning("ing"); + after_reasoning.push_reasoning(DecodedText::unattributed("ing")); after_reasoning.push_text("!"); output.append(after_reasoning); @@ -151,11 +160,67 @@ mod tests { output.events, vec![ UnifiedParserEvent::Text("hello world".to_string()), - UnifiedParserEvent::Reasoning("thinking".to_string()), + UnifiedParserEvent::Reasoning(DecodedText::unattributed("thinking")), UnifiedParserEvent::Text("!".to_string()), ] ); } + + #[test] + fn unified_parser_output_reasoning_coalescing_rebases_attributions() { + let mut output = UnifiedParserOutput::default(); + output.push_reasoning(DecodedText { + text: "think".to_string(), + attributions: [TokenAttribution { + token_id: 1, + anchor: TokenAnchor::Visible { byte_offset: 0 }, + }] + .into_iter() + .collect(), + }); + output.push_reasoning(DecodedText { + text: "ing".to_string(), + attributions: [TokenAttribution { + token_id: 2, + anchor: TokenAnchor::Visible { byte_offset: 0 }, + }] + .into_iter() + .collect(), + }); + // A zero-width-token-only piece is kept: the tokens count as reasoning. + output.push_reasoning(DecodedText { + text: String::new(), + attributions: [TokenAttribution { + token_id: 3, + anchor: TokenAnchor::ZeroWidth { byte_offset: 0 }, + }] + .into_iter() + .collect(), + }); + + assert_eq!( + output.events, + vec![UnifiedParserEvent::Reasoning(DecodedText { + text: "thinking".to_string(), + attributions: [ + TokenAttribution { + token_id: 1, + anchor: TokenAnchor::Visible { byte_offset: 0 }, + }, + TokenAttribution { + token_id: 2, + anchor: TokenAnchor::Visible { byte_offset: 5 }, + }, + TokenAttribution { + token_id: 3, + anchor: TokenAnchor::ZeroWidth { byte_offset: 8 }, + }, + ] + .into_iter() + .collect(), + })] + ); + } } /// Incremental parser that extracts reasoning and tool-call events from assistant output. @@ -186,7 +251,7 @@ pub trait UnifiedParser: Send { } /// Feed one decoded text delta into the parser, appending committed output into `output`. - fn parse_into(&mut self, delta: &str, output: &mut UnifiedParserOutput) -> Result<()>; + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()>; /// Flush any buffered parser state at end of stream. fn finish(&mut self) -> Result { diff --git a/rust/src/tokenizer/src/incremental/attribution.rs b/rust/src/tokenizer/src/incremental/attribution.rs index 358d83f421..8425c7f41d 100644 --- a/rust/src/tokenizer/src/incremental/attribution.rs +++ b/rust/src/tokenizer/src/incremental/attribution.rs @@ -67,9 +67,63 @@ impl DecodedText { } } + /// Return true when both the text and the attributions are empty. + pub fn is_empty(&self) -> bool { + self.text.is_empty() && self.attributions.is_empty() + } + + /// Take the decoded text and attributions, leaving `self` empty. + pub fn take(&mut self) -> Self { + take(self) + } + + /// Clear the text and attributions, keeping the allocations. + pub fn clear(&mut self) { + self.text.clear(); + self.attributions.clear(); + } + + /// Split off and return the prefix `[0, len)`; the suffix stays in `self`, + /// with its token anchors rebased by `-len`. + /// + /// Anchors follow the first-byte ownership rule: a `Visible` token moves to + /// the prefix when its first byte lies inside the prefix, so a token + /// straddling the split belongs to the span holding its first byte. A + /// `ZeroWidth` token moves to the prefix when its offset is `<= len`, + /// matching the `take_ready` cutoff rule and letting + /// `drain_prefix(text.len())` take trailing zero-width tokens. Relative + /// order across anchor kinds is preserved. + /// + /// Panics if `len` is not a char boundary, via `String::drain`. + #[must_use] + pub fn drain_prefix(&mut self, len: usize) -> DecodedText { + let text: String = self.text.drain(..len).collect(); + let len = offset_as_u32(len); + + let mut prefix = DecodedText { + text, + attributions: SmallVec::new(), + }; + for attribution in take(&mut self.attributions) { + let in_prefix = match attribution.anchor { + TokenAnchor::Visible { byte_offset } => byte_offset < len, + TokenAnchor::ZeroWidth { byte_offset } => byte_offset <= len, + }; + if in_prefix { + prefix.attributions.push(attribution); + } else { + self.attributions.push(TokenAttribution { + token_id: attribution.token_id, + anchor: attribution.anchor.offset_by(-i64::from(len)), + }); + } + } + prefix + } + /// Append another decoded fragment, rebasing its token anchors. pub fn append(&mut self, other: Self) { - if self.text.is_empty() && self.attributions.is_empty() { + if self.is_empty() { *self = other; return; } diff --git a/rust/src/tokenizer/src/incremental/attribution/tests.rs b/rust/src/tokenizer/src/incremental/attribution/tests.rs index 26b4acfefc..cd39a1c03c 100644 --- a/rust/src/tokenizer/src/incremental/attribution/tests.rs +++ b/rust/src/tokenizer/src/incremental/attribution/tests.rs @@ -169,6 +169,145 @@ fn decoded_text_append_preserves_zero_width_only_receiver() { ); } +#[test] +fn decoded_text_drain_prefix_moves_anchors_before_boundary() { + let mut dt = decoded("ab.cd", &[1, 2, 3], &[visible(0), visible(1), visible(3)]); + + let prefix = dt.drain_prefix(2); + + assert_eq!(prefix, decoded("ab", &[1, 2], &[visible(0), visible(1)])); + assert_eq!(dt, decoded(".cd", &[3], &[visible(1)])); +} + +#[test] +fn decoded_text_drain_prefix_keeps_straddling_token_in_prefix() { + // One token decodes to "ab"; a split inside it keeps the token in the + // prefix, the span holding its first byte. + let mut dt = decoded("ab", &[1], &[visible(0)]); + + let prefix = dt.drain_prefix(1); + + assert_eq!(prefix, decoded("a", &[1], &[visible(0)])); + assert_eq!(dt, decoded("b", &[], &[])); +} + +#[test] +fn decoded_text_drain_prefix_visible_anchor_at_boundary_goes_to_suffix() { + let mut dt = decoded("ab", &[1, 2], &[visible(0), visible(1)]); + + let prefix = dt.drain_prefix(1); + + assert_eq!(prefix, decoded("a", &[1], &[visible(0)])); + assert_eq!(dt, decoded("b", &[2], &[visible(0)])); +} + +#[test] +fn decoded_text_drain_prefix_zero_width_at_boundary_goes_to_prefix() { + let mut dt = decoded("ab", &[1, 2], &[zero_width(1), visible(1)]); + + let prefix = dt.drain_prefix(1); + + assert_eq!(prefix, decoded("a", &[1], &[zero_width(1)])); + assert_eq!(dt, decoded("b", &[2], &[visible(0)])); +} + +#[test] +fn decoded_text_drain_prefix_shared_anchor_group_moves_together() { + // A byte-fallback group jointly decoding one character shares one anchor; + // the whole group stays on whichever side holds that first byte. + let mut dt = decoded( + "a你", + &[1, 2, 3, 4], + &[visible(0), visible(1), visible(1), visible(1)], + ); + + let prefix = dt.drain_prefix(1); + + assert_eq!(prefix, decoded("a", &[1], &[visible(0)])); + assert_eq!( + dt, + decoded("你", &[2, 3, 4], &[visible(0), visible(0), visible(0)]) + ); + + let prefix = dt.drain_prefix(3); + + assert_eq!( + prefix, + decoded("你", &[2, 3, 4], &[visible(0), visible(0), visible(0)]) + ); + assert_eq!(dt, decoded("", &[], &[])); +} + +#[test] +fn decoded_text_drain_prefix_zero_drains_leading_zero_width() { + let mut dt = decoded("ab", &[1, 2], &[zero_width(0), visible(0)]); + + let prefix = dt.drain_prefix(0); + + assert_eq!(prefix, decoded("", &[1], &[zero_width(0)])); + assert_eq!(dt, decoded("ab", &[2], &[visible(0)])); +} + +#[test] +fn decoded_text_drain_prefix_full_length_takes_trailing_zero_width() { + let mut dt = decoded("ab", &[1, 2, 3], &[visible(0), visible(1), zero_width(2)]); + + let prefix = dt.drain_prefix(2); + + assert_eq!( + prefix, + decoded("ab", &[1, 2, 3], &[visible(0), visible(1), zero_width(2)]) + ); + assert_eq!(dt, decoded("", &[], &[])); +} + +#[test] +fn decoded_text_drain_prefix_preserves_order_across_anchor_kinds() { + let mut dt = decoded( + "ab", + &[1, 2, 3, 4], + &[visible(0), zero_width(1), visible(1), zero_width(2)], + ); + + let prefix = dt.drain_prefix(1); + + assert_eq!(prefix, decoded("a", &[1, 2], &[visible(0), zero_width(1)])); + assert_eq!(dt, decoded("b", &[3, 4], &[visible(0), zero_width(1)])); +} + +#[test] +fn decoded_text_drain_prefix_then_append_reconstructs_original() { + // ASCII, a byte-fallback group, and mid and trailing zero-width tokens. + let original = decoded( + "a你b", + &[1, 2, 3, 4, 5, 6, 7], + &[ + visible(0), + visible(1), + visible(1), + visible(1), + zero_width(4), + visible(4), + zero_width(5), + ], + ); + + for len in [0, 1, 4, 5] { + let mut rest = original.clone(); + let mut prefix = rest.drain_prefix(len); + prefix.append(rest); + assert_eq!(prefix, original, "split at byte {len}"); + } +} + +#[test] +#[should_panic(expected = "is_char_boundary")] +fn decoded_text_drain_prefix_panics_off_char_boundary() { + let mut dt = decoded("a你", &[1, 2], &[visible(0), visible(1)]); + + let _ = dt.drain_prefix(2); +} + struct AttributionCase<'a> { tokenizer: &'a dyn Tokenizer, prompt_token_ids: Vec,