diff --git a/pyproject.toml b/pyproject.toml index 906702fa075..a4c7ff487a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,10 +129,9 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer "tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py", "docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html", "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*", - "rust/src/chat/src/renderer/deepseek_v32/fixtures/*", - "rust/src/parser/src/tool/gemma4.rs", "rust/src/parser/src/unified/gemma4.rs", + "rust/src/chat/src/renderer/deepseek_v32/fixtures/*", "rust/src/parser/**", "rust/src/text/src/output/decoded.rs", - "rust/src/tokenizer/src/incremental.rs", "rust/src/parser/src/reasoning/tests.rs"] + "rust/src/tokenizer/src/incremental.rs"] ignore-hidden = false [tool.typos.default] diff --git a/rust/Cargo.lock b/rust/Cargo.lock index ee4d24622be..38f1c257756 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -2220,7 +2220,7 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" [[package]] name = "llm-multimodal" version = "1.7.1" -source = "git+https://github.com/smg-project/llm-multimodal?rev=5390032d6dc8a3e6fdc83acd320260367eb4b9b5#5390032d6dc8a3e6fdc83acd320260367eb4b9b5" +source = "git+https://github.com/smg-project/llm-multimodal?rev=15adba5e025d8636ba4a334fb379b1371f6196a1#15adba5e025d8636ba4a334fb379b1371f6196a1" dependencies = [ "anyhow", "base64 0.22.1", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 42bda825cf8..09f55cf07cd 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -59,7 +59,7 @@ indexmap = "2.13.0" indicatif = "0.18.4" itertools = "0.14.0" libc = "0.2.177" -llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "5390032d6dc8a3e6fdc83acd320260367eb4b9b5", default-features = false, features = ["native-tls"] } +llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "15adba5e025d8636ba4a334fb379b1371f6196a1", default-features = false, features = ["native-tls"] } mimalloc = "0.1.52" minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] } minijinja-contrib = { version = "2.0", features = ["pycompat"] } diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index d4999d07cd0..583a5fb936e 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -20,7 +20,7 @@ use crate::output::{ use crate::renderer::hf::{HfChatRenderer, MultimodalRenderInfo}; use crate::renderer::{ DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, HarmonyChatRenderer, - InklingChatRenderer, + InklingChatRenderer, KimiK3ChatRenderer, }; use crate::request::ChatRequest; use crate::{DynChatOutputProcessor, RendererSelection}; @@ -73,6 +73,7 @@ impl HfChatBackend { RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new()), RendererSelection::Harmony => Arc::new(HarmonyChatRenderer::new()?), RendererSelection::Inkling => Arc::new(InklingChatRenderer::new(tokenizer.clone())?), + RendererSelection::KimiK3 => Arc::new(KimiK3ChatRenderer::new(tokenizer.clone())), }; info!( diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 9c38c40be4e..0d538a24ba1 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -33,7 +33,8 @@ pub use parser::tool::{ToolParser, ToolParserError, ToolParserFactory}; pub use renderer::hf::ChatTemplateContentFormatOption; pub use renderer::{ ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, - HarmonyChatRenderer, InklingChatRenderer, RenderedPrompt, RendererSelection, + HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, RenderedPrompt, + RendererSelection, }; pub use request::{ ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool, @@ -353,6 +354,12 @@ mod tests { .unwrap(); } + #[test] + fn validate_parser_overrides_accepts_explicit_kimi_k3() { + let selection = ParserSelection::Explicit("kimi_k3".to_string()); + validate_parser_overrides(&selection, &selection).unwrap(); + } + #[test] fn validate_parser_overrides_accepts_auto_and_none() { validate_parser_overrides(&ParserSelection::Auto, &ParserSelection::None).unwrap(); @@ -366,7 +373,7 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, inkling, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); } #[test] @@ -377,6 +384,6 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, inkling, kimi, kimi_k2, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); + expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, inkling, kimi, kimi_k2, kimi_k3, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); } } diff --git a/rust/src/chat/src/parser/reasoning/mod.rs b/rust/src/chat/src/parser/reasoning/mod.rs index ca025268b04..9a1e37533b4 100644 --- a/rust/src/chat/src/parser/reasoning/mod.rs +++ b/rust/src/chat/src/parser/reasoning/mod.rs @@ -27,6 +27,7 @@ pub mod names { pub const GLM45: &str = "glm45"; pub const KIMI: &str = "kimi"; pub const KIMI_K2: &str = "kimi_k2"; + pub const KIMI_K3: &str = "kimi_k3"; pub const MINIMAX_M2: &str = "minimax_m2"; pub const MINIMAX_M3: &str = "minimax_m3"; pub const NEMOTRON_V3: &str = "nemotron_v3"; @@ -68,6 +69,7 @@ impl ReasoningParserFactory { .register_parser::(names::GLM45) .register_parser::(names::KIMI) .register_parser::(names::KIMI_K2) + .register_unified_dummy(names::KIMI_K3) .register_parser::(names::MINIMAX_M2) .register_parser::(names::MINIMAX_M3) .register_parser::(names::NEMOTRON_V3) diff --git a/rust/src/chat/src/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs index 78ccde8a4b3..e42b45a4862 100644 --- a/rust/src/chat/src/parser/tool/mod.rs +++ b/rust/src/chat/src/parser/tool/mod.rs @@ -33,6 +33,7 @@ pub mod names { // also routes to `Internlm2ToolParser` despite the version-agnostic name. pub const INTERNLM: &str = "internlm"; pub const KIMI_K2: &str = "kimi_k2"; + pub const KIMI_K3: &str = "kimi_k3"; pub const LLAMA3_JSON: &str = "llama3_json"; pub const LLAMA4_JSON: &str = "llama4_json"; pub const MINIMAX_M2: &str = "minimax_m2"; @@ -78,6 +79,7 @@ impl ToolParserFactory { .register_parser::(names::HY_V3) .register_parser::(names::INTERNLM) .register_parser::(names::KIMI_K2) + .register_unified_dummy(names::KIMI_K3) .register_parser::(names::LLAMA3_JSON) .register_parser::(names::LLAMA4_JSON) .register_parser::(names::MINIMAX_M2) diff --git a/rust/src/chat/src/parser/unified.rs b/rust/src/chat/src/parser/unified.rs index 2536bf78c25..6246733db28 100644 --- a/rust/src/chat/src/parser/unified.rs +++ b/rust/src/chat/src/parser/unified.rs @@ -5,7 +5,9 @@ use std::sync::LazyLock; -pub use vllm_parser::unified::{Gemma4UnifiedParser, InklingUnifiedParser, UnifiedParser}; +pub use vllm_parser::unified::{ + Gemma4UnifiedParser, InklingUnifiedParser, KimiK3UnifiedParser, UnifiedParser, +}; use vllm_tokenizer::DynTokenizer; use crate::parser::ParserFactory; @@ -15,6 +17,7 @@ use crate::request::ChatTool; pub mod names { pub const GEMMA4: &str = "gemma4"; pub const INKLING: &str = "inkling"; + pub const KIMI_K3: &str = "kimi_k3"; } /// Constructor signature for one registered unified parser implementation. @@ -39,11 +42,14 @@ impl UnifiedParserFactory { factory.register_parser::(names::GEMMA4); factory.register_parser::(names::INKLING); + factory.register_parser::(names::KIMI_K3); factory .register_pattern("gemma-4", names::GEMMA4) .register_pattern("gemma4", names::GEMMA4) - .register_pattern("inkling", names::INKLING); + .register_pattern("inkling", names::INKLING) + .register_pattern("kimi-k3", names::KIMI_K3) + .register_pattern("kimi_k3", names::KIMI_K3); factory } @@ -121,4 +127,20 @@ mod tests { ); factory.create(names::INKLING, &[], Arc::new(inkling_tokenizer())).unwrap(); } + + #[test] + fn factory_registers_kimi_k3() { + let factory = UnifiedParserFactory::new(); + let tokenizer = TestTokenizer::new() + .with_regular_token("<|open|>", 1001) + .with_regular_token("<|close|>", 1002) + .with_regular_token("<|sep|>", 1003); + + assert!(factory.contains(names::KIMI_K3)); + assert_eq!( + factory.resolve_name_for_model("moonshotai/Kimi-K3"), + Some(names::KIMI_K3) + ); + factory.create(names::KIMI_K3, &[], Arc::new(tokenizer)).unwrap(); + } } diff --git a/rust/src/chat/src/renderer/deepseek_v32/tests.rs b/rust/src/chat/src/renderer/deepseek_v32/tests.rs index 76796edd096..7960927219b 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/tests.rs @@ -59,7 +59,7 @@ fn fixture_request(input_name: &str) -> ChatRequest { fn deepseek_fixture_options() -> FixtureRequestOptions { FixtureRequestOptions { - enable_thinking: true, + enable_thinking: Some(true), no_generation_prompt_when_last_assistant: true, } } diff --git a/rust/src/chat/src/renderer/deepseek_v4/tests.rs b/rust/src/chat/src/renderer/deepseek_v4/tests.rs index 73380cef260..9068d460ee6 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/tests.rs @@ -27,7 +27,7 @@ fn fixture_request(input_name: &str) -> ChatRequest { fn deepseek_fixture_options() -> FixtureRequestOptions { FixtureRequestOptions { - enable_thinking: true, + enable_thinking: Some(true), no_generation_prompt_when_last_assistant: true, } } diff --git a/rust/src/chat/src/renderer/harmony/tests.rs b/rust/src/chat/src/renderer/harmony/tests.rs index 301d9707e71..4849c932233 100644 --- a/rust/src/chat/src/renderer/harmony/tests.rs +++ b/rust/src/chat/src/renderer/harmony/tests.rs @@ -22,7 +22,7 @@ fn fixture_request(input_name: &str) -> ChatRequest { fixture_chat_request( &fixture_path(input_name), FixtureRequestOptions { - enable_thinking: false, + enable_thinking: Some(false), no_generation_prompt_when_last_assistant: false, }, ) diff --git a/rust/src/chat/src/renderer/inkling/tests.rs b/rust/src/chat/src/renderer/inkling/tests.rs index 57e08ac0158..1c4cd7236db 100644 --- a/rust/src/chat/src/renderer/inkling/tests.rs +++ b/rust/src/chat/src/renderer/inkling/tests.rs @@ -105,7 +105,7 @@ fn fixture_request(name: &str) -> ChatRequest { fn inkling_fixture_options() -> FixtureRequestOptions { FixtureRequestOptions { - enable_thinking: false, + enable_thinking: Some(false), no_generation_prompt_when_last_assistant: false, } } diff --git a/rust/src/chat/src/renderer/kimi_k3/encoding.rs b/rust/src/chat/src/renderer/kimi_k3/encoding.rs new file mode 100644 index 00000000000..189d68c101c --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/encoding.rs @@ -0,0 +1,608 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Kimi K3 XTML prompt renderer. +//! +//! Port of Moonshot remote-code `encoding_k3.py::build_chat_segments()`. + +use std::collections::HashMap; + +use serde_json::{Map, Value, json}; +use vllm_tokenizer::Tokenizer; + +use crate::error::{Error, Result}; +use crate::request::{ + ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice, +}; +use crate::{AssistantContentBlock, AssistantToolCall}; + +pub(super) const OPEN: &str = "<|open|>"; +pub(super) const CLOSE: &str = "<|close|>"; +pub(super) const SEP: &str = "<|sep|>"; +pub(super) const END_OF_MSG: &str = "<|end_of_msg|>"; +pub(super) const IMAGE_PLACEHOLDER: &str = "<|media_pad|>"; + +const DEFAULT_THINKING_EFFORT: &str = "max"; +const VALID_THINKING_EFFORTS: &[&str] = &["low", "high", "max"]; + +/// K3 prompt encoder preserving Python's per-segment tokenization boundaries. +pub(super) struct K3TokenWriter<'a> { + tokenizer: &'a dyn Tokenizer, + token_ids: Vec, +} + +impl<'a> K3TokenWriter<'a> { + pub(super) fn new(tokenizer: &'a dyn Tokenizer) -> Self { + Self { + tokenizer, + token_ids: Vec::new(), + } + } + + /// Encode one trusted segment with normal added-token recognition. + pub(super) fn control(&mut self, text: &str) -> Result<()> { + if !text.is_empty() { + self.token_ids.extend(self.tokenizer.encode(text, false)?); + } + Ok(()) + } + + /// Encode one literal segment while bypassing every added-token matcher. + pub(super) fn ordinary(&mut self, text: &str) -> Result<()> { + if !text.is_empty() { + self.token_ids.extend(self.tokenizer.encode_ordinary(text)?); + } + Ok(()) + } + + pub(super) fn finish(self) -> Vec { + self.token_ids + } +} + +/// Render and tokenize one chat request using K3's segment-aware contract. +pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -> Result> { + let thinking = thinking_enabled(request)?; + let thinking_effort = thinking.then(|| thinking_effort(request)).transpose()?; + let tools = request_tools(request); + let mut out = K3TokenWriter::new(tokenizer); + + if !tools.is_empty() { + write_tool_declare(&mut out, tools, false)?; + } + + if let Some(effort) = thinking_effort { + // Preserve the checkpoint's literal guidance text: it still names + // `medium`, although the validator above no longer accepts it. + write_internal_system( + &mut out, + "thinking-effort", + &format!( + "`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}`." + ), + )?; + } + + // Track prior assistant tool-call ids for tool-result reordering / naming. + let mut tool_call_id_index: HashMap = HashMap::new(); + let mut pending_tool_run: Vec<(usize, ChatMessage)> = Vec::new(); + + let flush_tool_run = |out: &mut K3TokenWriter<'_>, + run: &mut Vec<(usize, ChatMessage)>, + id_index: &HashMap| + -> Result<()> { + if run.is_empty() { + return Ok(()); + } + + let mut resolved = Vec::with_capacity(run.len()); + let mut unresolved = false; + for (offset, message) in run.drain(..) { + let ChatMessage::ToolResponse { + content, + tool_call_id, + } = message + else { + unreachable!("pending tool run only holds tool responses"); + }; + match id_index.get(&tool_call_id) { + Some(&(position, ref name)) => { + resolved.push((position, offset, content, Some(name.clone()))); + } + None => { + unresolved = true; + resolved.push((usize::MAX, offset, content, None)); + } + } + } + + if unresolved { + // Preserve original order when the run cannot be fully matched. + resolved.sort_by_key(|item| item.1); + } else { + resolved.sort_by_key(|item| (item.0, item.1)); + } + + for (xtml_index, (_, _, content, name)) in resolved.into_iter().enumerate() { + let tool_name = name.as_deref().ok_or_else(|| { + Error::ChatTemplate( + "Kimi K3 tool messages need a resolvable tool name: \ + carry a matching tool_call_id against a preceding \ + assistant tool_call" + .to_string(), + ) + })?; + write_tool_message(out, tool_name, xtml_index + 1, &content)?; + } + Ok(()) + }; + + for (message_index, message) in request.messages.iter().enumerate() { + match message { + ChatMessage::ToolResponse { .. } => { + pending_tool_run.push((message_index, message.clone())); + continue; + } + _ => { + flush_tool_run(&mut out, &mut pending_tool_run, &tool_call_id_index)?; + } + } + + match message { + // Python: role=system with a `tools` field renders a dynamic + // tool-declare (`## New Tools Available`). Map that to Developer + // messages that carry tools (OpenAI "developer" / system-tools). + ChatMessage::Developer { + content, + tools: Some(local_tools), + } if !local_tools.is_empty() => { + write_tool_declare(&mut out, local_tools, true)?; + if !content_is_empty(content) { + write_role_message(&mut out, "system", None, content)?; + } + } + ChatMessage::System { content } | ChatMessage::Developer { content, .. } => { + write_role_message(&mut out, "system", None, content)?; + } + ChatMessage::User { content } => { + write_role_message(&mut out, "user", None, content)?; + } + ChatMessage::Assistant { content } => { + tool_call_id_index.clear(); + let mut call_position = 0usize; + for block in content { + if let AssistantContentBlock::ToolCall(call) = block { + call_position += 1; + if !call.id.is_empty() { + tool_call_id_index + .entry(call.id.clone()) + .or_insert((call_position, call.name.clone())); + } + } + } + + write_assistant_message(&mut out, content, thinking)?; + } + ChatMessage::ToolResponse { .. } => unreachable!("handled above"), + } + } + flush_tool_run(&mut out, &mut pending_tool_run, &tool_call_id_index)?; + + match &request.tool_choice { + ChatToolChoice::Required => { + write_internal_system( + &mut out, + "tool-choice", + "The system is invoked with `tool_choice=required`.\n\ + You MUST call tools in the next message.", + )?; + } + // Emit only when tools are present: Rust defaults tool_choice to None + // for tool-free requests, which must not inject a tool-choice message. + ChatToolChoice::None if !request.tools.is_empty() => { + write_internal_system( + &mut out, + "tool-choice", + "The system is invoked with `tool_choice=none`.\n\ + You MUST NOT call any tools in the next message.", + )?; + } + ChatToolChoice::None | ChatToolChoice::Auto | ChatToolChoice::Function { .. } => {} + } + + write_response_format(&mut out, request)?; + + if request.chat_options.add_generation_prompt() { + write_open_tag(&mut out, "message", &[("role", "assistant")])?; + write_open_tag(&mut out, if thinking { "think" } else { "response" }, &[])?; + } + + Ok(out.finish()) +} + +fn request_tools(request: &ChatRequest) -> &[ChatTool] { + // Declare tools whenever the request carries them. K3 tool-declare is + // independent of tool_choice; tool_choice only injects control messages. + request.tools.as_slice() +} + +fn thinking_enabled(request: &ChatRequest) -> Result { + if let Some(thinking) = request.parse_template_bool("thinking")? { + return Ok(thinking); + } + if let Some(enable_thinking) = request.parse_template_bool("enable_thinking")? { + return Ok(enable_thinking); + } + Ok(request + .chat_options + .reasoning_effort + .map(|effort| effort != crate::request::ReasoningEffort::None) + .unwrap_or(true)) +} + +fn thinking_effort(request: &ChatRequest) -> Result { + let effort = if let Some(value) = request.chat_options.template_kwargs.get("thinking_effort") { + value.as_str().ok_or_else(|| { + Error::ChatTemplate(format!( + "template kwarg `thinking_effort` must be a string, got {value}" + )) + })? + } else if let Some(effort) = request.chat_options.reasoning_effort { + effort.as_str() + } else if let Some(value) = request.chat_options.template_kwargs.get("reasoning_effort") { + value.as_str().ok_or_else(|| { + Error::ChatTemplate(format!( + "template kwarg `reasoning_effort` must be a string, got {value}" + )) + })? + } else { + DEFAULT_THINKING_EFFORT + }; + + if !VALID_THINKING_EFFORTS.contains(&effort) { + return Err(Error::ChatTemplate(format!( + "unsupported thinking_effort={effort:?}; supported values are `low`, `high`, and `max`" + ))); + } + Ok(effort.to_string()) +} + +fn content_is_empty(content: &ChatContent) -> bool { + match content { + ChatContent::Text(text) => text.is_empty(), + ChatContent::Parts(parts) => parts.iter().all(|part| match part { + ChatContentPart::Text { text } => text.is_empty(), + ChatContentPart::ImageUrl { .. } + | ChatContentPart::VideoUrl { .. } + | ChatContentPart::InputAudio { .. } + | ChatContentPart::AudioUrl { .. } => false, + }), + } +} + +fn write_tool_declare( + out: &mut K3TokenWriter<'_>, + tools: &[ChatTool], + dynamic: bool, +) -> Result<()> { + let mut specs = Vec::with_capacity(tools.len()); + for tool in tools { + let mut function = Map::new(); + function.insert( + "description".to_string(), + Value::String(tool.description.clone().unwrap_or_default()), + ); + function.insert("name".to_string(), Value::String(tool.name.clone())); + function.insert("parameters".to_string(), sort_json(&tool.parameters)); + specs.push(json!({ + "function": Value::Object(function), + "type": "function", + })); + } + let payload = compact_json(&sort_json(&Value::Array(specs)))?; + + 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\ + {payload}\n\ + ```" + ) + } else { + format!( + "# Tools\n\ + Here are the available tools, described in JSONSchema.\n\n\ + ```json\n\ + {payload}\n\ + ```" + ) + }; + + write_internal_system(out, "tool-declare", &body) +} + +fn write_internal_system( + out: &mut K3TokenWriter<'_>, + message_type: &str, + body: &str, +) -> Result<()> { + write_open_tag( + out, + "message", + &[("role", "system"), ("type", message_type)], + )?; + out.ordinary(body.trim())?; + write_close_tag(out, "message")?; + out.control(END_OF_MSG) +} + +fn write_role_message( + out: &mut K3TokenWriter<'_>, + role: &str, + name: Option<&str>, + content: &ChatContent, +) -> Result<()> { + let mut attrs = vec![("role", role.to_string())]; + if let Some(name) = name.filter(|name| !name.is_empty()) { + attrs.push(("name", name.to_string())); + } + let attr_refs: Vec<(&str, &str)> = attrs.iter().map(|(k, v)| (*k, v.as_str())).collect(); + write_open_tag(out, "message", &attr_refs)?; + write_content(out, content)?; + write_close_tag(out, "message")?; + out.control(END_OF_MSG) +} + +fn write_tool_message( + out: &mut K3TokenWriter<'_>, + tool_name: &str, + index: usize, + content: &ChatContent, +) -> Result<()> { + let index_str = index.to_string(); + write_open_tag( + out, + "message", + &[("role", "tool"), ("tool", tool_name), ("index", &index_str)], + )?; + write_content(out, content)?; + write_close_tag(out, "message")?; + out.control(END_OF_MSG) +} + +fn write_assistant_message( + out: &mut K3TokenWriter<'_>, + content: &[AssistantContentBlock], + thinking: bool, +) -> Result<()> { + write_open_tag(out, "message", &[("role", "assistant")])?; + + let mut reasoning = String::new(); + let mut response = String::new(); + let mut tool_calls = Vec::new(); + for block in content { + match block { + AssistantContentBlock::Reasoning { text } => reasoning.push_str(text), + AssistantContentBlock::Text { text } => response.push_str(text), + AssistantContentBlock::ToolCall(call) => tool_calls.push(call), + } + } + + // The think channel is structural: in thinking mode every assistant + // message carries open/close tags even when there is no reasoning content. + // In non-thinking mode the channel is dropped entirely. + if thinking { + write_open_tag(out, "think", &[])?; + if !reasoning.trim().is_empty() { + out.ordinary(&reasoning)?; + } + write_close_tag(out, "think")?; + } + + write_open_tag(out, "response", &[])?; + out.ordinary(&response)?; + write_close_tag(out, "response")?; + + if !tool_calls.is_empty() { + write_open_tag(out, "tools", &[])?; + for (index, tool_call) in tool_calls.into_iter().enumerate() { + write_assistant_tool_call(out, tool_call, index + 1)?; + } + write_close_tag(out, "tools")?; + } + + write_close_tag(out, "message")?; + out.control(END_OF_MSG) +} + +fn write_assistant_tool_call( + out: &mut K3TokenWriter<'_>, + tool_call: &AssistantToolCall, + index: usize, +) -> Result<()> { + let index_str = index.to_string(); + write_open_tag( + out, + "call", + &[ + ("tool", tool_call.name.as_str()), + ("index", index_str.as_str()), + ], + )?; + + let (args, json_block) = normalize_tool_arguments(&tool_call.arguments)?; + if let Some(raw) = json_block { + write_open_tag(out, "json", &[("type", "object")])?; + out.ordinary(&raw)?; + write_close_tag(out, "json")?; + } else { + for (key, value) in args { + let typ = xtml_type(&value); + write_open_tag(out, "argument", &[("key", key.as_str()), ("type", typ)])?; + out.ordinary(&xtml_value(&value))?; + write_close_tag(out, "argument")?; + } + } + + write_close_tag(out, "call") +} + +fn write_content(out: &mut K3TokenWriter<'_>, content: &ChatContent) -> Result<()> { + match content { + ChatContent::Text(text) => write_text_with_images(out, text), + ChatContent::Parts(parts) => { + for part in parts { + match part { + ChatContentPart::Text { text } => write_text_with_images(out, text)?, + ChatContentPart::ImageUrl { .. } => out.control(IMAGE_PLACEHOLDER)?, + ChatContentPart::VideoUrl { .. } => { + return Err(Error::UnsupportedMultimodalContent("video_url")); + } + ChatContentPart::InputAudio { .. } => { + return Err(Error::UnsupportedMultimodalContent("input_audio")); + } + ChatContentPart::AudioUrl { .. } => { + return Err(Error::UnsupportedMultimodalContent("audio_url")); + } + } + } + Ok(()) + } + } +} + +fn write_text_with_images(out: &mut K3TokenWriter<'_>, text: &str) -> Result<()> { + // Placeholder expansion is left as the literal K3 image token; multimodal + // preprocessing can replace it once image prompts are known. + out.ordinary(text) +} + +fn write_response_format(out: &mut K3TokenWriter<'_>, request: &ChatRequest) -> Result<()> { + let Some(rf) = request.chat_options.response_format.as_ref() else { + return Ok(()); + }; + + let rf_type = rf.get("type").and_then(Value::as_str).or_else(|| rf.as_str()).unwrap_or(""); + + match rf_type { + "json_object" => { + write_internal_system( + 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.", + )?; + } + "json_schema" => { + let schema = extract_response_schema(rf); + let schema_json = compact_json(&sort_json(&schema.unwrap_or(Value::Null)))?; + write_internal_system( + out, + "response-format", + &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_json}\n\ + ```" + ), + )?; + } + _ => {} + } + Ok(()) +} + +fn extract_response_schema(response_format: &Value) -> Option { + let json_schema = response_format.get("json_schema")?; + if let Some(schema) = json_schema.get("schema") { + return Some(schema.clone()); + } + if let Some(schema) = json_schema.get("json_schema") { + return Some(schema.clone()); + } + Some(json_schema.clone()) +} + +fn normalize_tool_arguments(arguments: &str) -> Result<(Map, Option)> { + let trimmed = arguments.trim(); + if trimmed.is_empty() { + return Ok((Map::new(), None)); + } + match serde_json::from_str::(trimmed) { + Ok(Value::Object(map)) => Ok((map, None)), + Ok(_) => Err(Error::ChatTemplate( + "Kimi K3 tool call arguments must be a JSON object".to_string(), + )), + Err(_) => Ok((Map::new(), Some(arguments.to_string()))), + } +} + +fn xtml_type(value: &Value) -> &'static str { + match value { + Value::Bool(_) => "boolean", + Value::Null => "null", + Value::Number(_) => "number", + Value::String(_) => "string", + Value::Object(_) => "object", + Value::Array(_) => "array", + } +} + +fn xtml_value(value: &Value) -> String { + match value { + Value::String(text) => text.clone(), + other => compact_json(other).unwrap_or_else(|_| other.to_string()), + } +} + +fn write_open_tag(out: &mut K3TokenWriter<'_>, tag: &str, attrs: &[(&str, &str)]) -> Result<()> { + out.control(OPEN)?; + out.ordinary(tag)?; + for (key, value) in attrs { + out.ordinary(&format!(" {key}"))?; + out.ordinary("=\"")?; + out.ordinary(&escape_attr_value(value))?; + out.ordinary("\"")?; + } + out.control(SEP) +} + +fn write_close_tag(out: &mut K3TokenWriter<'_>, tag: &str) -> Result<()> { + out.control(CLOSE)?; + out.ordinary(tag)?; + out.control(SEP) +} + +fn escape_attr_value(value: &str) -> String { + value.replace('&', "&").replace('"', """) +} + +fn compact_json(value: &Value) -> Result { + serde_json::to_string(value).map_err(|error| Error::ChatTemplate(error.to_string())) +} + +fn sort_json(value: &Value) -> Value { + match value { + Value::Array(items) => Value::Array(items.iter().map(sort_json).collect()), + Value::Object(map) => { + let mut sorted = Map::new(); + let mut keys = map.keys().collect::>(); + keys.sort(); + for key in keys { + sorted.insert(key.clone(), sort_json(&map[key])); + } + Value::Object(sorted) + } + _ => value.clone(), + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_input.json new file mode 100644 index 00000000000..569201f3b97 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_input.json @@ -0,0 +1,15 @@ +{ + "messages": [ + { + "role": "user", + "content": "json pls" + } + ], + "add_generation_prompt": true, + "response_format": { + "type": "json_object" + }, + "template_kwargs": { + "thinking": false + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_output.txt new file mode 100644 index 00000000000..57ce4a68761 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_output.txt @@ -0,0 +1,2 @@ +<|open|>message role="user"<|sep|>json pls<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="response-format"<|sep|>The system is invoked with `response_format=json_object`. +Your response must be raw JSON data without markdown code blocks (```json) or any additional formatting.<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>response<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_input.json new file mode 100644 index 00000000000..61ecebc268b --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_input.json @@ -0,0 +1,62 @@ +{ + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "weather", + "parameters": { + "type": "object", + "properties": { + "city": { + "type": "string" + }, + "days": { + "type": "number" + } + }, + "required": [ + "city" + ] + } + } + } + ], + "messages": [ + { + "role": "user", + "content": "hi" + }, + { + "role": "developer", + "tools": [ + { + "type": "function", + "function": { + "name": "calc", + "description": "calculator", + "parameters": { + "type": "object", + "properties": { + "x": { + "type": "number" + } + }, + "required": [ + "x" + ] + } + } + } + ] + }, + { + "role": "user", + "content": "use calc" + } + ], + "add_generation_prompt": true, + "template_kwargs": { + "thinking": true + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_output.txt new file mode 100644 index 00000000000..edec2bebf2a --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_output.txt @@ -0,0 +1,14 @@ +<|open|>message role="system" type="tool-declare"<|sep|># Tools +Here are the available tools, described in JSONSchema. + +```json +[{"function":{"description":"weather","name":"get_weather","parameters":{"properties":{"city":{"type":"string"},"days":{"type":"number"}},"required":["city"],"type":"object"}},"type":"function"}] +```<|close|>message<|sep|><|end_of_msg|><|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`. +Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>hi<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="tool-declare"<|sep|>## New Tools Available +The system dynamically extends the toolset via lazy-loading. +You have access to all existing and extended tools. +Here are the specs for the extended tools. + +```json +[{"function":{"description":"calculator","name":"calc","parameters":{"properties":{"x":{"type":"number"}},"required":["x"],"type":"object"}},"type":"function"}] +```<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>use calc<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_input.json new file mode 100644 index 00000000000..b7cf40c9b8b --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_input.json @@ -0,0 +1,31 @@ +{ + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "see this" + }, + { + "type": "image_url", + "image_url": "https://example.com/a.png" + } + ] + }, + { + "role": "assistant", + "content": "old answer", + "reasoning_content": "old reasoning" + }, + { + "role": "user", + "content": "continue" + } + ], + "add_generation_prompt": true, + "template_kwargs": { + "thinking": true, + "thinking_effort": "high" + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_output.txt new file mode 100644 index 00000000000..03582f1180a --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_output.txt @@ -0,0 +1,2 @@ +<|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`. +Now the system is invoked with `thinking_effort=high`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>see this<|media_pad|><|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>old reasoning<|close|>think<|sep|><|open|>response<|sep|>old answer<|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>continue<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_input.json new file mode 100644 index 00000000000..3a64f8dd0bb --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_input.json @@ -0,0 +1,87 @@ +{ + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "weather", + "parameters": { + "type": "object", + "properties": { + "city": { + "type": "string" + }, + "days": { + "type": "number" + } + }, + "required": [ + "city" + ] + } + } + }, + { + "type": "function", + "function": { + "name": "calc", + "description": "calculator", + "parameters": { + "type": "object", + "properties": { + "x": { + "type": "number" + } + }, + "required": [ + "x" + ] + } + } + } + ], + "messages": [ + { + "role": "user", + "content": "Hangzhou weather and calc" + }, + { + "role": "assistant", + "content": "I'll check.", + "reasoning_content": "Need tools.", + "tool_calls": [ + { + "id": "get_weather:0", + "type": "function", + "function": { + "name": "get_weather", + "arguments": "{\"city\":\"Hangzhou\",\"days\":1}" + } + }, + { + "id": "calc:1", + "type": "function", + "function": { + "name": "calc", + "arguments": "{\"x\":2}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "calc:1", + "content": "2" + }, + { + "role": "tool", + "tool_call_id": "get_weather:0", + "content": "{\"city\":\"Hangzhou\",\"condition\":\"rain\"}" + } + ], + "add_generation_prompt": true, + "tool_choice": "required", + "template_kwargs": { + "thinking": true + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_output.txt new file mode 100644 index 00000000000..1c24810072a --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_output.txt @@ -0,0 +1,8 @@ +<|open|>message role="system" type="tool-declare"<|sep|># Tools +Here are the available tools, described in JSONSchema. + +```json +[{"function":{"description":"weather","name":"get_weather","parameters":{"properties":{"city":{"type":"string"},"days":{"type":"number"}},"required":["city"],"type":"object"}},"type":"function"},{"function":{"description":"calculator","name":"calc","parameters":{"properties":{"x":{"type":"number"}},"required":["x"],"type":"object"}},"type":"function"}] +```<|close|>message<|sep|><|end_of_msg|><|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`. +Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>Hangzhou weather and calc<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>Need tools.<|close|>think<|sep|><|open|>response<|sep|>I'll check.<|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool="get_weather" index="1"<|sep|><|open|>argument key="city" type="string"<|sep|>Hangzhou<|close|>argument<|sep|><|open|>argument key="days" type="number"<|sep|>1<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool="calc" index="2"<|sep|><|open|>argument key="x" type="number"<|sep|>2<|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|>{"city":"Hangzhou","condition":"rain"}<|close|>message<|sep|><|end_of_msg|><|open|>message role="tool" tool="calc" index="2"<|sep|>2<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="tool-choice"<|sep|>The system is invoked with `tool_choice=required`. +You MUST call tools in the next message.<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/mod.rs b/rust/src/chat/src/renderer/kimi_k3/mod.rs new file mode 100644 index 00000000000..09bf356f5d6 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/mod.rs @@ -0,0 +1,39 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Native Kimi K3 XTML chat renderer. + +mod encoding; +#[cfg(test)] +mod tests; + +use vllm_text::Prompt; +use vllm_text::tokenizer::DynTokenizer; + +use super::{ChatRenderer, RenderedPrompt, request_template_kwargs}; +use crate::Result; +use crate::request::ChatRequest; + +/// Dedicated Kimi K3 XTML renderer. +#[derive(Clone)] +pub struct KimiK3ChatRenderer { + tokenizer: DynTokenizer, +} + +impl KimiK3ChatRenderer { + /// Create a Kimi K3 renderer. + pub fn new(tokenizer: DynTokenizer) -> Self { + Self { tokenizer } + } +} + +impl ChatRenderer for KimiK3ChatRenderer { + fn render(&self, request: &ChatRequest) -> Result { + request.validate()?; + + Ok(RenderedPrompt { + prompt: Prompt::TokenIds(encoding::render_request(request, self.tokenizer.as_ref())?), + effective_template_kwargs: request_template_kwargs(request), + }) + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/tests.rs b/rust/src/chat/src/renderer/kimi_k3/tests.rs new file mode 100644 index 00000000000..7bab90ddcf7 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/tests.rs @@ -0,0 +1,267 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Golden fixtures generated from HF remote-code `encoding_k3.py`. + +use std::path::PathBuf; +use std::sync::Arc; + +use expect_test::{expect, expect_file}; +use serde_json::json; +use vllm_text::Prompt; +use vllm_text::tokenizer::DynTokenizer; +use vllm_tokenizer::Tokenizer; +use vllm_tokenizer::test_utils::TestTokenizer; + +use super::KimiK3ChatRenderer; +use crate::AssistantContentBlock; +use crate::ChatRenderer; +use crate::renderer::kimi_k3::encoding::{CLOSE, END_OF_MSG, IMAGE_PLACEHOLDER, OPEN, SEP}; +use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; +use crate::request::{ChatContentPart, ChatMessage, GenerationPromptMode, ReasoningEffort}; + +const OPEN_ID: u32 = 256; +const CLOSE_ID: u32 = 257; +const SEP_ID: u32 = 258; +const END_OF_MSG_ID: u32 = 259; +const MEDIA_ID: u32 = 260; + +fn test_tokenizer() -> TestTokenizer { + TestTokenizer::new() + .with_special_token(OPEN, OPEN_ID) + .with_special_token(CLOSE, CLOSE_ID) + .with_special_token(SEP, SEP_ID) + .with_special_token(END_OF_MSG, END_OF_MSG_ID) + .with_special_token(IMAGE_PLACEHOLDER, MEDIA_ID) +} + +fn render_token_ids(request: &crate::request::ChatRequest, tokenizer: DynTokenizer) -> Vec { + let prompt = KimiK3ChatRenderer::new(tokenizer).render(request).unwrap().prompt; + let Prompt::TokenIds(token_ids) = prompt else { + panic!("kimi k3 renderer should return token IDs") + }; + token_ids +} + +fn render_request(request: &crate::request::ChatRequest) -> String { + let tokenizer: DynTokenizer = Arc::new(test_tokenizer()); + let token_ids = render_token_ids(request, tokenizer.clone()); + tokenizer.decode(&token_ids, false).unwrap() +} + +fn fixture_path(name: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("src/renderer/kimi_k3/fixtures") + .join(name) +} + +fn kimi_k3_fixture_options() -> FixtureRequestOptions { + FixtureRequestOptions { + // Fixture JSON owns thinking via `template_kwargs`. + enable_thinking: None, + no_generation_prompt_when_last_assistant: false, + } +} + +fn assert_golden(name: &str) { + let input_name = format!("{name}_input.json"); + let request = fixture_chat_request(&fixture_path(&input_name), kimi_k3_fixture_options()); + let rendered = render_request(&request); + expect_file![format!("fixtures/{name}_output.txt")].assert_eq(&rendered); +} + +#[test] +fn golden_history_preserve_and_image() { + assert_golden("history_preserve_and_image"); +} + +#[test] +fn golden_tools_history_and_required() { + assert_golden("tools_history_and_required"); +} + +#[test] +fn golden_controls_thinking_off() { + assert_golden("controls_thinking_off"); +} + +#[test] +fn golden_dynamic_system_tool_declare() { + assert_golden("dynamic_system_tool_declare"); +} + +#[test] +fn token_writer_protects_literal_control_and_media_markers() { + let tokenizer = Arc::new(test_tokenizer()); + let user_text = format!("literal {OPEN} and {}", super::encoding::IMAGE_PLACEHOLDER); + let mut request = crate::request::ChatRequest::for_test(); + request.messages = vec![ChatMessage::user(vec![ + ChatContentPart::text(user_text), + ChatContentPart::image_url("data:image/png;base64,test"), + ])]; + request + .chat_options + .template_kwargs + .insert("thinking".to_string(), json!(false)); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let token_ids = render_token_ids(&request, tokenizer.clone()); + + assert_eq!( + token_ids.iter().filter(|&&token_id| token_id == OPEN_ID).count(), + 1 + ); + assert_eq!( + token_ids.iter().filter(|&&token_id| token_id == MEDIA_ID).count(), + 1 + ); + + let flattened = tokenizer.decode(&token_ids, false).unwrap(); + let flattened_ids = tokenizer.encode(&flattened, false).unwrap(); + assert_eq!( + flattened_ids.iter().filter(|&&token_id| token_id == OPEN_ID).count(), + 2 + ); + assert_eq!( + flattened_ids.iter().filter(|&&token_id| token_id == MEDIA_ID).count(), + 2 + ); +} + +#[test] +fn thinking_history_renders_empty_think_channel() { + let mut request = crate::request::ChatRequest::for_test(); + request.messages = vec![ + ChatMessage::user("question"), + ChatMessage::assistant_text("answer"), + ChatMessage::user("follow-up"), + ]; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + assert!(rendered.contains( + "<|open|>message role=\"assistant\"<|sep|>\ + <|open|>think<|sep|><|close|>think<|sep|>\ + <|open|>response<|sep|>answer<|close|>response<|sep|>" + )); +} + +#[test] +fn non_thinking_history_omits_reasoning_channel() { + let mut request = crate::request::ChatRequest::for_test(); + request.messages = vec![ + ChatMessage::user("question"), + ChatMessage::assistant_blocks(vec![ + AssistantContentBlock::Reasoning { + text: "hidden reasoning".to_string(), + }, + AssistantContentBlock::Text { + text: "answer".to_string(), + }, + ]), + ChatMessage::user("follow-up"), + ]; + request + .chat_options + .template_kwargs + .insert("thinking".to_string(), json!(false)); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + assert!(!rendered.contains("hidden reasoning")); + assert!(!rendered.contains("<|open|>think<|sep|>")); + assert!(rendered.contains( + "<|open|>message role=\"assistant\"<|sep|>\ + <|open|>response<|sep|>answer<|close|>response<|sep|>" + )); +} + +#[test] +fn defaults_thinking_effort_to_max() { + let rendered = render_request(&crate::request::ChatRequest::for_test()); + + expect![[r#"<|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`. +Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>test<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>"#]] + .assert_eq(&rendered); +} + +#[test] +fn translates_standard_thinking_kwargs() { + let mut request = crate::request::ChatRequest::for_test(); + request + .chat_options + .template_kwargs + .insert("enable_thinking".to_string(), json!(true)); + request + .chat_options + .template_kwargs + .insert("reasoning_effort".to_string(), json!("high")); + + let rendered = render_request(&request); + + assert!(rendered.contains("thinking_effort=high")); + assert!(rendered.ends_with("<|open|>think<|sep|>")); +} + +#[test] +fn native_k3_kwargs_take_precedence() { + let mut request = crate::request::ChatRequest::for_test(); + request.chat_options.template_kwargs.extend([ + ("thinking".to_string(), json!(true)), + ("enable_thinking".to_string(), json!(false)), + ("thinking_effort".to_string(), json!("low")), + ("reasoning_effort".to_string(), json!("high")), + ]); + + let rendered = render_request(&request); + + assert!(rendered.contains("thinking_effort=low")); + assert!(!rendered.contains("thinking_effort=high")); +} + +#[test] +fn standard_none_disables_thinking() { + let mut request = crate::request::ChatRequest::for_test(); + request.chat_options.template_kwargs.extend([ + ("enable_thinking".to_string(), json!(false)), + ("reasoning_effort".to_string(), json!("none")), + ]); + + let rendered = render_request(&request); + + assert!(!rendered.contains("type=\"thinking-effort\"")); + assert!(rendered.ends_with("<|open|>response<|sep|>")); +} + +#[test] +fn typed_none_disables_thinking() { + let mut request = crate::request::ChatRequest::for_test(); + request.chat_options.reasoning_effort = Some(ReasoningEffort::None); + + let rendered = render_request(&request); + + assert!(!rendered.contains("type=\"thinking-effort\"")); + assert!(rendered.ends_with("<|open|>response<|sep|>")); +} + +#[test] +fn rejects_removed_medium_thinking_effort() { + let mut request = crate::request::ChatRequest::for_test(); + request + .chat_options + .template_kwargs + .insert("thinking_effort".to_string(), json!("medium")); + + let error = KimiK3ChatRenderer::new(Arc::new(test_tokenizer())) + .render(&request) + .unwrap_err(); + + expect![[r#" + ChatTemplate( + "unsupported thinking_effort=\"medium\"; supported values are `low`, `high`, and `max`", + ) + "#]] + .assert_debug_eq(&error); +} diff --git a/rust/src/chat/src/renderer/mod.rs b/rust/src/chat/src/renderer/mod.rs index 4d9c1581a01..a83953f3318 100644 --- a/rust/src/chat/src/renderer/mod.rs +++ b/rust/src/chat/src/renderer/mod.rs @@ -15,6 +15,7 @@ pub mod deepseek_v4; pub mod harmony; pub mod hf; mod inkling; +mod kimi_k3; mod selection; #[cfg(test)] mod test_utils; @@ -23,6 +24,7 @@ pub use deepseek_v4::DeepSeekV4ChatRenderer; pub use deepseek_v32::DeepSeekV32ChatRenderer; pub use harmony::HarmonyChatRenderer; pub use inkling::InklingChatRenderer; +pub use kimi_k3::KimiK3ChatRenderer; pub use selection::RendererSelection; /// Rendered chat prompt submitted to the text backend. diff --git a/rust/src/chat/src/renderer/selection.rs b/rust/src/chat/src/renderer/selection.rs index 711c0d87ddf..5b6642c4c03 100644 --- a/rust/src/chat/src/renderer/selection.rs +++ b/rust/src/chat/src/renderer/selection.rs @@ -26,6 +26,8 @@ pub enum RendererSelection { Harmony, /// Force the Inkling native token renderer. Inkling, + /// Force the Kimi K3 XTML renderer. + KimiK3, } impl RendererSelection { @@ -37,6 +39,7 @@ impl RendererSelection { pub const HF_LITERAL: &str = "hf"; pub const INKLING_LITERAL: &str = "inkling"; pub const INKLING_MODEL_TYPE: &str = "inkling_mm_model"; + pub const KIMI_K3_LITERAL: &str = "kimi_k3"; /// Resolve the renderer selection using the given model type string, if /// it's `Auto`. @@ -47,6 +50,7 @@ impl RendererSelection { Self::DEEPSEEK_V4_LITERAL => Self::DeepSeekV4, Self::GPT_OSS_MODEL_TYPE => Self::Harmony, Self::INKLING_MODEL_TYPE => Self::Inkling, + Self::KIMI_K3_LITERAL => Self::KimiK3, _ => Self::Hf, }, selection => selection, @@ -70,6 +74,8 @@ impl FromStr for RendererSelection { Ok(Self::Harmony) } else if value.eq_ignore_ascii_case(Self::INKLING_LITERAL) { Ok(Self::Inkling) + } else if value.eq_ignore_ascii_case(Self::KIMI_K3_LITERAL) { + Ok(Self::KimiK3) } else { Err(format!( "unknown renderer `{value}` (expected one of: {})", @@ -88,6 +94,7 @@ impl fmt::Display for RendererSelection { Self::DeepSeekV4 => f.write_str(Self::DEEPSEEK_V4_LITERAL), Self::Harmony => f.write_str(Self::HARMONY_LITERAL), Self::Inkling => f.write_str(Self::INKLING_LITERAL), + Self::KimiK3 => f.write_str(Self::KIMI_K3_LITERAL), } } } @@ -114,7 +121,7 @@ mod tests { fn renderer_selection_expected_error_message() { let err = RendererSelection::from_str("unknown").unwrap_err(); expect_test::expect![ - "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling)" + "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling, kimi_k3)" ] .assert_eq(&err); } diff --git a/rust/src/chat/src/renderer/test_utils.rs b/rust/src/chat/src/renderer/test_utils.rs index 08eeca8176d..0a4edc81379 100644 --- a/rust/src/chat/src/renderer/test_utils.rs +++ b/rust/src/chat/src/renderer/test_utils.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::collections::HashMap; use std::fs; use std::path::Path; @@ -16,8 +17,9 @@ use crate::request::{ /// Options for constructing a [`ChatRequest`] from a fixture file. #[derive(Debug, Clone, Copy)] pub(crate) struct FixtureRequestOptions { - /// Whether to set the template kwarg `[enable_]thinking=true`. - pub enable_thinking: bool, + /// Optional thinking toggle applied only when the fixture does not already + /// set `thinking` / `enable_thinking` in [`FixtureRequest::template_kwargs`]. + pub enable_thinking: Option, /// Whether fixtures ending in an assistant message should omit the /// trailing generation prompt. pub no_generation_prompt_when_last_assistant: bool, @@ -46,6 +48,15 @@ pub(crate) struct FixtureRequest { messages: Vec, add_generation_prompt: Option, reasoning_effort: Option, + /// Standard response format passed to model-specific renderers. + #[serde(default)] + response_format: Option, + /// Extra chat-template kwargs (thinking, preserve_thinking, …). + #[serde(default)] + template_kwargs: HashMap, + /// When omitted, defaults to `auto` if tools are present, otherwise `none`. + #[serde(default)] + tool_choice: Option, } impl FixtureFile { @@ -57,6 +68,9 @@ impl FixtureFile { messages, add_generation_prompt: None, reasoning_effort: None, + response_format: None, + template_kwargs: HashMap::new(), + tool_choice: None, }, } } @@ -69,6 +83,7 @@ pub(crate) enum FixtureMessage { content: FixtureContent, }, Developer { + #[serde(default)] content: FixtureContent, #[serde(default)] tools: Vec, @@ -98,6 +113,12 @@ pub(crate) enum FixtureContent { Parts(Vec), } +impl Default for FixtureContent { + fn default() -> Self { + Self::Text(String::new()) + } +} + #[derive(Debug, Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] pub(crate) enum FixtureContentPart { @@ -136,6 +157,13 @@ struct FixtureToolCallFunction { impl FixtureRequest { fn into_chat_request(self, options: FixtureRequestOptions) -> ChatRequest { + let tools = to_chat_tools(&self.tools); + let tool_choice = self.tool_choice.unwrap_or(if tools.is_empty() { + ChatToolChoice::None + } else { + ChatToolChoice::Auto + }); + let mut request = ChatRequest { request_id: "renderer-fixture".to_string(), messages: self @@ -144,12 +172,8 @@ impl FixtureRequest { .enumerate() .map(|(index, message)| fixture_message_to_chat_message(index, message)) .collect(), - tools: to_chat_tools(&self.tools), - tool_choice: if self.tools.is_empty() { - ChatToolChoice::None - } else { - ChatToolChoice::Auto - }, + tools, + tool_choice, ..ChatRequest::for_test() }; @@ -162,9 +186,15 @@ impl FixtureRequest { request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; } request.chat_options.reasoning_effort = self.reasoning_effort; - if options.enable_thinking { - for key in ["thinking", "enable_thinking"] { - request.chat_options.template_kwargs.insert(key.to_string(), Value::Bool(true)); + request.chat_options.response_format = self.response_format; + request.chat_options.template_kwargs.extend(self.template_kwargs); + + // Options supply a default thinking toggle only when the fixture did not. + if let Some(thinking) = options.enable_thinking { + let kwargs = &mut request.chat_options.template_kwargs; + if !kwargs.contains_key("thinking") && !kwargs.contains_key("enable_thinking") { + kwargs.insert("thinking".to_string(), Value::Bool(thinking)); + kwargs.insert("enable_thinking".to_string(), Value::Bool(thinking)); } } diff --git a/rust/src/chat/src/request.rs b/rust/src/chat/src/request.rs index eb72557c7d2..56a764f8987 100644 --- a/rust/src/chat/src/request.rs +++ b/rust/src/chat/src/request.rs @@ -397,6 +397,10 @@ pub struct ChatOptions { /// Effort level exposed to chat templates for reasoning models. pub reasoning_effort: Option, + /// Standard response format available to model-specific renderers. + #[serde(default)] + pub response_format: Option, + /// Additional keyword arguments exposed to the chat template. pub template_kwargs: HashMap, } @@ -407,6 +411,7 @@ impl Default for ChatOptions { generation_prompt_mode: GenerationPromptMode::StartNewAssistant, chat_template: None, reasoning_effort: None, + response_format: None, template_kwargs: HashMap::new(), } } diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index ff15aa929ef..9cd670d29d0 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -212,6 +212,23 @@ impl RoundtripCase { } } + /// Kimi K3 XTML tool/reasoning channels (native renderer + unified parser). + /// + /// Needs HF tokenizer files under `HF_HOME` (`tiktoken.model` + + /// `tokenizer_config.json`). Weights are not required for this text-level + /// roundtrip. + fn kimi_k3() -> Self { + Self { + model_id: "moonshotai/Kimi-K3", + assistant_stop_suffix: "<|end_of_msg|>", + tool_call_parser: ParserSelection::Auto, + reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::Toggleable { default: true }, + json_fmt: compact_json_fmt(), + sort_json_keys: false, + } + } + /// SeedOSS with `` / `` reasoning tags. fn seed_oss() -> Self { Self { @@ -255,7 +272,7 @@ impl RoundtripCase { fn gpt_oss() -> Self { Self { model_id: "openai/gpt-oss-20b", - assistant_stop_suffix: "", // not applicable for token-id cases + assistant_stop_suffix: "", tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Always { value: true }, @@ -268,7 +285,7 @@ impl RoundtripCase { fn inkling() -> Self { Self { model_id: "thinkingmachines/Inkling", - assistant_stop_suffix: "", + assistant_stop_suffix: "<|content_model_end_sampling|>", tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Always { value: true }, @@ -312,6 +329,8 @@ roundtrip_tests! { nemotron_v3 => [reasoning_and_content], gemma4 => [tool_call_mix], // Gemma4 strips reasoning in history if there's no tool call kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history + // K3 drops plain-assistant reasoning in history; tool-call turns keep it. + kimi_k3 => [tool_call_mix], gpt_oss => [tool_call_mix], // Harmony strips reasoning in history if there's no tool call inkling => [reasoning_and_content, tool_call_mix], } @@ -670,14 +689,23 @@ fn decoded_completion_stream( .collect() } Prompt::TokenIds(token_ids) => { - ensure!( - assistant_stop_suffix.is_empty(), - "token-id roundtrip cases do not support text stop suffixes" - ); + let body = if assistant_stop_suffix.is_empty() { + token_ids.as_slice() + } else { + let stop_token_ids = tokenizer + .encode(assistant_stop_suffix, false) + .context("failed to encode token-id completion stop suffix")?; + token_ids.strip_suffix(stop_token_ids.as_slice()).with_context(|| { + format!( + "token-id completion did not end with {:?}: {:?}", + assistant_stop_suffix, token_ids + ) + })? + }; incremental_decode_chunks( tokenizer, &prompt_token_ids, - token_ids, + body, TOKEN_COMPLETION_CHUNK_TOKENS, )? } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index b0926fb24ac..bfece8b6092 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -666,7 +666,7 @@ fn serve_args_reject_unknown_renderer_value() { .unwrap_err(); expect![[r#" - error: invalid value 'definitely_missing' for '--tokenizer-mode ': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling) + error: invalid value 'definitely_missing' for '--tokenizer-mode ': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling, kimi_k3) For more information, try '--help'. "#]] diff --git a/rust/src/parser/src/unified/kimi_k3.rs b/rust/src/parser/src/unified/kimi_k3.rs new file mode 100644 index 00000000000..7e716225885 --- /dev/null +++ b/rust/src/parser/src/unified/kimi_k3.rs @@ -0,0 +1,1149 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Unified parser for the Kimi K3 (XTML) chat format. +//! +//! Original Python implementations: +//! - `vllm/reasoning/kimi_k3_reasoning_parser.py` +//! - `vllm/tool_parsers/kimi_k3_tool_parser.py` +//! +//! K3 wraps one assistant message into XTML channels built from the dedicated +//! special tokens `<|open|>`, `<|close|>`, and `<|sep|>`: +//! +//! ```text +//! <|open|>think<|sep|>reasoning<|close|>think<|sep|> +//! <|open|>response<|sep|>visible answer<|close|>response<|sep|> +//! <|open|>tools<|sep|> +//! <|open|>call tool="get_weather" index="1"<|sep|> +//! <|open|>argument key="city" type="string"<|sep|>Hangzhou<|close|>argument<|sep|> +//! <|close|>call<|sep|> +//! <|close|>tools<|sep|> +//! <|close|>message<|sep|> +//! ``` +//! +//! In chat serving the generation prompt ends with `<|open|>think<|sep|>` +//! (thinking) or `<|open|>response<|sep|>` (instruct), so the model output +//! starts *inside* that channel without re-emitting the open tag; +//! [`UnifiedParser::initialize`] detects this from the prompt token IDs. +//! +//! Argument decoding mirrors the renderer's type tagging (inverse encoding): +//! `type="string"` values pass the raw text through, other types are +//! JSON-decoded, and a raw `json` block is passed through unmodified. Attribute +//! values reverse the renderer escaping (`"` before `&`). +//! +//! Known limitation (shared with the Python parser): string argument and +//! response bodies are emitted raw, so a value that literally contains +//! `<|close|>argument<|sep|>` or `<|close|>response<|sep|>` is +//! indistinguishable from a real closing marker. + +mod structural_tag; + +pub use structural_tag::KimiK3StructuralTagBuilder; + +use serde_json::{Map, Value}; +use vllm_tokenizer::DynTokenizer; +use winnow::ascii::{multispace0 as ws0, multispace1 as ws1}; +use winnow::combinator::{alt, delimited, eof, preceded, repeat, seq, terminated}; +use winnow::error::{ContextError, ErrMode, ModalResult, StrContext}; +use winnow::prelude::*; +use winnow::stream::Partial; +use winnow::token::{literal, rest, take_till, take_until, take_while}; + +use self::structural_tag::KIMI_K3_STRUCTURAL_TAG_BUILDER; +use super::{Result, UnifiedParser, UnifiedParserOutput, token_id}; +use crate::tool::{StructuralTagBuilder, Tool, ToolCallDelta}; +use crate::unified::parsing_failed; +use crate::utils::{MarkerScanState, parse_buffered_event, safe_text_len_mul, take_until_marker}; + +const OPEN: &str = "<|open|>"; +const SEP: &str = "<|sep|>"; +const END_OF_MSG: &str = "<|end_of_msg|>"; + +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 TOOLS_OPEN: &str = "<|open|>tools<|sep|>"; +const TOOLS_CLOSE: &str = "<|close|>tools<|sep|>"; +const MESSAGE_CLOSE: &str = "<|close|>message<|sep|>"; +const CALL_OPEN: &str = "<|open|>call"; +const CALL_CLOSE: &str = "<|close|>call<|sep|>"; +const ARG_OPEN: &str = "<|open|>argument"; +const ARG_CLOSE: &str = "<|close|>argument<|sep|>"; +const JSON_OPEN: &str = "<|open|>json"; +const JSON_CLOSE: &str = "<|close|>json<|sep|>"; + +const IDLE_MARKERS: &[&str] = &[ + THINK_OPEN, + RESPONSE_OPEN, + TOOLS_OPEN, + MESSAGE_CLOSE, + END_OF_MSG, +]; +const REASONING_MARKERS: &[&str] = &[THINK_CLOSE, END_OF_MSG]; +const RESPONSE_MARKERS: &[&str] = &[RESPONSE_CLOSE, TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]; +const EPILOGUE_MARKERS: &[&str] = &[TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]; +const TOOLS_MARKERS: &[&str] = &[CALL_OPEN, TOOLS_CLOSE, MESSAGE_CLOSE, END_OF_MSG]; + +/// Channel tags are a couple of text tokens; longer `<|open|>…<|sep|>` spans in +/// the prompt tail (attribute-bearing message opens, message bodies) never name +/// a generation channel. +const MAX_PREFILL_TAG_TOKENS: usize = 8; + +type KimiK3Input<'i> = Partial<&'i str>; + +#[derive(Debug, Clone, PartialEq, Eq)] +enum KimiK3Event { + Text { + len: usize, + }, + Reasoning { + len: usize, + }, + /// Structural noise consumed without emitting anything. + Skip, + ThinkOpen, + ThinkClose, + ResponseOpen, + ResponseClose, + ToolsOpen, + ToolsClose, + /// The assistant message closed; everything after it is ignored. + MessageEnd, + CallOpen { + name: String, + index: Option, + }, + CallComplete { + arguments: String, + }, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +enum KimiK3Mode { + /// Before any channel opens: channel opens are expected, but raw text + /// falls through as visible text so marker-free output still streams. + #[default] + Idle, + /// Inside the `think` channel. + Reasoning, + /// Inside the `response` channel. + Response, + /// After the `response` (or `tools`) channel closed: only a `tools` + /// channel or the message close may follow, and noise is never content. + Epilogue, + /// Inside the `tools` channel, between `call` blocks. + Tools, + /// Inside one `call` block, buffering its body until the close marker. + Call { + name: String, + index: Option, + scan: MarkerScanState, + }, + /// After the message closed: ignore the rest (EOS leakage guard). + Done, +} + +/// Unified parser for Kimi K3 XTML think / response / tools channels. +pub struct KimiK3UnifiedParser { + buffer: String, + mode: KimiK3Mode, + /// Parser-provided tool-call IDs (`{tool}:{zero_based_index}`) by tool + /// index; its length is also the count of emitted calls. + call_ids: Vec, + tokenizer: DynTokenizer, + open_token_id: u32, + sep_token_id: u32, +} + +impl KimiK3UnifiedParser { + /// Create a Kimi K3 parser. + pub fn new(_tools: &[Tool], tokenizer: DynTokenizer) -> Result { + let open_token_id = token_id(tokenizer.as_ref(), OPEN)?; + let sep_token_id = token_id(tokenizer.as_ref(), SEP)?; + + Ok(Self { + buffer: String::new(), + mode: KimiK3Mode::default(), + call_ids: Vec::new(), + tokenizer, + open_token_id, + sep_token_id, + }) + } + + /// Detect the prefilled generation channel from the prompt tail. + /// + /// `add_generation_prompt` ends the prompt with `<|open|>think<|sep|>` or + /// `<|open|>response<|sep|>`, so generation starts inside that channel and + /// never re-emits the open tag. Locate the last `<|open|>…<|sep|>` pair and + /// decode the short tag between the two structural tokens. + fn initialize_mode(&mut self, prompt_token_ids: &[u32]) { + self.mode = KimiK3Mode::Idle; + + let Some(sep_pos) = prompt_token_ids.iter().rposition(|&id| id == self.sep_token_id) else { + return; + }; + let Some(open_pos) = + prompt_token_ids[..sep_pos].iter().rposition(|&id| id == self.open_token_id) + else { + return; + }; + let tag_ids = &prompt_token_ids[open_pos + 1..sep_pos]; + if tag_ids.is_empty() || tag_ids.len() > MAX_PREFILL_TAG_TOKENS { + return; + } + let Ok(tag) = self.tokenizer.decode(tag_ids, /* skip_special_tokens */ false) else { + return; + }; + self.mode = match tag.trim() { + "think" => KimiK3Mode::Reasoning, + "response" => KimiK3Mode::Response, + _ => KimiK3Mode::Idle, + }; + } + + fn apply_event(&mut self, event: KimiK3Event, output: &mut UnifiedParserOutput) -> Result<()> { + match event { + KimiK3Event::Text { len } => output.push_text(self.buffer[..len].to_string()), + KimiK3Event::Reasoning { len } => { + output.push_reasoning(self.buffer[..len].to_string()); + } + KimiK3Event::Skip => {} + KimiK3Event::ThinkOpen => self.mode = KimiK3Mode::Reasoning, + KimiK3Event::ThinkClose => self.mode = KimiK3Mode::Idle, + KimiK3Event::ResponseOpen => self.mode = KimiK3Mode::Response, + KimiK3Event::ResponseClose => self.mode = KimiK3Mode::Epilogue, + KimiK3Event::ToolsOpen => self.mode = KimiK3Mode::Tools, + KimiK3Event::ToolsClose => self.mode = KimiK3Mode::Epilogue, + KimiK3Event::MessageEnd => self.mode = KimiK3Mode::Done, + KimiK3Event::CallOpen { name, index } => { + self.mode = KimiK3Mode::Call { + name, + index, + scan: MarkerScanState::default(), + }; + } + KimiK3Event::CallComplete { arguments } => { + let mode = std::mem::replace(&mut self.mode, KimiK3Mode::Tools); + let KimiK3Mode::Call { name, index, .. } = mode else { + return Err(parsing_failed!( + "Kimi K3 call completion without an active tool call" + )); + }; + // An empty/garbage call block without a tool name is dropped, + // matching the Python parser. + if name.is_empty() { + return Ok(()); + } + + let tool_index = self.call_ids.len(); + self.call_ids.push(tool_call_id_for(&name, index.as_deref())); + output.push_call(ToolCallDelta { + tool_index, + name: Some(name), + arguments, + }); + } + } + Ok(()) + } + + fn reset_state(&mut self) -> String { + self.mode = KimiK3Mode::Idle; + self.call_ids.clear(); + std::mem::take(&mut self.buffer) + } +} + +impl UnifiedParser for KimiK3UnifiedParser { + fn create(tools: &[Tool], tokenizer: DynTokenizer) -> Result> + where + Self: Sized + 'static, + { + Self::new(tools, tokenizer).map(|parser| Box::new(parser) as Box) + } + + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> { + self.buffer.clear(); + self.call_ids.clear(); + self.initialize_mode(prompt_token_ids); + Ok(()) + } + + fn preserve_special_tokens(&self) -> bool { + true + } + + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + Some(&KIMI_K3_STRUCTURAL_TAG_BUILDER) + } + + fn tool_call_id(&self, tool_index: usize) -> Option<&str> { + self.call_ids.get(tool_index).map(String::as_str) + } + + fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> { + self.buffer.push_str(chunk); + + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| { + parse_next_kimi_k3_event(input, &mut self.mode) + })? { + self.apply_event(event, output)?; + self.buffer.drain(..consumed_len); + } + + Ok(()) + } + + fn finish(&mut self) -> Result { + let mut output = UnifiedParserOutput::default(); + + match &self.mode { + KimiK3Mode::Idle | KimiK3Mode::Response => { + output.push_text(std::mem::take(&mut self.buffer)); + } + KimiK3Mode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)), + KimiK3Mode::Epilogue | KimiK3Mode::Done => self.buffer.clear(), + // A tools channel truncated between complete calls loses only its + // closing markers; keep the calls already emitted. + KimiK3Mode::Tools if self.buffer.is_empty() => {} + KimiK3Mode::Tools | KimiK3Mode::Call { .. } => { + return Err(parsing_failed!("incomplete Kimi K3 tool call")); + } + } + + // Keep call_ids so tool_call_id() stays available after the stream ends. + self.mode = KimiK3Mode::Idle; + Ok(output) + } + + fn reset(&mut self) -> String { + self.reset_state() + } +} + +/// Build the API-side tool-call ID from the XTML one-based `index` attribute. +/// +/// The ID uses the zero-based call ordinal; XTML's message index stays +/// one-based when rendering tool result messages. +fn tool_call_id_for(name: &str, index: Option<&str>) -> String { + match index { + None => name.to_string(), + Some(raw) => match raw.parse::() { + Ok(one_based) => format!("{name}:{}", one_based - 1), + Err(_) => format!("{name}:{raw}"), + }, + } +} + +/// Parse one Kimi K3 event from buffered streaming input. +fn parse_next_kimi_k3_event( + input: &mut KimiK3Input<'_>, + mode: &mut KimiK3Mode, +) -> ModalResult { + match mode { + KimiK3Mode::Idle => parse_idle_event(input), + KimiK3Mode::Reasoning => parse_reasoning_event(input), + KimiK3Mode::Response => parse_response_event(input), + KimiK3Mode::Epilogue => parse_epilogue_event(input), + KimiK3Mode::Tools => parse_tools_event(input), + KimiK3Mode::Call { scan, .. } => call_body_event(input, scan), + KimiK3Mode::Done => parse_done_event(input), + } +} + +/// Parse an event while waiting for the next channel open. +fn parse_idle_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + literal(THINK_OPEN).value(KimiK3Event::ThinkOpen), + literal(RESPONSE_OPEN).value(KimiK3Event::ResponseOpen), + literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen), + message_end_event, + safe_idle_text_event, + )) + .parse_next(input) +} + +/// Parse an event inside the `think` channel. +fn parse_reasoning_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + literal(THINK_CLOSE).value(KimiK3Event::ThinkClose), + // `<|end_of_msg|>` can reach the parser under `ignore_eos` or + // `include_stop_str_in_output`; never leak it into reasoning. + literal(END_OF_MSG).value(KimiK3Event::MessageEnd), + safe_reasoning_event, + )) + .parse_next(input) +} + +/// Parse an event inside the `response` channel. +fn parse_response_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + literal(RESPONSE_CLOSE).value(KimiK3Event::ResponseClose), + // The response body also implicitly ends at a `tools` channel. + literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen), + message_end_event, + safe_response_text_event, + )) + .parse_next(input) +} + +/// Parse an event after the response channel closed. +fn parse_epilogue_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen), + message_end_event, + skip_epilogue_noise_event, + )) + .parse_next(input) +} + +/// Parse an event inside the `tools` channel, between `call` blocks. +fn parse_tools_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + call_open_event, + literal(TOOLS_CLOSE).value(KimiK3Event::ToolsClose), + // Defensive: an unterminated tools channel still ends with the message. + message_end_event, + skip_tools_noise_event, + )) + .parse_next(input) +} + +/// Parse a message close or end-of-message marker. +fn message_end_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt((literal(MESSAGE_CLOSE), literal(END_OF_MSG))) + .value(KimiK3Event::MessageEnd) + .parse_next(input) +} + +/// Ignore everything after the assistant message closed. +fn parse_done_event(input: &mut KimiK3Input<'_>) -> ModalResult { + rest.value(KimiK3Event::Skip).parse_next(input) +} + +/// Parse safe text while waiting for the next channel marker. +fn safe_idle_text_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, IDLE_MARKERS).map(|len| KimiK3Event::Text { len }) +} + +/// Parse safe reasoning before the think close marker. +fn safe_reasoning_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, REASONING_MARKERS).map(|len| KimiK3Event::Reasoning { len }) +} + +/// Parse safe response text before the next channel marker. +fn safe_response_text_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, RESPONSE_MARKERS).map(|len| KimiK3Event::Text { len }) +} + +/// Skip non-content noise after the response channel closed. +fn skip_epilogue_noise_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, EPILOGUE_MARKERS).map(|_| KimiK3Event::Skip) +} + +/// Skip non-content noise between `call` blocks. +fn skip_tools_noise_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, TOOLS_MARKERS).map(|_| KimiK3Event::Skip) +} + +/// Parse a `call` open tag into its tool name and one-based index. +fn call_open_event(input: &mut KimiK3Input<'_>) -> ModalResult { + let (attrs,) = seq!( + _: literal(CALL_OPEN), + take_until(0.., SEP), + _: literal(SEP), + ) + .parse_next(input)?; + let attrs = parse_tag_attrs(attrs)?; + + Ok(KimiK3Event::CallOpen { + name: attr_value(&attrs, "tool").unwrap_or_default().to_string(), + index: attr_value(&attrs, "index") + .filter(|index| !index.is_empty()) + .map(str::to_string), + }) +} + +/// Parse the buffered call body through `<|close|>call<|sep|>` into a +/// completed call. +fn call_body_event( + input: &mut KimiK3Input<'_>, + scan: &mut MarkerScanState, +) -> ModalResult { + let (body,) = seq!( + take_until_marker(CALL_CLOSE, scan), + _: literal(CALL_CLOSE), + ) + .parse_next(input)?; + let arguments = parse_call_arguments(body)?; + + Ok(KimiK3Event::CallComplete { arguments }) +} + +/// Parse a complete `call` body into the OpenAI-style arguments JSON string. +/// +/// The body is either one raw `json` block (passed through unmodified) or a +/// sequence of typed `argument` blocks (converted per their `type` tags). +fn parse_call_arguments(body: &str) -> ModalResult { + let mut input = body; + terminated( + delimited(ws0, alt((json_block_arguments, typed_arguments)), ws0), + eof, + ) + .parse_next(&mut input) + .map_err(|_| xtml_error("Kimi K3 call body")) +} + +/// Parse one raw `json` argument block, passing its body through unmodified. +fn json_block_arguments(input: &mut &str) -> ModalResult { + let (raw,) = seq!( + _: literal(JSON_OPEN), + _: take_until(0.., SEP), // attrs (`type="object"`), unused on decode + _: literal(SEP), + take_until(0.., JSON_CLOSE), + _: literal(JSON_CLOSE), + ) + .parse_next(input)?; + Ok(raw.to_string()) +} + +/// Parse typed `argument` blocks into a serialized JSON object, preserving +/// argument order. +fn typed_arguments(input: &mut &str) -> ModalResult { + let pairs: Vec<(String, Value)> = + repeat(0.., terminated(argument_block, ws0)).parse_next(input)?; + let arguments = pairs.into_iter().collect::>(); + serde_json::to_string(&arguments).map_err(|_| xtml_error("Kimi K3 arguments")) +} + +/// Parse one typed `argument` block into its key/value pair. +fn argument_block(input: &mut &str) -> ModalResult<(String, Value)> { + let (attrs, raw) = seq!( + _: literal(ARG_OPEN), + take_until(0.., SEP), + _: literal(SEP), + take_until(0.., ARG_CLOSE), + _: literal(ARG_CLOSE), + ) + .parse_next(input)?; + let attrs = parse_tag_attrs(attrs)?; + + let key = attr_value(&attrs, "key").unwrap_or_default().to_string(); + let arg_type = attr_value(&attrs, "type").unwrap_or("string"); + Ok((key, decode_argument_value(arg_type, raw))) +} + +/// Decode one typed argument value per its XTML `type` tag. +/// +/// `string` values pass the raw text through (the renderer emits them +/// unescaped); other types are JSON-decoded, falling back to the raw text on +/// malformed payloads so one quirky value does not fail the whole call. +fn decode_argument_value(arg_type: &str, raw: &str) -> Value { + if arg_type == "string" { + return Value::String(raw.to_string()); + } + serde_json::from_str(raw).unwrap_or_else(|_| Value::String(raw.to_string())) +} + +/// Parse a complete XTML attribute string like ` tool="get_weather" index="1"`. +fn parse_tag_attrs(attrs: &str) -> ModalResult> { + let mut input = attrs; + terminated(repeat(0.., preceded(ws1, tag_attr)), (ws0, eof)) + .parse_next(&mut input) + .map_err(|_| xtml_error("XTML tag attributes")) +} + +/// Parse one XTML `key="value"` attribute pair. +fn tag_attr(input: &mut &str) -> ModalResult<(String, String)> { + seq!( + take_while(1.., |char: char| char.is_alphanumeric() || char == '_').map(str::to_string), + _: literal("=\""), + take_till(0.., '"').map(unescape_attr_value), + _: literal("\""), + ) + .parse_next(input) +} + +/// Reverse XTML attribute escaping: `"` first, then `&` (the inverse +/// of the encode order). +fn unescape_attr_value(value: &str) -> String { + value.replace(""", "\"").replace("&", "&") +} + +/// Look up one parsed attribute value by key. +fn attr_value<'a>(attrs: &'a [(String, String)], key: &str) -> Option<&'a str> { + attrs.iter().find(|(name, _)| name == key).map(|(_, value)| value.as_str()) +} + +/// Build a cut error for determinably malformed XTML structure. +fn xtml_error(label: &'static str) -> ErrMode { + let mut error = ContextError::new(); + error.push(StrContext::Label(label)); + ErrMode::Cut(error) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use serde_json::{Value, json}; + use thiserror_ext::AsReport; + use vllm_tokenizer::Tokenizer as _; + use vllm_tokenizer::test_utils::TestTokenizer; + + use super::{ + END_OF_MSG, KimiK3UnifiedParser, OPEN, RESPONSE_CLOSE, RESPONSE_OPEN, SEP, THINK_CLOSE, + THINK_OPEN, TOOLS_CLOSE, TOOLS_OPEN, + }; + use crate::tool::ToolCallDelta; + use crate::unified::{ + UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput, + }; + + const OPEN_ID: u32 = 256; + const CLOSE_ID: u32 = 257; + const SEP_ID: u32 = 258; + const END_OF_MSG_ID: u32 = 259; + + fn tokenizer() -> TestTokenizer { + TestTokenizer::new() + .with_special_token(OPEN, OPEN_ID) + .with_special_token("<|close|>", CLOSE_ID) + .with_special_token(SEP, SEP_ID) + .with_special_token(END_OF_MSG, END_OF_MSG_ID) + } + + trait UnifiedParserTestExt { + fn parse_chunk(&mut self, chunk: &str) -> super::Result; + fn parse_complete(&mut self, text: &str) -> super::Result; + } + + impl UnifiedParserTestExt for KimiK3UnifiedParser { + fn parse_chunk(&mut self, chunk: &str) -> super::Result { + let mut output = UnifiedParserOutput::default(); + self.parse_into(chunk, &mut output)?; + Ok(output) + } + + fn parse_complete(&mut self, text: &str) -> super::Result { + let mut output = self.parse_chunk(text)?; + output.append(self.finish()?); + Ok(output) + } + } + + trait UnifiedOutputTestExt { + fn normal_text(&self) -> String; + fn reasoning_text(&self) -> String; + fn calls(&self) -> Vec; + } + + impl UnifiedOutputTestExt for UnifiedParserOutput { + fn normal_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Text(text) => Some(text.as_str()), + _ => None, + }) + .collect() + } + + fn reasoning_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Reasoning(text) => Some(text.as_str()), + _ => None, + }) + .collect() + } + + fn calls(&self) -> Vec { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::ToolCall(call) => Some(call.clone()), + _ => None, + }) + .collect() + } + } + + fn test_parser() -> KimiK3UnifiedParser { + KimiK3UnifiedParser::new(&[], Arc::new(tokenizer())).unwrap() + } + + fn collect_stream(parser: &mut KimiK3UnifiedParser, chunks: &[&str]) -> UnifiedParserOutput { + let mut output = UnifiedParserOutput::default(); + for chunk in chunks { + output.append(parser.parse_chunk(chunk).unwrap()); + } + output.append(parser.finish().unwrap()); + output + } + + /// Split `text` into small chunks to stress marker-split handling. + fn char_chunks(text: &str, size: usize) -> Vec { + let chars: Vec = text.chars().collect(); + chars.chunks(size).map(|chunk| chunk.iter().collect()).collect() + } + + fn arg(key: &str, arg_type: &str, value: &str) -> String { + format!( + "{OPEN}argument key=\"{key}\" type=\"{arg_type}\"{SEP}{value}<|close|>argument{SEP}" + ) + } + + fn call(attrs: &str, body: &str) -> String { + format!("{OPEN}call {attrs}{SEP}{body}<|close|>call{SEP}") + } + + fn thinking_output(reasoning: &str, response: &str, tools_body: &str) -> String { + let mut output = format!("{THINK_OPEN}{reasoning}{THINK_CLOSE}"); + output.push_str(&format!("{RESPONSE_OPEN}{response}{RESPONSE_CLOSE}")); + if !tools_body.is_empty() { + output.push_str(&format!("{TOOLS_OPEN}{tools_body}{TOOLS_CLOSE}")); + } + output.push_str("<|close|>message<|sep|>"); + output + } + + fn first_call(output: &UnifiedParserOutput) -> ToolCallDelta { + output.calls().first().expect("expected one tool call").clone() + } + + #[test] + fn kimi_k3_create_requires_structural_tokens() { + let error = match KimiK3UnifiedParser::new(&[], Arc::new(TestTokenizer::new())) { + Ok(_) => panic!("expected missing token error"), + Err(error) => error, + }; + + assert!(matches!( + error, + UnifiedParserError::MissingToken { token } if token == OPEN + )); + } + + #[test] + fn kimi_k3_parses_reasoning_response_and_typed_tool_call() { + let body = [ + arg("city", "string", "Hangzhou"), + arg("days", "number", "1.5"), + arg("detailed", "boolean", "true"), + arg("filters", "object", r#"{"kind":"rain"}"#), + arg("hours", "array", "[8,20]"), + ] + .concat(); + let text = thinking_output( + "Need the weather tool.", + "I'll check.", + &call("tool=\"get_weather\" index=\"1\"", &body), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(output.reasoning_text(), "Need the weather tool."); + assert_eq!(output.normal_text(), "I'll check."); + let call = first_call(&output); + assert_eq!(call.tool_index, 0); + assert_eq!(call.name.as_deref(), Some("get_weather")); + assert_eq!( + serde_json::from_str::(&call.arguments).unwrap(), + json!({ + "city": "Hangzhou", + "days": 1.5, + "detailed": true, + "filters": { "kind": "rain" }, + "hours": [8, 20], + }) + ); + assert_eq!(parser.tool_call_id(0), Some("get_weather:0")); + } + + #[test] + fn kimi_k3_arguments_preserve_order_and_number_formatting() { + let body = [ + arg("y", "number", "1.0"), + arg("x", "number", "2"), + arg("items", "array", r#"["left","right"]"#), + ] + .concat(); + let text = thinking_output("t", "", &call("tool=\"add\" index=\"1\"", &body)); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!( + first_call(&output).arguments, + r#"{"y":1.0,"x":2,"items":["left","right"]}"# + ); + } + + #[test] + fn kimi_k3_streaming_splits_markers_across_chunks() { + let text = thinking_output( + "step by step", + "the answer", + &call("tool=\"calc\" index=\"1\"", &arg("x", "number", "42")), + ); + + for size in [1, 3, 7] { + let chunks = char_chunks(&text, size); + let chunk_refs: Vec<&str> = chunks.iter().map(String::as_str).collect(); + let output = collect_stream(&mut test_parser(), &chunk_refs); + + assert_eq!(output.reasoning_text(), "step by step", "chunk size {size}"); + assert_eq!(output.normal_text(), "the answer", "chunk size {size}"); + assert_eq!(first_call(&output).name.as_deref(), Some("calc")); + assert_eq!(first_call(&output).arguments, r#"{"x":42}"#); + } + } + + #[test] + fn kimi_k3_streaming_emits_text_incrementally() { + let mut parser = test_parser(); + let prompt = tokenizer().encode("<|open|>response<|sep|>", false).unwrap(); + parser.initialize(&prompt).unwrap(); + + let first = parser.parse_chunk("Hel").unwrap(); + assert_eq!(first.normal_text(), "Hel"); + + let second = parser.parse_chunk("lo<|close|>resp").unwrap(); + assert_eq!(second.normal_text(), "lo"); + + let mut output = parser.parse_chunk("onse<|sep|>").unwrap(); + output.append(parser.finish().unwrap()); + assert_eq!(output.normal_text(), ""); + } + + #[test] + fn kimi_k3_initialize_think_prefill_starts_in_reasoning() { + let mut parser = test_parser(); + let prompt = tokenizer() + .encode( + "<|open|>message role=\"assistant\"<|sep|><|open|>think<|sep|>", + false, + ) + .unwrap(); + parser.initialize(&prompt).unwrap(); + + let output = parser + .parse_complete(&format!( + "reasoning{THINK_CLOSE}{RESPONSE_OPEN}answer{RESPONSE_CLOSE}<|close|>message{SEP}" + )) + .unwrap(); + + assert_eq!(output.reasoning_text(), "reasoning"); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn kimi_k3_initialize_response_prefill_starts_in_response() { + let mut parser = test_parser(); + let prompt = tokenizer() + .encode( + "<|open|>message role=\"assistant\"<|sep|><|open|>response<|sep|>", + false, + ) + .unwrap(); + parser.initialize(&prompt).unwrap(); + + let output = parser + .parse_complete(&format!("answer{RESPONSE_CLOSE}<|close|>message{SEP}")) + .unwrap(); + + assert_eq!(output.normal_text(), "answer"); + assert!(output.reasoning_text().is_empty()); + } + + #[test] + fn kimi_k3_initialize_message_open_prefill_starts_idle() { + let mut parser = test_parser(); + let prompt = + tokenizer().encode("<|open|>message role=\"assistant\"<|sep|>", false).unwrap(); + parser.initialize(&prompt).unwrap(); + + let output = parser + .parse_complete(&format!("{THINK_OPEN}reason{THINK_CLOSE}{RESPONSE_OPEN}hi")) + .unwrap(); + + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "hi"); + } + + #[test] + fn kimi_k3_plain_text_falls_through_as_text() { + let output = collect_stream(&mut test_parser(), &["plain ", "answer"]); + + assert_eq!(output.normal_text(), "plain answer"); + assert!(output.reasoning_text().is_empty()); + assert!(output.calls().is_empty()); + } + + #[test] + fn kimi_k3_tool_call_waits_for_close_marker() { + let mut parser = test_parser(); + let mut output = UnifiedParserOutput::default(); + + let argument = arg("x", "number", "1"); + for chunk in [ + TOOLS_OPEN, + "<|open|>call tool=\"calc\" index=\"1\"<|sep|>", + argument.as_str(), + ] { + output.append(parser.parse_chunk(chunk).unwrap()); + assert!(output.calls().is_empty()); + } + + output.append(parser.parse_chunk("<|close|>call<|sep|>").unwrap()); + + assert_eq!(first_call(&output).name.as_deref(), Some("calc")); + assert_eq!(first_call(&output).arguments, r#"{"x":1}"#); + } + + #[test] + fn kimi_k3_parses_multiple_tool_calls() { + let tools_body = format!( + "{}{}", + call( + "tool=\"get_weather\" index=\"1\"", + &arg("city", "string", "SF") + ), + call("tool=\"get_time\" index=\"2\"", ""), + ); + let text = thinking_output("t", "r", &tools_body); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + let calls = output.calls(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].tool_index, 0); + assert_eq!(calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(calls[0].arguments, r#"{"city":"SF"}"#); + assert_eq!(calls[1].tool_index, 1); + assert_eq!(calls[1].name.as_deref(), Some("get_time")); + assert_eq!(calls[1].arguments, "{}"); + assert_eq!(parser.tool_call_id(0), Some("get_weather:0")); + assert_eq!(parser.tool_call_id(1), Some("get_time:1")); + } + + #[test] + fn kimi_k3_json_block_arguments_pass_through_raw() { + // Spacing and key order must survive unmodified: raw `json` blocks are + // not validated or normalized. + let raw = r#"{"b": 1, "a": [2 , 3]}"#; + let body = format!("{OPEN}json type=\"object\"{SEP}{raw}<|close|>json{SEP}"); + let text = thinking_output("t", "", &call("tool=\"run\" index=\"1\"", &body)); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(first_call(&output).arguments, raw); + } + + #[test] + fn kimi_k3_string_argument_passes_raw_text_through() { + let value = "line one\nline two {\"not\": \"json\"} & "; + let text = thinking_output( + "t", + "", + &call( + "tool=\"write\" index=\"1\"", + &arg("content", "string", value), + ), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!( + serde_json::from_str::(&first_call(&output).arguments).unwrap(), + json!({ "content": value }) + ); + } + + #[test] + fn kimi_k3_malformed_typed_argument_falls_back_to_raw_text() { + let text = thinking_output( + "t", + "", + &call( + "tool=\"calc\" index=\"1\"", + &arg("x", "number", "not a number"), + ), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(first_call(&output).arguments, r#"{"x":"not a number"}"#); + } + + #[test] + fn kimi_k3_attribute_values_are_unescaped() { + let text = thinking_output( + "t", + "", + &call( + "tool=\"a"b&c\" index=\"1\"", + &arg("key", "string", "value"), + ), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(first_call(&output).name.as_deref(), Some("a\"b&c")); + } + + #[test] + fn kimi_k3_call_without_tool_name_is_dropped() { + let tools_body = format!( + "{}{}", + call("index=\"1\"", &arg("x", "number", "1")), + call("tool=\"real\" index=\"2\"", ""), + ); + let text = thinking_output("t", "", &tools_body); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + let calls = output.calls(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].tool_index, 0); + assert_eq!(calls[0].name.as_deref(), Some("real")); + assert_eq!(parser.tool_call_id(0), Some("real:1")); + } + + #[test] + fn kimi_k3_tool_call_id_follows_index_attribute() { + let tools_body = format!( + "{}{}{}", + call("tool=\"first\" index=\"3\"", ""), + call("tool=\"second\"", ""), + call("tool=\"third\" index=\"x\"", ""), + ); + let text = thinking_output("t", "", &tools_body); + + let mut parser = test_parser(); + parser.parse_complete(&text).unwrap(); + + assert_eq!(parser.tool_call_id(0), Some("first:2")); + assert_eq!(parser.tool_call_id(1), Some("second")); + assert_eq!(parser.tool_call_id(2), Some("third:x")); + } + + #[test] + fn kimi_k3_ignores_output_after_message_close() { + let mut parser = test_parser(); + let output = parser + .parse_complete(&format!( + "{RESPONSE_OPEN}answer{RESPONSE_CLOSE}<|close|>message{SEP}junk{END_OF_MSG}" + )) + .unwrap(); + + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn kimi_k3_epilogue_noise_is_not_content() { + let text = format!( + "{RESPONSE_OPEN}answer{RESPONSE_CLOSE}\n{TOOLS_OPEN}{}{TOOLS_CLOSE}\n<|close|>message{SEP}", + call("tool=\"calc\" index=\"1\"", ""), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(output.normal_text(), "answer"); + assert_eq!(output.calls().len(), 1); + } + + #[test] + fn kimi_k3_finish_flushes_unclosed_reasoning() { + let mut parser = test_parser(); + let mut output = parser.parse_chunk(&format!("{THINK_OPEN}still thinking")).unwrap(); + output.append(parser.finish().unwrap()); + + assert_eq!(output.reasoning_text(), "still thinking"); + assert!(output.normal_text().is_empty()); + } + + #[test] + fn kimi_k3_finish_flushes_partial_marker_as_text() { + let mut parser = test_parser(); + let mut output = parser.parse_chunk("answer<|clo").unwrap(); + output.append(parser.finish().unwrap()); + + assert_eq!(output.normal_text(), "answer<|clo"); + } + + #[test] + fn kimi_k3_finish_fails_mid_tool_call() { + let mut parser = test_parser(); + parser + .parse_chunk(&format!( + "{TOOLS_OPEN}<|open|>call tool=\"calc\" index=\"1\"<|sep|>{}", + arg("x", "number", "1") + )) + .unwrap(); + + let error = parser.finish().unwrap_err(); + + assert!(error.to_report_string().contains("incomplete Kimi K3 tool call")); + } + + #[test] + fn kimi_k3_finish_after_truncated_tools_keeps_complete_calls() { + let mut parser = test_parser(); + let mut output = parser + .parse_chunk(&format!( + "{TOOLS_OPEN}{}", + call("tool=\"calc\" index=\"1\"", &arg("x", "number", "1")) + )) + .unwrap(); + output.append(parser.finish().unwrap()); + + assert_eq!(first_call(&output).name.as_deref(), Some("calc")); + } + + #[test] + fn kimi_k3_malformed_call_attributes_fail_fast() { + let mut parser = test_parser(); + let error = parser + .parse_chunk(&format!("{TOOLS_OPEN}<|open|>call garbage attrs<|sep|>")) + .unwrap_err(); + + assert!(error.to_report_string().contains("XTML tag attributes")); + } + + #[test] + fn kimi_k3_empty_response_channel_emits_nothing() { + let text = thinking_output("t", "", &call("tool=\"calc\" index=\"1\"", "")); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert!(output.normal_text().is_empty()); + assert_eq!(output.reasoning_text(), "t"); + assert_eq!(output.calls().len(), 1); + } + + #[test] + fn kimi_k3_reset_returns_buffered_text() { + let mut parser = test_parser(); + let prompt = tokenizer().encode("<|open|>response<|sep|>", false).unwrap(); + parser.initialize(&prompt).unwrap(); + parser.parse_chunk("answer<|close|>resp").unwrap(); + + let raw = parser.reset(); + + assert_eq!(raw, "<|close|>resp"); + } +} diff --git a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs new file mode 100644 index 00000000000..d3515ad2dc3 --- /dev/null +++ b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs @@ -0,0 +1,503 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Structural-tag grammar for Kimi K3 XTML tool calls. + +use serde_json::{Map, Value}; +use xgrammar_structural_tag::Result; +use xgrammar_structural_tag::builders::{ + StructuralTagBuilder, StructuralTagContext, StructuralTagOptions, +}; +use xgrammar_structural_tag::format::{Format, JsonSchemaFormat, StructuralTag, TagFormat}; +use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam, function_parameters}; + +use super::{ + ARG_CLOSE, CALL_CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN, RESPONSE_CLOSE, + RESPONSE_OPEN, SEP, THINK_CLOSE, TOOLS_CLOSE, TOOLS_OPEN, +}; + +pub(super) static KIMI_K3_STRUCTURAL_TAG_BUILDER: KimiK3StructuralTagBuilder = + KimiK3StructuralTagBuilder; + +const XTML_TYPES: &[&str] = &["string", "number", "boolean", "null", "object", "array"]; + +/// Kimi K3 XTML structural-tag builder. +#[derive(Debug, Clone, Copy, Default)] +pub struct KimiK3StructuralTagBuilder; + +impl StructuralTagBuilder for KimiK3StructuralTagBuilder { + fn build(&self, ctx: StructuralTagContext<'_>) -> Result { + let mut elements = response_prefix(ctx.options.reasoning); + + let tools = match ctx.tool_choice { + // Serving lowering filters empty tools before calling the builder, + // while direct `build_structural_tag(..., auto, ...)` calls do not. + BuilderToolChoice::Auto if ctx.function_tools.is_empty() => None, + BuilderToolChoice::Auto => Some(Format::optional(tools_channel( + ctx.function_tools, + ctx.options, + ))), + BuilderToolChoice::Forced | BuilderToolChoice::Required => { + Some(tools_channel(ctx.function_tools, ctx.options)) + } + }; + if let Some(tools) = tools { + elements.push(tools); + } + elements.push(Format::optional(Format::const_string(MESSAGE_CLOSE))); + + Ok(StructuralTag::new(Format::sequence(elements))) + } +} + +fn response_prefix(reasoning: bool) -> Vec { + let mut elements = Vec::new(); + if reasoning { + elements.push(Format::tag( + "", + Format::any_text_excluding(&[THINK_CLOSE, END_OF_MSG]), + THINK_CLOSE, + )); + elements.push(Format::const_string(RESPONSE_OPEN)); + } else { + elements.push(Format::optional(Format::const_string(RESPONSE_OPEN))); + } + elements.push(Format::tag( + "", + Format::any_text_excluding(&[RESPONSE_CLOSE, TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]), + RESPONSE_CLOSE, + )); + elements +} + +fn tools_channel(tools: &[FunctionToolParam], options: StructuralTagOptions) -> Format { + let calls = tools.iter().map(|tool| call_tag(tool, options)).collect(); + Format::tag( + TOOLS_OPEN, + Format::tags_with_separator(calls, "", true, false), + TOOLS_CLOSE, + ) +} + +fn call_tag(tool: &FunctionToolParam, options: StructuralTagOptions) -> TagFormat { + let parameters = function_parameters(&tool.function); + let call_body = Format::or(vec![ + typed_arguments(¶meters, options), + raw_json_arguments(¶meters, options), + ]); + + TagFormat::new( + format!( + "{OPEN}call tool=\"{}\" index=\"", + escape_attr_value(&tool.function.name) + ), + Format::sequence(vec![ + Format::regex("[1-9][0-9]*"), + Format::const_string(format!("\"{SEP}")), + call_body, + ]), + CALL_CLOSE, + ) +} + +fn typed_arguments(parameters: &Value, options: StructuralTagOptions) -> Format { + let Some(schema) = parameters.as_object() else { + return if parameters == &Value::Bool(false) { + Format::const_string("") + } else { + Format::star(permissive_argument()) + }; + }; + let Some(properties) = schema.get("properties").and_then(Value::as_object) else { + return Format::star(permissive_argument()); + }; + if properties.is_empty() { + return Format::star(permissive_argument()); + } + + let root_defs = root_definitions(schema); + let arguments = properties + .iter() + .flat_map(|(key, schema)| argument_tags(key, schema, &root_defs, options)) + .map(Format::Tag) + .collect::>(); + let arguments = match arguments.as_slice() { + [argument] => argument.clone(), + _ => Format::or(arguments), + }; + // Keep typed arguments order-agnostic and non-unique, but do not allow an + // empty call when the root schema declares required properties. + if schema + .get("required") + .and_then(Value::as_array) + .is_some_and(|required| !required.is_empty()) + { + Format::plus(arguments) + } else { + Format::star(arguments) + } +} + +fn argument_tags( + key: &str, + schema: &Value, + root_defs: &Map, + options: StructuralTagOptions, +) -> Vec { + let types = schema_types(schema); + types + .into_iter() + .map(|xtml_type| { + let content = if xtml_type == "string" { + string_argument_content(schema) + } else { + json_schema( + attach_root_definitions(&narrow_schema_type(schema, xtml_type), root_defs), + options, + ) + }; + TagFormat::new( + format!( + "{OPEN}argument key=\"{}\" type=\"{xtml_type}\"{SEP}", + escape_attr_value(key) + ), + content, + ARG_CLOSE, + ) + }) + .collect() +} + +fn schema_types(schema: &Value) -> Vec<&'static str> { + let Some(schema) = schema.as_object() else { + return XTML_TYPES.to_vec(); + }; + let mut types = Vec::new(); + match schema.get("type") { + Some(Value::String(value)) => push_schema_type(&mut types, value), + Some(Value::Array(values)) => { + for value in values.iter().filter_map(Value::as_str) { + push_schema_type(&mut types, value); + } + } + _ => {} + } + if types.is_empty() + && let Some(value) = schema.get("const") + { + push_value_type(&mut types, value); + } + if types.is_empty() + && let Some(values) = schema.get("enum").and_then(Value::as_array) + { + for value in values { + push_value_type(&mut types, value); + } + } + if types.is_empty() { + XTML_TYPES.to_vec() + } else { + types + } +} + +fn push_value_type(types: &mut Vec<&'static str>, value: &Value) { + let xtml_type = match value { + Value::String(_) => "string", + Value::Number(_) => "number", + Value::Bool(_) => "boolean", + Value::Null => "null", + Value::Object(_) => "object", + Value::Array(_) => "array", + }; + if !types.contains(&xtml_type) { + types.push(xtml_type); + } +} + +fn narrow_schema_type(schema: &Value, xtml_type: &str) -> Value { + let Some(mut schema) = schema.as_object().cloned() else { + return schema.clone(); + }; + let json_type = if xtml_type == "number" && explicitly_integer_only(&schema) { + "integer" + } else { + xtml_type + }; + schema.insert("type".to_string(), Value::String(json_type.to_string())); + Value::Object(schema) +} + +fn explicitly_integer_only(schema: &Map) -> bool { + match schema.get("type") { + Some(Value::String(value)) => value == "integer", + Some(Value::Array(values)) => { + let values = values.iter().filter_map(Value::as_str).collect::>(); + values.contains(&"integer") && !values.contains(&"number") + } + _ => false, + } +} + +fn push_schema_type(types: &mut Vec<&'static str>, json_type: &str) { + let xtml_type = match json_type { + "string" => Some("string"), + "integer" | "number" => Some("number"), + "boolean" => Some("boolean"), + "null" => Some("null"), + "object" => Some("object"), + "array" => Some("array"), + _ => None, + }; + if let Some(xtml_type) = xtml_type + && !types.contains(&xtml_type) + { + types.push(xtml_type); + } +} + +fn string_argument_content(schema: &Value) -> Format { + let Some(schema) = schema.as_object() else { + return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]); + }; + let values = schema + .get("enum") + .and_then(Value::as_array) + .cloned() + .or_else(|| schema.get("const").cloned().map(|value| vec![value])); + let Some(values) = values else { + return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]); + }; + if values.is_empty() + || values.len() > 256 + || values + .iter() + .any(|value| value.as_str().is_none_or(|value| value.contains("<|"))) + { + return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]); + } + let values = values.iter().filter_map(Value::as_str).collect::>(); + match values.as_slice() { + [value] => Format::const_string(*value), + _ => Format::or(values.into_iter().map(Format::const_string).collect()), + } +} + +fn raw_json_arguments(parameters: &Value, options: StructuralTagOptions) -> Format { + Format::tag( + format!("{JSON_OPEN} type=\"object\"{SEP}"), + json_schema(parameters.clone(), options), + JSON_CLOSE, + ) +} + +fn permissive_argument() -> Format { + let key = Format::regex(r#"(?:[^<\"&]|&(?:amp|quot);|<[^|])*"#); + let alternatives = XTML_TYPES + .iter() + .map(|xtml_type| { + Format::sequence(vec![ + key.clone(), + Format::const_string(format!("\" type=\"{xtml_type}\"{SEP}")), + if *xtml_type == "string" { + Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]) + } else { + Format::json_schema(Value::Bool(true)) + }, + ]) + }) + .collect(); + Format::tag( + format!("{OPEN}argument key=\""), + Format::or(alternatives), + ARG_CLOSE, + ) +} + +fn json_schema(schema: Value, options: StructuralTagOptions) -> Format { + Format::JsonSchema( + JsonSchemaFormat::new(schema) + .with_any_order(options.any_order) + .with_max_whitespace_cnt(options.max_whitespace_cnt), + ) +} + +fn root_definitions(schema: &Map) -> Map { + ["$defs", "definitions"] + .into_iter() + .filter_map(|key| schema.get(key).map(|value| (key.to_string(), value.clone()))) + .collect() +} + +fn attach_root_definitions(schema: &Value, root_defs: &Map) -> Value { + let Some(mut schema) = schema.as_object().cloned() else { + return schema.clone(); + }; + for (key, value) in root_defs { + schema.entry(key.clone()).or_insert_with(|| value.clone()); + } + Value::Object(schema) +} + +fn escape_attr_value(value: &str) -> String { + value.replace('&', "&").replace('"', """) +} + +#[cfg(test)] +mod tests { + use expect_test::expect; + use serde_json::json; + use xgrammar_structural_tag::builders::StructuralTagOptions; + use xgrammar_structural_tag::{ + FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag, + }; + + use super::KimiK3StructuralTagBuilder; + + fn tool(name: &str, parameters: serde_json::Value) -> ToolParam { + ToolParam::Function(FunctionToolParam::new( + FunctionDefinition::new(name).with_parameters(parameters), + )) + } + + #[test] + fn required_structural_tag_matches_xtml_channels() { + let tools = vec![tool( + "get_weather", + json!({ + "$defs": { + "place": { "type": "object", "properties": { "city": { "type": "string" } } } + }, + "type": "object", + "properties": { + "unit": { "type": "string", "enum": ["celsius", "fahrenheit"] }, + "place": { "$ref": "#/$defs/place", "type": "object" } + }, + "required": ["place"] + }), + )]; + let tag = build_structural_tag( + KimiK3StructuralTagBuilder, + &tools, + ToolChoice::required(), + StructuralTagOptions::default().with_reasoning(false), + ) + .unwrap(); + + expect![[r##"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"optional","content":{"type":"const_string","value":"<|open|>response<|sep|>"}},{"type":"tag","begin":"","content":{"type":"any_text","excludes":["<|close|>response<|sep|>","<|open|>tools<|sep|>","<|close|>message<|sep|>","<|end_of_msg|>"]},"end":"<|close|>response<|sep|>"},{"type":"tag","begin":"<|open|>tools<|sep|>","content":{"type":"tags_with_separator","tags":[{"begin":"<|open|>call tool=\"get_weather\" index=\"","content":{"type":"sequence","elements":[{"type":"regex","pattern":"[1-9][0-9]*"},{"type":"const_string","value":"\"<|sep|>"},{"type":"or","elements":[{"type":"plus","content":{"type":"or","elements":[{"type":"tag","begin":"<|open|>argument key=\"unit\" type=\"string\"<|sep|>","content":{"type":"or","elements":[{"type":"const_string","value":"celsius"},{"type":"const_string","value":"fahrenheit"}]},"end":"<|close|>argument<|sep|>"},{"type":"tag","begin":"<|open|>argument key=\"place\" type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$ref":"#/$defs/place","type":"object","$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}}},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>argument<|sep|>"}]}},{"type":"tag","begin":"<|open|>json type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}},"type":"object","properties":{"unit":{"type":"string","enum":["celsius","fahrenheit"]},"place":{"$ref":"#/$defs/place","type":"object"}},"required":["place"]},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>json<|sep|>"}]}]},"end":"<|close|>call<|sep|>"}],"separator":"","at_least_one":true,"stop_after_first":false},"end":"<|close|>tools<|sep|>"},{"type":"optional","content":{"type":"const_string","value":"<|close|>message<|sep|>"}}]}}"##]].assert_eq(&tag.to_json_string().unwrap()); + } + + #[test] + fn typed_arguments_require_one_tag_only_for_nonempty_required() { + let required = super::typed_arguments( + &json!({ + "type": "object", + "properties": { "query": { "type": "string" } }, + "required": ["query"] + }), + StructuralTagOptions::default(), + ); + let optional = super::typed_arguments( + &json!({ + "type": "object", + "properties": { "query": { "type": "string" } } + }), + StructuralTagOptions::default(), + ); + let empty_required = super::typed_arguments( + &json!({ + "type": "object", + "properties": { "query": { "type": "string" } }, + "required": [] + }), + StructuralTagOptions::default(), + ); + + assert_eq!(serde_json::to_value(required).unwrap()["type"], "plus"); + assert_eq!(serde_json::to_value(optional).unwrap()["type"], "star"); + assert_eq!( + serde_json::to_value(empty_required).unwrap()["type"], + "star" + ); + } + + #[test] + fn reasoning_grammar_starts_inside_prefilled_think_channel() { + let tools = vec![tool("ping", json!({ "type": "object", "properties": {} }))]; + let tag = build_structural_tag( + KimiK3StructuralTagBuilder, + &tools, + ToolChoice::auto(), + StructuralTagOptions::default().with_reasoning(true), + ) + .unwrap(); + let value = serde_json::to_value(tag).unwrap(); + + assert_eq!( + value["format"]["elements"][0]["end"], + "<|close|>think<|sep|>" + ); + assert_eq!( + value["format"]["elements"][1]["value"], + "<|open|>response<|sep|>" + ); + assert_eq!(value["format"]["elements"][3]["type"], "optional"); + } + + #[test] + fn forced_choice_keeps_only_the_named_tool() { + let tools = vec![ + tool("search", json!({ "type": "object" })), + tool("lookup", json!({ "type": "object" })), + ]; + let tag = build_structural_tag( + KimiK3StructuralTagBuilder, + &tools, + ToolChoice::function("lookup"), + StructuralTagOptions::default().with_reasoning(false), + ) + .unwrap() + .to_json_string() + .unwrap(); + + assert!(tag.contains("lookup")); + assert!(!tag.contains("search")); + } + + #[test] + fn union_argument_content_matches_its_xtml_type() { + let tools = vec![tool( + "set_count", + json!({ + "type": "object", + "properties": { + "count": { "type": ["integer", "null"] } + } + }), + )]; + let tag = build_structural_tag( + KimiK3StructuralTagBuilder, + &tools, + ToolChoice::required(), + StructuralTagOptions::default(), + ) + .unwrap() + .to_json_string() + .unwrap(); + + assert!(tag.contains(r#"type=\"number\""#)); + assert!(tag.contains(r#""json_schema":{"type":"integer"}"#), "{tag}"); + assert!(tag.contains(r#"type=\"null\""#)); + assert!(tag.contains(r#""json_schema":{"type":"null"}"#), "{tag}"); + } + + #[test] + fn unsafe_string_enum_falls_back_as_a_whole() { + let format = super::string_argument_content(&json!({ + "type": "string", + "enum": ["safe", "<|unsafe"] + })); + + assert_eq!(serde_json::to_value(format).unwrap()["type"], "any_text"); + } +} diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index 4700211ae93..ae5524759df 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -6,10 +6,12 @@ mod combined; mod gemma4; mod inkling; +mod kimi_k3; pub use combined::CombinedParser; pub use gemma4::Gemma4UnifiedParser; pub use inkling::InklingUnifiedParser; +pub use kimi_k3::{KimiK3StructuralTagBuilder, KimiK3UnifiedParser}; use thiserror::Error; use thiserror_ext::Macro; use vllm_tokenizer::DynTokenizer; diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index d75f3fab162..2c4050f8f1e 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -89,6 +89,15 @@ pub(super) fn prepare_chat_request( )?; let template_kwargs = request.chat_template_kwargs.unwrap_or_default(); + let response_format = + request.response_format.as_ref().map(serde_json::to_value).transpose().map_err( + |error| { + ApiError::invalid_request( + format!("failed to serialize response_format: {error}"), + Some("response_format"), + ) + }, + )?; let include_usage = (request.stream_options.as_ref()) .and_then(|options| options.include_usage) @@ -150,6 +159,7 @@ pub(super) fn prepare_chat_request( generation_prompt_mode, chat_template: request.chat_template, reasoning_effort: request.reasoning_effort, + response_format, template_kwargs, }, tools: convert_tools(request.tools)?, @@ -400,6 +410,7 @@ fn convert_tool_choice(tool_choice: Option<&ToolChoice>) -> Result &'static str { let model_type = config.effective_model_type(); match model_type { - Some("kimi" | "kimi_k2" | "kimi_k25" | "deepseek_v3") => KIMI_PATTERN, + Some("kimi" | "kimi_k2" | "kimi_k25" | "kimi_k3" | "deepseek_v3") => KIMI_PATTERN, _ => CL100K_BASE_PATTERN, } } @@ -817,6 +817,7 @@ mod tests { #[test] fn tiktoken_detects_kimi_pattern_from_model_type() { let kimi = config_json!({ "model_type": "kimi_k25" }); + let kimi_k3 = config_json!({ "model_type": "kimi_k3" }); let baseten_kimi = config_json!({ "model_type": "deepseek_v3" }); let nested_kimi = config_json!({ "model_type": "composite_wrapper", @@ -830,6 +831,7 @@ mod tests { let missing = config_json!({ "text_config": {} }); assert_eq!(detect_bpe_pattern(&kimi), KIMI_PATTERN); + assert_eq!(detect_bpe_pattern(&kimi_k3), KIMI_PATTERN); assert_eq!(detect_bpe_pattern(&baseten_kimi), KIMI_PATTERN); assert_eq!(detect_bpe_pattern(&nested_kimi), CL100K_BASE_PATTERN); assert_eq!(detect_bpe_pattern(&generic), CL100K_BASE_PATTERN);