From 63dab72ee57ead37d96055fa33490e2dbae79814 Mon Sep 17 00:00:00 2001 From: lightseek-bot <243258330+lightseek-bot@users.noreply.github.com> Date: Sun, 26 Jul 2026 20:25:21 +0000 Subject: [PATCH 1/5] feat(kimi-k3): add K3 support Signed-off-by: lightseek-bot <243258330+lightseek-bot@users.noreply.github.com> Co-authored-by: Keyang Ru Co-authored-by: Chen Hongtao <56470055+chenht2022@users.noreply.github.com> --- crates/multimodal/src/registry/kimi_k25.rs | 35 +- crates/multimodal/src/vision/processor.rs | 10 +- crates/reasoning_parser/Cargo.toml | 1 + crates/reasoning_parser/src/factory.rs | 20 +- crates/reasoning_parser/src/lib.rs | 5 +- .../reasoning_parser/src/parsers/kimi_k3.rs | 632 ++++++++ crates/reasoning_parser/src/parsers/mod.rs | 2 + crates/tokenizer/src/encoders/kimi_k3_xtml.rs | 1118 ++++++++++++++ crates/tokenizer/src/encoders/mod.rs | 1 + crates/tokenizer/src/kimi_k2_tokenizer.rs | 25 +- crates/tokenizer/src/tiktoken.rs | 73 +- .../fixtures/kimi_k3/k3_render_fixtures.json | 1349 +++++++++++++++++ crates/tokenizer/tests/kimi_k3_renderer.rs | 209 +++ crates/tool_parser/src/factory.rs | 9 +- crates/tool_parser/src/lib.rs | 4 +- crates/tool_parser/src/parsers/kimi_k3.rs | 666 ++++++++ crates/tool_parser/src/parsers/mod.rs | 2 + .../src/routers/grpc/utils/chat_utils.rs | 116 +- 18 files changed, 4243 insertions(+), 34 deletions(-) create mode 100644 crates/reasoning_parser/src/parsers/kimi_k3.rs create mode 100644 crates/tokenizer/src/encoders/kimi_k3_xtml.rs create mode 100644 crates/tokenizer/tests/fixtures/kimi_k3/k3_render_fixtures.json create mode 100644 crates/tokenizer/tests/kimi_k3_renderer.rs create mode 100644 crates/tool_parser/src/parsers/kimi_k3.rs diff --git a/crates/multimodal/src/registry/kimi_k25.rs b/crates/multimodal/src/registry/kimi_k25.rs index 2abbff858..327784d55 100644 --- a/crates/multimodal/src/registry/kimi_k25.rs +++ b/crates/multimodal/src/registry/kimi_k25.rs @@ -28,11 +28,14 @@ impl ModelProcessorSpec for KimiK25VisionSpec { } fn matches(&self, metadata: &ModelMetadata) -> bool { + // Kimi-K3 reuses K2.5's MoonViT vision stack and `<|media_pad|>` + // placeholder (media_placeholder_token_id 163605), so it shares this + // spec. let id = metadata.model_id.to_ascii_lowercase(); - id.contains("kimi") && id.contains("k2") + id.contains("kimi") && (id.contains("k2") || id.contains("k3")) || metadata .config_model_type() - .is_some_and(|mt| mt == "kimi_k25") + .is_some_and(|mt| mt == "kimi_k25" || mt == "kimi_k3") } fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult { @@ -120,6 +123,34 @@ mod tests { assert_eq!(spec.name(), "kimi_k25"); } + #[test] + fn kimi_k3_matches_model_id_and_model_type() { + let tokenizer = TestTokenizer::new(&[("<|media_pad|>", 163605)]); + // Match by model_id containing kimi + k3. + let config = json!({ + "model_type": "kimi_k3", + "media_placeholder_token_id": 163605 + }); + let metadata = ModelMetadata { + model_id: "moonshotai/Kimi-K3", + tokenizer: &tokenizer, + config: &config, + }; + let registry = ModelRegistry::new(); + let spec = registry + .lookup(&metadata) + .expect("kimi_k3 -> kimi_k25 spec"); + assert_eq!(spec.name(), "kimi_k25"); + + // Also match by model_type alone (id without a k3 hint). + let metadata_by_type = ModelMetadata { + model_id: "internal/checkpoint-final", + tokenizer: &tokenizer, + config: &config, + }; + assert!(registry.lookup(&metadata_by_type).is_some()); + } + #[test] fn kimi_k25_prompt_replacements() { let tokenizer = TestTokenizer::new(&[("<|media_pad|>", 163605)]); diff --git a/crates/multimodal/src/vision/processor.rs b/crates/multimodal/src/vision/processor.rs index 27654ae87..1569bc244 100644 --- a/crates/multimodal/src/vision/processor.rs +++ b/crates/multimodal/src/vision/processor.rs @@ -325,7 +325,7 @@ impl VisionProcessorRegistry { Box::new(super::processors::Llama4VisionProcessor::new()), ); - // Register Kimi-K2.5 Vision + // Register Kimi-K2.5 Vision (also used by Kimi-K3, same MoonViT stack) registry.register( "kimi-k2", Box::new(super::processors::KimiK25Processor::new()), @@ -334,6 +334,14 @@ impl VisionProcessorRegistry { "kimi_k2", Box::new(super::processors::KimiK25Processor::new()), ); + registry.register( + "kimi-k3", + Box::new(super::processors::KimiK25Processor::new()), + ); + registry.register( + "kimi_k3", + Box::new(super::processors::KimiK25Processor::new()), + ); registry } diff --git a/crates/reasoning_parser/Cargo.toml b/crates/reasoning_parser/Cargo.toml index 71ad479d4..ef1068d0d 100644 --- a/crates/reasoning_parser/Cargo.toml +++ b/crates/reasoning_parser/Cargo.toml @@ -14,6 +14,7 @@ categories = ["parsing", "text-processing"] [dependencies] parking_lot = { workspace = true } +regex = "1.12" thiserror.workspace = true tokio = { workspace = true, features = ["sync"] } diff --git a/crates/reasoning_parser/src/factory.rs b/crates/reasoning_parser/src/factory.rs index 0e9c2b57f..2e89b4d59 100644 --- a/crates/reasoning_parser/src/factory.rs +++ b/crates/reasoning_parser/src/factory.rs @@ -7,7 +7,7 @@ use parking_lot::RwLock; use crate::{ parsers::{ BaseReasoningParser, CohereCmdParser, DeepSeekR1Parser, Glm45Parser, InklingParser, - KimiParser, MiniMaxParser, NanoV3Parser, PassthroughParser, Qwen3Parser, + KimiK3Parser, KimiParser, MiniMaxParser, NanoV3Parser, PassthroughParser, Qwen3Parser, QwenThinkingParser, Step3Parser, }, traits::{ParserConfig, ReasoningParser, DEFAULT_MAX_BUFFER_SIZE}, @@ -184,6 +184,9 @@ impl ParserFactory { Box::new(BaseReasoningParser::new(config).with_model_type("kimi_thinking".to_string())) }); + // Kimi K3 XTML think channel (structural <|open|>/<|close|>/<|sep|> tokens). + registry.register_parser("kimi_k3", || Box::new(KimiK3Parser::new())); + registry.register_pattern("deepseek-r1", "deepseek_r1"); registry.register_pattern("deepseek-v3.1", "deepseek_v31"); registry.register_pattern("deepseek-v3-1", "deepseek_v31"); @@ -196,6 +199,10 @@ impl ParserFactory { registry.register_pattern("glm-5", "glm45"); // GLM-5.x reuse glm45 reasoning format registry.register_pattern("kimi-k2-thinking", "kimi_thinking"); registry.register_pattern("kimi-k2.5", "kimi_k25"); + // K3 patterns must precede the generic "kimi" pattern below, since + // "kimi-k3".contains("kimi") and first-substring-match wins. + registry.register_pattern("kimi-k3", "kimi_k3"); + registry.register_pattern("kimi_k3", "kimi_k3"); registry.register_pattern("kimi", "kimi"); // legacy: Kimi-K2-Instruct with unicode tokens registry.register_pattern("step3", "step3"); registry.register_pattern("minimax", "minimax"); @@ -280,6 +287,17 @@ mod tests { assert_eq!(parser.model_type(), "kimi"); } + #[test] + fn test_factory_creates_kimi_k3() { + let factory = ParserFactory::new(); + // Kimi-K3 ids must resolve to the dedicated kimi_k3 parser, not the + // legacy "kimi" parser (whose pattern is a substring of "kimi-k3"). + assert_eq!(factory.create("moonshotai/Kimi-K3").model_type(), "kimi_k3"); + assert_eq!(factory.create("Kimi_K3").model_type(), "kimi_k3"); + // The legacy kimi pattern still resolves plain Kimi-K2 ids. + assert_eq!(factory.create("kimi-chat").model_type(), "kimi"); + } + #[test] fn test_factory_creates_inkling() { let factory = ParserFactory::new(); diff --git a/crates/reasoning_parser/src/lib.rs b/crates/reasoning_parser/src/lib.rs index 58a090518..a04a5f160 100644 --- a/crates/reasoning_parser/src/lib.rs +++ b/crates/reasoning_parser/src/lib.rs @@ -4,8 +4,9 @@ pub mod traits; pub use factory::{ParserFactory, ParserRegistry}; pub use parsers::{ - BaseReasoningParser, CohereCmdParser, DeepSeekR1Parser, Glm45Parser, InklingParser, KimiParser, - MiniMaxParser, NanoV3Parser, PassthroughParser, Qwen3Parser, QwenThinkingParser, Step3Parser, + BaseReasoningParser, CohereCmdParser, DeepSeekR1Parser, Glm45Parser, InklingParser, + KimiK3Parser, KimiParser, MiniMaxParser, NanoV3Parser, PassthroughParser, Qwen3Parser, + QwenThinkingParser, Step3Parser, }; pub use traits::{ ParseError, ParserConfig, ParserResult, ReasoningParser, DEFAULT_MAX_BUFFER_SIZE, diff --git a/crates/reasoning_parser/src/parsers/kimi_k3.rs b/crates/reasoning_parser/src/parsers/kimi_k3.rs new file mode 100644 index 000000000..771a93b43 --- /dev/null +++ b/crates/reasoning_parser/src/parsers/kimi_k3.rs @@ -0,0 +1,632 @@ +// Kimi K3 reasoning parser for the XTML tag format. +// +// K3 uses structural special tokens <|open|>, <|close|>, <|sep|> to delimit +// channels. This parser extracts the `think` channel into `reasoning_text` and +// delivers the remainder (with `response`/`message` wrapper tokens removed) as +// `normal_text`. The `tools` channel is preserved verbatim for the tool parser. +// +// Ported from the Kimi-K3 reference reasoning parser (XTML `think` channel). + +use regex::Regex; + +use crate::traits::{ParseError, ParserResult, ReasoningParser, DEFAULT_MAX_BUFFER_SIZE}; + +/// Literal marker strings (used for streaming overlap detection). +const THINK_OPEN: &str = "<|open|>think<|sep|>"; +const THINK_CLOSE: &str = "<|close|>think<|sep|>"; +const RESPONSE_OPEN: &str = "<|open|>response<|sep|>"; +const RESPONSE_CLOSE: &str = "<|close|>response<|sep|>"; +const MESSAGE_OPEN: &str = "<|open|>message<|sep|>"; +const MESSAGE_CLOSE: &str = "<|close|>message<|sep|>"; + +/// Reasoning parser for the Kimi K3 (XTML) chat format. +/// +/// Extracts the `think` channel into `reasoning_text` and strips the +/// `response`/`message` wrapper tokens from `normal_text` by substitution, +/// preserving any `tools` channel for downstream parsing. +#[derive(Debug, Clone)] +pub struct KimiK3Parser { + think_open_re: Regex, + think_close_re: Regex, + response_open_re: Regex, + response_close_re: Regex, + message_open_re: Regex, + message_close_re: Regex, + + // Shared state + in_reasoning: bool, + + // Streaming accumulation + buffer: String, + reason_phase_ended: bool, + /// Byte offset in `buffer` where post-think content begins. + content_tail_start: usize, + /// Safe-to-emit reasoning already returned to the caller. + emitted_reasoning: String, + /// Safe-to-emit content already returned to the caller. + emitted_content: String, +} + +impl KimiK3Parser { + /// Create a new `KimiK3Parser` with pre-compiled tolerant XTML regexes. + #[expect( + clippy::expect_used, + reason = "regex patterns are hardcoded valid literals; failure indicates a programming error" + )] + pub fn new() -> Self { + let open = r"<\|open\|>"; + let close = r"<\|close\|>"; + let sep = r"<\|sep\|>"; + + Self { + think_open_re: Regex::new(&format!(r"{open}\s*think\s*{sep}")) + .expect("valid think-open regex"), + think_close_re: Regex::new(&format!(r"{close}\s*think\s*{sep}")) + .expect("valid think-close regex"), + response_open_re: Regex::new(&format!(r"{open}\s*response\s*{sep}")) + .expect("valid response-open regex"), + response_close_re: Regex::new(&format!(r"{close}\s*response\s*{sep}")) + .expect("valid response-close regex"), + message_open_re: Regex::new(&format!(r"{open}\s*message\s*{sep}")) + .expect("valid message-open regex"), + message_close_re: Regex::new(&format!(r"{close}\s*message\s*{sep}")) + .expect("valid message-close regex"), + + in_reasoning: false, + buffer: String::new(), + reason_phase_ended: false, + content_tail_start: 0, + emitted_reasoning: String::new(), + emitted_content: String::new(), + } + } + + /// Strip the `response` and `message` wrapper markers + /// (`<|open|>response<|sep|>`, `<|close|>response<|sep|>`, + /// `<|open|>message<|sep|>`, `<|close|>message<|sep|>`) from `text` via + /// substitution. + /// + /// Unlike a structured extraction that pulls only the response body, + /// this approach removes only the wrapper markers, preserving any + /// `<|open|>tools<|sep|>…` channel that follows for the tool parser. + fn strip_content_wrapper(&self, text: &str) -> String { + let text = self.response_open_re.replace_all(text, ""); + let text = self.response_close_re.replace_all(&text, ""); + let text = self.message_open_re.replace_all(&text, ""); + self.message_close_re.replace_all(&text, "").into_owned() + } + + /// Return the prefix of accumulated reasoning text that is safe to stream. + /// + /// Skips past the think-open marker if present, then holds back any + /// partial marker suffix that could complete as a think-open or think-close. + fn reasoning_text_ready_to_emit<'a>(&self, text: &'a str) -> &'a str { + let after_open = if let Some(m) = self.think_open_re.find(text) { + &text[m.end()..] + } else { + text + }; + + let overlap = Self::compute_overlap(after_open, &[THINK_OPEN, THINK_CLOSE]); + if overlap > 0 { + &after_open[..after_open.len() - overlap] + } else { + after_open + } + } + + /// Return the prefix of post-reasoning content that is safe to stream. + /// + /// Strips the response-open prefix, removes complete response-close and + /// message-close markers, then holds back any partial marker suffix. + fn content_ready_to_emit(&self, text: &str) -> String { + // Strip response-open prefix (first occurrence) + let text = if let Some(m) = self.response_open_re.find(text) { + &text[m.end()..] + } else { + text + }; + + // Remove all complete response-close, message-open, and message-close markers + let text = self.response_close_re.replace_all(text, ""); + let text = self.message_open_re.replace_all(&text, ""); + let text = self.message_close_re.replace_all(&text, ""); + + // Hold back any partial marker suffix + let text_str: &str = &text; + let overlap = Self::compute_overlap( + text_str, + &[RESPONSE_OPEN, RESPONSE_CLOSE, MESSAGE_OPEN, MESSAGE_CLOSE], + ); + if overlap > 0 { + text_str[..text_str.len() - overlap].to_owned() + } else { + text_str.to_owned() + } + } + + /// Compute the maximum length of a suffix of `text` that is also a + /// non-empty prefix of any marker in `markers`. + fn compute_overlap(text: &str, markers: &[&str]) -> usize { + let mut overlap = 0usize; + for marker in markers { + let max_check = marker.len().saturating_sub(1).min(text.len()); + for n in (1..=max_check).rev() { + if text.ends_with(&marker[..n]) { + overlap = overlap.max(n); + break; + } + } + } + overlap + } +} + +impl Default for KimiK3Parser { + fn default() -> Self { + Self::new() + } +} + +impl ReasoningParser for KimiK3Parser { + fn detect_and_parse_reasoning(&mut self, text: &str) -> Result { + // Guard against oversized single-shot input, mirroring the streaming + // entry point (`BaseReasoningParser` and `InklingParser` apply the same + // guard to their non-streaming path). + if text.len() > DEFAULT_MAX_BUFFER_SIZE { + return Err(ParseError::BufferOverflow(text.len())); + } + + let m_open = self.think_open_re.find(text); + let reasoning_start = m_open.map(|m| m.end()).unwrap_or(0); + + // No think channel markers at all. + if m_open.is_none() && self.think_close_re.find(text).is_none() { + if self.in_reasoning { + // Generation prefix consumed the open marker; think not yet closed. + // + // This diverges from the reference non-streaming extraction + // path, which is stateless and always treats marker-free text as + // plain content. SMG instead tracks `in_reasoning` explicitly + // (set via `mark_reasoning_started()` when the caller's + // generation prompt already pre-fills K3's thinking-mode + // prefix `<|open|>think<|sep|>`). So if the model's output is + // truncated before a close marker appears, the entire output + // IS reasoning — classify it all as reasoning here rather + // than losing it as content. + return Ok(ParserResult::reasoning(text.to_owned())); + } + // Plain content with no think channel. + return Ok(ParserResult::normal(self.strip_content_wrapper(text))); + } + + // Search for the close marker starting after the open (or from the start). + let m_close = self.think_close_re.find_at(text, reasoning_start); + if let Some(m_close) = m_close { + let reasoning = text[reasoning_start..m_close.start()].to_owned(); + let normal = self.strip_content_wrapper(&text[m_close.end()..]); + self.in_reasoning = false; + Ok(ParserResult::new(normal, reasoning)) + } else { + // Think opened but not yet closed — still in reasoning. + self.in_reasoning = true; + Ok(ParserResult::reasoning(text[reasoning_start..].to_owned())) + } + } + + fn parse_reasoning_streaming_incremental( + &mut self, + text: &str, + ) -> Result { + // Guard against unbounded buffer growth, mirroring `BaseReasoningParser`. + let buffered_size = self.buffer.len() + text.len(); + if buffered_size > DEFAULT_MAX_BUFFER_SIZE { + return Err(ParseError::BufferOverflow(buffered_size)); + } + + self.buffer.push_str(text); + + // Without a prefilled think channel, the stream begins in normal content. + // Hold only a possible split think-open marker; otherwise route the + // accumulated text downstream so structural tool tokens reach the tool parser. + if !self.in_reasoning + && !self.reason_phase_ended + && self.think_open_re.find(&self.buffer).is_none() + { + if THINK_OPEN.starts_with(self.buffer.as_str()) { + return Ok(ParserResult::default()); + } + self.reason_phase_ended = true; + self.content_tail_start = 0; + } + + if self.reason_phase_ended { + // Content phase: emit new content from the post-think tail. + let content_tail = &self.buffer[self.content_tail_start..]; + let current_safe = self.content_ready_to_emit(content_tail); + // Invariant: `emitted_content` is always a prefix of `current_safe`; + // on the impossible mismatch, emit nothing rather than double-emit. + let new_content = if current_safe.starts_with(self.emitted_content.as_str()) { + current_safe[self.emitted_content.len()..].to_owned() + } else { + String::new() + }; + self.emitted_content = current_safe; + return Ok(ParserResult { + normal_text: new_content, + reasoning_text: String::new(), + }); + } + + // Check whether the think-close marker now completes in the buffer. + if let Some(m_close) = self.think_close_re.find(&self.buffer) { + // Transition: reasoning phase ends. + self.reason_phase_ended = true; + self.content_tail_start = m_close.end(); + self.in_reasoning = false; + + // Final reasoning delta. + let r_start = self + .think_open_re + .find(&self.buffer) + .map(|m| m.end()) + .unwrap_or(0); + let full_reasoning = &self.buffer[r_start..m_close.start()]; + // Invariant: `emitted_reasoning` is always a prefix of `full_reasoning`; + // on the impossible mismatch, emit nothing rather than double-emit. + let reasoning_delta = if full_reasoning.starts_with(self.emitted_reasoning.as_str()) { + full_reasoning[self.emitted_reasoning.len()..].to_owned() + } else { + String::new() + }; + + // Initial content delta (safe prefix of the content tail). + let content_tail = &self.buffer[m_close.end()..]; + let content_safe = self.content_ready_to_emit(content_tail); + self.emitted_content.clone_from(&content_safe); + + Ok(ParserResult { + normal_text: content_safe, + reasoning_text: reasoning_delta, + }) + } else { + // Still in reasoning — emit the safe prefix of accumulated reasoning. + // The tool parser is gated on `is_in_reasoning()`, so the flag must + // stay set until the think-close transition above clears it. + self.in_reasoning = true; + let current_safe = self.reasoning_text_ready_to_emit(&self.buffer).to_owned(); + // Invariant: `emitted_reasoning` is always a prefix of `current_safe`; + // on the impossible mismatch, emit nothing rather than double-emit. + let reasoning_delta = if current_safe.starts_with(self.emitted_reasoning.as_str()) { + current_safe[self.emitted_reasoning.len()..].to_owned() + } else { + String::new() + }; + self.emitted_reasoning = current_safe; + Ok(ParserResult { + normal_text: String::new(), + reasoning_text: reasoning_delta, + }) + } + } + + fn reset(&mut self) { + self.in_reasoning = false; + self.buffer.clear(); + self.reason_phase_ended = false; + self.content_tail_start = 0; + self.emitted_reasoning.clear(); + self.emitted_content.clear(); + } + + fn model_type(&self) -> &str { + "kimi_k3" + } + + fn requires_special_tokens(&self) -> bool { + true + } + + fn is_in_reasoning(&self) -> bool { + self.in_reasoning + } + + fn mark_reasoning_started(&mut self) { + self.in_reasoning = true; + } + + fn mark_think_start_stripped(&mut self) { + // For K3 streaming, the absence of a think-open marker in the buffer + // is already handled by treating reasoning_start as 0. No additional + // state is needed. + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const OPEN: &str = "<|open|>"; + const CLOSE: &str = "<|close|>"; + const SEP: &str = "<|sep|>"; + fn think_open() -> String { + format!("{OPEN}think{SEP}") + } + fn think_close() -> String { + format!("{CLOSE}think{SEP}") + } + fn response_open() -> String { + format!("{OPEN}response{SEP}") + } + fn response_close() -> String { + format!("{CLOSE}response{SEP}") + } + fn message_open() -> String { + format!("{OPEN}message{SEP}") + } + fn message_close() -> String { + format!("{CLOSE}message{SEP}") + } + + // ------------------------------------------------------------------------- + // Non-streaming tests + // ------------------------------------------------------------------------- + + #[test] + fn model_type_and_special_tokens() { + let p = KimiK3Parser::new(); + assert_eq!(p.model_type(), "kimi_k3"); + assert!(p.requires_special_tokens()); + } + + #[test] + fn extract_with_xtml_tags() { + let mut p = KimiK3Parser::new(); + let input = format!( + "{}step{}{}answer", + think_open(), + think_close(), + response_open() + ); + let r = p.detect_and_parse_reasoning(&input).unwrap(); + assert_eq!(r.reasoning_text, "step"); + assert_eq!(r.normal_text, "answer"); + } + + #[test] + fn extract_with_generation_prefix_consumed() { + // open think marker absent (was the prompt prefix) + let mut p = KimiK3Parser::new(); + p.mark_reasoning_started(); + let input = format!("step{}{}answer", think_close(), response_open()); + let r = p.detect_and_parse_reasoning(&input).unwrap(); + assert_eq!(r.reasoning_text, "step"); + assert_eq!(r.normal_text, "answer"); + } + + #[test] + fn strips_response_wrapper_with_close_and_keeps_tools_channel() { + let mut p = KimiK3Parser::new(); + let input = format!( + "{}step{}{}answer{}{}tools{}CALLS{}tools{}", + think_open(), + think_close(), + response_open(), + response_close(), + OPEN, + SEP, + CLOSE, + SEP + ); + let r = p.detect_and_parse_reasoning(&input).unwrap(); + assert_eq!(r.reasoning_text, "step"); + // response wrapper removed; tools channel preserved for the tool parser + assert_eq!( + r.normal_text, + format!("answer{OPEN}tools{SEP}CALLS{CLOSE}tools{SEP}") + ); + } + + #[test] + fn strips_message_wrapper_like_response_wrapper() { + // The `message` wrapper is stripped symmetrically (open + close), just + // like the `response` wrapper, leaving only the answer body. + let mut p = KimiK3Parser::new(); + let input = format!( + "{}r{}{}answer{}", + think_open(), + think_close(), + message_open(), + message_close() + ); + let r = p.detect_and_parse_reasoning(&input).unwrap(); + assert_eq!(r.reasoning_text, "r"); + assert_eq!(r.normal_text, "answer"); + } + + #[test] + fn no_think_channel_is_all_content() { + let mut p = KimiK3Parser::new(); + let r = p.detect_and_parse_reasoning("just an answer").unwrap(); + assert_eq!(r.reasoning_text, ""); + assert_eq!(r.normal_text, "just an answer"); + } + + #[test] + fn marker_free_text_while_in_reasoning_is_all_reasoning() { + // Once the caller has told us we're already in the think channel + // (e.g. the generation prompt pre-filled `<|open|>think<|sep|>`), + // truncated output with no markers at all must come back as + // reasoning in full, not be misclassified as content. + let mut p = KimiK3Parser::new(); + p.mark_reasoning_started(); + let text = "still thinking, no markers"; + let r = p.detect_and_parse_reasoning(text).unwrap(); + assert_eq!(r.reasoning_text, text); + assert_eq!(r.normal_text, ""); + } + + #[test] + fn detect_and_parse_reasoning_buffer_overflow_is_guarded() { + let mut p = KimiK3Parser::new(); + let oversized = "a".repeat(DEFAULT_MAX_BUFFER_SIZE + 1); + let result = p.detect_and_parse_reasoning(&oversized); + assert!(matches!( + result, + Err(ParseError::BufferOverflow(size)) if size == DEFAULT_MAX_BUFFER_SIZE + 1 + )); + } + + // ------------------------------------------------------------------------- + // Streaming tests + // ------------------------------------------------------------------------- + + #[test] + fn streaming_thinking_disabled_routes_tools_to_normal_content() { + let mut p = KimiK3Parser::new(); + let chunks = [ + "<|close|>response", + "<|sep|>", + "<|open|>tools", + "<|sep|>", + r#"<|open|>call tool="kvv_walle_case" index="1""#, + "<|sep|>", + "<|close|>call", + "<|sep|>", + "<|close|>tools", + "<|sep|>", + ]; + + let mut reasoning = String::new(); + let mut normal = String::new(); + for chunk in chunks { + let result = p.parse_reasoning_streaming_incremental(chunk).unwrap(); + reasoning.push_str(&result.reasoning_text); + normal.push_str(&result.normal_text); + } + + assert_eq!(reasoning, ""); + assert_eq!( + normal, + concat!( + "<|open|>tools<|sep|>", + r#"<|open|>call tool="kvv_walle_case" index="1"<|sep|>"#, + "<|close|>call<|sep|><|close|>tools<|sep|>" + ) + ); + assert!(!p.is_in_reasoning()); + } + + #[test] + fn streaming_split_open_marker_held_back() { + let mut p = KimiK3Parser::new(); + assert_eq!( + p.parse_reasoning_streaming_incremental("<|open|>") + .unwrap() + .reasoning_text, + "" + ); + assert_eq!( + p.parse_reasoning_streaming_incremental("think") + .unwrap() + .reasoning_text, + "" + ); + let r = p + .parse_reasoning_streaming_incremental("<|sep|>step") + .unwrap(); + assert_eq!(r.reasoning_text, "step"); + } + + #[test] + fn streaming_split_close_hands_content_downstream() { + let mut p = KimiK3Parser::new(); + p.parse_reasoning_streaming_incremental("<|open|>think<|sep|>step") + .unwrap(); + let partial = p + .parse_reasoning_streaming_incremental("<|close|>") + .unwrap(); + assert_eq!(partial.reasoning_text, ""); // partial close held back + assert_eq!(partial.normal_text, ""); + let closed = p + .parse_reasoning_streaming_incremental("think<|sep|><|open|>response<|sep|>answer") + .unwrap(); + assert_eq!(closed.reasoning_text, ""); + assert_eq!(closed.normal_text, "answer"); // response-open stripped + } + + #[test] + fn streaming_split_message_wrapper_markers_do_not_leak_into_content() { + // Analogous to `streaming_split_close_hands_content_downstream`, but + // the marker split during the post-reasoning content phase is the + // `message` wrapper rather than the think-close marker. + let mut p = KimiK3Parser::new(); + // Reach the content phase in one shot. + p.parse_reasoning_streaming_incremental("<|open|>think<|sep|>step<|close|>think<|sep|>") + .unwrap(); + + // Split message-open marker across chunks. + let partial_open = p + .parse_reasoning_streaming_incremental("<|open|>mess") + .unwrap(); + assert_eq!(partial_open.normal_text, ""); // partial open marker held back + let after_open = p + .parse_reasoning_streaming_incremental("age<|sep|>answer") + .unwrap(); + assert_eq!(after_open.normal_text, "answer"); // message-open stripped, no leak + + // Split message-close marker across chunks. + let partial_close = p + .parse_reasoning_streaming_incremental(" more<|clo") + .unwrap(); + assert_eq!(partial_close.normal_text, " more"); // partial close marker held back + let after_close = p + .parse_reasoning_streaming_incremental("se|>message<|sep|> tail") + .unwrap(); + assert_eq!(after_close.normal_text, " tail"); // message-close stripped, no leak + } + + #[test] + fn streaming_maintains_in_reasoning_flag() { + let mut p = KimiK3Parser::new(); + // While accumulating reasoning content, the parser reports in-reasoning. + p.parse_reasoning_streaming_incremental("<|open|>think<|sep|>step") + .unwrap(); + assert!(p.is_in_reasoning()); + // A partial close marker is held back but we are still in reasoning. + p.parse_reasoning_streaming_incremental("<|close|>") + .unwrap(); + assert!(p.is_in_reasoning()); + // Completing the think-close transition clears the flag. + p.parse_reasoning_streaming_incremental("think<|sep|><|open|>response<|sep|>answer") + .unwrap(); + assert!(!p.is_in_reasoning()); + } + + #[test] + fn streaming_buffer_overflow_is_guarded() { + let mut p = KimiK3Parser::new(); + let oversized = "a".repeat(DEFAULT_MAX_BUFFER_SIZE + 1); + let result = p.parse_reasoning_streaming_incremental(&oversized); + assert!(matches!( + result, + Err(ParseError::BufferOverflow(size)) if size == DEFAULT_MAX_BUFFER_SIZE + 1 + )); + } + + #[test] + fn reset_clears_state() { + let mut p = KimiK3Parser::new(); + // Put the parser into reasoning mode via streaming. + p.parse_reasoning_streaming_incremental("<|open|>think<|sep|>step") + .unwrap(); + assert!(p.is_in_reasoning() || !p.buffer.is_empty()); + p.reset(); + assert!(!p.is_in_reasoning()); + assert!(p.buffer.is_empty()); + // After reset, a fresh non-streaming parse should work correctly. + let r = p.detect_and_parse_reasoning("just content").unwrap(); + assert_eq!(r.normal_text, "just content"); + assert_eq!(r.reasoning_text, ""); + } +} diff --git a/crates/reasoning_parser/src/parsers/mod.rs b/crates/reasoning_parser/src/parsers/mod.rs index 447942577..22baf14aa 100644 --- a/crates/reasoning_parser/src/parsers/mod.rs +++ b/crates/reasoning_parser/src/parsers/mod.rs @@ -4,6 +4,7 @@ pub mod deepseek_r1; pub mod glm45; pub mod inkling; pub mod kimi; +pub mod kimi_k3; pub mod minimax; pub mod nano_v3; pub mod passthrough; @@ -16,6 +17,7 @@ pub use deepseek_r1::DeepSeekR1Parser; pub use glm45::Glm45Parser; pub use inkling::InklingParser; pub use kimi::KimiParser; +pub use kimi_k3::KimiK3Parser; pub use minimax::MiniMaxParser; pub use nano_v3::NanoV3Parser; pub use passthrough::PassthroughParser; diff --git a/crates/tokenizer/src/encoders/kimi_k3_xtml.rs b/crates/tokenizer/src/encoders/kimi_k3_xtml.rs new file mode 100644 index 000000000..25afd8e42 --- /dev/null +++ b/crates/tokenizer/src/encoders/kimi_k3_xtml.rs @@ -0,0 +1,1118 @@ +//! Kimi-K3 XTML chat-template renderer. +//! +//! Ported from the upstream Python reference `encoding_k3.py::build_chat_segments` +//! (entry point). Unlike the Kimi-K2.5 encoder, Kimi-K3 ships **no Jinja chat +//! template**: its prompt is rendered entirely in Python into XTML using the +//! `<|open|>` / `<|close|>` / `<|sep|>` / `<|end_of_msg|>` control tokens. This +//! module reproduces that rendering so SMG's gRPC path can build the prompt +//! itself, dispatched via `Renderer::KimiK3Xtml`. +//! +//! # Encoding fidelity limitation +//! +//! The Python reference distinguishes *structural markers* (encoded as tiktoken +//! special-token IDs) from *user/tool text* (encoded as ordinary BPE) by +//! emitting a list of `EncodeSegment { text, allow_special }`. This port instead +//! concatenates every segment's `text` into a flat `String`. A later single +//! `encode` therefore treats all `<|...|>` occurrences uniformly — a known +//! limitation for user/tool text that literally contains control-token strings +//! (e.g. a user typing `<|open|>`). Matching the flat-`String` contract of the +//! sibling renderers (`kimi_k25_tools`, `deepseek_v32`) is intentional here; +//! segment-aware encoding is out of scope. +//! +//! # Deferred (TODO) +//! +//! The following `build_chat_segments` branches are intentionally not ported yet +//! (not exercised by the current gRPC path / golden fixtures): +//! - multimodal `image_prompts` substitution (image parts render as the bare +//! `<|kimi_image_placeholder|>` marker; a separate multimodal task will wire +//! real image prompts through). + +use anyhow::{anyhow, Result}; +use serde_json::{Map, Value}; + +use crate::chat_template::ChatTemplateParams; + +const OPEN_TOKEN: &str = "<|open|>"; +const CLOSE_TOKEN: &str = "<|close|>"; +const SEP_TOKEN: &str = "<|sep|>"; +const END_OF_MSG_TOKEN: &str = "<|end_of_msg|>"; +const IMAGE_PLACEHOLDER: &str = "<|kimi_image_placeholder|>"; + +/// Render a Kimi-K3 chat prompt to an XTML `String`. +/// +/// `params.tools` (when non-empty) produces the leading `tool-declare` system +/// message. `params.add_generation_prompt` appends the assistant generation +/// prompt tail. Thinking mode is resolved from `template_kwargs["thinking"]` +/// then `params.thinking`, defaulting to `true` to match the Python +/// `build_chat_segments(thinking=True)` default; it selects `think` vs +/// `response` for both the generation-prompt tail and (per-turn) any prior +/// assistant reasoning. +pub fn apply_kimi_k3_xtml(messages: &[Value], params: &ChatTemplateParams) -> Result { + // Re-sort tool results by tool_call_id, then normalize each message + // (deep-sort tool schemas, coerce tool-call arguments) — both side-effect + // free, mirroring the Python entry point. + let reordered = normalize_xtml_tool_result_messages(messages); + let normalized: Vec = reordered.iter().map(normalize_message).collect(); + + // Top-level tools declaration is deep-sorted before compact JSON encoding. + let tools_sorted: Option = params + .tools + .filter(|t| !t.is_empty()) + .map(|t| deep_sort(&Value::Array(t.to_vec()))); + + let thinking = params + .template_kwargs + .and_then(|k| k.get("thinking")) + .and_then(Value::as_bool) + .or(params.thinking) + .unwrap_or(true); + + let mut out = String::new(); + + if let Some(tools) = &tools_sorted { + push_tool_declare(&mut out, tools, false)?; + } + + // Effort directive (`thinking-effort` internal system message). Only emitted + // while thinking is on, mirroring the Python reference (both the validation + // and the emit are gated on `thinking`). + // + // Precedence: + // 1. An explicit `thinking_effort` (via `chat_template_kwargs`) wins and is + // validated strictly — a provided-but-unsupported value is a hard error, + // matching the reference `assert thinking_effort in _VALID_THINKING_EFFORTS`. + // 2. Otherwise the OpenAI-standard top-level `reasoning_effort` level is + // used, but only when it names a supported K3 effort (`low`/`high`/`max`). + // `minimal`/`none` already switch thinking off upstream and `medium` has + // no K3 equivalent, so such values emit no directive rather than erroring + // on an otherwise-valid OpenAI field. + // + // Absent both, no directive is emitted and the model applies its intrinsic + // `max` default — matching the reference, which injects nothing when + // `thinking_effort` is unset. `preserve_thinking` (which some callers send + // alongside the effort) is not read by the reference renderer and is ignored. + if thinking { + if let Some(effort_val) = params + .template_kwargs + .and_then(|k| k.get("thinking_effort")) + .filter(|v| !v.is_null()) + { + match effort_val.as_str() { + Some(effort @ ("low" | "high" | "max")) => { + push_thinking_effort(&mut out, effort); + } + _ => { + return Err(anyhow!( + "Unsupported thinking_effort={effort_val}; \ + supported values are [\"high\", \"low\", \"max\"]." + )); + } + } + } else if let Some(effort) = params + .template_kwargs + .and_then(|k| k.get("reasoning_effort")) + .and_then(Value::as_str) + .filter(|e| matches!(*e, "low" | "high" | "max")) + { + push_thinking_effort(&mut out, effort); + } + } + + // Tracks the most recent assistant `tool_calls` so tool messages can resolve + // a missing name by position, mirroring the Python module-local state. + let mut current_tool_calls: Option<&Vec> = None; + let mut tool_index: usize = 0; + + for message in &normalized { + let obj = match message.as_object() { + Some(o) => o, + None => continue, + }; + let role = obj.get("role").and_then(Value::as_str).unwrap_or(""); + match role { + "user" => { + let mut attrs = vec![("role", "user".to_string())]; + if let Some(name) = nonempty_name(obj) { + attrs.push(("name", name)); + } + push_open_tag(&mut out, "message", &attrs); + push_content(&mut out, obj.get("content")); + push_close_tag(&mut out, "message"); + push_end_of_msg(&mut out); + } + "system" => { + if let Some(tools) = obj + .get("tools") + .filter(|t| t.as_array().is_some_and(|a| !a.is_empty())) + { + // Dynamic tool declaration; already deep-sorted by + // `normalize_message`. + push_tool_declare(&mut out, tools, true)?; + } else { + let mut attrs = vec![("role", "system".to_string())]; + if let Some(name) = nonempty_name(obj) { + attrs.push(("name", name)); + } + push_open_tag(&mut out, "message", &attrs); + push_content(&mut out, obj.get("content")); + push_close_tag(&mut out, "message"); + push_end_of_msg(&mut out); + } + } + "tool" => { + tool_index += 1; + let mut tool_name = obj + .get("tool") + .and_then(Value::as_str) + .or_else(|| obj.get("name").and_then(Value::as_str)) + .map(str::to_string); + if tool_name.is_none() { + if let Some(tcs) = current_tool_calls { + if tool_index <= tcs.len() { + let tc = &tcs[tool_index - 1]; + let fnobj = function_or_self(tc); + tool_name = fnobj + .get("name") + .and_then(Value::as_str) + .map(str::to_string); + } + } + } + let tool_name = tool_name.ok_or_else(|| { + anyhow!( + "Kimi K3 tool messages need a resolvable tool name: carry `tool`/`name`, \ + or match a preceding assistant tool_call by order." + ) + })?; + push_open_tag( + &mut out, + "message", + &[ + ("role", "tool".to_string()), + ("tool", tool_name), + ("index", tool_index.to_string()), + ], + ); + push_content(&mut out, obj.get("content")); + push_close_tag(&mut out, "message"); + push_end_of_msg(&mut out); + } + "assistant" => { + current_tool_calls = obj.get("tool_calls").and_then(Value::as_array); + tool_index = 0; + let mut attrs = vec![("role", "assistant".to_string())]; + if let Some(name) = nonempty_name(obj) { + attrs.push(("name", name)); + } + push_open_tag(&mut out, "message", &attrs); + render_assistant_segments(&mut out, obj, thinking)?; + push_close_tag(&mut out, "message"); + push_end_of_msg(&mut out); + } + // Unknown roles produce no output, matching the Python loop's lack + // of a matching branch. + _ => {} + } + } + + // Post-loop `tool_choice` / `response_format` internal system messages, + // emitted after the conversation and before the generation-prompt tail to + // mirror `build_chat_segments`. + match params + .template_kwargs + .and_then(|kwargs| kwargs.get("tool_choice")) + .and_then(Value::as_str) + { + Some("required") => push_internal_system_message( + &mut out, + "tool-choice", + "The system is invoked with `tool_choice=required`.\n\ + You MUST call tools in the next message.", + ), + Some("none") => push_internal_system_message( + &mut out, + "tool-choice", + "The system is invoked with `tool_choice=none`.\n\ + You MUST NOT call any tools in the next message.", + ), + _ => {} + } + + if let Some(response_format) = params + .template_kwargs + .and_then(|kwargs| kwargs.get("response_format")) + { + push_response_format(&mut out, response_format, params.template_kwargs)?; + } + + if params.add_generation_prompt { + push_open_tag(&mut out, "message", &[("role", "assistant".to_string())]); + push_open_tag(&mut out, if thinking { "think" } else { "response" }, &[]); + } + + Ok(out) +} + +// --------------------------------------------------------------------------- +// Rendering helpers (segment text emitted straight into the output String) +// --------------------------------------------------------------------------- + +fn escape_attr_value(value: &str) -> String { + value.replace('&', "&").replace('"', """) +} + +fn push_attr(out: &mut String, key: &str, value: &str) { + out.push(' '); + out.push_str(key); + out.push_str("=\""); + out.push_str(&escape_attr_value(value)); + out.push('"'); +} + +fn push_open_tag(out: &mut String, tag: &str, attrs: &[(&str, String)]) { + out.push_str(OPEN_TOKEN); + out.push_str(tag); + for (key, value) in attrs { + push_attr(out, key, value); + } + out.push_str(SEP_TOKEN); +} + +fn push_close_tag(out: &mut String, tag: &str) { + out.push_str(CLOSE_TOKEN); + out.push_str(tag); + out.push_str(SEP_TOKEN); +} + +fn push_end_of_msg(out: &mut String) { + out.push_str(END_OF_MSG_TOKEN); +} + +/// Emit an internal `role="system"` message with the given `type` and a +/// (stripped) body, mirroring the Python `_internal_system_message` helper. +fn push_internal_system_message(out: &mut String, message_type: &str, body: &str) { + push_open_tag( + out, + "message", + &[ + ("role", "system".to_string()), + ("type", message_type.to_string()), + ], + ); + out.push_str(body.trim()); + push_close_tag(out, "message"); + push_end_of_msg(out); +} + +/// Render `content` (string, or an OpenAI content-part array) into `out`. +/// +/// Image parts emit the bare `<|kimi_image_placeholder|>` marker; real +/// `image_prompts` substitution is deferred (see module docs). +fn push_content(out: &mut String, content: Option<&Value>) { + match content { + Some(Value::String(s)) => out.push_str(s), + Some(Value::Array(parts)) => { + for part in parts { + let ty = part.get("type").and_then(Value::as_str); + if matches!(ty, Some("image") | Some("image_url")) { + out.push_str(IMAGE_PLACEHOLDER); + } else if let Some(text) = part.get("text").and_then(Value::as_str) { + out.push_str(text); + } + } + } + _ => {} + } +} + +fn push_tool_declare(out: &mut String, tools: &Value, dynamic: bool) -> Result<()> { + let compact = json_compact(tools)?; + let body = if dynamic { + format!( + "## New Tools Available\n\ + The system dynamically extends the toolset via lazy-loading.\n\ + You have access to all existing and extended tools.\n\ + Here are the specs for the extended tools.\n\n\ + ```json\n{compact}\n```" + ) + } else { + format!( + "# Tools\n\ + Here are the available tools, described in JSONSchema.\n\n\ + ```json\n{compact}\n```" + ) + }; + push_open_tag( + out, + "message", + &[ + ("role", "system".to_string()), + ("type", "tool-declare".to_string()), + ], + ); + out.push_str(&body); + push_close_tag(out, "message"); + push_end_of_msg(out); + Ok(()) +} + +/// Emit the `thinking-effort` internal system message. +/// +/// Mirrors the Python reference `_internal_system_message("thinking-effort", …)`: +/// a `role="system" type="thinking-effort"` message whose (stripped) body states +/// the requested effort. `effort` is pre-validated to `low`/`high`/`max`; the +/// body text intentionally reproduces the reference verbatim (including its +/// mention of `medium`, which the validation set does not accept). +fn push_thinking_effort(out: &mut String, effort: &str) { + let body = format!( + concat!( + "`thinking_effort` guides on how much to think in your ", + "thinking channel (not including the response channel), ", + "supported values include `low`, `medium`, `high`, and `max`.\n", + "Now the system is invoked with `thinking_effort={effort}`.", + ), + effort = effort, + ); + push_internal_system_message(out, "thinking-effort", &body); +} + +/// Emit the `response-format` internal system message for `json_object` / +/// `json_schema`, mirroring the Python reference. The schema is resolved from an +/// explicit `response_schema` template kwarg, falling back to +/// `response_format.json_schema.schema`; it is deep-sorted then compacted so the +/// output is stable and byte-identical to the reference renderer. +fn push_response_format( + out: &mut String, + response_format: &Value, + template_kwargs: Option<&std::collections::HashMap>, +) -> Result<()> { + let response_type = response_format + .get("type") + .and_then(Value::as_str) + .or_else(|| response_format.as_str()); + match response_type { + Some("json_object") => push_internal_system_message( + out, + "response-format", + "The system is invoked with `response_format=json_object`.\n\ + Your response must be raw JSON data without markdown code \ + blocks (```json) or any additional formatting.", + ), + Some("json_schema") => { + let response_schema = template_kwargs + .and_then(|kwargs| kwargs.get("response_schema")) + .or_else(|| { + response_format + .get("json_schema") + .and_then(|schema| schema.get("schema")) + }); + let schema = response_schema + .map(deep_sort) + .map(|schema| json_compact(&schema)) + .transpose()? + .unwrap_or_else(|| "null".to_string()); + let body = format!( + "The system is invoked with `response_format=json_schema`.\n\ + Your response must be raw JSON data without markdown code \ + blocks (```json) or any additional formatting.\n\ + The JSON data must match the following schema:\n\ + ```json\n{schema}\n```" + ); + push_internal_system_message(out, "response-format", &body); + } + _ => {} + } + Ok(()) +} + +fn render_assistant_segments( + out: &mut String, + msg: &Map, + thinking: bool, +) -> Result<()> { + // The `` channel is structural: in thinking mode every assistant + // message carries the open/close tags even when there is no reasoning content + // to fill in. In non-thinking mode the channel is dropped entirely. Mirrors + // the Python `_render_assistant_segments(..., thinking)`. + if thinking { + // `reasoning_content or reasoning` — Python truthiness picks the first + // non-falsy value, falling back to `reasoning` otherwise. + let reasoning = msg + .get("reasoning_content") + .filter(|v| is_truthy(v)) + .or_else(|| msg.get("reasoning")); + push_open_tag(out, "think", &[]); + if let Some(rc) = reasoning { + let rc_str = plain_string(rc); + if !rc_str.trim().is_empty() { + out.push_str(&rc_str); + } + } + push_close_tag(out, "think"); + } + + push_open_tag(out, "response", &[]); + push_content(out, msg.get("content")); + push_close_tag(out, "response"); + + if let Some(tool_calls) = msg + .get("tool_calls") + .and_then(Value::as_array) + .filter(|a| !a.is_empty()) + { + push_open_tag(out, "tools", &[]); + for (idx, tool_call) in tool_calls.iter().enumerate() { + let index = idx + 1; + let fnobj = function_or_self(tool_call); + let name = fnobj + .get("name") + .and_then(Value::as_str) + .ok_or_else(|| anyhow!("Kimi K3 tool call is missing a function name"))?; + push_open_tag( + out, + "call", + &[("tool", name.to_string()), ("index", index.to_string())], + ); + let json_block = fnobj.get("_xtml_json_block").and_then(Value::as_str); + if let Some(block) = json_block { + push_open_tag(out, "json", &[("type", "object".to_string())]); + out.push_str(block); + push_close_tag(out, "json"); + } else if let Some(args) = fnobj.get("arguments").and_then(Value::as_object) { + for (key, value) in args { + push_open_tag( + out, + "argument", + &[("key", key.clone()), ("type", xtml_type(value))], + ); + out.push_str(&xtml_value(value)); + push_close_tag(out, "argument"); + } + } + push_close_tag(out, "call"); + } + push_close_tag(out, "tools"); + } + + Ok(()) +} + +/// `tool_call.get("function", tool_call)`: use the `function` object when +/// present, otherwise treat the tool-call object itself as the function shape +/// (arguments are attached at the top level in that case). +fn function_or_self(tool_call: &Value) -> &Value { + match tool_call.get("function") { + Some(f) if f.is_object() => f, + _ => tool_call, + } +} + +fn nonempty_name(obj: &Map) -> Option { + obj.get("name") + .and_then(Value::as_str) + .filter(|s| !s.is_empty()) + .map(str::to_string) +} + +// --------------------------------------------------------------------------- +// XTML value typing +// --------------------------------------------------------------------------- + +fn xtml_type(value: &Value) -> String { + match value { + Value::Bool(_) => "boolean", + Value::Null => "null", + Value::Number(_) => "number", + Value::String(_) => "string", + Value::Object(_) => "object", + Value::Array(_) => "array", + } + .to_string() +} + +/// `_xtml_value`: strings pass through verbatim; everything else is JSON-encoded. +/// +/// The Python reference uses `json.dumps(value, ensure_ascii=False)` with +/// Python's default `(", ", ": ")` separators, whereas `serde_json` emits +/// compact `(",", ":")` separators — so object/array argument values differ in +/// separator spacing. Scalar (number/bool/null) values are identical. Argument +/// values in practice (and in the golden fixtures) are strings, so this only +/// affects rarely-seen structured argument values. +fn xtml_value(value: &Value) -> String { + match value { + Value::String(s) => s.clone(), + other => serde_json::to_string(other).unwrap_or_default(), + } +} + +/// `str(value)` for text emission: strings verbatim, everything else JSON. +fn plain_string(value: &Value) -> String { + match value { + Value::String(s) => s.clone(), + other => serde_json::to_string(other).unwrap_or_default(), + } +} + +fn is_truthy(value: &Value) -> bool { + match value { + Value::Null => false, + Value::Bool(b) => *b, + Value::String(s) => !s.is_empty(), + Value::Number(n) => n.as_f64().is_none_or(|f| f != 0.0), + Value::Array(a) => !a.is_empty(), + Value::Object(o) => !o.is_empty(), + } +} + +/// `str(scalar)` used for tool-call-id keys. +fn scalar_str(value: &Value) -> String { + match value { + Value::String(s) => s.clone(), + Value::Bool(true) => "True".to_string(), + Value::Bool(false) => "False".to_string(), + Value::Null => "None".to_string(), + other => other.to_string(), + } +} + +// --------------------------------------------------------------------------- +// JSON helpers +// --------------------------------------------------------------------------- + +/// Compact JSON with `serde_json`'s `(",", ":")` separators and no non-ASCII +/// escaping — matching `json.dumps(..., ensure_ascii=False, separators=(",", ":"))`. +/// Callers deep-sort the value beforehand so object keys serialize in sorted +/// order despite the crate's `preserve_order` feature. +fn json_compact(value: &Value) -> Result { + serde_json::to_string(value).map_err(Into::into) +} + +fn empty_object() -> Value { + Value::Object(Map::new()) +} + +/// Recursively sort object keys (ascending); array order is preserved. +/// +/// Required because the crate enables `serde_json`'s `preserve_order` feature, +/// so serialization would otherwise keep the caller's insertion order. +fn deep_sort(value: &Value) -> Value { + match value { + Value::Object(map) => { + let mut entries: Vec<(&String, &Value)> = map.iter().collect(); + entries.sort_by_key(|entry| entry.0); + let mut sorted = Map::with_capacity(entries.len()); + for (key, val) in entries { + sorted.insert(key.clone(), deep_sort(val)); + } + Value::Object(sorted) + } + Value::Array(arr) => Value::Array(arr.iter().map(deep_sort).collect()), + other => other.clone(), + } +} + +// --------------------------------------------------------------------------- +// Message normalization (side-effect free — inputs are never mutated) +// --------------------------------------------------------------------------- + +/// Coerce a tool call's `arguments` into `(object, optional_raw_json_block)`. +/// +/// Mirrors Python's `normalize_tool_arguments`, but is infallible: malformed +/// argument types that the Python reference would `raise` on instead degrade to +/// empty arguments (or, for unparsable strings, a raw JSON block), so a single +/// bad tool call cannot abort the whole render. +fn normalize_tool_arguments(arguments: Option<&Value>) -> (Value, Option) { + match arguments { + None | Some(Value::Null) => (empty_object(), None), + Some(Value::Object(m)) => (Value::Object(m.clone()), None), + Some(Value::String(s)) => { + if s.trim().is_empty() { + return (empty_object(), None); + } + match serde_json::from_str::(s) { + Ok(Value::Object(m)) => (Value::Object(m), None), + // Non-object JSON (Python raises); keep the raw text as a block. + Ok(_) => (empty_object(), Some(s.clone())), + // Unparsable (Python's `except: return {}, arguments`). + Err(_) => (empty_object(), Some(s.clone())), + } + } + // Non-str/non-dict (Python raises TypeError); drop to empty arguments. + Some(_) => (empty_object(), None), + } +} + +/// Deep-sort a message's `tools` and normalize any `tool_calls[].arguments`. +fn normalize_message(message: &Value) -> Value { + let obj = match message.as_object() { + Some(o) => o, + None => return message.clone(), + }; + let mut normalized = obj.clone(); + + if let Some(tools) = normalized.get("tools") { + if !tools.is_null() { + let sorted = deep_sort(tools); + normalized.insert("tools".to_string(), sorted); + } + } + + let tool_calls = match normalized.get("tool_calls").and_then(Value::as_array) { + Some(tc) if !tc.is_empty() => tc.clone(), + _ => return Value::Object(normalized), + }; + + let mut normalized_calls: Vec = Vec::with_capacity(tool_calls.len()); + for tool_call in &tool_calls { + let tc_obj = match tool_call.as_object() { + Some(o) => o, + None => { + normalized_calls.push(tool_call.clone()); + continue; + } + }; + let mut tc = tc_obj.clone(); + match tc.get("function").and_then(|f| f.as_object()).cloned() { + Some(mut fnmap) => { + let (args, json_block) = normalize_tool_arguments(fnmap.get("arguments")); + fnmap.insert("arguments".to_string(), args); + apply_json_block(&mut fnmap, json_block); + tc.insert("function".to_string(), Value::Object(fnmap)); + } + None => { + let (args, json_block) = normalize_tool_arguments(tc.get("arguments")); + tc.insert("arguments".to_string(), args); + apply_json_block(&mut tc, json_block); + } + } + normalized_calls.push(Value::Object(tc)); + } + normalized.insert("tool_calls".to_string(), Value::Array(normalized_calls)); + Value::Object(normalized) +} + +fn apply_json_block(map: &mut Map, json_block: Option) { + match json_block { + None => { + map.remove("_xtml_json_block"); + } + Some(block) => { + map.insert("_xtml_json_block".to_string(), Value::String(block)); + } + } +} + +/// Map assistant `tool_calls[].id` to `(1-based position, function name)`. +/// +/// Every entry advances the position (even an id-less one); duplicate ids keep +/// their first occurrence. +fn tool_call_id_index(tool_calls: &Value) -> Vec<(String, usize, Option)> { + let mut index: Vec<(String, usize, Option)> = Vec::new(); + let arr = match tool_calls.as_array() { + Some(a) => a, + None => return index, + }; + for (pos0, tool_call) in arr.iter().enumerate() { + let position = pos0 + 1; + let tc_obj = match tool_call.as_object() { + Some(o) => o, + None => continue, + }; + let call_id = match tc_obj.get("id") { + Some(v) if !v.is_null() => v, + _ => continue, + }; + let key = scalar_str(call_id); + if index.iter().any(|(k, _, _)| k == &key) { + continue; + } + let name = match tc_obj.get("function").and_then(|f| f.as_object()) { + Some(f) => f.get("name").and_then(Value::as_str).map(str::to_string), + None => tc_obj + .get("name") + .and_then(Value::as_str) + .map(str::to_string), + }; + index.push((key, position, name)); + } + index +} + +fn lookup_index<'a>( + index: &'a [(String, usize, Option)], + key: &str, +) -> Option<&'a (String, usize, Option)> { + index.iter().find(|(k, _, _)| k == key) +} + +/// Re-sort each run of consecutive `tool` messages into the most recent +/// assistant `tool_calls` order, matching by `tool_call_id == tool_calls[].id`. +/// +/// A fully-matched run is sorted by the matched 1-based position (stable on +/// original offset) and each matched message's `tool`/`name` is rewritten to the +/// authoritative call name. A run that cannot be fully matched is left +/// untouched. Side-effect free: matched messages are shallow-copied; every other +/// message is cloned through unchanged. Re-running is idempotent. +fn normalize_xtml_tool_result_messages(messages: &[Value]) -> Vec { + let mut output: Vec = Vec::with_capacity(messages.len()); + let mut current_index: Vec<(String, usize, Option)> = Vec::new(); + let n = messages.len(); + let mut i = 0; + + while i < n { + let message = &messages[i]; + let role = role_of(message); + + if role == Some("assistant") { + current_index = match message.get("tool_calls") { + Some(v) if v.as_array().is_some_and(|a| !a.is_empty()) => tool_call_id_index(v), + _ => Vec::new(), + }; + output.push(message.clone()); + i += 1; + continue; + } + + if role != Some("tool") { + output.push(message.clone()); + i += 1; + continue; + } + + // Gather a run of consecutive tool messages. + let mut run: Vec<(Option, usize, &Value, Option)> = Vec::new(); + let mut unresolved = false; + let mut offset = 0; + while i < n && role_of(&messages[i]) == Some("tool") { + let tool_message = &messages[i]; + let call_id = tool_message + .get("tool_call_id") + .or_else(|| tool_message.get("id")) + .filter(|v| !v.is_null()); + let matched = call_id + .map(scalar_str) + .and_then(|k| lookup_index(¤t_index, &k).cloned()); + match matched { + None => { + unresolved = true; + run.push((None, offset, tool_message, None)); + } + Some((_, position, name)) => { + run.push((Some(position), offset, tool_message, name)); + } + } + offset += 1; + i += 1; + } + + if unresolved { + for item in &run { + output.push(item.2.clone()); + } + } else { + run.sort_by_key(|item| (item.0, item.1)); + for (_, _, tool_message, name) in run { + match name { + None => output.push(tool_message.clone()), + Some(nm) => { + let mut resolved = tool_message.as_object().cloned().unwrap_or_default(); + resolved.insert("tool".to_string(), Value::String(nm.clone())); + if resolved.contains_key("name") { + resolved.insert("name".to_string(), Value::String(nm)); + } + output.push(Value::Object(resolved)); + } + } + } + } + } + + output +} + +fn role_of(message: &Value) -> Option<&str> { + message + .as_object() + .and_then(|o| o.get("role")) + .and_then(Value::as_str) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use serde_json::json; + + use super::*; + + fn params(thinking: Option, tools: Option<&[Value]>) -> ChatTemplateParams<'_> { + ChatTemplateParams { + add_generation_prompt: true, + tools, + thinking, + ..Default::default() + } + } + + fn params_kw( + thinking: Option, + template_kwargs: &HashMap, + add_generation_prompt: bool, + ) -> ChatTemplateParams<'_> { + ChatTemplateParams { + add_generation_prompt, + thinking, + template_kwargs: Some(template_kwargs), + ..Default::default() + } + } + + #[test] + fn thinking_defaults_on_when_unspecified() { + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let rendered = apply_kimi_k3_xtml(&messages, ¶ms(None, None)).unwrap(); + assert!( + rendered.ends_with("<|open|>think<|sep|>"), + "got: {rendered}" + ); + } + + #[test] + fn deep_sort_orders_nested_keys() { + let value = json!({"b": 1, "a": {"d": 2, "c": 3}}); + let sorted = json_compact(&deep_sort(&value)).unwrap(); + assert_eq!(sorted, r#"{"a":{"c":3,"d":2},"b":1}"#); + } + + #[test] + fn attr_values_are_escaped() { + let mut out = String::new(); + push_attr(&mut out, "name", "a&b\"c"); + assert_eq!(out, " name=\"a&b"c\""); + } + + #[test] + fn string_arguments_render_verbatim() { + let messages = vec![json!({ + "role": "assistant", + "content": "", + "tool_calls": [{ + "function": {"name": "f", "arguments": "{\"k\": \"v\"}"} + }] + })]; + let rendered = apply_kimi_k3_xtml(&messages, ¶ms(Some(true), None)).unwrap(); + assert!( + rendered.contains( + "<|open|>argument key=\"k\" type=\"string\"<|sep|>v<|close|>argument<|sep|>" + ), + "got: {rendered}" + ); + } + + // --- Effort directive: reasoning_effort -> thinking_effort bridge --------- + + #[test] + fn reasoning_effort_emits_thinking_effort_directive() { + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let kwargs = HashMap::from([("reasoning_effort".to_string(), json!("high"))]); + let rendered = apply_kimi_k3_xtml(&messages, ¶ms_kw(None, &kwargs, true)).unwrap(); + assert!( + rendered.contains("<|open|>message role=\"system\" type=\"thinking-effort\"<|sep|>"), + "missing thinking-effort message: {rendered}" + ); + assert!( + rendered.contains( + "Now the system is invoked with `thinking_effort=high`.<|close|>message<|sep|>" + ), + "wrong effort level: {rendered}" + ); + } + + #[test] + fn explicit_thinking_effort_wins_over_reasoning_effort() { + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let kwargs = HashMap::from([ + ("reasoning_effort".to_string(), json!("high")), + ("thinking_effort".to_string(), json!("low")), + ]); + let rendered = apply_kimi_k3_xtml(&messages, ¶ms_kw(None, &kwargs, true)).unwrap(); + assert!( + rendered.contains("Now the system is invoked with `thinking_effort=low`."), + "explicit thinking_effort should win: {rendered}" + ); + assert!( + !rendered.contains("thinking_effort=high"), + "reasoning_effort should not leak when overridden: {rendered}" + ); + } + + #[test] + fn unsupported_reasoning_effort_emits_no_directive() { + // `medium` is a valid OpenAI reasoning_effort but has no K3 equivalent: + // it must be ignored (no directive), never error like an explicit value. + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let kwargs = HashMap::from([("reasoning_effort".to_string(), json!("medium"))]); + let rendered = apply_kimi_k3_xtml(&messages, ¶ms_kw(None, &kwargs, true)).unwrap(); + assert!( + !rendered.contains("type=\"thinking-effort\""), + "medium must not emit a directive: {rendered}" + ); + } + + #[test] + fn unsupported_explicit_thinking_effort_errors() { + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let kwargs = HashMap::from([("thinking_effort".to_string(), json!("medium"))]); + assert!(apply_kimi_k3_xtml(&messages, ¶ms_kw(None, &kwargs, true)).is_err()); + } + + #[test] + fn effort_directive_suppressed_when_thinking_off() { + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let kwargs = HashMap::from([("reasoning_effort".to_string(), json!("low"))]); + let rendered = + apply_kimi_k3_xtml(&messages, ¶ms_kw(Some(false), &kwargs, true)).unwrap(); + assert!( + !rendered.contains("type=\"thinking-effort\""), + "no effort directive when thinking is off: {rendered}" + ); + } + + #[test] + fn no_effort_directive_when_unspecified() { + // Byte-parity with the reference: absent both keys, nothing is injected + // (the model applies its intrinsic `max` default). + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let kwargs = HashMap::new(); + let rendered = apply_kimi_k3_xtml(&messages, ¶ms_kw(None, &kwargs, true)).unwrap(); + assert!( + !rendered.contains("type=\"thinking-effort\""), + "no directive should be injected by default: {rendered}" + ); + } + + // --- Structural channel ------------------------------------------ + + #[test] + fn think_channel_is_structural_when_thinking() { + // Assistant history message with no reasoning still carries empty tags. + let messages = vec![ + json!({"role": "user", "content": "Hi"}), + json!({"role": "assistant", "content": "ok"}), + ]; + let kwargs = HashMap::new(); + let rendered = + apply_kimi_k3_xtml(&messages, ¶ms_kw(Some(true), &kwargs, false)).unwrap(); + assert!( + rendered.contains( + "<|open|>think<|sep|><|close|>think<|sep|>\ + <|open|>response<|sep|>ok<|close|>response<|sep|>" + ), + "empty think channel must still be emitted: {rendered}" + ); + } + + #[test] + fn think_channel_dropped_when_not_thinking() { + let messages = vec![ + json!({"role": "user", "content": "Hi"}), + json!({"role": "assistant", "content": "ok"}), + ]; + let kwargs = HashMap::new(); + let rendered = + apply_kimi_k3_xtml(&messages, ¶ms_kw(Some(false), &kwargs, false)).unwrap(); + assert!( + !rendered.contains("<|open|>think<|sep|>"), + "think channel must be dropped in non-thinking mode: {rendered}" + ); + assert!( + rendered.contains("<|open|>response<|sep|>ok<|close|>response<|sep|>"), + "response channel must still render: {rendered}" + ); + } + + // --- tool_choice / response_format internal messages --------------------- + + #[test] + fn tool_choice_required_emits_internal_message() { + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let kwargs = HashMap::from([("tool_choice".to_string(), json!("required"))]); + let rendered = apply_kimi_k3_xtml(&messages, ¶ms_kw(None, &kwargs, true)).unwrap(); + assert!( + rendered.contains( + "<|open|>message role=\"system\" type=\"tool-choice\"<|sep|>\ + The system is invoked with `tool_choice=required`.\n\ + You MUST call tools in the next message.<|close|>message<|sep|>" + ), + "got: {rendered}" + ); + } + + #[test] + fn tool_choice_none_emits_internal_message() { + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let kwargs = HashMap::from([("tool_choice".to_string(), json!("none"))]); + let rendered = apply_kimi_k3_xtml(&messages, ¶ms_kw(None, &kwargs, true)).unwrap(); + assert!( + rendered.contains( + "The system is invoked with `tool_choice=none`.\n\ + You MUST NOT call any tools in the next message." + ), + "got: {rendered}" + ); + } + + #[test] + fn tool_choice_auto_emits_nothing() { + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let kwargs = HashMap::from([("tool_choice".to_string(), json!("auto"))]); + let rendered = apply_kimi_k3_xtml(&messages, ¶ms_kw(None, &kwargs, true)).unwrap(); + assert!( + !rendered.contains("type=\"tool-choice\""), + "got: {rendered}" + ); + } + + #[test] + fn response_format_json_object_emits_internal_message() { + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let kwargs = HashMap::from([( + "response_format".to_string(), + json!({"type": "json_object"}), + )]); + let rendered = apply_kimi_k3_xtml(&messages, ¶ms_kw(None, &kwargs, true)).unwrap(); + assert!( + rendered.contains( + "<|open|>message role=\"system\" type=\"response-format\"<|sep|>\ + The system is invoked with `response_format=json_object`." + ), + "got: {rendered}" + ); + } + + #[test] + fn response_format_json_schema_emits_sorted_schema() { + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let kwargs = HashMap::from([( + "response_format".to_string(), + json!({ + "type": "json_schema", + "json_schema": { + "name": "x", + "schema": {"type": "object", "properties": {"a": {"type": "string"}}} + } + }), + )]); + let rendered = apply_kimi_k3_xtml(&messages, ¶ms_kw(None, &kwargs, true)).unwrap(); + assert!( + rendered.contains("The system is invoked with `response_format=json_schema`."), + "got: {rendered}" + ); + assert!( + rendered.contains( + "```json\n{\"properties\":{\"a\":{\"type\":\"string\"}},\"type\":\"object\"}\n```" + ), + "schema must be deep-sorted and compacted: {rendered}" + ); + } +} diff --git a/crates/tokenizer/src/encoders/mod.rs b/crates/tokenizer/src/encoders/mod.rs index a20d380e0..b1283d246 100644 --- a/crates/tokenizer/src/encoders/mod.rs +++ b/crates/tokenizer/src/encoders/mod.rs @@ -1,3 +1,4 @@ pub mod deepseek_v32; pub mod deepseek_v4; pub mod kimi_k25_tools; +pub mod kimi_k3_xtml; diff --git a/crates/tokenizer/src/kimi_k2_tokenizer.rs b/crates/tokenizer/src/kimi_k2_tokenizer.rs index 7f24dc5aa..8e710a458 100644 --- a/crates/tokenizer/src/kimi_k2_tokenizer.rs +++ b/crates/tokenizer/src/kimi_k2_tokenizer.rs @@ -31,7 +31,7 @@ pub(crate) const KIMI_K2_PATTERN: &str = r"[\p{Han}]+|[^\r\n\p{L}\p{N}]?[\p{Lu}\ /// `tokenization_kimi` (via `auto_map`, `tokenizer_class`, etc.). Callers pass /// the parsed JSON so we don't re-read the file the tiktoken loader already /// parsed. Fallback: read sibling `config.json` and check `model_type` ∈ -/// `{kimi_k2, kimi_k25}`. +/// `{kimi_k2, kimi_k25, kimi_k3}`. pub(crate) fn matches(tokenizer_config: Option<&Value>, dir: &Path) -> bool { if tokenizer_config.is_some_and(value_mentions_kimi_tokenizer) { return true; @@ -77,7 +77,10 @@ fn read_json(path: &Path) -> Option { fn model_config_is_kimi(config: &Value) -> bool { let model_type = config.get("model_type").and_then(Value::as_str); - matches!(model_type, Some("kimi_k2") | Some("kimi_k25")) + matches!( + model_type, + Some("kimi_k2") | Some("kimi_k25") | Some("kimi_k3") + ) } fn value_mentions_kimi_tokenizer(value: &Value) -> bool { @@ -219,6 +222,24 @@ mod tests { assert!(matches(tokenizer_config(dir.path()).as_ref(), dir.path())); } + #[test] + fn matches_via_model_type_kimi_k3() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("tiktoken.model"), MINIMAL_TIKTOKEN_MODEL).unwrap(); + std::fs::write( + dir.path().join("tokenizer_config.json"), + r#"{ "added_tokens_decoder": {} }"#, + ) + .unwrap(); + std::fs::write( + dir.path().join("config.json"), + r#"{ "model_type": "kimi_k3" }"#, + ) + .unwrap(); + + assert!(matches(tokenizer_config(dir.path()).as_ref(), dir.path())); + } + #[test] fn substring_does_not_falsely_match_kimi() { // Names that *contain* "tokenization_kimi" as a substring but aren't diff --git a/crates/tokenizer/src/tiktoken.rs b/crates/tokenizer/src/tiktoken.rs index 754c75346..071088bc3 100644 --- a/crates/tokenizer/src/tiktoken.rs +++ b/crates/tokenizer/src/tiktoken.rs @@ -17,7 +17,7 @@ use crate::{ load_chat_template_from_file, ChatTemplateContentFormat, ChatTemplateParams, ChatTemplateState, ThinkingKeyName, ThinkingToggle, }, - encoders::kimi_k25_tools::apply_kimi_k25_tools, + encoders::{kimi_k25_tools::apply_kimi_k25_tools, kimi_k3_xtml::apply_kimi_k3_xtml}, factory::discover_chat_template_in_dir, kimi_k2_tokenizer, traits::{Decoder, Encoder, Encoding, SpecialTokens, TokenIdType, Tokenizer as TokenizerTrait}, @@ -27,6 +27,9 @@ use crate::{ enum Renderer { Jinja, KimiK25Tools, + /// Kimi-K3 renders its prompt entirely in the native XTML encoder (no Jinja + /// chat template). See [`crate::encoders::kimi_k3_xtml`]. + KimiK3Xtml, } /// Regex pattern for cl100k_base tokenization. @@ -551,6 +554,7 @@ impl TokenizerTrait for TiktokenTokenizer { match self.renderer { Renderer::Jinja => self.chat_template.apply(messages, params), Renderer::KimiK25Tools => apply_kimi_k25_tools(&self.chat_template, messages, params), + Renderer::KimiK3Xtml => apply_kimi_k3_xtml(messages, ¶ms), } } @@ -559,18 +563,31 @@ impl TokenizerTrait for TiktokenTokenizer { } fn thinking_toggle(&self) -> ThinkingToggle { - self.chat_template.thinking_toggle() + match self.renderer { + // K3 has no Jinja template; its native renderer defaults thinking + // ON (Python `build_chat_segments(thinking=True)`). + Renderer::KimiK3Xtml => ThinkingToggle::DefaultOn, + _ => self.chat_template.thinking_toggle(), + } } fn thinking_key_name(&self) -> Option { - self.chat_template.thinking_key_name() + match self.renderer { + Renderer::KimiK3Xtml => Some(ThinkingKeyName::Thinking), + _ => self.chat_template.thinking_key_name(), + } } fn eos_token_ids(&self) -> &[TokenIdType] { &self.eos_token_ids } fn think_in_prefill(&self) -> bool { - self.chat_template.think_in_prefill() + match self.renderer { + // K3's generation-prompt tail opens `` when thinking is on + // (the default), so completions start mid-reasoning. + Renderer::KimiK3Xtml => true, + _ => self.chat_template.think_in_prefill(), + } } fn set_chat_template(&mut self, template: String) -> Result<()> { @@ -603,14 +620,16 @@ fn detect_renderer_from_config(dir: &Path) -> Renderer { return Renderer::Jinja; } }; - let is_kimi = value + let arch_strs: Vec<&str> = value .get("architectures") .and_then(|v| v.as_array()) - .is_some_and(|a| { - a.iter() - .any(|v| v.as_str() == Some("KimiK25ForConditionalGeneration")) - }); - if is_kimi { + .map(|a| a.iter().filter_map(|v| v.as_str()).collect()) + .unwrap_or_default(); + if arch_strs.contains(&"KimiK3ForConditionalGeneration") { + tracing::debug!(?path, "selected KimiK3Xtml chat-template renderer"); + return Renderer::KimiK3Xtml; + } + if arch_strs.contains(&"KimiK25ForConditionalGeneration") { tracing::debug!(?path, "selected KimiK25Tools chat-template renderer"); return Renderer::KimiK25Tools; } @@ -637,6 +656,40 @@ mod tests { dir } + #[test] + fn test_detect_renderer_from_config() { + let k3 = write_minimal_tiktoken_dir( + "{}", + Some(r#"{"architectures": ["KimiK3ForConditionalGeneration"]}"#), + ); + assert!(matches!( + detect_renderer_from_config(k3.path()), + Renderer::KimiK3Xtml + )); + + let k25 = write_minimal_tiktoken_dir( + "{}", + Some(r#"{"architectures": ["KimiK25ForConditionalGeneration"]}"#), + ); + assert!(matches!( + detect_renderer_from_config(k25.path()), + Renderer::KimiK25Tools + )); + + let other = + write_minimal_tiktoken_dir("{}", Some(r#"{"architectures": ["LlamaForCausalLM"]}"#)); + assert!(matches!( + detect_renderer_from_config(other.path()), + Renderer::Jinja + )); + + let missing = write_minimal_tiktoken_dir("{}", None); + assert!(matches!( + detect_renderer_from_config(missing.path()), + Renderer::Jinja + )); + } + #[test] fn test_tiktoken_creation() { let tokenizer = TiktokenTokenizer::new(TiktokenModel::Cl100kBase).unwrap(); diff --git a/crates/tokenizer/tests/fixtures/kimi_k3/k3_render_fixtures.json b/crates/tokenizer/tests/fixtures/kimi_k3/k3_render_fixtures.json new file mode 100644 index 000000000..c62183b9d --- /dev/null +++ b/crates/tokenizer/tests/fixtures/kimi_k3/k3_render_fixtures.json @@ -0,0 +1,1349 @@ +{ + "plain_user_thinking": { + "text": "<|open|>message role=\"user\"<|sep|>Hi<|close|>message<|sep|><|end_of_msg|><|open|>message role=\"assistant\"<|sep|><|open|>think<|sep|>", + "segments": [ + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "user", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "Hi", + "allow_special": false + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|end_of_msg|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "assistant", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "think", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + } + ] + }, + "system_user_thinking": { + "text": "<|open|>message role=\"system\"<|sep|>You are helpful<|close|>message<|sep|><|end_of_msg|><|open|>message role=\"user\"<|sep|>Hi<|close|>message<|sep|><|end_of_msg|><|open|>message role=\"assistant\"<|sep|><|open|>think<|sep|>", + "segments": [ + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "system", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "You are helpful", + "allow_special": false + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|end_of_msg|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "user", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "Hi", + "allow_special": false + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|end_of_msg|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "assistant", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "think", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + } + ] + }, + "plain_user_no_thinking": { + "text": "<|open|>message role=\"user\"<|sep|>Hi<|close|>message<|sep|><|end_of_msg|><|open|>message role=\"assistant\"<|sep|><|open|>response<|sep|>", + "segments": [ + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "user", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "Hi", + "allow_special": false + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|end_of_msg|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "assistant", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "response", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + } + ] + }, + "assistant_prior_turn": { + "text": "<|open|>message role=\"user\"<|sep|>Hi<|close|>message<|sep|><|end_of_msg|><|open|>message role=\"assistant\"<|sep|><|open|>think<|sep|><|close|>think<|sep|><|open|>response<|sep|>Hello!<|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|><|open|>message role=\"user\"<|sep|>Bye<|close|>message<|sep|><|end_of_msg|><|open|>message role=\"assistant\"<|sep|><|open|>think<|sep|>", + "segments": [ + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "user", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "Hi", + "allow_special": false + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|end_of_msg|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "assistant", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "think", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "think", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "response", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "Hello!", + "allow_special": false + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "response", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|end_of_msg|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "user", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "Bye", + "allow_special": false + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|end_of_msg|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "assistant", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "think", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + } + ] + }, + "with_tools": { + "text": "<|open|>message role=\"system\" type=\"tool-declare\"<|sep|># Tools\nHere are the available tools, described in JSONSchema.\n\n```json\n[{\"function\":{\"name\":\"get_weather\",\"parameters\":{\"properties\":{\"city\":{\"type\":\"string\"}},\"type\":\"object\"}},\"type\":\"function\"}]\n```<|close|>message<|sep|><|end_of_msg|><|open|>message role=\"user\"<|sep|>weather in Paris?<|close|>message<|sep|><|end_of_msg|><|open|>message role=\"assistant\"<|sep|><|open|>think<|sep|>", + "segments": [ + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "system", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": " type", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "tool-declare", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "# Tools\nHere are the available tools, described in JSONSchema.\n\n```json\n[{\"function\":{\"name\":\"get_weather\",\"parameters\":{\"properties\":{\"city\":{\"type\":\"string\"}},\"type\":\"object\"}},\"type\":\"function\"}]\n```", + "allow_special": false + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|end_of_msg|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "user", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "weather in Paris?", + "allow_special": false + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|end_of_msg|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "assistant", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "think", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + } + ] + }, + "assistant_tool_call_then_result": { + "text": "<|open|>message role=\"system\" type=\"tool-declare\"<|sep|># Tools\nHere are the available tools, described in JSONSchema.\n\n```json\n[{\"function\":{\"name\":\"get_weather\",\"parameters\":{\"properties\":{\"city\":{\"type\":\"string\"}},\"type\":\"object\"}},\"type\":\"function\"}]\n```<|close|>message<|sep|><|end_of_msg|><|open|>message role=\"user\"<|sep|>weather?<|close|>message<|sep|><|end_of_msg|><|open|>message role=\"assistant\"<|sep|><|open|>think<|sep|><|close|>think<|sep|><|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"city\" type=\"string\"<|sep|>Paris<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|><|open|>message role=\"tool\" tool=\"get_weather\" index=\"1\"<|sep|>sunny<|close|>message<|sep|><|end_of_msg|><|open|>message role=\"assistant\"<|sep|><|open|>think<|sep|>", + "segments": [ + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "system", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": " type", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "tool-declare", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "# Tools\nHere are the available tools, described in JSONSchema.\n\n```json\n[{\"function\":{\"name\":\"get_weather\",\"parameters\":{\"properties\":{\"city\":{\"type\":\"string\"}},\"type\":\"object\"}},\"type\":\"function\"}]\n```", + "allow_special": false + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|end_of_msg|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "user", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "weather?", + "allow_special": false + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|end_of_msg|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "assistant", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "think", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "think", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "response", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "response", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "tools", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "call", + "allow_special": false + }, + { + "text": " tool", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "get_weather", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": " index", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "1", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "argument", + "allow_special": false + }, + { + "text": " key", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "city", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": " type", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "string", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "Paris", + "allow_special": false + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "argument", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "call", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "tools", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|end_of_msg|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "tool", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": " tool", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "get_weather", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": " index", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "1", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "sunny", + "allow_special": false + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|end_of_msg|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "assistant", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "think", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + } + ] + }, + "thinking_effort_low": { + "text": "<|open|>message role=\"system\" type=\"thinking-effort\"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`.\nNow the system is invoked with `thinking_effort=low`.<|close|>message<|sep|><|end_of_msg|><|open|>message role=\"user\"<|sep|>Hi<|close|>message<|sep|><|end_of_msg|><|open|>message role=\"assistant\"<|sep|><|open|>think<|sep|>", + "segments": [ + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "system", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": " type", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "thinking-effort", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`.\nNow the system is invoked with `thinking_effort=low`.", + "allow_special": false + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|end_of_msg|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "user", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "Hi", + "allow_special": false + }, + { + "text": "<|close|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|end_of_msg|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "message", + "allow_special": false + }, + { + "text": " role", + "allow_special": false + }, + { + "text": "=\"", + "allow_special": false + }, + { + "text": "assistant", + "allow_special": false + }, + { + "text": "\"", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + }, + { + "text": "<|open|>", + "allow_special": true + }, + { + "text": "think", + "allow_special": false + }, + { + "text": "<|sep|>", + "allow_special": true + } + ] + } +} diff --git a/crates/tokenizer/tests/kimi_k3_renderer.rs b/crates/tokenizer/tests/kimi_k3_renderer.rs new file mode 100644 index 000000000..080e39cb2 --- /dev/null +++ b/crates/tokenizer/tests/kimi_k3_renderer.rs @@ -0,0 +1,209 @@ +//! Golden tests for the Kimi-K3 XTML chat-template renderer. +//! +//! Each case builds the equivalent messages/tools/params in Rust and asserts +//! the rendered `String` equals the `text` field of the corresponding entry in +//! `tests/fixtures/kimi_k3/k3_render_fixtures.json` byte-for-byte. Those +//! fixtures are the authoritative expected outputs produced by the upstream +//! Python `encoding_k3.py::build_chat_segments`. + +#![allow(clippy::expect_used, clippy::unwrap_used, clippy::panic)] + +use std::{collections::HashMap, fs}; + +use llm_tokenizer::{ + chat_template::ChatTemplateParams, encoders::kimi_k3_xtml::apply_kimi_k3_xtml, + traits::Tokenizer as TokenizerTrait, TiktokenTokenizer, +}; +use serde_json::{json, Value}; +use tempfile::TempDir; + +const MIN_TIKTOKEN_MODEL: &str = "aGVsbG8= 0\n"; + +/// Load a single fixture's expected `text` by case name. +fn fixture_text(case: &str) -> String { + let raw = fs::read_to_string( + std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests/fixtures/kimi_k3/k3_render_fixtures.json"), + ) + .expect("k3 fixtures must exist"); + let value: Value = serde_json::from_str(&raw).expect("fixtures must be valid JSON"); + value + .get(case) + .and_then(|c| c.get("text")) + .and_then(Value::as_str) + .unwrap_or_else(|| panic!("fixture case `{case}` missing text")) + .to_string() +} + +fn render(messages: &[Value], tools: Option<&[Value]>, thinking: bool) -> String { + let params = ChatTemplateParams { + add_generation_prompt: true, + tools, + thinking: Some(thinking), + ..Default::default() + }; + apply_kimi_k3_xtml(messages, ¶ms).expect("k3 render should succeed") +} + +fn get_weather_tools() -> Vec { + vec![json!({ + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}} + } + } + })] +} + +#[test] +fn plain_user_thinking() { + let messages = vec![json!({"role": "user", "content": "Hi"})]; + assert_eq!( + render(&messages, None, true), + fixture_text("plain_user_thinking") + ); +} + +#[test] +fn system_user_thinking() { + let messages = vec![ + json!({"role": "system", "content": "You are helpful"}), + json!({"role": "user", "content": "Hi"}), + ]; + assert_eq!( + render(&messages, None, true), + fixture_text("system_user_thinking") + ); +} + +#[test] +fn plain_user_no_thinking() { + let messages = vec![json!({"role": "user", "content": "Hi"})]; + assert_eq!( + render(&messages, None, false), + fixture_text("plain_user_no_thinking") + ); +} + +#[test] +fn assistant_prior_turn() { + let messages = vec![ + json!({"role": "user", "content": "Hi"}), + json!({"role": "assistant", "content": "Hello!"}), + json!({"role": "user", "content": "Bye"}), + ]; + assert_eq!( + render(&messages, None, true), + fixture_text("assistant_prior_turn") + ); +} + +#[test] +fn with_tools() { + let messages = vec![json!({"role": "user", "content": "weather in Paris?"})]; + let tools = get_weather_tools(); + assert_eq!( + render(&messages, Some(&tools), true), + fixture_text("with_tools") + ); +} + +#[test] +fn assistant_tool_call_then_result() { + let messages = vec![ + json!({"role": "user", "content": "weather?"}), + json!({ + "role": "assistant", + "content": "", + "tool_calls": [{ + "function": {"name": "get_weather", "arguments": {"city": "Paris"}} + }] + }), + json!({"role": "tool", "tool": "get_weather", "content": "sunny"}), + ]; + let tools = get_weather_tools(); + assert_eq!( + render(&messages, Some(&tools), true), + fixture_text("assistant_tool_call_then_result") + ); +} + +#[test] +fn thinking_effort_low() { + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let template_kwargs = HashMap::from([("thinking_effort".to_string(), json!("low"))]); + let params = ChatTemplateParams { + add_generation_prompt: true, + thinking: Some(true), + template_kwargs: Some(&template_kwargs), + ..Default::default() + }; + let rendered = apply_kimi_k3_xtml(&messages, ¶ms).expect("k3 render should succeed"); + assert_eq!(rendered, fixture_text("thinking_effort_low")); +} + +#[test] +fn thinking_effort_ignored_when_thinking_off() { + // Reference gates both validation and emission on `thinking`, so an effort + // provided while thinking is off produces no thinking-effort message. + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let template_kwargs = HashMap::from([("thinking_effort".to_string(), json!("low"))]); + let params = ChatTemplateParams { + add_generation_prompt: true, + thinking: Some(false), + template_kwargs: Some(&template_kwargs), + ..Default::default() + }; + let rendered = apply_kimi_k3_xtml(&messages, ¶ms).expect("k3 render should succeed"); + assert!(!rendered.contains("thinking-effort"), "got: {rendered}"); + assert_eq!(rendered, fixture_text("plain_user_no_thinking")); +} + +#[test] +fn thinking_effort_invalid_is_rejected() { + // Mirrors the reference `assert thinking_effort in _VALID_THINKING_EFFORTS`; + // `medium` is described in the body text but not an accepted value. + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let template_kwargs = HashMap::from([("thinking_effort".to_string(), json!("medium"))]); + let params = ChatTemplateParams { + add_generation_prompt: true, + thinking: Some(true), + template_kwargs: Some(&template_kwargs), + ..Default::default() + }; + assert!(apply_kimi_k3_xtml(&messages, ¶ms).is_err()); +} + +/// End-to-end: a tokenizer loaded from a K3 directory (no chat template at all) +/// must load successfully, detect the K3 renderer, and render XTML through +/// `apply_chat_template`. +#[test] +fn tokenizer_loads_and_renders_k3_without_chat_template() { + let dir = TempDir::new().unwrap(); + fs::write(dir.path().join("tiktoken.model"), MIN_TIKTOKEN_MODEL).unwrap(); + fs::write( + dir.path().join("config.json"), + r#"{"architectures": ["KimiK3ForConditionalGeneration"]}"#, + ) + .unwrap(); + fs::write(dir.path().join("tokenizer_config.json"), "{}").unwrap(); + // Note: intentionally NO chat_template.json / .jinja in this directory. + + let tok = TiktokenTokenizer::from_dir(dir.path()).expect("K3 tokenizer should load"); + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let rendered = tok + .apply_chat_template( + &messages, + ChatTemplateParams { + add_generation_prompt: true, + thinking: Some(true), + ..Default::default() + }, + ) + .expect("K3 render should succeed"); + + assert_eq!(rendered, fixture_text("plain_user_thinking")); +} diff --git a/crates/tool_parser/src/factory.rs b/crates/tool_parser/src/factory.rs index 695e05ee0..06466d83c 100644 --- a/crates/tool_parser/src/factory.rs +++ b/crates/tool_parser/src/factory.rs @@ -10,8 +10,9 @@ use tokio::sync::Mutex; use crate::{ parsers::{ CohereParser, DeepSeek31Parser, DeepSeekDsmlParser, DeepSeekParser, Glm4MoeParser, - InklingParser, JsonParser, KimiK2Parser, LlamaParser, MinimaxM2Parser, MistralParser, - PassthroughParser, PythonicParser, QwenParser, QwenXmlParser, SarashinaParser, Step3Parser, + InklingParser, JsonParser, KimiK2Parser, KimiK3Parser, LlamaParser, MinimaxM2Parser, + MistralParser, PassthroughParser, PythonicParser, QwenParser, QwenXmlParser, + SarashinaParser, Step3Parser, }, traits::ToolParser, }; @@ -332,6 +333,7 @@ impl ParserFactory { || Box::new(KimiK2Parser::new()), KimiK2Parser::build_structural_tag, ); + registry.register_parser("kimi_k3", || Box::new(KimiK3Parser::new())); registry.register_parser_with_structural_tag( "inkling", || Box::new(InklingParser::new()), @@ -416,6 +418,9 @@ impl ParserFactory { registry.map_model("kimi-k2*", "kimik2"); registry.map_model("Kimi-K2*", "kimik2"); registry.map_model("moonshot*/Kimi-K2*", "kimik2"); + registry.map_model("kimi-k3*", "kimi_k3"); + registry.map_model("Kimi-K3*", "kimi_k3"); + registry.map_model("moonshot*/Kimi-K3*", "kimi_k3"); // Inkling models use TML JSON tool calls. registry.map_model("inkling*", "inkling"); diff --git a/crates/tool_parser/src/lib.rs b/crates/tool_parser/src/lib.rs index c33677d45..2bb98a646 100644 --- a/crates/tool_parser/src/lib.rs +++ b/crates/tool_parser/src/lib.rs @@ -18,8 +18,8 @@ mod tests; pub use factory::{ParserFactory, PooledParser, ToolConstraint}; pub use parsers::{ CohereParser, DeepSeek31Parser, DeepSeekDsmlParser, DeepSeekParser, Glm4MoeParser, - InklingParser, JsonParser, KimiK2Parser, LlamaParser, MinimaxM2Parser, MistralParser, - PythonicParser, QwenParser, Step3Parser, + InklingParser, JsonParser, KimiK2Parser, KimiK3Parser, LlamaParser, MinimaxM2Parser, + MistralParser, PythonicParser, QwenParser, Step3Parser, }; pub use traits::ToolParser; pub use types::{FunctionCall, StreamingParseResult, ToolCall}; diff --git a/crates/tool_parser/src/parsers/kimi_k3.rs b/crates/tool_parser/src/parsers/kimi_k3.rs new file mode 100644 index 000000000..192695ee1 --- /dev/null +++ b/crates/tool_parser/src/parsers/kimi_k3.rs @@ -0,0 +1,666 @@ +//! Kimi-K3 (XTML) tool-call parser +//! +//! Ports the Kimi-K3 reference tool-call parser: turns the generated XTML +//! `response` and `tools` channels back into plain `content` and `ToolCall`s. +//! +//! # Format +//! +//! K3 assistant tool calls live in a nested `tools` channel, emitted after a +//! sibling `response` channel that is unwrapped into content: +//! ```text +//! <|open|>response<|sep|> CONTENT <|close|>response<|sep|> +//! <|open|>tools<|sep|> +//! <|open|>call tool="NAME" index="N"<|sep|> +//! <|open|>argument key="K" type="T"<|sep|>VALUE<|close|>argument<|sep|> +//! <|close|>call<|sep|> +//! <|close|>tools<|sep|> +//! ``` +//! `<|open|>`, `<|close|>`, `<|sep|>` are literal strings in the detokenized +//! text (not special tokens from this parser's point of view). Markers +//! tolerate optional whitespace within them (e.g. `<|open|> tools <|sep|>`) +//! as defense in depth; this is a no-op on clean input. Bodies use +//! non-greedy matching so each block stops at its own first closing marker. +//! +//! # Argument decoding +//! - `type="string"` -> the value is raw text, used as-is (no unescaping). +//! - any other type -> the value is JSON-decoded; on JSON error, falls back +//! to the raw string rather than erroring mid-stream. +//! +//! # Attribute escaping +//! Attribute values (`tool=`, `index=`, `key=`, `type=`) are escaped on the +//! encode side (`&` -> `&`, `"` -> `"`); decoding reverses that in +//! order (`"` -> `"` then `&` -> `&`). +//! +//! # Tool-call id +//! SMG's [`ToolCall`] has no `id` field (one is assigned later by +//! `model_gateway`), so this parser only produces `name` + `arguments` (+ +//! content). The XTML `index` attribute is still parsed to drive +//! [`ToolCallItem::tool_index`] (`index - 1`, falling back to an ordinal +//! counter when the attribute is missing or unparsable). +//! +//! Known limitation (inherited from the reference): because string argument +//! and response bodies are emitted raw, a value that literally contains +//! `<|close|>argument<|sep|>` or `<|close|>response<|sep|>` is +//! indistinguishable from a real closing marker. + +use std::collections::HashMap; + +use async_trait::async_trait; +use openai_protocol::common::Tool; +use regex::Regex; +use serde_json::Value; + +use crate::{ + errors::ParserResult, + traits::ToolParser, + types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem}, +}; + +const TOOLS_OPEN: &str = "<|open|>tools<|sep|>"; +const RESPONSE_OPEN: &str = "<|open|>response<|sep|>"; +const RESPONSE_CLOSE: &str = "<|close|>response<|sep|>"; + +/// A decoded `<|open|>call ...<|sep|>...<|close|>call<|sep|>` block. +struct DecodedCall { + name: String, + /// Compact JSON object string. + arguments: String, + /// Zero-based ordinal derived from the XTML `index` attribute + /// (`index - 1`); `None` when the attribute is missing or unparsable, + /// in which case the caller falls back to an ordinal counter. + tool_index: Option, +} + +/// Kimi-K3 XTML tool-call parser. +/// +/// Handles both non-streaming (`parse_complete`) and streaming +/// (`parse_incremental`) extraction of the `tools` channel, and unwraps the +/// sibling `response` channel into `content`. +pub struct KimiK3Parser { + /// Matches `<|open|>tools<|sep|>` (tolerant of inner whitespace). + tools_open_re: Regex, + /// Matches `<|close|>tools<|sep|>`. + tools_close_re: Regex, + /// Matches `<|open|>response<|sep|>`. + response_open_re: Regex, + /// Matches `<|close|>response<|sep|>`. + response_close_re: Regex, + /// Matches a stray `<|close|>message<|sep|>` (stripped from content). + message_close_re: Regex, + /// Matches one `call` block, capturing `attrs` and `body`. + call_re: Regex, + /// Matches one `argument` block, capturing `attrs` and `val`. + arg_re: Regex, + /// Matches one `key="value"` attribute pair. + attr_re: Regex, + /// Matches a complete, unwrapped `response` channel, capturing `c`. + response_re: Regex, + + /// Accumulates every chunk seen so far (plays the role of the streaming + /// `current_text`, which the engine rebuilds from scratch each call). + buffer: String, + /// Byte offset into `buffer` up to which response content has already + /// been emitted. + sent_content_idx: usize, + /// Number of tool calls already emitted. + sent_tool_call_count: usize, +} + +impl KimiK3Parser { + /// Create a new Kimi-K3 parser. + #[expect( + clippy::expect_used, + reason = "regex patterns are compile-time string literals" + )] + pub fn new() -> Self { + Self { + tools_open_re: Regex::new(r"<\|open\|>\s*tools\s*<\|sep\|>") + .expect("valid regex"), + tools_close_re: Regex::new(r"<\|close\|>\s*tools\s*<\|sep\|>") + .expect("valid regex"), + response_open_re: Regex::new(r"<\|open\|>\s*response\s*<\|sep\|>") + .expect("valid regex"), + response_close_re: Regex::new(r"<\|close\|>\s*response\s*<\|sep\|>") + .expect("valid regex"), + message_close_re: Regex::new(r"<\|close\|>\s*message\s*<\|sep\|>") + .expect("valid regex"), + call_re: Regex::new( + r"(?s)<\|open\|>\s*call\s+(?P.*?)<\|sep\|>(?P.*?)<\|close\|>\s*call\s*<\|sep\|>", + ) + .expect("valid regex"), + arg_re: Regex::new( + r"(?s)<\|open\|>\s*argument\s+(?P.*?)<\|sep\|>(?P.*?)<\|close\|>\s*argument\s*<\|sep\|>", + ) + .expect("valid regex"), + attr_re: Regex::new(r#"(?P\w+)="(?P[^"]*)""#).expect("valid regex"), + response_re: Regex::new( + r"(?s)<\|open\|>\s*response\s*<\|sep\|>(?P.*?)<\|close\|>\s*response\s*<\|sep\|>", + ) + .expect("valid regex"), + buffer: String::new(), + sent_content_idx: 0, + sent_tool_call_count: 0, + } + } + + /// Parse the `key="value"` attribute pairs in `s`, unescaping each value + /// (`"` -> `"` then `&` -> `&`, the reverse of the encode order). + fn attrs(&self, s: &str) -> HashMap { + self.attr_re + .captures_iter(s) + .map(|m| { + let key = m.name("k").map_or("", |g| g.as_str()).to_string(); + let value = m + .name("v") + .map_or("", |g| g.as_str()) + .replace(""", "\"") + .replace("&", "&"); + (key, value) + }) + .collect() + } + + /// Decode one `call` block (its `attrs` segment and `body`) into a + /// [`DecodedCall`]. Each argument is re-typed per its `type=` tag: + /// strings pass through raw, everything else is JSON-decoded (falling + /// back to the raw string on malformed JSON, so a partial stream never + /// errors). Returns `None` when no tool name is present. + fn decode_call(&self, attrs: &str, body: &str) -> Option { + let call_attrs = self.attrs(attrs); + let tool_name = call_attrs.get("tool").cloned().unwrap_or_default(); + if tool_name.is_empty() { + return None; + } + let tool_index = call_attrs + .get("index") + .and_then(|s| s.parse::().ok()) + .and_then(|n| n.checked_sub(1)) + .and_then(|n| usize::try_from(n).ok()); + + let mut arguments = serde_json::Map::new(); + for arg_match in self.arg_re.captures_iter(body) { + let arg_attrs_text = arg_match.name("attrs").map_or("", |g| g.as_str()); + let arg_attrs = self.attrs(arg_attrs_text); + let key = arg_attrs.get("key").cloned().unwrap_or_default(); + let arg_type = arg_attrs.get("type").map_or("string", String::as_str); + let raw_value = arg_match.name("val").map_or("", |g| g.as_str()); + + let value = if arg_type == "string" { + Value::String(raw_value.to_string()) + } else { + serde_json::from_str::(raw_value) + .unwrap_or_else(|_| Value::String(raw_value.to_string())) + }; + arguments.insert(key, value); + } + + let arguments_str = serde_json::to_string(&arguments).unwrap_or_else(|_| "{}".to_string()); + + Some(DecodedCall { + name: tool_name, + arguments: arguments_str, + tool_index, + }) + } + + /// Strip XTML response/message markers from generated response text. + /// + /// In chat serving, `<|open|>response<|sep|>` is often part of the + /// prompt generation prefix, so the model output may only contain the + /// body plus `<|close|>response<|sep|>`. Handles both that + /// consumed-prefix shape and a complete + /// `<|open|>response<|sep|>...<|close|>response<|sep|>` wrapper. + fn strip_response_content(&self, text: &str) -> Option { + let stripped = if let Some(m_open) = self.response_open_re.find(text) { + if let Some(m_close) = self.response_close_re.find_at(text, m_open.end()) { + text[m_open.end()..m_close.start()].to_string() + } else { + text[m_open.end()..].to_string() + } + } else { + self.response_close_re.replace_all(text, "").into_owned() + }; + let stripped = self + .message_close_re + .replace_all(&stripped, "") + .into_owned(); + if stripped.is_empty() { + None + } else { + Some(stripped) + } + } + + /// Compute non-streaming content: prefer the unwrapped `response` + /// channel; else fall back to stripping markers from `before` (the text + /// preceding the `tools` channel, or the whole output when there is no + /// `tools` channel at all). + fn content(&self, model_output: &str, before: &str) -> Option { + if let Some(m) = self.response_re.captures(model_output) { + let c = m.name("c").map_or("", |g| g.as_str()); + return if c.is_empty() { + None + } else { + Some(c.to_string()) + }; + } + self.strip_response_content(before) + } + + /// Compute the streaming-safe slice of response content, updating + /// `sent_content_idx`. This is what keeps split markers from leaking: + /// content is only released up to the start of the next recognized + /// marker (or up to the point where a partial marker might still be + /// growing at the tail of `current_text`). + fn extract_response_content(&mut self, current_text: &str) -> Option { + let m_open = self.response_open_re.find(current_text); + let body_start = m_open.map_or(0, |m| m.end()); + + let tools_start = self + .tools_open_re + .find_at(current_text, body_start) + .map(|m| m.start()); + let response_end = self + .response_close_re + .find_at(current_text, body_start) + .map(|m| m.start()); + + let sendable_idx = match (tools_start, response_end) { + (Some(a), Some(b)) => a.min(b), + (Some(a), None) | (None, Some(a)) => a, + (None, None) => { + let overlap = partial_tag_overlap(current_text, RESPONSE_OPEN) + .max(partial_tag_overlap(current_text, RESPONSE_CLOSE)) + .max(partial_tag_overlap(current_text, TOOLS_OPEN)); + current_text.len() - overlap + } + }; + + if sendable_idx <= body_start { + return None; + } + if self.sent_content_idx < body_start { + self.sent_content_idx = body_start; + } + if sendable_idx <= self.sent_content_idx { + return None; + } + + let content = current_text[self.sent_content_idx..sendable_idx].to_string(); + self.sent_content_idx = sendable_idx; + if content.is_empty() { + None + } else { + Some(content) + } + } + + /// Decode every complete `call` block in `section` (a slice starting + /// just past the `tools`-open marker), skipping blocks with no tool + /// name. + fn decode_calls_in_section(&self, section: &str) -> Vec { + let mut calls = Vec::new(); + for m in self.call_re.captures_iter(section) { + let attrs_text = m.name("attrs").map_or("", |g| g.as_str()); + let body = m.name("body").map_or("", |g| g.as_str()); + if let Some(decoded) = self.decode_call(attrs_text, body) { + calls.push(decoded); + } + } + calls + } +} + +/// Length of the longest prefix of `tag` that `text` ends with (up to +/// `tag.len() - 1`, since a full match would have been caught by the marker +/// regexes already). Used to hold back a possibly-still-growing partial +/// marker at the tail of the streamed text instead of leaking it as content. +fn partial_tag_overlap(text: &str, tag: &str) -> usize { + let max_len = text.len().min(tag.len().saturating_sub(1)); + for n in (1..=max_len).rev() { + if text.ends_with(&tag[..n]) { + return n; + } + } + 0 +} + +impl Default for KimiK3Parser { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl ToolParser for KimiK3Parser { + async fn parse_complete(&self, output: &str) -> ParserResult<(String, Vec)> { + let Some(m_open) = self.tools_open_re.find(output) else { + let content = self.content(output, output); + return Ok((content.unwrap_or_default(), vec![])); + }; + + let before = &output[..m_open.start()]; + let start = m_open.end(); + let section_end = self + .tools_close_re + .find_at(output, start) + .map_or(output.len(), |m| m.start()); + let section = &output[start..section_end]; + + let calls = self + .decode_calls_in_section(section) + .into_iter() + .map(|decoded| ToolCall { + function: FunctionCall { + name: decoded.name, + arguments: decoded.arguments, + }, + }) + .collect(); + + let content = self.content(output, before); + Ok((content.unwrap_or_default(), calls)) + } + + async fn parse_incremental( + &mut self, + chunk: &str, + _tools: &[Tool], + ) -> ParserResult { + self.buffer.push_str(chunk); + let current_text = self.buffer.clone(); + + let content = self.extract_response_content(¤t_text); + + let Some(m_tools) = self.tools_open_re.find(¤t_text) else { + return Ok(StreamingParseResult { + normal_text: content.unwrap_or_default(), + calls: vec![], + }); + }; + + let section = ¤t_text[m_tools.end()..]; + let decoded_calls = self.decode_calls_in_section(section); + + if decoded_calls.len() <= self.sent_tool_call_count { + return Ok(StreamingParseResult { + normal_text: content.unwrap_or_default(), + calls: vec![], + }); + } + + let calls = decoded_calls + .iter() + .skip(self.sent_tool_call_count) + .enumerate() + .map(|(i, decoded)| ToolCallItem { + tool_index: decoded.tool_index.unwrap_or(self.sent_tool_call_count + i), + name: Some(decoded.name.clone()), + parameters: decoded.arguments.clone(), + }) + .collect(); + self.sent_tool_call_count = decoded_calls.len(); + + Ok(StreamingParseResult { + normal_text: content.unwrap_or_default(), + calls, + }) + } + + fn has_tool_markers(&self, text: &str) -> bool { + self.tools_open_re.is_match(text) + } + + fn reset(&mut self) { + self.buffer.clear(); + self.sent_content_idx = 0; + self.sent_tool_call_count = 0; + } +} + +#[cfg(test)] +mod tests { + use serde_json::Value; + + use super::*; + + const OPEN: &str = "<|open|>"; + const CLOSE: &str = "<|close|>"; + const SEP: &str = "<|sep|>"; + + fn _arg(key: &str, typ: &str, value: &str) -> String { + format!(r#"{OPEN}argument key="{key}" type="{typ}"{SEP}{value}{CLOSE}argument{SEP}"#) + } + + fn _call(tool: &str, index: i32, args: &[String]) -> String { + let body: String = args.concat(); + format!(r#"{OPEN}call tool="{tool}" index="{index}"{SEP}{body}{CLOSE}call{SEP}"#) + } + + fn _response(content: &str) -> String { + format!("{OPEN}response{SEP}{content}{CLOSE}response{SEP}") + } + + fn _tools(calls: &[String]) -> String { + let body: String = calls.concat(); + format!("{OPEN}tools{SEP}{body}{CLOSE}tools{SEP}") + } + + #[tokio::test] + async fn test_parse_complete_response_and_typed_arguments() { + let parser = KimiK3Parser::new(); + let input = _response("answer") + + &_tools(&[_call( + "calc", + 1, + &[ + _arg("x", "number", "1"), + _arg("flag", "boolean", "true"), + _arg("text", "string", "raw"), + ], + )]); + + let (content, calls) = parser.parse_complete(&input).await.unwrap(); + assert_eq!(content, "answer"); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].function.name, "calc"); + let args: Value = serde_json::from_str(&calls[0].function.arguments).unwrap(); + assert_eq!( + args, + serde_json::json!({"x": 1, "flag": true, "text": "raw"}) + ); + } + + #[tokio::test] + async fn test_parse_complete_unescapes_attributes() { + let parser = KimiK3Parser::new(); + let input = _tools(&[_call( + "a&b"c", + 1, + &[_arg("k&q", "string", "v")], + )]); + + let (_content, calls) = parser.parse_complete(&input).await.unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].function.name, "a&b\"c"); + let args: Value = serde_json::from_str(&calls[0].function.arguments).unwrap(); + assert_eq!(args, serde_json::json!({"k&q": "v"})); + } + + #[tokio::test] + async fn test_parse_complete_allows_less_than_in_attributes() { + let parser = KimiK3Parser::new(); + let input = _tools(&[_call("calc HashMap { - let kwargs_capacity = 1 + request.chat_template_kwargs.as_ref().map_or(0, |k| k.len()); + let kwargs_capacity = 3 + request.chat_template_kwargs.as_ref().map_or(0, |k| k.len()); let mut combined = HashMap::with_capacity(kwargs_capacity); if let Some(reasoning_effort) = &request.reasoning_effort { combined.insert( @@ -194,6 +198,16 @@ fn build_chat_template_kwargs(request: &ChatCompletionRequest) -> HashMap u /// * `history_count` - Number of tool calls in previous messages /// /// # Returns -/// A unique ID string. KimiK2 uses `functions.{name}:{global_index}`, others use `call_{uuid}` +/// A unique ID string: +/// - Kimi-K3 (XTML): `{name}:{tool_index}` — an opaque, per-message zero-based +/// ordinal. K3 never renders the id into the prompt and matches tool results +/// back to calls by opaque id equality scoped to the most recent assistant +/// message, so the id carries no `functions.` prefix and no history offset. +/// Mirrors the K3 reference decode parser (`{tool_name}:{xtml_index - 1}`). +/// - Kimi-K2: `functions.{name}:{history_count + tool_index}` (globally unique). +/// - others: `call_{24-char-uuid}`. pub(crate) fn generate_tool_call_id( model: &str, tool_name: &str, tool_index: usize, history_count: usize, ) -> String { - // Case-insensitive check without allocation (search for "kimi" substring) + // Case-insensitive substring checks without allocation. let is_kimi = model .as_bytes() .windows(4) // "kimi".len() .any(|window| window.eq_ignore_ascii_case(b"kimi")); - if is_kimi { - // KimiK2 format: functions.{name}:{global_index} - format!("functions.{}:{}", tool_name, history_count + tool_index) - } else { + if !is_kimi { // Standard OpenAI format: call_{24-char-uuid} - format!("call_{}", &Uuid::now_v7().simple().to_string()[..24]) + return format!("call_{}", &Uuid::now_v7().simple().to_string()[..24]); + } + + let is_k3 = model + .as_bytes() + .windows(2) // "k3".len() + .any(|window| window.eq_ignore_ascii_case(b"k3")); + + if is_k3 { + // Kimi-K3 (XTML) opaque format: {name}:{per-message zero-based ordinal}. + format!("{tool_name}:{tool_index}") + } else { + // Kimi-K2 format: functions.{name}:{global_index}. + format!("functions.{}:{}", tool_name, history_count + tool_index) } } @@ -1327,6 +1358,33 @@ mod tests { assert_eq!(kwargs.get("custom"), Some(&Value::Bool(true))); } + #[test] + fn tool_and_output_controls_are_forwarded_to_renderer() { + let request: ChatCompletionRequest = serde_json::from_value(json!({ + "model": "kimi-k3", + "messages": [{"role": "user", "content": "hello"}], + "tool_choice": "required", + "response_format": {"type": "json_object"} + })) + .unwrap(); + let kwargs = build_chat_template_kwargs(&request); + assert_eq!(kwargs.get(TOOL_CHOICE_KEY), Some(&json!("required"))); + assert_eq!( + kwargs.get(RESPONSE_FORMAT_KEY), + Some(&json!({"type": "json_object"})) + ); + + // Absent both, neither key is forwarded. + let bare: ChatCompletionRequest = serde_json::from_value(json!({ + "model": "kimi-k3", + "messages": [{"role": "user", "content": "hello"}] + })) + .unwrap(); + let kwargs = build_chat_template_kwargs(&bare); + assert!(!kwargs.contains_key(TOOL_CHOICE_KEY)); + assert!(!kwargs.contains_key(RESPONSE_FORMAT_KEY)); + } + /// End-to-end: run a real MMBench-shaped `[text, image]` message through the /// full SMG pipeline (`process_content_format` + the actual model chat /// template) and assert the rendered prompt places the image BEFORE the @@ -1390,4 +1448,38 @@ mod tests { "image must precede the question (vstart={vstart}, qpos={qpos}).\n--- rendered ---\n{rendered}" ); } + + #[test] + fn test_generate_tool_call_id_kimi_k3_opaque_format() { + // K3: `{name}:{per-message zero-based ordinal}` — no `functions.` prefix, + // no history offset (matches the K3 reference decode parser). + assert_eq!( + generate_tool_call_id("moonshotai/Kimi-K3", "get_weather", 0, 0), + "get_weather:0" + ); + assert_eq!( + generate_tool_call_id("kimi_k3", "get_weather", 1, 5), + "get_weather:1" + ); + } + + #[test] + fn test_generate_tool_call_id_kimi_k2_unchanged() { + // K2 keeps the globally unique `functions.{name}:{history + index}` form. + assert_eq!( + generate_tool_call_id("moonshotai/Kimi-K2-Instruct", "get_weather", 0, 0), + "functions.get_weather:0" + ); + assert_eq!( + generate_tool_call_id("Kimi-K2", "get_weather", 1, 2), + "functions.get_weather:3" + ); + } + + #[test] + fn test_generate_tool_call_id_non_kimi_uses_uuid() { + let id = generate_tool_call_id("gpt-4o", "get_weather", 0, 0); + assert!(id.starts_with("call_"), "got: {id}"); + assert!(!id.contains("get_weather"), "got: {id}"); + } } From 851f38927f3eee100ab07cb108ac239666d78b43 Mon Sep 17 00:00:00 2001 From: key4ng Date: Mon, 27 Jul 2026 10:34:55 -0700 Subject: [PATCH 2/5] fix(kimi-k3): address review feedback on renderer detection and tool indexing - Detect the KimiK3Xtml / KimiK25Tools renderers from config.json `model_type` when `architectures` is absent or unrecognized, matching the existing model_type-based tokenizer detection (is_kimi_tokenizer). - Register underscore-spelled K3 tool-parser model globs (kimi_k3*, Kimi_K3*, moonshot*/Kimi_K3*) for parity with the reasoning-parser factory. - Assign streaming tool indices by emission order instead of trusting the model-supplied XTML `index`, so duplicate/sparse/out-of-order indices can no longer collide; this matches the non-streaming path. Remove the now-unused DecodedCall index plumbing. - Document the intentional response-open prefix-strip asymmetry and the non-flushed partial-marker tail in the reasoning parser, and clarify in the XTML encoder why HTML-escaping control tokens is not the correct fix. Signed-off-by: key4ng --- .../reasoning_parser/src/parsers/kimi_k3.rs | 19 +++++- crates/tokenizer/src/encoders/kimi_k3_xtml.rs | 9 +++ crates/tokenizer/src/tiktoken.rs | 32 +++++++++- crates/tool_parser/src/factory.rs | 6 ++ crates/tool_parser/src/parsers/kimi_k3.rs | 61 ++++++++++++++----- 5 files changed, 109 insertions(+), 18 deletions(-) diff --git a/crates/reasoning_parser/src/parsers/kimi_k3.rs b/crates/reasoning_parser/src/parsers/kimi_k3.rs index 771a93b43..f19da6a34 100644 --- a/crates/reasoning_parser/src/parsers/kimi_k3.rs +++ b/crates/reasoning_parser/src/parsers/kimi_k3.rs @@ -119,8 +119,22 @@ impl KimiK3Parser { /// /// Strips the response-open prefix, removes complete response-close and /// message-close markers, then holds back any partial marker suffix. + /// + /// The held-back partial-marker suffix is intentionally never force-flushed: + /// the [`ReasoningParser`] trait has no finalize/end-of-stream hook, and the + /// only bytes ever withheld are a proper prefix of a control marker (e.g. + /// `<|clo`). A complete generation always ends on a whole marker, so nothing + /// is lost in practice; a stream truncated mid-marker-prefix is already + /// incomplete output, and dropping that dangling fragment is preferable to + /// leaking a partial control token as content. fn content_ready_to_emit(&self, text: &str) -> String { - // Strip response-open prefix (first occurrence) + // Strip the response-open marker as a prefix (first occurrence only). + // In the K3 grammar the response channel opens exactly once and its + // opening marker is a prefix of the content tail, so a slice from its + // end is correct and a second occurrence cannot legitimately appear. + // The response-close / message-open / message-close markers, by + // contrast, are terminators that may sit anywhere in the tail, so they + // are removed globally below. This asymmetry is deliberate. let text = if let Some(m) = self.response_open_re.find(text) { &text[m.end()..] } else { @@ -132,7 +146,8 @@ impl KimiK3Parser { let text = self.message_open_re.replace_all(&text, ""); let text = self.message_close_re.replace_all(&text, ""); - // Hold back any partial marker suffix + // Hold back any partial marker suffix (see the fn-level note above: + // withheld only while it could still complete into a control marker). let text_str: &str = &text; let overlap = Self::compute_overlap( text_str, diff --git a/crates/tokenizer/src/encoders/kimi_k3_xtml.rs b/crates/tokenizer/src/encoders/kimi_k3_xtml.rs index 25afd8e42..4e0dc0a28 100644 --- a/crates/tokenizer/src/encoders/kimi_k3_xtml.rs +++ b/crates/tokenizer/src/encoders/kimi_k3_xtml.rs @@ -19,6 +19,15 @@ //! sibling renderers (`kimi_k25_tools`, `deepseek_v32`) is intentional here; //! segment-aware encoding is out of scope. //! +//! HTML-escaping the control tokens in user/tool text is **not** the fix: the +//! reference emits that text as ordinary BPE without escaping, so the model was +//! trained on the raw bytes; rewriting `<|open|>` to `<|open|>` (or +//! similar) would both feed the model text it never saw and corrupt legitimate +//! content. The correct fix is to carry the reference's per-segment +//! `allow_special` flag through to the tokenizer so structural markers keep +//! their special-token ids while surrounding text is encoded as ordinary BPE — +//! a cross-cutting change across this renderer family, tracked as future work. +//! //! # Deferred (TODO) //! //! The following `build_chat_segments` branches are intentionally not ported yet diff --git a/crates/tokenizer/src/tiktoken.rs b/crates/tokenizer/src/tiktoken.rs index 071088bc3..4d78872cf 100644 --- a/crates/tokenizer/src/tiktoken.rs +++ b/crates/tokenizer/src/tiktoken.rs @@ -633,7 +633,21 @@ fn detect_renderer_from_config(dir: &Path) -> Renderer { tracing::debug!(?path, "selected KimiK25Tools chat-template renderer"); return Renderer::KimiK25Tools; } - Renderer::Jinja + // Fall back to `model_type` when `architectures` is absent or unrecognized. + // Some Kimi checkpoints omit the architecture entry but still carry a + // `model_type`, and `is_kimi_tokenizer` already keys off it — keep the + // renderer selection consistent with that tokenizer detection. + match value.get("model_type").and_then(|v| v.as_str()) { + Some("kimi_k3") => { + tracing::debug!(?path, "selected KimiK3Xtml renderer via model_type"); + Renderer::KimiK3Xtml + } + Some("kimi_k25") => { + tracing::debug!(?path, "selected KimiK25Tools renderer via model_type"); + Renderer::KimiK25Tools + } + _ => Renderer::Jinja, + } } #[cfg(test)] @@ -688,6 +702,22 @@ mod tests { detect_renderer_from_config(missing.path()), Renderer::Jinja )); + + // `model_type` fallback: checkpoints without a recognized architecture + // entry are still detected by their `model_type`, matching + // `is_kimi_tokenizer`'s detection. + let k3_model_type = write_minimal_tiktoken_dir("{}", Some(r#"{"model_type": "kimi_k3"}"#)); + assert!(matches!( + detect_renderer_from_config(k3_model_type.path()), + Renderer::KimiK3Xtml + )); + + let k25_model_type = + write_minimal_tiktoken_dir("{}", Some(r#"{"model_type": "kimi_k25"}"#)); + assert!(matches!( + detect_renderer_from_config(k25_model_type.path()), + Renderer::KimiK25Tools + )); } #[test] diff --git a/crates/tool_parser/src/factory.rs b/crates/tool_parser/src/factory.rs index 06466d83c..d9e586325 100644 --- a/crates/tool_parser/src/factory.rs +++ b/crates/tool_parser/src/factory.rs @@ -421,6 +421,12 @@ impl ParserFactory { registry.map_model("kimi-k3*", "kimi_k3"); registry.map_model("Kimi-K3*", "kimi_k3"); registry.map_model("moonshot*/Kimi-K3*", "kimi_k3"); + // Underscore spellings (e.g. `kimi_k3`, `Kimi_K3`) as used by some + // checkpoint ids; mirrors the reasoning-parser factory, whose glob + // matching is case-insensitive but still enumerates both separators. + registry.map_model("kimi_k3*", "kimi_k3"); + registry.map_model("Kimi_K3*", "kimi_k3"); + registry.map_model("moonshot*/Kimi_K3*", "kimi_k3"); // Inkling models use TML JSON tool calls. registry.map_model("inkling*", "inkling"); diff --git a/crates/tool_parser/src/parsers/kimi_k3.rs b/crates/tool_parser/src/parsers/kimi_k3.rs index 192695ee1..daf6ecf98 100644 --- a/crates/tool_parser/src/parsers/kimi_k3.rs +++ b/crates/tool_parser/src/parsers/kimi_k3.rs @@ -34,9 +34,13 @@ //! # Tool-call id //! SMG's [`ToolCall`] has no `id` field (one is assigned later by //! `model_gateway`), so this parser only produces `name` + `arguments` (+ -//! content). The XTML `index` attribute is still parsed to drive -//! [`ToolCallItem::tool_index`] (`index - 1`, falling back to an ordinal -//! counter when the attribute is missing or unparsable). +//! content). Streaming [`ToolCallItem::tool_index`] is a per-message +//! zero-based ordinal assigned in emission order, independent of the XTML +//! `index` attribute. Deriving it from the ordinal — rather than trusting the +//! model-supplied `index` — keeps streamed indices (and the ids later built +//! from them) unique and monotonic even when a model emits duplicate, sparse, +//! or out-of-order `index` values, and matches the non-streaming path, which +//! also indexes tool calls by their emission order. //! //! Known limitation (inherited from the reference): because string argument //! and response bodies are emitted raw, a value that literally contains @@ -61,14 +65,14 @@ const RESPONSE_OPEN: &str = "<|open|>response<|sep|>"; const RESPONSE_CLOSE: &str = "<|close|>response<|sep|>"; /// A decoded `<|open|>call ...<|sep|>...<|close|>call<|sep|>` block. +/// +/// The XTML `index` attribute is deliberately not captured: streamed tool +/// indices are assigned by emission order (see the module-level `Tool-call id` +/// docs), so a model-supplied `index` never influences the parser output. struct DecodedCall { name: String, /// Compact JSON object string. arguments: String, - /// Zero-based ordinal derived from the XTML `index` attribute - /// (`index - 1`); `None` when the attribute is missing or unparsable, - /// in which case the caller falls back to an ordinal counter. - tool_index: Option, } /// Kimi-K3 XTML tool-call parser. @@ -171,11 +175,6 @@ impl KimiK3Parser { if tool_name.is_empty() { return None; } - let tool_index = call_attrs - .get("index") - .and_then(|s| s.parse::().ok()) - .and_then(|n| n.checked_sub(1)) - .and_then(|n| usize::try_from(n).ok()); let mut arguments = serde_json::Map::new(); for arg_match in self.arg_re.captures_iter(body) { @@ -199,7 +198,6 @@ impl KimiK3Parser { Some(DecodedCall { name: tool_name, arguments: arguments_str, - tool_index, }) } @@ -394,7 +392,9 @@ impl ToolParser for KimiK3Parser { .skip(self.sent_tool_call_count) .enumerate() .map(|(i, decoded)| ToolCallItem { - tool_index: decoded.tool_index.unwrap_or(self.sent_tool_call_count + i), + // Per-message zero-based ordinal in emission order; the XTML + // `index` attribute is deliberately ignored (see module docs). + tool_index: self.sent_tool_call_count + i, name: Some(decoded.name.clone()), parameters: decoded.arguments.clone(), }) @@ -545,7 +545,7 @@ mod tests { } #[tokio::test] - async fn test_decode_call_missing_index_falls_back_to_ordinal() { + async fn test_decode_call_without_index_attribute() { let parser = KimiK3Parser::new(); let input = format!( r#"{OPEN}tools{SEP}{OPEN}call tool="calc"{SEP}{CLOSE}call{SEP}{CLOSE}tools{SEP}"# @@ -632,6 +632,25 @@ mod tests { assert_eq!(result.calls[1].tool_index, 1); } + #[tokio::test] + async fn test_streaming_duplicate_xtml_index_uses_ordinal() { + // A model that emits colliding, sparse, or out-of-order `index` + // attributes must still yield unique, monotonic streamed indices: the + // parser assigns them by emission order and ignores the attribute. + let mut parser = KimiK3Parser::new(); + let input = _tools(&[ + _call("first", 7, &[_arg("a", "number", "1")]), + _call("second", 7, &[_arg("b", "number", "2")]), + _call("third", 3, &[_arg("c", "number", "3")]), + ]); + + let result = parser.parse_incremental(&input, &[]).await.unwrap(); + assert_eq!(result.calls.len(), 3); + assert_eq!(result.calls[0].tool_index, 0); + assert_eq!(result.calls[1].tool_index, 1); + assert_eq!(result.calls[2].tool_index, 2); + } + #[tokio::test] async fn test_reset_clears_streaming_state() { let mut parser = KimiK3Parser::new(); @@ -663,4 +682,16 @@ mod tests { Some("kimi_k3".to_string()) ); } + + #[test] + fn test_factory_resolves_underscore_kimi_k3() { + let factory = crate::factory::ParserFactory::new(); + for id in ["kimi_k3", "Kimi_K3", "moonshotai/Kimi_K3"] { + assert_eq!( + factory.registry().resolve_model_to_parser(id), + Some("kimi_k3".to_string()), + "expected `{id}` to resolve to the kimi_k3 parser" + ); + } + } } From 28ab345dcdba2c771fc55d5358541f1b10715392 Mon Sep 17 00:00:00 2001 From: Keyang Ru Date: Mon, 27 Jul 2026 04:16:13 +0000 Subject: [PATCH 3/5] feat(tool-parser): add guided-decoding structural tag for kimi_k3 Register the kimi_k3 parser with a structural-tag builder so guided decoding (xgrammar) constrains Kimi-K3 XTML tool calls, mirroring the existing kimi_k2 support. This closes the gap that let K3 emit invalid or empty tool calls (schema-invalid args, empty {}, semantically empty). KimiK3Parser::build_structural_tag emits a triggered_tags grammar for K3 XTML: a single tools section whose content is a plus of call blocks (native parallel calls). Each call pins the declared tool name and index, and every argument is constrained from the JSON schema (Strict): required properties forced present in declared order, optionals wrapped in optional, and each type/value pinned per property type -- string bodies verbatim, integer as number, enums as literal alternates, escaped attribute values. Schemas without usable properties fall back to a permissive skeleton so the grammar is never infeasible. Like kimi_k2, stop_after_first is left unset so the trailing message/end-of-msg tokens after the tools section remain valid. Signed-off-by: Keyang Ru --- crates/tool_parser/src/factory.rs | 6 +- crates/tool_parser/src/parsers/kimi_k3.rs | 392 +++++++++++++++++++++- 2 files changed, 391 insertions(+), 7 deletions(-) diff --git a/crates/tool_parser/src/factory.rs b/crates/tool_parser/src/factory.rs index d9e586325..3c062d7dd 100644 --- a/crates/tool_parser/src/factory.rs +++ b/crates/tool_parser/src/factory.rs @@ -333,7 +333,11 @@ impl ParserFactory { || Box::new(KimiK2Parser::new()), KimiK2Parser::build_structural_tag, ); - registry.register_parser("kimi_k3", || Box::new(KimiK3Parser::new())); + registry.register_parser_with_structural_tag( + "kimi_k3", + || Box::new(KimiK3Parser::new()), + KimiK3Parser::build_structural_tag, + ); registry.register_parser_with_structural_tag( "inkling", || Box::new(InklingParser::new()), diff --git a/crates/tool_parser/src/parsers/kimi_k3.rs b/crates/tool_parser/src/parsers/kimi_k3.rs index daf6ecf98..513060fce 100644 --- a/crates/tool_parser/src/parsers/kimi_k3.rs +++ b/crates/tool_parser/src/parsers/kimi_k3.rs @@ -47,12 +47,12 @@ //! `<|close|>argument<|sep|>` or `<|close|>response<|sep|>` is //! indistinguishable from a real closing marker. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use async_trait::async_trait; use openai_protocol::common::Tool; use regex::Regex; -use serde_json::Value; +use serde_json::{json, Value}; use crate::{ errors::ParserResult, @@ -61,9 +61,15 @@ use crate::{ }; const TOOLS_OPEN: &str = "<|open|>tools<|sep|>"; +const TOOLS_CLOSE: &str = "<|close|>tools<|sep|>"; const RESPONSE_OPEN: &str = "<|open|>response<|sep|>"; const RESPONSE_CLOSE: &str = "<|close|>response<|sep|>"; +/// Bare XTML control markers, used to assemble structural-tag literals. +const OPEN: &str = "<|open|>"; +const CLOSE: &str = "<|close|>"; +const SEP: &str = "<|sep|>"; + /// A decoded `<|open|>call ...<|sep|>...<|close|>call<|sep|>` block. /// /// The XTML `index` attribute is deliberately not captured: streamed tool @@ -111,6 +117,63 @@ pub struct KimiK3Parser { } impl KimiK3Parser { + /// Build an xgrammar structural tag that constrains Kimi-K3 XTML tool + /// calls to the declared `tools`. + /// + /// The grammar mirrors the encoder (`kimi_k3_xtml.rs`): a tool call is + /// + /// ```text + /// <|open|>tools<|sep|> + /// <|open|>call tool="NAME" index="N"<|sep|> + /// <|open|>argument key="K" type="T"<|sep|>VALUE<|close|>argument<|sep|> + /// <|close|>call<|sep|> + /// <|close|>tools<|sep|> + /// ``` + /// + /// A `triggered_tags` format dispatches to the tools section once + /// `<|open|>tools<|sep|>` appears, then forces the framing above. All K3 + /// tool calls live in a single section, so its `content` is a `plus` of + /// call blocks (native parallel calls) and a lone section suffices — unlike + /// K2, which repeats a tag per call. Within a call, arguments are + /// constrained per the JSON schema (Strict mode): every `required` property + /// must appear, in the schema's declared order, followed by any optional + /// properties; each `type=` attribute and each value is pinned to the + /// property's type. Schemas without usable `properties` fall back to a + /// permissive skeleton so the grammar is never infeasible. + /// + /// `at_least_one` is wired to `tool_choice`: `true` for `"required"` (a + /// tool call must be emitted), `false` for `"auto"`. Mirroring K2, + /// `stop_after_first` is left unset so trailing tokens after the section + /// (`<|close|>message<|sep|><|end_of_msg|>`) remain valid. + pub fn build_structural_tag(tools: &[Tool], at_least_one: bool) -> Value { + let call_formats: Vec = tools + .iter() + .filter(|tool| !tool.function.name.is_empty()) + .map(|tool| build_call_format(&tool.function.name, &tool.function.parameters)) + .collect(); + + let call_choice = if call_formats.is_empty() { + permissive_call_format() + } else { + json!({ "type": "or", "elements": call_formats }) + }; + + let tools_section = json!({ + "begin": TOOLS_OPEN, + "content": { "type": "plus", "content": call_choice }, + "end": TOOLS_CLOSE, + }); + + json!({ + "format": { + "type": "triggered_tags", + "triggers": [TOOLS_OPEN], + "tags": [tools_section], + "at_least_one": at_least_one, + } + }) + } + /// Create a new Kimi-K3 parser. #[expect( clippy::expect_used, @@ -309,6 +372,191 @@ impl KimiK3Parser { } } +// --------------------------------------------------------------------------- +// Structural-tag construction (guided decoding) +// --------------------------------------------------------------------------- + +/// How one argument value is constrained in the structural tag. +enum ArgShape { + /// The `type=` attribute and value grammar are pinned from the schema. + Fixed { + type_attr: &'static str, + value: Value, + }, + /// The schema is ambiguous (union / missing / `anyOf`); allow any + /// well-formed `type=` attribute and any raw value up to the marker. + Permissive, +} + +/// Escape an attribute value the same way the encoder does (`&` -> `&`, +/// then `"` -> `"`), so grammar literals match the emitted bytes. +fn escape_attr(value: &str) -> String { + value.replace('&', "&").replace('"', """) +} + +/// Build the `sequence` grammar for a single `<|open|>call ...<|close|>call` +/// block for the tool `name` with the given JSON-schema `params`. +fn build_call_format(name: &str, params: &Value) -> Value { + let esc_name = escape_attr(name); + json!({ + "type": "sequence", + "elements": [ + { "type": "const_string", "value": format!("{OPEN}call tool=\"{esc_name}\" index=\"") }, + { "type": "regex", "pattern": "[1-9][0-9]*" }, + { "type": "const_string", "value": format!("\"{SEP}") }, + build_args_format(params), + { "type": "const_string", "value": format!("{CLOSE}call{SEP}") }, + ] + }) +} + +/// Build the argument-region grammar for a tool's parameter schema (Strict +/// mode): required properties in declared order, then optional properties. +/// Falls back to a permissive skeleton when the schema has no usable +/// `properties`. +fn build_args_format(params: &Value) -> Value { + let Some(properties) = params.get("properties").and_then(Value::as_object) else { + return permissive_args_format(); + }; + if properties.is_empty() { + return permissive_args_format(); + } + + let required: HashSet<&str> = params + .get("required") + .and_then(Value::as_array) + .map(|values| values.iter().filter_map(Value::as_str).collect()) + .unwrap_or_default(); + + let elements: Vec = properties + .iter() + .map(|(key, subschema)| { + let argument = build_argument_format(key, subschema); + if required.contains(key.as_str()) { + argument + } else { + json!({ "type": "optional", "content": argument }) + } + }) + .collect(); + + json!({ "type": "sequence", "elements": elements }) +} + +/// Build the grammar for one `<|open|>argument ...<|close|>argument` block, +/// pinning `key=`, `type=`, and the value grammar from the property schema. +fn build_argument_format(key: &str, subschema: &Value) -> Value { + let esc_key = escape_attr(key); + match classify_arg(subschema) { + ArgShape::Fixed { type_attr, value } => json!({ + "type": "sequence", + "elements": [ + { "type": "const_string", "value": format!("{OPEN}argument key=\"{esc_key}\" type=\"{type_attr}\"{SEP}") }, + value, + { "type": "const_string", "value": format!("{CLOSE}argument{SEP}") }, + ] + }), + ArgShape::Permissive => json!({ + "type": "sequence", + "elements": [ + { "type": "const_string", "value": format!("{OPEN}argument key=\"{esc_key}\" type=\"") }, + { "type": "any_text", "excludes": ["\"", "<|"] }, + { "type": "const_string", "value": format!("\"{SEP}") }, + { "type": "any_text", "excludes": [CLOSE, OPEN, SEP] }, + { "type": "const_string", "value": format!("{CLOSE}argument{SEP}") }, + ] + }), + } +} + +/// Map a property schema to its XTML `type=` attribute and value grammar. +/// K3 emits string bodies verbatim (no JSON quoting) while every other type is +/// compact JSON, and it derives `type=` from the *value*, so `integer` becomes +/// `number`. +fn classify_arg(subschema: &Value) -> ArgShape { + // A string `enum` constrains the raw body to one of the literals. + if let Some(values) = subschema.get("enum").and_then(Value::as_array) { + if !values.is_empty() && values.iter().all(Value::is_string) { + let options: Vec = values + .iter() + .filter_map(Value::as_str) + .map(|v| json!({ "type": "const_string", "value": v })) + .collect(); + return ArgShape::Fixed { + type_attr: "string", + value: json!({ "type": "or", "elements": options }), + }; + } + } + + match subschema.get("type").and_then(Value::as_str) { + // Verbatim body -> json_schema (which expects quotes) cannot be used. + Some("string") => ArgShape::Fixed { + type_attr: "string", + value: json!({ "type": "any_text", "excludes": [CLOSE, OPEN, SEP] }), + }, + Some("integer") | Some("number") => ArgShape::Fixed { + type_attr: "number", + value: json!({ "type": "json_schema", "json_schema": subschema }), + }, + Some("boolean") => ArgShape::Fixed { + type_attr: "boolean", + value: json!({ "type": "json_schema", "json_schema": { "type": "boolean" } }), + }, + Some("object") => ArgShape::Fixed { + type_attr: "object", + value: json!({ "type": "json_schema", "json_schema": subschema }), + }, + Some("array") => ArgShape::Fixed { + type_attr: "array", + value: json!({ "type": "json_schema", "json_schema": subschema }), + }, + Some("null") => ArgShape::Fixed { + type_attr: "null", + value: json!({ "type": "const_string", "value": "null" }), + }, + _ => ArgShape::Permissive, + } +} + +/// Permissive argument region: zero or more structurally-valid argument blocks +/// with unconstrained keys/types/values. Used when the schema is unusable. +fn permissive_args_format() -> Value { + json!({ "type": "star", "content": permissive_argument_format() }) +} + +/// One structurally-valid argument block with unconstrained key/type/value. +fn permissive_argument_format() -> Value { + json!({ + "type": "sequence", + "elements": [ + { "type": "const_string", "value": format!("{OPEN}argument key=\"") }, + { "type": "any_text", "excludes": ["\"", "<|"] }, + { "type": "const_string", "value": "\" type=\"" }, + { "type": "any_text", "excludes": ["\"", "<|"] }, + { "type": "const_string", "value": format!("\"{SEP}") }, + { "type": "any_text", "excludes": [CLOSE, OPEN, SEP] }, + { "type": "const_string", "value": format!("{CLOSE}argument{SEP}") }, + ] + }) +} + +/// Permissive call block used only when no declared tool name is usable. +fn permissive_call_format() -> Value { + json!({ + "type": "sequence", + "elements": [ + { "type": "const_string", "value": format!("{OPEN}call tool=\"") }, + { "type": "any_text", "excludes": ["\"", "<|"] }, + { "type": "const_string", "value": "\" index=\"" }, + { "type": "regex", "pattern": "[1-9][0-9]*" }, + { "type": "const_string", "value": format!("\"{SEP}") }, + permissive_args_format(), + { "type": "const_string", "value": format!("{CLOSE}call{SEP}") }, + ] + }) +} + /// Length of the longest prefix of `tag` that `text` ends with (up to /// `tag.len() - 1`, since a full match would have been caught by the marker /// regexes already). Used to hold back a possibly-still-growing partial @@ -424,10 +672,6 @@ mod tests { use super::*; - const OPEN: &str = "<|open|>"; - const CLOSE: &str = "<|close|>"; - const SEP: &str = "<|sep|>"; - fn _arg(key: &str, typ: &str, value: &str) -> String { format!(r#"{OPEN}argument key="{key}" type="{typ}"{SEP}{value}{CLOSE}argument{SEP}"#) } @@ -694,4 +938,140 @@ mod tests { ); } } + + #[test] + fn test_factory_registers_kimi_k3_structural_tag() { + let factory = crate::factory::ParserFactory::new(); + assert!(factory.registry().has_structural_tag("kimi_k3")); + } + + fn sample_tools() -> Vec { + serde_json::from_value(serde_json::json!([{ + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "days": {"type": "integer"}, + "units": {"type": "string", "enum": ["c", "f"]} + }, + "required": ["city", "days"] + } + } + }])) + .unwrap() + } + + #[test] + fn test_build_structural_tag_frames_tools_section() { + let tag = KimiK3Parser::build_structural_tag(&sample_tools(), true); + let format = &tag["format"]; + assert_eq!(format["type"], "triggered_tags"); + assert_eq!(format["triggers"][0], TOOLS_OPEN); + assert_eq!(format["at_least_one"], true); + // `stop_after_first` is intentionally unset (mirrors K2) so trailing + // `<|close|>message<|sep|><|end_of_msg|>` tokens stay valid. + assert!(format.get("stop_after_first").is_none()); + + let section = &format["tags"][0]; + assert_eq!(section["begin"], TOOLS_OPEN); + assert_eq!(section["end"], TOOLS_CLOSE); + // A tools section is one or more call blocks. + assert_eq!(section["content"]["type"], "plus"); + } + + #[test] + fn test_build_structural_tag_pins_call_framing_and_tool_name() { + let tag = KimiK3Parser::build_structural_tag(&sample_tools(), true); + let call = &tag["format"]["tags"][0]["content"]["content"]["elements"][0]; + assert_eq!(call["type"], "sequence"); + let elems = call["elements"].as_array().unwrap(); + assert_eq!( + elems[0]["value"], + "<|open|>call tool=\"get_weather\" index=\"" + ); + assert_eq!(elems[1]["type"], "regex"); + assert_eq!(elems[2]["value"], "\"<|sep|>"); + assert_eq!(elems[4]["value"], "<|close|>call<|sep|>"); + } + + #[test] + fn test_build_structural_tag_strict_arguments_by_type() { + let tag = KimiK3Parser::build_structural_tag(&sample_tools(), true); + let call = &tag["format"]["tags"][0]["content"]["content"]["elements"][0]; + let args = &call["elements"][3]; + assert_eq!(args["type"], "sequence"); + let arg_elems = args["elements"].as_array().unwrap(); + // city (required, string), days (required, number), units (optional, enum). + assert_eq!(arg_elems.len(), 3); + + // Required string: bare sequence, verbatim body via any_text. + assert_eq!(arg_elems[0]["type"], "sequence"); + assert_eq!( + arg_elems[0]["elements"][0]["value"], + "<|open|>argument key=\"city\" type=\"string\"<|sep|>" + ); + assert_eq!(arg_elems[0]["elements"][1]["type"], "any_text"); + + // Required integer: rendered as type="number", JSON-schema-constrained value. + assert_eq!( + arg_elems[1]["elements"][0]["value"], + "<|open|>argument key=\"days\" type=\"number\"<|sep|>" + ); + assert_eq!(arg_elems[1]["elements"][1]["type"], "json_schema"); + + // Optional string enum: wrapped in `optional`, value is an `or` of literals. + assert_eq!(arg_elems[2]["type"], "optional"); + let units = &arg_elems[2]["content"]; + assert_eq!( + units["elements"][0]["value"], + "<|open|>argument key=\"units\" type=\"string\"<|sep|>" + ); + assert_eq!(units["elements"][1]["type"], "or"); + assert_eq!(units["elements"][1]["elements"][0]["value"], "c"); + assert_eq!(units["elements"][1]["elements"][1]["value"], "f"); + } + + #[test] + fn test_build_structural_tag_escapes_attribute_values() { + let tools: Vec = serde_json::from_value(serde_json::json!([{ + "type": "function", + "function": { + "name": "a&b\"c", + "parameters": { + "type": "object", + "properties": {"k\"q": {"type": "string"}}, + "required": ["k\"q"] + } + } + }])) + .unwrap(); + let tag = KimiK3Parser::build_structural_tag(&tools, true); + let call = &tag["format"]["tags"][0]["content"]["content"]["elements"][0]; + assert_eq!( + call["elements"][0]["value"], + "<|open|>call tool=\"a&b"c\" index=\"" + ); + let arg = &call["elements"][3]["elements"][0]; + assert_eq!( + arg["elements"][0]["value"], + "<|open|>argument key=\"k"q\" type=\"string\"<|sep|>" + ); + } + + #[test] + fn test_build_structural_tag_falls_back_when_no_properties() { + let tools: Vec = serde_json::from_value(serde_json::json!([{ + "type": "function", + "function": {"name": "ping", "parameters": {"type": "object"}} + }])) + .unwrap(); + let tag = KimiK3Parser::build_structural_tag(&tools, false); + assert_eq!(tag["format"]["at_least_one"], false); + let call = &tag["format"]["tags"][0]["content"]["content"]["elements"][0]; + // Argument region degrades to a permissive `star` of argument blocks. + assert_eq!(call["elements"][3]["type"], "star"); + } } From c4754790d2bdf5f82d477f72b2d62f4efbebe3f8 Mon Sep 17 00:00:00 2001 From: key4ng Date: Mon, 27 Jul 2026 11:13:10 -0700 Subject: [PATCH 4/5] docs(tool-parser): clarify k3 structural-tag argument-order and marker-exclusion intent Correct the build_structural_tag / build_args_format docs: the argument grammar emits one slot per property in the schema's declared order (preserve_order is on) with required properties mandatory and optionals individually skippable in place -- not "required first, then optional" as previously described. Note that this declared order is a generation constraint only; decode_call parses arguments order-insensitively. Also document why the permissive attribute regions exclude the `<|` marker prefix while free-text value regions exclude only the complete OPEN/CLOSE/SEP markers -- a deliberate distinction (quote-delimited identifiers vs marker-delimited free text), not an inconsistency. Signed-off-by: key4ng --- crates/tool_parser/src/parsers/kimi_k3.rs | 29 ++++++++++++++++------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/crates/tool_parser/src/parsers/kimi_k3.rs b/crates/tool_parser/src/parsers/kimi_k3.rs index 513060fce..96a77e51b 100644 --- a/crates/tool_parser/src/parsers/kimi_k3.rs +++ b/crates/tool_parser/src/parsers/kimi_k3.rs @@ -135,11 +135,16 @@ impl KimiK3Parser { /// tool calls live in a single section, so its `content` is a `plus` of /// call blocks (native parallel calls) and a lone section suffices — unlike /// K2, which repeats a tag per call. Within a call, arguments are - /// constrained per the JSON schema (Strict mode): every `required` property - /// must appear, in the schema's declared order, followed by any optional - /// properties; each `type=` attribute and each value is pinned to the - /// property's type. Schemas without usable `properties` fall back to a - /// permissive skeleton so the grammar is never infeasible. + /// constrained per the JSON schema (Strict mode): properties are pinned in + /// the schema's declared order (the crate enables `serde_json`'s + /// `preserve_order`), each `required` property mandatory and each optional + /// property individually skippable in its declared slot; every `type=` + /// attribute and value is pinned to the property's type. This declared + /// order is a *generation* constraint only — `decode_call` parses arguments + /// order-insensitively, so a model constrained here still round-trips even + /// though the grammar admits just one ordering. Schemas without usable + /// `properties` fall back to a permissive skeleton so the grammar is never + /// infeasible. /// /// `at_least_one` is wired to `tool_choice`: `true` for `"required"` (a /// tool call must be emitted), `false` for `"auto"`. Mirroring K2, @@ -411,9 +416,10 @@ fn build_call_format(name: &str, params: &Value) -> Value { } /// Build the argument-region grammar for a tool's parameter schema (Strict -/// mode): required properties in declared order, then optional properties. -/// Falls back to a permissive skeleton when the schema has no usable -/// `properties`. +/// mode): one slot per property in the schema's declared order, `required` +/// properties mandatory and the rest wrapped in `optional` (individually +/// skippable, in place). Falls back to a permissive skeleton when the schema +/// has no usable `properties`. fn build_args_format(params: &Value) -> Value { let Some(properties) = params.get("properties").and_then(Value::as_object) else { return permissive_args_format(); @@ -526,6 +532,13 @@ fn permissive_args_format() -> Value { } /// One structurally-valid argument block with unconstrained key/type/value. +/// +/// The two exclusion conventions below are deliberate, not an oversight: +/// quote-delimited attribute tokens (`key=`, `type=`) exclude the `<|` marker +/// *prefix*, which blocks any partial or full control marker inside a short +/// identifier; the free-text value is marker-delimited and may legitimately +/// contain a lone `<` or `|`, so it excludes only the complete `OPEN`/`CLOSE`/ +/// `SEP` markers to avoid over-restricting real content. fn permissive_argument_format() -> Value { json!({ "type": "sequence", From 5eb832edb07dc536d31558448844675d94ec1d9d Mon Sep 17 00:00:00 2001 From: key4ng Date: Mon, 27 Jul 2026 11:20:07 -0700 Subject: [PATCH 5/5] test(reasoning-parser): use a neutral example tool name in k3 streaming test Signed-off-by: key4ng --- crates/reasoning_parser/src/parsers/kimi_k3.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/reasoning_parser/src/parsers/kimi_k3.rs b/crates/reasoning_parser/src/parsers/kimi_k3.rs index f19da6a34..4976c16dd 100644 --- a/crates/reasoning_parser/src/parsers/kimi_k3.rs +++ b/crates/reasoning_parser/src/parsers/kimi_k3.rs @@ -504,7 +504,7 @@ mod tests { "<|sep|>", "<|open|>tools", "<|sep|>", - r#"<|open|>call tool="kvv_walle_case" index="1""#, + r#"<|open|>call tool="get_weather" index="1""#, "<|sep|>", "<|close|>call", "<|sep|>", @@ -525,7 +525,7 @@ mod tests { normal, concat!( "<|open|>tools<|sep|>", - r#"<|open|>call tool="kvv_walle_case" index="1"<|sep|>"#, + r#"<|open|>call tool="get_weather" index="1"<|sep|>"#, "<|close|>call<|sep|><|close|>tools<|sep|>" ) );