From 36229e1337ae0e44e60075a11ebfde210c11d706 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Thu, 1 Oct 2026 12:21:50 +0800 Subject: [PATCH] [Rust Frontend] Token-aware marker parsing for Kimi K3 (#58358) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/Cargo.lock | 2 + rust/Cargo.toml | 1 + rust/src/chat/src/parser/unified.rs | 3 +- rust/src/parser/Cargo.toml | 2 + rust/src/parser/src/unified/kimi_k3.rs | 394 +++++++++++++-- rust/src/parser/src/unified/mod.rs | 9 + rust/src/parser/src/utils.rs | 197 +++++--- rust/src/parser/src/utils/marker.rs | 660 +++++++++++++++++++++++++ 8 files changed, 1135 insertions(+), 133 deletions(-) create mode 100644 rust/src/parser/src/utils/marker.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index d6143f941b..bd5a63e81d 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5964,8 +5964,10 @@ dependencies = [ name = "vllm-parser" version = "0.1.0" dependencies = [ + "auto_impl", "criterion", "easy-ext", + "educe", "expect-test", "futures", "openai-protocol", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 0002796d40..a2e6d0c1ed 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -31,6 +31,7 @@ async-openai = { version = "0.33.1", default-features = false, features = ["nati async-trait = "0.1.89" asynk-strim-attr = "0.1.0" auto_enums = { version = "0.8.9", features = ["tokio1"] } +auto_impl = "1.3.0" axum = "0.8.8" base64 = "0.22.1" bytemuck = { version = "1.25.0", features = ["extern_crate_alloc"] } diff --git a/rust/src/chat/src/parser/unified.rs b/rust/src/chat/src/parser/unified.rs index 5ab830150a..52cabaf162 100644 --- a/rust/src/chat/src/parser/unified.rs +++ b/rust/src/chat/src/parser/unified.rs @@ -141,7 +141,8 @@ mod tests { let tokenizer = TestTokenizer::new() .with_regular_token("<|open|>", 1001) .with_regular_token("<|close|>", 1002) - .with_regular_token("<|sep|>", 1003); + .with_regular_token("<|sep|>", 1003) + .with_regular_token("<|end_of_msg|>", 1004); assert!(factory.contains(names::KIMI_K3)); assert_eq!( diff --git a/rust/src/parser/Cargo.toml b/rust/src/parser/Cargo.toml index 25546cb1df..3086de03e9 100644 --- a/rust/src/parser/Cargo.toml +++ b/rust/src/parser/Cargo.toml @@ -8,7 +8,9 @@ license.workspace = true test-util = [] [dependencies] +auto_impl.workspace = true easy-ext.workspace = true +educe.workspace = true serde.workspace = true serde_json.workspace = true serde_with.workspace = true diff --git a/rust/src/parser/src/unified/kimi_k3.rs b/rust/src/parser/src/unified/kimi_k3.rs index 08ad91193c..d6e32463c1 100644 --- a/rust/src/parser/src/unified/kimi_k3.rs +++ b/rust/src/parser/src/unified/kimi_k3.rs @@ -26,6 +26,14 @@ //! starts *inside* that channel without re-emitting the open tag; //! [`UnifiedParser::initialize`] detects this from the prompt token IDs. //! +//! Channel markers are structural only when their `<|open|>` / `<|close|>` / +//! `<|sep|>` bytes come from the dedicated special tokens. The model can spell +//! the same text with ordinary BPE tokens (for example while quoting a +//! configuration snippet in its reasoning); such text stays content in the +//! current channel. The parser therefore expects token-attributed input from +//! the incremental detokenizer: fed unattributed text, it treats every marker +//! as content. Channel names such as `think` are ordinary text either way. +//! //! Argument decoding mirrors the renderer's type tagging (inverse encoding): //! `type="string"` values pass the raw text through, other types are //! JSON-decoded, and a raw `json` block is passed through unmodified. Attribute @@ -34,29 +42,33 @@ //! Known limitation (shared with the Python parser): string argument and //! response bodies are emitted raw, so a value that literally contains //! `<|close|>argument<|sep|>` or `<|close|>response<|sep|>` is -//! indistinguishable from a real closing marker. +//! indistinguishable from a real closing marker. With token attribution this +//! now only applies to `argument` and `json` blocks inside a call body, which +//! are still parsed from text. mod structural_tag; pub use structural_tag::KimiK3StructuralTagBuilder; use serde_json::{Map, Value}; -use vllm_tokenizer::{DecodedText, DynTokenizer}; +use vllm_tokenizer::{DecodedText, DynTokenizer, Tokenizer}; 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}; use winnow::prelude::*; -use winnow::stream::Partial; use winnow::token::{literal, rest, take_till, take_until, take_while}; use self::structural_tag::KIMI_K3_STRUCTURAL_TAG_BUILDER; -use super::{Result, UnifiedParser, UnifiedParserOutput, token_id}; +use super::{Result, UnifiedParser, UnifiedParserOutput, special_token}; use crate::output_grammar::{ self, BuiltOutputGrammar, OutputGrammarContext, visible_format_from_builder, }; use crate::tool::{Tool, ToolCallDelta}; use crate::unified::parsing_failed; -use crate::utils::{MarkerScanState, parse_buffered_event, safe_text_len_mul, take_until_marker}; +use crate::utils::{ + Attributed, AttributedExt as _, Marker, MarkerScanState, SpecialToken, + parse_buffered_event_attributed, safe_text_len_mul, take_until_marker, +}; const OPEN: &str = "<|open|>"; const CLOSE: &str = "<|close|>"; @@ -70,31 +82,82 @@ const RESPONSE_CLOSE: &str = "<|close|>response<|sep|>"; const TOOLS_OPEN: &str = "<|open|>tools<|sep|>"; const TOOLS_CLOSE: &str = "<|close|>tools<|sep|>"; const MESSAGE_CLOSE: &str = "<|close|>message<|sep|>"; -const CALL_OPEN: &str = "<|open|>call"; const CALL_CLOSE: &str = "<|close|>call<|sep|>"; const ARG_OPEN: &str = "<|open|>argument"; const ARG_CLOSE: &str = "<|close|>argument<|sep|>"; const JSON_OPEN: &str = "<|open|>json"; const JSON_CLOSE: &str = "<|close|>json<|sep|>"; -const IDLE_MARKERS: &[&str] = &[ - THINK_OPEN, - RESPONSE_OPEN, - TOOLS_OPEN, - MESSAGE_CLOSE, - END_OF_MSG, -]; -const REASONING_MARKERS: &[&str] = &[THINK_CLOSE, END_OF_MSG]; -const RESPONSE_MARKERS: &[&str] = &[RESPONSE_CLOSE, TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]; -const EPILOGUE_MARKERS: &[&str] = &[TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]; -const TOOLS_MARKERS: &[&str] = &[CALL_OPEN, TOOLS_CLOSE, MESSAGE_CLOSE, END_OF_MSG]; - /// Channel tags are a couple of text tokens; longer `<|open|>…<|sep|>` spans in /// the prompt tail (attribute-bearing message opens, message bodies) never name /// a generation channel. const MAX_PREFILL_TAG_TOKENS: usize = 8; -type KimiK3Input<'i> = Partial<&'i str>; +type KimiK3Input<'i> = Attributed<'i, Markers>; + +/// The K3 structural special tokens, resolved once from the active tokenizer. +struct SpecialTokens { + open: SpecialToken, + close: SpecialToken, + sep: SpecialToken, + end_of_msg: SpecialToken, +} + +impl SpecialTokens { + fn new(tokenizer: &dyn Tokenizer) -> Result { + Ok(Self { + open: special_token(tokenizer, OPEN)?, + close: special_token(tokenizer, CLOSE)?, + sep: special_token(tokenizer, SEP)?, + end_of_msg: special_token(tokenizer, END_OF_MSG)?, + }) + } +} + +/// The K3 channel markers, with their special-token segments guarded by the +/// active tokenizer's IDs. +/// +/// One definition serves both the `alt` branches that consume a marker and the +/// safe-text scans that stop in front of one, so the two can never disagree on +/// whether a spelling is structural. +#[derive(Debug)] +struct Markers { + think_open: Marker, + think_close: Marker, + response_open: Marker, + response_close: Marker, + tools_open: Marker, + tools_close: Marker, + message_close: Marker, + end_of_msg: Marker, + /// `<|open|>call`; the tag attributes and the terminating `<|sep|>` follow. + call_open: Marker, + call_close: Marker, + sep: Marker, +} + +impl Markers { + fn new(tokens: &SpecialTokens) -> Self { + let open = + |tag: &str| Marker::special(&tokens.open).then_text(tag).then_special(&tokens.sep); + let close = + |tag: &str| Marker::special(&tokens.close).then_text(tag).then_special(&tokens.sep); + + Self { + think_open: open("think"), + think_close: close("think"), + response_open: open("response"), + response_close: close("response"), + tools_open: open("tools"), + tools_close: close("tools"), + message_close: close("message"), + end_of_msg: Marker::special(&tokens.end_of_msg), + call_open: Marker::special(&tokens.open).then_text("call"), + call_close: close("call"), + sep: Marker::special(&tokens.sep), + } + } +} #[derive(Debug, Clone, PartialEq, Eq)] enum KimiK3Event { @@ -148,26 +211,25 @@ enum KimiK3Mode { pub struct KimiK3UnifiedParser { buffer: DecodedText, mode: KimiK3Mode, + markers: Markers, /// Number of calls emitted in the current response. emitted_call_count: usize, tokenizer: DynTokenizer, - open_token_id: u32, - sep_token_id: u32, + special_tokens: SpecialTokens, } impl KimiK3UnifiedParser { - /// Create a Kimi K3 parser. + /// Create a Kimi K3 parser expecting token-attributed input. pub fn new(_tools: &[Tool], tokenizer: DynTokenizer) -> Result { - let open_token_id = token_id(tokenizer.as_ref(), OPEN)?; - let sep_token_id = token_id(tokenizer.as_ref(), SEP)?; + let special_tokens = SpecialTokens::new(tokenizer.as_ref())?; Ok(Self { buffer: DecodedText::default(), mode: KimiK3Mode::default(), + markers: Markers::new(&special_tokens), emitted_call_count: 0, tokenizer, - open_token_id, - sep_token_id, + special_tokens, }) } @@ -180,11 +242,14 @@ impl KimiK3UnifiedParser { fn initialize_mode(&mut self, prompt_token_ids: &[u32]) { self.mode = KimiK3Mode::Idle; - let Some(sep_pos) = prompt_token_ids.iter().rposition(|&id| id == self.sep_token_id) else { + let Some(sep_pos) = + prompt_token_ids.iter().rposition(|&id| id == self.special_tokens.sep.id) + else { return; }; - let Some(open_pos) = - prompt_token_ids[..sep_pos].iter().rposition(|&id| id == self.open_token_id) + let Some(open_pos) = prompt_token_ids[..sep_pos] + .iter() + .rposition(|&id| id == self.special_tokens.open.id) else { return; }; @@ -291,9 +356,11 @@ impl UnifiedParser for KimiK3UnifiedParser { 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.text, |input| { - parse_next_kimi_k3_event(input, &mut self.mode) - })? { + while let Some((event, consumed_len)) = + parse_buffered_event_attributed(&self.buffer, &self.markers, |input| { + parse_next_kimi_k3_event(input, &mut self.mode) + })? + { let piece = self.buffer.drain_prefix(consumed_len); self.apply_event(event, piece, output)?; } @@ -347,10 +414,11 @@ fn parse_next_kimi_k3_event( /// Parse an event while waiting for the next channel open. fn parse_idle_event(input: &mut KimiK3Input<'_>) -> ModalResult { + let markers = input.markers(); alt(( - literal(THINK_OPEN).value(KimiK3Event::ThinkOpen), - literal(RESPONSE_OPEN).value(KimiK3Event::ResponseOpen), - literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen), + markers.think_open.value(KimiK3Event::ThinkOpen), + markers.response_open.value(KimiK3Event::ResponseOpen), + markers.tools_open.value(KimiK3Event::ToolsOpen), message_end_event, safe_idle_text_event, )) @@ -359,11 +427,12 @@ fn parse_idle_event(input: &mut KimiK3Input<'_>) -> ModalResult { /// Parse an event inside the `think` channel. fn parse_reasoning_event(input: &mut KimiK3Input<'_>) -> ModalResult { + let markers = input.markers(); alt(( - literal(THINK_CLOSE).value(KimiK3Event::ThinkClose), + markers.think_close.value(KimiK3Event::ThinkClose), // `<|end_of_msg|>` can reach the parser under `ignore_eos` or // `include_stop_str_in_output`; never leak it into reasoning. - literal(END_OF_MSG).value(KimiK3Event::MessageEnd), + markers.end_of_msg.value(KimiK3Event::MessageEnd), safe_reasoning_event, )) .parse_next(input) @@ -371,10 +440,11 @@ fn parse_reasoning_event(input: &mut KimiK3Input<'_>) -> ModalResult) -> ModalResult { + let markers = input.markers(); alt(( - literal(RESPONSE_CLOSE).value(KimiK3Event::ResponseClose), + markers.response_close.value(KimiK3Event::ResponseClose), // The response body also implicitly ends at a `tools` channel. - literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen), + markers.tools_open.value(KimiK3Event::ToolsOpen), message_end_event, safe_response_text_event, )) @@ -383,8 +453,9 @@ fn parse_response_event(input: &mut KimiK3Input<'_>) -> ModalResult /// Parse an event after the response channel closed. fn parse_epilogue_event(input: &mut KimiK3Input<'_>) -> ModalResult { + let markers = input.markers(); alt(( - literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen), + markers.tools_open.value(KimiK3Event::ToolsOpen), message_end_event, skip_epilogue_noise_event, )) @@ -393,9 +464,10 @@ fn parse_epilogue_event(input: &mut KimiK3Input<'_>) -> ModalResult /// Parse an event inside the `tools` channel, between `call` blocks. fn parse_tools_event(input: &mut KimiK3Input<'_>) -> ModalResult { + let markers = input.markers(); alt(( call_open_event, - literal(TOOLS_CLOSE).value(KimiK3Event::ToolsClose), + markers.tools_close.value(KimiK3Event::ToolsClose), // Defensive: an unterminated tools channel still ends with the message. message_end_event, skip_tools_noise_event, @@ -405,7 +477,8 @@ fn parse_tools_event(input: &mut KimiK3Input<'_>) -> ModalResult { /// Parse a message close or end-of-message marker. fn message_end_event(input: &mut KimiK3Input<'_>) -> ModalResult { - alt((literal(MESSAGE_CLOSE), literal(END_OF_MSG))) + let markers = input.markers(); + alt((&markers.message_close, &markers.end_of_msg)) .value(KimiK3Event::MessageEnd) .parse_next(input) } @@ -417,35 +490,82 @@ 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(|_| KimiK3Event::Text) + let markers = input.markers(); + safe_text_len_mul( + input, + &[ + &markers.think_open, + &markers.response_open, + &markers.tools_open, + &markers.message_close, + &markers.end_of_msg, + ], + ) + .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(|_| KimiK3Event::Reasoning) + let markers = input.markers(); + safe_text_len_mul(input, &[&markers.think_close, &markers.end_of_msg]) + .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(|_| KimiK3Event::Text) + let markers = input.markers(); + safe_text_len_mul( + input, + &[ + &markers.response_close, + &markers.tools_open, + &markers.message_close, + &markers.end_of_msg, + ], + ) + .map(|_| KimiK3Event::Text) } /// Skip non-content noise after the response channel closed. fn skip_epilogue_noise_event(input: &mut KimiK3Input<'_>) -> ModalResult { - safe_text_len_mul(input, EPILOGUE_MARKERS).map(|_| KimiK3Event::Skip) + let markers = input.markers(); + safe_text_len_mul( + input, + &[ + &markers.tools_open, + &markers.message_close, + &markers.end_of_msg, + ], + ) + .map(|_| KimiK3Event::Skip) } /// Skip non-content noise between `call` blocks. fn skip_tools_noise_event(input: &mut KimiK3Input<'_>) -> ModalResult { - safe_text_len_mul(input, TOOLS_MARKERS).map(|_| KimiK3Event::Skip) + let markers = input.markers(); + safe_text_len_mul( + input, + &[ + &markers.call_open, + &markers.tools_close, + &markers.message_close, + &markers.end_of_msg, + ], + ) + .map(|_| KimiK3Event::Skip) } /// Parse a `call` open tag into its tool name and one-based index. +/// +/// The attributes end at the structural `<|sep|>`; an ordinary spelling of +/// `<|sep|>` inside an attribute value stays attribute text. fn call_open_event(input: &mut KimiK3Input<'_>) -> ModalResult { + let markers = input.markers(); + let mut scan = MarkerScanState::default(); let (attrs,) = seq!( - _: literal(CALL_OPEN), - take_until(0.., SEP), - _: literal(SEP), + _: &markers.call_open, + take_until_marker(&markers.sep, &mut scan), + _: &markers.sep, ) .parse_next(input)?; let attrs = parse_tag_attrs(attrs)?; @@ -464,9 +584,10 @@ fn call_body_event( input: &mut KimiK3Input<'_>, scan: &mut MarkerScanState, ) -> ModalResult { + let markers = input.markers(); let (body,) = seq!( - take_until_marker(CALL_CLOSE, scan), - _: literal(CALL_CLOSE), + take_until_marker(&markers.call_close, scan), + _: &markers.call_close, ) .parse_next(input)?; let arguments = parse_call_arguments(body)?; @@ -614,9 +735,11 @@ mod tests { } impl UnifiedParserTestExt for KimiK3UnifiedParser { + /// Parse `chunk` as the detokenizer would deliver it: markers spelled in + /// full are structural special tokens, everything else ordinary text. fn parse_chunk(&mut self, chunk: &str) -> super::Result { let mut output = UnifiedParserOutput::default(); - self.parse_into(DecodedText::unattributed(chunk), &mut output)?; + self.parse_into(decode_ids(&structural(chunk)), &mut output)?; Ok(output) } @@ -665,14 +788,41 @@ mod tests { } } + /// The parser as constructed in production: markers must be spelled by + /// their special tokens. fn test_parser() -> KimiK3UnifiedParser { KimiK3UnifiedParser::new(&[], Arc::new(tokenizer())).unwrap() } + /// Run token IDs through the real incremental decoder to obtain attributed + /// input, exactly as the frontend does. + fn decode_ids(ids: &[u32]) -> DecodedText { + let tokenizer = tokenizer(); + let mut decoder = tokenizer.create_decode_stream(&[], false, 0); + for &id in ids { + decoder.push_token(id).unwrap(); + } + decoder.flush(None).unwrap().1 + } + + /// Encode with the special tokens, so markers are structural. + fn structural(text: &str) -> Vec { + tokenizer().encode(text, false).unwrap() + } + + /// Encode as ordinary bytes only, so marker spellings are content. + fn ordinary(text: &str) -> Vec { + tokenizer().encode_ordinary(text).unwrap() + } + + /// Stream the structural encoding of the concatenated `chunks`, split at the + /// chunk boundaries, so a marker split across chunks keeps its token anchor. fn collect_stream(parser: &mut KimiK3UnifiedParser, chunks: &[&str]) -> UnifiedParserOutput { + let mut rest = decode_ids(&structural(&chunks.concat())); let mut output = UnifiedParserOutput::default(); for chunk in chunks { - output.append(parser.parse_chunk(chunk).unwrap()); + let piece = rest.drain_prefix(chunk.len()); + parser.parse_into(piece, &mut output).unwrap(); } output.append(parser.finish().unwrap()); output @@ -1136,4 +1286,140 @@ mod tests { assert_eq!(raw, "<|close|>resp"); } + + // --- Ordinary marker spellings ------------------------------------------ + + fn start_in_reasoning(parser: &mut KimiK3UnifiedParser) { + parser.initialize(&structural(THINK_OPEN)).unwrap(); + } + + fn parse_decoded( + parser: &mut KimiK3UnifiedParser, + decoded: DecodedText, + ) -> UnifiedParserOutput { + let mut output = UnifiedParserOutput::default(); + parser.parse_into(decoded, &mut output).unwrap(); + output.append(parser.finish().unwrap()); + output + } + + #[test] + fn kimi_k3_markers_spell_the_documented_strings() { + let special_tokens = super::SpecialTokens::new(&tokenizer()).unwrap(); + let markers = super::Markers::new(&special_tokens); + + assert_eq!(markers.think_open.as_str(), THINK_OPEN); + assert_eq!(markers.think_close.as_str(), THINK_CLOSE); + assert_eq!(markers.response_open.as_str(), RESPONSE_OPEN); + assert_eq!(markers.response_close.as_str(), RESPONSE_CLOSE); + assert_eq!(markers.tools_open.as_str(), TOOLS_OPEN); + assert_eq!(markers.tools_close.as_str(), TOOLS_CLOSE); + assert_eq!(markers.message_close.as_str(), super::MESSAGE_CLOSE); + assert_eq!(markers.call_close.as_str(), super::CALL_CLOSE); + assert_eq!(markers.end_of_msg.as_str(), END_OF_MSG); + } + + #[test] + fn kimi_k3_ordinary_close_marker_text_stays_in_reasoning() { + // The production incident: the model spells `<|close|>think<|sep|>` with + // ordinary tokens while reasoning, then really closes the channel later. + let mut ids = ordinary(THINK_CLOSE); + ids.extend(structural(&format!( + " after{THINK_CLOSE}{RESPONSE_OPEN}answer{RESPONSE_CLOSE}<|close|>message{SEP}" + ))); + let decoded = decode_ids(&ids); + assert!( + decoded.text.starts_with(THINK_CLOSE), + "same bytes as the structural marker" + ); + + let mut parser = test_parser(); + start_in_reasoning(&mut parser); + let output = parse_decoded(&mut parser, decoded); + + assert_eq!(output.reasoning_text(), format!("{THINK_CLOSE} after")); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn kimi_k3_ordinary_open_marker_text_stays_in_reasoning() { + let mut ids = structural("config: "); + ids.extend(ordinary(THINK_OPEN)); + ids.extend(structural(&format!( + "{THINK_CLOSE}{RESPONSE_OPEN}ok{RESPONSE_CLOSE}" + ))); + + let mut parser = test_parser(); + start_in_reasoning(&mut parser); + let output = parse_decoded(&mut parser, decode_ids(&ids)); + + assert_eq!(output.reasoning_text(), format!("config: {THINK_OPEN}")); + assert_eq!(output.normal_text(), "ok"); + } + + #[test] + fn kimi_k3_attributed_streaming_splits_at_every_char_boundary() { + // A lookalike close inside reasoning, then genuine markers and a call. + let mut ids = structural(THINK_OPEN); + ids.extend(ordinary(THINK_CLOSE)); + ids.extend(structural(&format!( + " step{THINK_CLOSE}{RESPONSE_OPEN}the answer{RESPONSE_CLOSE}{TOOLS_OPEN}{}{TOOLS_CLOSE}<|close|>message{SEP}", + call("tool=\"calc\" index=\"1\"", &arg("x", "number", "42")), + ))); + let full = decode_ids(&ids); + + for split in (0..=full.text.len()).filter(|&index| full.text.is_char_boundary(index)) { + let mut rest = full.clone(); + let head = rest.drain_prefix(split); + let mut parser = test_parser(); + let mut output = UnifiedParserOutput::default(); + parser.parse_into(head, &mut output).unwrap(); + parser.parse_into(rest, &mut output).unwrap(); + output.append(parser.finish().unwrap()); + + assert_eq!( + output.reasoning_text(), + format!("{THINK_CLOSE} step"), + "split {split}" + ); + assert_eq!(output.normal_text(), "the answer", "split {split}"); + assert_eq!( + first_call(&output).arguments, + r#"{"x":42}"#, + "split {split}" + ); + } + } + + #[test] + fn kimi_k3_incomplete_utf8_before_structural_close_still_closes() { + // A stray incomplete byte right before the close marker: the decoder + // anchors the special token at its own spelling, so the boundary holds. + let mut ids = vec![0xe4]; + ids.extend(structural(&format!( + "{THINK_CLOSE}{RESPONSE_OPEN}ok{RESPONSE_CLOSE}" + ))); + + let mut parser = test_parser(); + start_in_reasoning(&mut parser); + let output = parse_decoded(&mut parser, decode_ids(&ids)); + + assert_eq!(output.reasoning_text(), "\u{FFFD}"); + assert_eq!(output.normal_text(), "ok"); + } + + #[test] + fn kimi_k3_ordinary_sep_text_inside_call_attributes_stays_attribute_text() { + let mut ids = structural(&format!("{TOOLS_OPEN}{OPEN}call tool=\"a")); + ids.extend(ordinary(SEP)); + ids.extend(structural(&format!( + "b\" index=\"1\"{SEP}{}<|close|>call{SEP}{TOOLS_CLOSE}", + arg("x", "number", "1") + ))); + + let output = parse_decoded(&mut test_parser(), decode_ids(&ids)); + + assert_eq!(first_call(&output).name, Some(format!("a{SEP}b"))); + assert_eq!(first_call(&output).arguments, r#"{"x":1}"#); + } } diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index e65f2de087..c621075a02 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -21,6 +21,7 @@ use vllm_tokenizer::{DecodedText, DynTokenizer}; use crate::output_grammar::{self, BuiltOutputGrammar, OutputGrammarContext}; use crate::reasoning::ReasoningError; use crate::tool::{Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput}; +use crate::utils::SpecialToken; /// Result alias for unified parser operations. pub type Result = std::result::Result; @@ -288,3 +289,11 @@ fn token_id(tokenizer: &dyn vllm_tokenizer::Tokenizer, token: &str) -> Result Result { + Ok(SpecialToken { + id: token_id(tokenizer, token)?, + text: token.to_string(), + }) +} diff --git a/rust/src/parser/src/utils.rs b/rust/src/parser/src/utils.rs index e8ef70a861..1324e297fb 100644 --- a/rust/src/parser/src/utils.rs +++ b/rust/src/parser/src/utils.rs @@ -3,14 +3,23 @@ //! Shared helpers for streaming parsers. +use std::fmt::Debug; + +use vllm_tokenizer::DecodedText; use winnow::Parser; use winnow::error::{ContextError, ErrMode, ModalResult, Needed, StrContext, StrContextValue}; -use winnow::stream::{FindSlice, Offset, Partial, Stream}; +use winnow::stream::{Partial, Stream}; use crate::tool::{Result, ToolParserError}; +pub(crate) mod marker; pub(crate) mod recursion; +pub use marker::{ + Attributed, AttributedExt, Marker, MarkerLike, MarkerStream, SpecialToken, + attributed_with_markers, +}; + /// Return the byte length of the longest proper prefix of `token` that is also /// a suffix of `buffer`. /// @@ -48,48 +57,81 @@ pub fn partial_prefix_len(buffer: &str, token: &str) -> usize { 0 } +/// Where the next marker may be, in the coordinate system of +/// [`MarkerStream::offset`]. +enum Scan { + /// A complete marker starts at this offset. + Found(usize), + /// A marker may start at this offset but needs more input to decide; text + /// from here on must be held back. + Pending(usize), + /// Every candidate before this offset (the end of the input) was rejected. + Clear(usize), +} + +/// Scan for the earliest of `markers` at or after `from`. +/// +/// Candidates come from the stream ([`MarkerStream::next_candidate`]): text +/// occurrences for plain markers, visible anchors for markers guarded by a +/// special token. Each candidate is decided by probing the marker's own parser +/// on a cloned cursor, so the scanner and the consuming `alt` branch can never +/// disagree; a rejected candidate (an ordinary lookalike) is skipped and the +/// scan continues. +fn scan_markers<'i, I: MarkerStream<'i>, M: MarkerLike>( + input: &I, + markers: &[M], + from: usize, +) -> ModalResult { + let start = input.offset(); + let mut from = from.max(start); + while let Some(at) = input.next_candidate(markers, from) { + for marker in markers { + let mut probe = input.clone(); + probe.next_slice(at - start); + match marker.as_marker().parse_next(&mut probe) { + Ok(_) => return Ok(Scan::Found(at)), + Err(ErrMode::Backtrack(_)) => {} + Err(ErrMode::Incomplete(_)) => return Ok(Scan::Pending(at)), + Err(error @ ErrMode::Cut(_)) => return Err(error), + } + } + from = at + 1; + } + Ok(Scan::Clear(start + input.eof_offset())) +} + /// Parse a safe text run before the next marker. /// This is the single-marker variant of [`safe_text_len_mul`]. /// /// Returns the text length in bytes, and advances the input. -pub fn safe_text_len(input: &mut Partial<&str>, marker: &str) -> ModalResult { - let text = **input; - if text.is_empty() { - return incomplete(); - } - - if let Some(start_idx) = text.find(marker) { - input.next_slice(start_idx); - return Ok(start_idx); - } - - let keep_len = partial_prefix_len(text, marker); - let emit_len = text.len().saturating_sub(keep_len); - if emit_len == 0 { - return incomplete(); - } - - input.next_slice(emit_len); - Ok(emit_len) +pub fn safe_text_len<'i, I: MarkerStream<'i>, M: MarkerLike>( + input: &mut I, + marker: M, +) -> ModalResult { + safe_text_len_mul(input, &[marker]) } /// Parse a safe text run before the earliest next marker. /// This is the multi-marker variant of [`safe_text_len`]. /// +/// Text that may still grow into a marker is held back: a partial spelling at +/// the end of the input for plain markers, and only a genuine special token's +/// incomplete marker for guarded ones. +/// /// Returns the text length in bytes, and advances the input. -pub fn safe_text_len_mul(input: &mut Partial<&str>, markers: &[&str]) -> ModalResult { - let text = **input; - if text.is_empty() { +pub fn safe_text_len_mul<'i, I: MarkerStream<'i>, M: MarkerLike>( + input: &mut I, + markers: &[M], +) -> ModalResult { + if input.eof_offset() == 0 { return incomplete(); } - if let Some(start_idx) = find_slice_mul(text, markers) { - input.next_slice(start_idx); - return Ok(start_idx); - } - - let keep_len = markers.iter().map(|marker| partial_prefix_len(text, marker)).max().unwrap_or(0); - let emit_len = text.len().saturating_sub(keep_len); + let start = input.offset(); + let stop = match scan_markers(input, markers, start)? { + Scan::Found(at) | Scan::Pending(at) | Scan::Clear(at) => at, + }; + let emit_len = stop - start; if emit_len == 0 { return incomplete(); } @@ -98,23 +140,11 @@ pub fn safe_text_len_mul(input: &mut Partial<&str>, markers: &[&str]) -> ModalRe Ok(emit_len) } -#[inline(always)] -fn find_slice_mul(text: &str, markers: &[&str]) -> Option { - let range = match markers { - // Use the fast specialized `winnow::stream::FindSlice` impl for 1-3 markers. - [first] => text.find_slice(*first), - [first, second] => text.find_slice((*first, *second)), - [first, second, third] => text.find_slice((*first, *second, *third)), - // Fall back to a linear scan for 4+ markers. - _ => return markers.iter().filter_map(|marker| text.find(marker)).min(), - }; - range.map(|range| range.start) -} - /// Streaming scan state for a buffered marker search [`take_until_marker`], /// so that we don't have to rescan the whole buffered prefix when resuming. #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct MarkerScanState { + /// Offset relative to the cursor at call time from which the scan resumes. scan_start: usize, } @@ -128,55 +158,47 @@ impl MarkerScanState { /// /// This is the streaming-buffered variant of `winnow::token::take_until(0.., /// marker)`: it returns the slice before `marker` and leaves `marker` for the -/// caller to consume. On incomplete input, it stores the earliest byte offset -/// that can still match `marker` and returns `Incomplete` without consuming -/// input, so the next parse can avoid rescanning the whole buffered prefix. +/// caller to consume. On incomplete input, it stores the earliest offset that +/// can still start `marker` and returns `Incomplete` without consuming input, +/// so the next parse can avoid rescanning the whole buffered prefix. /// /// Use this for outer parser states that keep the full buffered input across /// chunks while waiting for a closing marker. Plain `take_until` is still a /// better fit for one-shot parsers over a complete body, and for `1..` cases /// where an empty slice before the marker should be rejected. -pub fn take_until_marker<'i, 'a>( - marker: &'a str, +pub fn take_until_marker<'i, 'a, I: MarkerStream<'i>, M: MarkerLike + 'a>( + marker: M, state: &'a mut MarkerScanState, -) -> impl Parser, &'i str, ErrMode> + 'a { - move |input: &mut Partial<&'i str>| take_until_marker_(input, marker, state) +) -> impl Parser> + 'a { + move |input: &mut I| take_until_marker_(input, &marker, state) } -fn take_until_marker_<'i>( - input: &mut Partial<&'i str>, - marker: &str, +fn take_until_marker_<'i, I: MarkerStream<'i>, M: MarkerLike>( + input: &mut I, + marker: &M, state: &mut MarkerScanState, ) -> ModalResult<&'i str> { - debug_assert!(!marker.is_empty()); + debug_assert!(!marker.as_marker().as_str().is_empty()); - let text = **input; - if text.is_empty() { + if input.eof_offset() == 0 { return incomplete(); } - // Normal updates store a char boundary; this keeps stale or misused state from panicking. - let scan_start = floor_char_boundary(text, state.scan_start); - - if let Some(offset) = text[scan_start..].find(marker) { - let marker_start = scan_start + offset; - let body = &text[..marker_start]; - input.next_slice(marker_start); - state.reset(); - return Ok(body); + let start = input.offset(); + match scan_markers( + input, + std::slice::from_ref(marker), + start + state.scan_start, + )? { + Scan::Found(at) => { + state.reset(); + Ok(input.next_slice(at - start)) + } + Scan::Pending(at) | Scan::Clear(at) => { + state.scan_start = at - start; + incomplete() + } } - - let keep_len = partial_prefix_len(text, marker); - state.scan_start = text.len() - keep_len; - incomplete() -} - -fn floor_char_boundary(text: &str, index: usize) -> usize { - let mut index = index.min(text.len()); - while !text.is_char_boundary(index) { - index -= 1; - } - index } /// Streaming lexical state for a top-level JSON object. @@ -394,7 +416,26 @@ pub fn parse_buffered_event( buffer: &str, parse: impl FnOnce(&mut Partial<&str>) -> ModalResult, ) -> Result> { - let mut input = Partial::new(buffer); + parse_buffered_stream_event(Partial::new(buffer), parse) +} + +/// Parse one event from a buffered [`DecodedText`], for token-aware parsers. +/// +/// Same contract as [`parse_buffered_event`]; the input also carries the +/// buffer's token anchors and the parser's `markers` ([`Attributed`]). +pub fn parse_buffered_event_attributed( + buffer: &DecodedText, + markers: &M, + parse: impl FnOnce(&mut Attributed<'_, M>) -> ModalResult, +) -> Result> { + parse_buffered_stream_event(attributed_with_markers(buffer, markers), parse) +} + +fn parse_buffered_stream_event<'i, I: MarkerStream<'i>, E>( + mut input: I, + parse: impl FnOnce(&mut I) -> ModalResult, +) -> Result> { + let buffer = input.remaining(); let checkpoint = input.checkpoint(); let event = match parse(&mut input) { Ok(event) => event, diff --git a/rust/src/parser/src/utils/marker.rs b/rust/src/parser/src/utils/marker.rs new file mode 100644 index 0000000000..cc36f8e6bf --- /dev/null +++ b/rust/src/parser/src/utils/marker.rs @@ -0,0 +1,660 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Token-aware marker matching for streaming parsers. +//! +//! Model protocols delimit channels with dedicated special tokens such as +//! `<|close|>` and `<|sep|>`. Ordinary BPE tokens can decode to the very same +//! spelling, so matching a marker by text alone mistakes model-written content +//! for structure. The helpers here keep text as the parse axis and consult the +//! generated-token attribution carried by [`DecodedText`] only where a marker's +//! bytes must have come from a special token. +//! +//! - [`Attributed`] is a winnow input over one `DecodedText` buffer: the text +//! plus its buffer-local token anchors, composed from winnow's own stream +//! wrappers so every existing combinator keeps working. +//! - [`Marker`] is a fixed marker string whose special-token segments are +//! *guarded*: they match only when the expected special token produced them. +//! A marker without guards matches by spelling alone, which is exactly what a +//! plain `&str` marker is; [`MarkerRef`] is the borrowed view shared by both. +//! One marker definition is both the parser that consumes the marker and the +//! definition the safe-text scanner in the parent module stops in front of, +//! so the two never disagree. +//! - [`MarkerStream`] abstracts over `Partial<&str>` (text-only parsers) and +//! [`Attributed`], letting one generic scanner serve both. + +use std::fmt::Debug; +use std::ops::Range; + +use auto_impl::auto_impl; +use educe::Educe; +use vllm_tokenizer::{DecodedText, TokenAnchor, TokenAttribution}; +use winnow::Parser; +use winnow::error::{ContextError, ErrMode, ModalResult}; +use winnow::stream::{ + Compare, FindSlice, LocatingSlice, Location, Partial, Stateful, Stream, StreamIsPartial, +}; +use winnow::token::literal; + +use super::partial_prefix_len; + +/// Streaming winnow input over one [`DecodedText`] buffer. +/// +/// `LocatingSlice` reports the cursor's byte offset from the buffer start, which +/// is the coordinate system of the anchors; `Partial` keeps the incomplete-input +/// semantics streaming parsers rely on; `Stateful` carries the immutable +/// [`AttributedState`]. winnow forwards `Compare`, `FindSlice`, `Location`, and +/// `UpdateSlice` through all three, so `literal`, `alt`, `seq!`, `take_until`, +/// and `rest` work unchanged. +pub type Attributed<'i, M = ()> = Stateful>, AttributedState<'i, M>>; + +/// Immutable state of an [`Attributed`] input. +#[derive(Debug, Educe)] +#[educe(Clone, Copy)] +pub struct AttributedState<'i, M = ()> { + /// Token anchors of the buffer, sorted by byte offset. + anchors: &'i [TokenAttribution], + /// The parser's markers; see [`AttributedExt::markers`]. + markers: &'i M, +} + +#[easy_ext::ext(AttributedExt)] +pub impl<'i, M> Attributed<'i, M> { + /// The parser's markers, reachable from every grammar function through the + /// input instead of an extra parameter. + fn markers(&self) -> &'i M { + self.state.markers + } +} + +/// Create an [`Attributed`] input over `buffer`, carrying `markers`. +pub fn attributed_with_markers<'i, M>( + buffer: &'i DecodedText, + markers: &'i M, +) -> Attributed<'i, M> { + Stateful { + input: Partial::new(LocatingSlice::new(buffer.text.as_str())), + state: AttributedState { + anchors: &buffer.attributions, + markers, + }, + } +} + +/// A `&str`-sliced streaming input that can say where markers may begin. +/// +/// Implemented for `Partial<&str>` (text-only parsers, no anchors) and +/// [`Attributed`]. +pub trait MarkerStream<'i>: + Stream + StreamIsPartial + Clone + for<'m> Compare<&'m str> +{ + /// Remaining input from the cursor. + fn remaining(&self) -> &'i str; + + /// Byte offset of the cursor in the coordinate system of + /// [`Self::anchors`] and of resumable scan state. + /// + /// `Partial<&str>` tracks no origin and reports 0, so its offsets are + /// relative to the cursor at the time of the call; [`Attributed`] reports + /// the offset from the buffer start. + fn offset(&self) -> usize; + + /// Token anchors of the buffer, sorted by byte offset; empty for text-only + /// input. + fn anchors(&self) -> &'i [TokenAttribution]; + + /// Earliest offset at or after `from` where one of `markers` may begin. + /// + /// A guarded marker can only begin at a visible anchor of its first special + /// token, so text-only input never yields one. Any other marker begins at a + /// complete occurrence of its spelling, or at a proper prefix of it that + /// ends the input. + fn next_candidate(&self, markers: &[M], from: usize) -> Option { + let start = self.offset(); + let from = from.max(start); + let markers = || markers.iter().map(MarkerLike::as_marker); + + let anchored = markers().any(|marker| marker.first_token_id().is_some()).then(|| { + visible_in(self.anchors(), from..usize::MAX).find_map(|(at, token_id)| { + markers().any(|marker| marker.first_token_id() == Some(token_id)).then_some(at) + }) + }); + let by_text = text_candidate( + self.remaining(), + markers() + .filter(|marker| marker.first_token_id().is_none()) + .map(MarkerRef::as_str), + from - start, + ) + .map(|at| start + at); + + match (anchored.flatten(), by_text) { + (Some(anchored), Some(by_text)) => Some(anchored.min(by_text)), + (anchored, by_text) => anchored.or(by_text), + } + } +} + +impl<'i> MarkerStream<'i> for Partial<&'i str> { + fn remaining(&self) -> &'i str { + **self + } + + fn offset(&self) -> usize { + 0 + } + + fn anchors(&self) -> &'i [TokenAttribution] { + &[] + } +} + +impl<'i, M: Debug> MarkerStream<'i> for Attributed<'i, M> { + fn remaining(&self) -> &'i str { + ****self + } + + fn offset(&self) -> usize { + self.current_token_start() + } + + fn anchors(&self) -> &'i [TokenAttribution] { + self.state.anchors + } +} + +fn anchor_offset(anchor: TokenAnchor) -> usize { + match anchor { + TokenAnchor::Visible { byte_offset } | TokenAnchor::ZeroWidth { byte_offset } => { + byte_offset as usize + } + } +} + +/// Visible anchors within `range`, as `(byte_offset, token_id)`, in generation +/// order. +fn visible_in( + anchors: &[TokenAttribution], + range: Range, +) -> impl Iterator + '_ { + let first = + anchors.partition_point(|attribution| anchor_offset(attribution.anchor) < range.start); + anchors[first..] + .iter() + .take_while(move |attribution| anchor_offset(attribution.anchor) < range.end) + .filter_map(|attribution| match attribution.anchor { + TokenAnchor::Visible { byte_offset } => { + Some((byte_offset as usize, attribution.token_id)) + } + TokenAnchor::ZeroWidth { .. } => None, + }) +} + +/// Return whether `range` is exactly the visible span of the single token +/// `token_id`: one visible anchor at `range.start` with that ID and no other +/// visible anchor inside. Zero-width records are ignored. +/// +/// Special tokens never merge with neighbouring text, so this is the identity +/// test for a special token's spelling. Anchors beyond the available input do +/// not exist yet, so a partially arrived genuine spelling passes and is left to +/// `literal` to report as incomplete. +fn is_single_token(anchors: &[TokenAttribution], range: Range, token_id: u32) -> bool { + let mut visible = visible_in(anchors, range.clone()); + visible.next() == Some((range.start, token_id)) && visible.next().is_none() +} + +/// Earliest text position at or after `from` where one of `spellings` occurs +/// in full, or where a proper prefix of one of them ends the input. +fn text_candidate<'m>( + text: &str, + spellings: impl Iterator + Clone, + from: usize, +) -> Option { + // Round up: a marker starts on a char boundary at or after `from`, and a + // rejected candidate inside a multi-byte char must not be found again. + let from = text.ceil_char_boundary(from); + let tail = &text[from..]; + if tail.is_empty() { + return None; + } + if let Some(at) = find_earliest(tail, spellings.clone()) { + return Some(from + at); + } + let keep_len = spellings.map(|spelling| partial_prefix_len(tail, spelling)).max().unwrap_or(0); + (keep_len > 0).then(|| text.len() - keep_len) +} + +/// Earliest occurrence of any spelling in `text`. +#[inline(always)] +fn find_earliest<'m>(text: &str, mut spellings: impl Iterator) -> Option { + let first = spellings.next()?; + let second = spellings.next(); + let third = second.and_then(|_| spellings.next()); + let fourth = third.and_then(|_| spellings.next()); + // Use the fast specialized `FindSlice` impls for 1-3 spellings, and fall + // back to a linear scan for 4+. + let range = match (second, third, fourth) { + (None, _, _) => text.find_slice(first), + (Some(second), None, _) => text.find_slice((first, second)), + (Some(second), Some(third), None) => text.find_slice((first, second, third)), + (Some(second), Some(third), Some(fourth)) => { + return [first, second, third, fourth] + .into_iter() + .chain(spellings) + .filter_map(|spelling| text.find(spelling)) + .min(); + } + }; + range.map(|range| range.start) +} + +/// A special token's spelling and its ID in the active tokenizer, resolved +/// together so a marker cannot pair one token's spelling with another's ID. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SpecialToken { + pub text: String, + pub id: u32, +} + +/// A fixed marker string whose special-token segments are guarded by token +/// identity. +/// +/// Build it once from tokenizer IDs, then use it both as the parser that +/// consumes the marker (it implements [`Parser`] by reference) and as the +/// definition the safe-text scanner stops in front of. +/// +/// ```ignore +/// let think_close = Marker::special(&close).then_text("think").then_special(&sep); +/// ``` +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Marker { + text: String, + guards: Vec, +} + +/// One special-token segment of a marker. +#[derive(Clone, Debug, PartialEq, Eq)] +struct Guard { + /// Byte range of the spelling within the marker text. + range: Range, + token_id: u32, +} + +impl Marker { + /// A marker matched by spelling alone. + pub fn text(text: impl Into) -> Self { + Self { + text: text.into(), + guards: Vec::new(), + } + } + + /// A marker beginning with the special token `token`. + pub fn special(token: &SpecialToken) -> Self { + Self::text(String::new()).then_special(token) + } + + /// Append ordinary text, matched by spelling. + #[must_use] + pub fn then_text(mut self, text: &str) -> Self { + self.text.push_str(text); + self + } + + /// Append the special token `token`. + #[must_use] + pub fn then_special(mut self, token: &SpecialToken) -> Self { + let start = self.text.len(); + self.text.push_str(&token.text); + self.guards.push(Guard { + range: start..self.text.len(), + token_id: token.id, + }); + self + } + + /// Drop the token-identity guards, so the marker matches by spelling alone, + /// for callers whose input carries no token attribution. + #[must_use] + pub fn without_guards(mut self) -> Self { + self.guards.clear(); + self + } + + /// The full spelling of the marker. + pub fn as_str(&self) -> &str { + &self.text + } +} + +/// Anything the scanning helpers accept as a marker, viewed through +/// [`MarkerRef`]. +/// +/// A plain `str` is a marker without guards, so text-only parsers keep passing +/// `&[&str]` (or `&String`); token-aware parsers pass `&Marker`. +#[auto_impl(&)] +pub trait MarkerLike { + /// Borrow the marker's spelling and guards. + fn as_marker(&self) -> MarkerRef<'_>; +} + +impl MarkerLike for str { + fn as_marker(&self) -> MarkerRef<'_> { + MarkerRef { + text: self, + guards: &[], + } + } +} + +impl MarkerLike for String { + fn as_marker(&self) -> MarkerRef<'_> { + self.as_str().as_marker() + } +} + +impl MarkerLike for Marker { + fn as_marker(&self) -> MarkerRef<'_> { + MarkerRef { + text: &self.text, + guards: &self.guards, + } + } +} + +/// Borrowed view of a marker: its spelling and guarded segments. +/// +/// A plain `&str` is a marker without guards; see [`MarkerLike`]. +#[derive(Clone, Copy, Debug)] +pub struct MarkerRef<'m> { + text: &'m str, + guards: &'m [Guard], +} + +impl<'m> MarkerRef<'m> { + /// The full spelling of the marker. + pub fn as_str(self) -> &'m str { + self.text + } + + /// The special token the marker begins with, if it begins with one. + pub fn first_token_id(self) -> Option { + self.guards + .first() + .filter(|guard| guard.range.start == 0) + .map(|guard| guard.token_id) + } + + /// Parse the marker at the cursor, returning its spelling. + /// + /// Every guard whose first byte is available is checked before the spelling: + /// when those bytes do not come from the expected special token the parser + /// backtracks at once instead of asking for more input, so an ordinary + /// partial lookalike streams as content immediately. A genuine special token + /// split by output holdback still reports incomplete, because its anchor + /// arrived with its first byte. Nothing is consumed on failure. + pub fn parse_next<'i, I: MarkerStream<'i>>(self, input: &mut I) -> ModalResult<&'i str> { + let start = input.offset(); + let available = input.eof_offset(); + let anchors = input.anchors(); + for guard in self.guards.iter().take_while(|guard| guard.range.start < available) { + let range = start + guard.range.start..start + guard.range.end; + if !is_single_token(anchors, range, guard.token_id) { + return Err(ErrMode::Backtrack(ContextError::new())); + } + } + literal(self.text).parse_next(input) + } +} + +impl<'i, I: MarkerStream<'i>> Parser> for &Marker { + fn parse_next(&mut self, input: &mut I) -> ModalResult<&'i str> { + self.as_marker().parse_next(input) + } +} + +#[cfg(test)] +mod tests { + use vllm_tokenizer::{DecodedText, TokenAnchor, TokenAttribution}; + use winnow::Parser; + use winnow::error::ErrMode; + use winnow::stream::{Partial, Stream}; + + use super::{ + Attributed, Marker, MarkerLike, MarkerStream, SpecialToken, attributed_with_markers, + }; + use crate::utils::safe_text_len_mul; + + const CLOSE_ID: u32 = 256; + const SEP_ID: u32 = 258; + + /// Build attributed text from `(piece, token_id)` pairs, one visible anchor + /// at the first byte of each piece. + fn pieces(pieces: &[(&str, u32)]) -> DecodedText { + let mut decoded = DecodedText::default(); + for &(piece, token_id) in pieces { + decoded.attributions.push(TokenAttribution { + token_id, + anchor: TokenAnchor::Visible { + byte_offset: decoded.text.len() as u32, + }, + }); + decoded.text.push_str(piece); + } + decoded + } + + fn token(text: &str, id: u32) -> SpecialToken { + SpecialToken { + text: text.to_string(), + id, + } + } + + /// An input over `buffer` without parser markers. + fn attributed(buffer: &DecodedText) -> Attributed<'_> { + attributed_with_markers(buffer, &()) + } + + fn close() -> Marker { + Marker::special(&token("<|close|>", CLOSE_ID)) + } + + fn think_close() -> Marker { + close().then_text("think").then_special(&token("<|sep|>", SEP_ID)) + } + + /// `<|close|>` spelled by ordinary tokens. + fn ordinary_close() -> Vec<(&'static str, u32)> { + vec![("<", 60), ("|", 124), ("close", 900), ("|", 124), (">", 62)] + } + + #[test] + fn special_marker_accepts_single_token_with_expected_id() { + let buffer = pieces(&[("<|close|>", CLOSE_ID), ("x", 1)]); + let mut input = attributed(&buffer); + + let matched = (&close()).parse_next(&mut input).unwrap(); + + assert_eq!(matched, "<|close|>"); + assert_eq!(input.remaining(), "x"); + } + + #[test] + fn special_marker_rejects_ordinary_tokens_with_same_spelling() { + let mut fixture = ordinary_close(); + fixture.push(("x", 1)); + let buffer = pieces(&fixture); + let mut input = attributed(&buffer); + + let error = (&close()).parse_next(&mut input).unwrap_err(); + + assert!(matches!(error, ErrMode::Backtrack(_))); + assert_eq!(input.remaining(), "<|close|>x", "nothing consumed"); + } + + #[test] + fn special_marker_rejects_single_ordinary_token_with_wrong_id() { + let buffer = pieces(&[("<|close|>", 900)]); + let mut input = attributed(&buffer); + + let error = (&close()).parse_next(&mut input).unwrap_err(); + + assert!(matches!(error, ErrMode::Backtrack(_))); + } + + #[test] + fn special_marker_reports_incomplete_for_split_genuine_spelling() { + // Output holdback delivered only the first bytes; the anchor came with them. + let buffer = pieces(&[("<|clo", CLOSE_ID)]); + let mut input = attributed(&buffer); + + let error = (&close()).parse_next(&mut input).unwrap_err(); + + assert!(matches!(error, ErrMode::Incomplete(_))); + } + + #[test] + fn special_marker_rejects_partial_lookalike_without_waiting() { + let buffer = pieces(&[("<", 60), ("|", 124), ("clo", 901)]); + let mut input = attributed(&buffer); + + let error = (&close()).parse_next(&mut input).unwrap_err(); + + assert!(matches!(error, ErrMode::Backtrack(_))); + } + + #[test] + fn special_marker_ignores_zero_width_records_inside_span() { + let mut buffer = pieces(&[("<|close|>", CLOSE_ID)]); + buffer.attributions.push(TokenAttribution { + token_id: 7, + anchor: TokenAnchor::ZeroWidth { byte_offset: 3 }, + }); + let mut input = attributed(&buffer); + + (&close()).parse_next(&mut input).unwrap(); + } + + #[test] + fn marker_without_guards_matches_spelling_only() { + let buffer = pieces(&ordinary_close()); + let marker = close().without_guards(); + + (&marker).parse_next(&mut attributed(&buffer)).unwrap(); + (&marker).parse_next(&mut Partial::new("<|close|>")).unwrap(); + } + + #[test] + fn guarded_marker_never_matches_text_only_input() { + let error = (&close()).parse_next(&mut Partial::new("<|close|>")).unwrap_err(); + + assert!(matches!(error, ErrMode::Backtrack(_))); + } + + #[test] + fn marker_requires_every_special_segment() { + let marker = think_close(); + let genuine = pieces(&[("<|close|>", CLOSE_ID), ("think", 5), ("<|sep|>", SEP_ID)]); + assert_eq!( + (&marker).parse_next(&mut attributed(&genuine)).unwrap(), + "<|close|>think<|sep|>" + ); + + // Real close, ordinary separator: the strict rule keeps it as content. + let mixed = pieces(&[("<|close|>", CLOSE_ID), ("think", 5), ("<|sep|>", 902)]); + let mut input = attributed(&mixed); + assert!(matches!( + (&marker).parse_next(&mut input).unwrap_err(), + ErrMode::Backtrack(_) + )); + assert_eq!( + input.remaining(), + "<|close|>think<|sep|>", + "nothing consumed" + ); + } + + #[test] + fn marker_first_token_id_requires_leading_special() { + assert_eq!(think_close().as_marker().first_token_id(), Some(CLOSE_ID)); + assert_eq!(Marker::text("").as_marker().first_token_id(), None); + assert_eq!( + Marker::text("x") + .then_special(&token("<|sep|>", SEP_ID)) + .as_marker() + .first_token_id(), + None + ); + assert_eq!("".as_marker().first_token_id(), None); + } + + #[test] + fn attributed_next_candidate_uses_anchors_for_guarded_markers() { + let mut fixture = vec![("say ", 1)]; + fixture.extend(ordinary_close()); // lookalike at 4 + fixture.extend([("think", 5), ("<|sep|>", 903), (" then ", 2)]); + let genuine_at = fixture.iter().map(|(piece, _)| piece.len()).sum::(); + fixture.extend([("<|close|>", CLOSE_ID), ("think", 5), ("<|sep|>", SEP_ID)]); + let buffer = pieces(&fixture); + let input = attributed(&buffer); + + let marker = think_close(); + assert_eq!(input.next_candidate(&[&marker], 0), Some(genuine_at)); + assert_eq!(input.next_candidate(&[&marker], genuine_at + 1), None); + } + + #[test] + fn attributed_next_candidate_falls_back_to_text_for_unguarded_markers() { + let buffer = pieces(&[("a", 1), ("", 2), ("b", 3)]); + let input = attributed(&buffer); + + assert_eq!(input.next_candidate(&[""], 0), Some(1)); + assert_eq!(input.next_candidate(&[""], 0), None); + // Partial spelling at the end of the input is a candidate too. + let buffer = pieces(&[("a", 1), (""], 0), Some(1)); + } + + #[test] + fn text_next_candidate_matches_previous_scan_semantics() { + let input = Partial::new("hello<|tool"); + + assert_eq!( + input.next_candidate(&["<|tool_call>", "<|channel>thought\n"], 0), + Some("hello".len()) + ); + assert_eq!(input.next_candidate(&["<|tool_call>"], 18), None); + assert_eq!( + Partial::new("hello<|tool_call>") + .next_candidate(&["<|tool_call>", ""], 0), + Some(5) + ); + } + + #[test] + fn scan_skips_rejected_text_candidate_starting_with_multibyte_char() { + // A marker whose leading text is found by spelling but whose guard then + // fails: the scan must move past the whole char, not retry inside it. + let marker = Marker::text("é").then_special(&token("<|sep|>", SEP_ID)); + let buffer = pieces(&[("aé", 1), ("<", 60), ("|sep|>", 904), ("b", 2)]); + let mut input = attributed(&buffer); + + let len = safe_text_len_mul(&mut input, &[&marker]).unwrap(); + + assert_eq!(len, buffer.text.len()); + } + + #[test] + fn attributed_offsets_follow_the_cursor() { + let buffer = pieces(&[("ab", 1), ("<|close|>", CLOSE_ID)]); + let mut input = attributed(&buffer); + input.next_slice(2); + + assert_eq!(input.offset(), 2); + assert_eq!(input.remaining(), "<|close|>"); + (&close()).parse_next(&mut input).unwrap(); + assert_eq!(input.offset(), 11); + } +}