iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Rust Frontend] Token-aware marker parsing for Kimi K3 (#58358)

Browse Source 
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Bugen Zhao
2026-10-01 04:21:50 +00:00
committed by GitHub
co-authored by Claude Opus 5.5
parent ec5ef92c87
commit 36229e1337
8 changed files with 1135 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/Cargo.lock
Generated
+2
View File
@@ -5964,8 +5964,10 @@ dependencies = [
name = "vllm-parser"
version = "0.1.0"
dependencies = [
"auto_impl",
"criterion",
"easy-ext",
"educe",
"expect-test",
"futures",
"openai-protocol",
rust/Cargo.toml
+1
View File
@@ -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"] }
rust/src/chat/src/parser/unified.rs
+2 -1
View File
@@ -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!(
rust/src/parser/Cargo.toml
+2
View File
@@ -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
rust/src/parser/src/unified/kimi_k3.rs
+340 -54
View File
@@ -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<Self> {
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<Self> {
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<KimiK3Event> {
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<KimiK3Event> {
/// Parse an event inside the `think` channel.
fn parse_reasoning_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {
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<KimiK3Event
/// Parse an event inside the `response` channel.
fn parse_response_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {
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<KimiK3Event>
/// Parse an event after the response channel closed.
fn parse_epilogue_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {
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<KimiK3Event>
/// Parse an event inside the `tools` channel, between `call` blocks.
fn parse_tools_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {
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<KimiK3Event> {
/// Parse a message close or end-of-message marker.
fn message_end_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {
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<KimiK3Event> {
/// Parse safe text while waiting for the next channel marker.
fn safe_idle_text_event(input: &mut KimiK3Input<'_>) -> ModalResult<KimiK3Event> {
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<KimiK3Event> {
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<KimiK3Event> {
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<KimiK3Event> {
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<KimiK3Event> {
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<KimiK3Event> {
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<KimiK3Event> {
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<UnifiedParserOutput> {
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<u32> {
tokenizer().encode(text, false).unwrap()
}
/// Encode as ordinary bytes only, so marker spellings are content.
fn ordinary(text: &str) -> Vec<u32> {
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}"#);
}
}
rust/src/parser/src/unified/mod.rs
+9
View File
@@ -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<T> = std::result::Result<T, UnifiedParserError>;
@@ -288,3 +289,11 @@ fn token_id(tokenizer: &dyn vllm_tokenizer::Tokenizer, token: &str) -> Result<u3
token: token.to_string(),
})
}
/// Resolves `token` to a [`SpecialToken`], or an error if it's not found.
fn special_token(tokenizer: &dyn vllm_tokenizer::Tokenizer, token: &str) -> Result<SpecialToken> {
Ok(SpecialToken {
id: token_id(tokenizer, token)?,
text: token.to_string(),
})
}
rust/src/parser/src/utils.rs
+119 -78
View File
@@ -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<Scan> {
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<usize> {
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<usize> {
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<usize> {
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<usize> {
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<usize> {
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<Partial<&'i str>, &'i str, ErrMode<ContextError>> + 'a {
move |input: &mut Partial<&'i str>| take_until_marker_(input, marker, state)
) -> impl Parser<I, &'i str, ErrMode<ContextError>> + '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<E>(
buffer: &str,
parse: impl FnOnce(&mut Partial<&str>) -> ModalResult<E>,
) -> Result<Option<(E, usize)>> {
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<M: Debug, E>(
buffer: &DecodedText,
markers: &M,
parse: impl FnOnce(&mut Attributed<'_, M>) -> ModalResult<E>,
) -> Result<Option<(E, usize)>> {
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<E>,
) -> Result<Option<(E, usize)>> {
let buffer = input.remaining();
let checkpoint = input.checkpoint();
let event = match parse(&mut input) {
Ok(event) => event,
rust/src/parser/src/utils/marker.rs
+660
View File
@@ -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<Partial<LocatingSlice<&'i str>>, 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<Slice = &'i str> + 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<M: MarkerLike>(&self, markers: &[M], from: usize) -> Option<usize> {
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<usize>,
) -> impl Iterator<Item = (usize, u32)> + '_ {
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<usize>, 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<Item = &'m str> + Clone,
from: usize,
) -> Option<usize> {
// 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<Item = &'m str>) -> Option<usize> {
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<Guard>,
}
/// 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<usize>,
token_id: u32,
}
impl Marker {
/// A marker matched by spelling alone.
pub fn text(text: impl Into<String>) -> 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<u32> {
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<I, &'i str, ErrMode<ContextError>> 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("</think>").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!("</think>".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::<usize>();
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), ("</think>", 2), ("b", 3)]);
let input = attributed(&buffer);
assert_eq!(input.next_candidate(&["</think>"], 0), Some(1));
assert_eq!(input.next_candidate(&["</thinking>"], 0), None);
// Partial spelling at the end of the input is a candidate too.
let buffer = pieces(&[("a", 1), ("</thi", 2)]);
let input = attributed(&buffer);
assert_eq!(input.next_candidate(&["</think>"], 0), Some(1));
}
#[test]
fn text_next_candidate_matches_previous_scan_semantics() {
let input = Partial::new("hello<not_marker><|tool");
assert_eq!(
input.next_candidate(&["<|tool_call>", "<|channel>thought\n"], 0),
Some("hello<not_marker>".len())
);
assert_eq!(input.next_candidate(&["<|tool_call>"], 18), None);
assert_eq!(
Partial::new("hello<channel|><|tool_call>")
.next_candidate(&["<|tool_call>", "<channel|>"], 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);
}
}