From e1f068a92719b1157c7e2893b6e79cc147829397 Mon Sep 17 00:00:00 2001 From: Indrajit Bhosale Date: Fri, 21 Aug 2026 12:31:34 -0700 Subject: [PATCH 01/10] K3 structual tag defualt Signed-off-by: Indrajit Bhosale --- .../src/tool_calling/structural_tag/dsml.rs | 1 + .../src/tool_calling/structural_tag/format.rs | 4 + .../tool_calling/structural_tag/kimi_k2.rs | 1 + .../tool_calling/structural_tag/kimi_k3.rs | 478 +++++++++++++++++- .../structural_tag/triggered_tags.rs | 1 + 5 files changed, 465 insertions(+), 20 deletions(-) diff --git a/parsers/v1/src/tool_calling/structural_tag/dsml.rs b/parsers/v1/src/tool_calling/structural_tag/dsml.rs index 71439cc5c..2f1f99a17 100644 --- a/parsers/v1/src/tool_calling/structural_tag/dsml.rs +++ b/parsers/v1/src/tool_calling/structural_tag/dsml.rs @@ -97,6 +97,7 @@ pub(crate) fn build_dsml_tool_calls( tags: vec![block_tag], at_least_one: outer_at_least_one, stop_after_first: ctx.stop_after_first(), + excludes: vec![], }), })) } diff --git a/parsers/v1/src/tool_calling/structural_tag/format.rs b/parsers/v1/src/tool_calling/structural_tag/format.rs index 036120fa8..85e6c1e5f 100644 --- a/parsers/v1/src/tool_calling/structural_tag/format.rs +++ b/parsers/v1/src/tool_calling/structural_tag/format.rs @@ -88,6 +88,8 @@ pub struct TriggeredTagsFormat { pub tags: Vec, pub at_least_one: bool, pub stop_after_first: bool, + #[serde(skip_serializing_if = "Vec::is_empty")] + pub excludes: Vec, } /// `tags_with_separator`: a tag sequence with fixed separators and no free text. @@ -227,6 +229,7 @@ mod tests { tags: vec![], at_least_one: false, stop_after_first: false, + excludes: vec![], }; let value = serde_json::to_value(&format).unwrap(); @@ -240,6 +243,7 @@ mod tests { tags: vec![], at_least_one: true, stop_after_first: false, + excludes: vec![], }; let value = serde_json::to_value(&format).unwrap(); diff --git a/parsers/v1/src/tool_calling/structural_tag/kimi_k2.rs b/parsers/v1/src/tool_calling/structural_tag/kimi_k2.rs index f1eb62d6a..80b2186c6 100644 --- a/parsers/v1/src/tool_calling/structural_tag/kimi_k2.rs +++ b/parsers/v1/src/tool_calling/structural_tag/kimi_k2.rs @@ -101,6 +101,7 @@ pub(crate) fn build_kimi_k2( }], at_least_one: outer_at_least_one, stop_after_first: ctx.stop_after_first(), + excludes: vec![], }) } else { Format::Sequence(SequenceFormat { diff --git a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs index 673e7ec5f..3f79555a4 100644 --- a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs +++ b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs @@ -9,13 +9,15 @@ //! format so named and required tool choices can be constrained without //! changing what the K3 parser expects. +use std::collections::HashSet; + use serde_json::{Map, Value}; use super::builder::{ToolCallFormatBuildContext, resolve_tools_to_include}; use super::format::{ AnyTextFormat, ConstStringFormat, Format, JsonSchemaFormat, JsonSchemaStyle, OptionalFormat, OrFormat, RegexFormat, SequenceFormat, StarFormat, StructuralTag, TagFormat, - TagsWithSeparatorFormat, + TagsWithSeparatorFormat, TriggeredTagsFormat, }; use crate::tool_calling::ToolDefinition; @@ -24,8 +26,11 @@ const CLOSE: &str = "<|close|>"; const SEP: &str = "<|sep|>"; const RESPONSE_OPEN: &str = "<|open|>response<|sep|>"; const RESPONSE_CLOSE: &str = "<|close|>response<|sep|>"; +const THINK_OPEN: &str = "<|open|>think<|sep|>"; +const THINK_CLOSE: &str = "<|close|>think<|sep|>"; const TOOLS_OPEN: &str = "<|open|>tools<|sep|>"; const TOOLS_CLOSE: &str = "<|close|>tools<|sep|>"; +const CALL_OPEN: &str = "<|open|>call"; const CALL_CLOSE: &str = "<|close|>call<|sep|>"; const ARGUMENT_CLOSE: &str = "<|close|>argument<|sep|>"; const MESSAGE_CLOSE: &str = "<|close|>message<|sep|>"; @@ -56,6 +61,219 @@ fn one_of(elements: Vec) -> Format { } } +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +enum JsonType { + String, + Number, + Integer, + Boolean, + Array, + Object, + Null, +} + +const JSON_TYPES: [JsonType; 7] = [ + JsonType::String, + JsonType::Number, + JsonType::Integer, + JsonType::Boolean, + JsonType::Array, + JsonType::Object, + JsonType::Null, +]; + +impl JsonType { + fn from_name(value: &str) -> Option { + match value { + "string" => Some(Self::String), + "number" => Some(Self::Number), + "integer" => Some(Self::Integer), + "boolean" => Some(Self::Boolean), + "array" => Some(Self::Array), + "object" => Some(Self::Object), + "null" => Some(Self::Null), + _ => None, + } + } + + fn for_value(value: &Value) -> Self { + match value { + Value::Null => Self::Null, + Value::Bool(_) => Self::Boolean, + Value::Number(number) if number.is_i64() || number.is_u64() => Self::Integer, + Value::Number(_) => Self::Number, + Value::String(_) => Self::String, + Value::Array(_) => Self::Array, + Value::Object(_) => Self::Object, + } + } + + fn xtml_name(self) -> &'static str { + match self { + Self::String => "string", + Self::Number | Self::Integer => "number", + Self::Boolean => "boolean", + Self::Array => "array", + Self::Object => "object", + Self::Null => "null", + } + } +} + +fn resolve_local_ref<'a>(reference: &str, root: &'a Value) -> Option<&'a Value> { + let mut value = root; + for raw_part in reference.strip_prefix("#/")?.split('/') { + let part = raw_part.replace("~1", "/").replace("~0", "~"); + value = value.get(&part)?; + } + matches!(value, Value::Bool(_) | Value::Object(_)).then_some(value) +} + +fn schema_types(schema: &Value, root: &Value, seen_refs: &HashSet) -> Vec { + if let Some(allowed) = schema.as_bool() { + return if allowed { + JSON_TYPES.to_vec() + } else { + Vec::new() + }; + } + let Some(schema) = schema.as_object() else { + return JSON_TYPES.to_vec(); + }; + + if let Some(reference) = schema.get("$ref").and_then(Value::as_str) + && !seen_refs.contains(reference) + && let Some(target) = resolve_local_ref(reference, root) + { + let mut nested_seen = seen_refs.clone(); + nested_seen.insert(reference.to_string()); + return schema_types(target, root, &nested_seen); + } + + if let Some(schema_type) = schema.get("type") { + if let Some(schema_type) = schema_type.as_str() { + return JsonType::from_name(schema_type) + .map(|value| vec![value]) + .unwrap_or_else(|| JSON_TYPES.to_vec()); + } + if let Some(schema_types) = schema_type.as_array() { + return JSON_TYPES + .into_iter() + .filter(|candidate| { + schema_types.iter().any(|value| { + value + .as_str() + .and_then(JsonType::from_name) + .is_some_and(|value| value == *candidate) + }) + }) + .collect(); + } + } + + for keyword in ["anyOf", "oneOf"] { + if let Some(options) = schema.get(keyword).and_then(Value::as_array) { + let option_types: HashSet<_> = options + .iter() + .filter(|option| matches!(option, Value::Bool(_) | Value::Object(_))) + .flat_map(|option| schema_types(option, root, seen_refs)) + .collect(); + return JSON_TYPES + .into_iter() + .filter(|candidate| option_types.contains(candidate)) + .collect(); + } + } + + if let Some(options) = schema.get("allOf").and_then(Value::as_array) { + let all_types: HashSet<_> = JSON_TYPES.into_iter().collect(); + let mut constrained = options + .iter() + .filter(|option| matches!(option, Value::Bool(_) | Value::Object(_))) + .map(|option| { + schema_types(option, root, seen_refs) + .into_iter() + .collect::>() + }) + .filter(|types| types != &all_types); + if let Some(mut result) = constrained.next() { + if result.contains(&JsonType::Number) { + result.insert(JsonType::Integer); + } + for mut types in constrained { + if types.contains(&JsonType::Number) { + types.insert(JsonType::Integer); + } + result.retain(|value| types.contains(value)); + } + if result.contains(&JsonType::Number) { + result.remove(&JsonType::Integer); + } + return JSON_TYPES + .into_iter() + .filter(|candidate| result.contains(candidate)) + .collect(); + } + } + + if let Some(value) = schema.get("const") { + return vec![JsonType::for_value(value)]; + } + if let Some(values) = schema.get("enum").and_then(Value::as_array) { + let enum_types: HashSet<_> = values.iter().map(JsonType::for_value).collect(); + return JSON_TYPES + .into_iter() + .filter(|candidate| enum_types.contains(candidate)) + .collect(); + } + + const OBJECT_KEYWORDS: [&str; 9] = [ + "additionalProperties", + "dependentRequired", + "dependentSchemas", + "maxProperties", + "minProperties", + "patternProperties", + "properties", + "propertyNames", + "required", + ]; + const ARRAY_KEYWORDS: [&str; 8] = [ + "contains", + "items", + "maxContains", + "maxItems", + "minContains", + "minItems", + "prefixItems", + "uniqueItems", + ]; + const STRING_KEYWORDS: [&str; 4] = ["format", "maxLength", "minLength", "pattern"]; + const NUMBER_KEYWORDS: [&str; 5] = [ + "exclusiveMaximum", + "exclusiveMinimum", + "maximum", + "minimum", + "multipleOf", + ]; + if OBJECT_KEYWORDS.iter().any(|key| schema.contains_key(*key)) { + vec![JsonType::Object] + } else if ARRAY_KEYWORDS.iter().any(|key| schema.contains_key(*key)) { + vec![JsonType::Array] + } else if STRING_KEYWORDS.iter().any(|key| schema.contains_key(*key)) { + vec![JsonType::String] + } else if NUMBER_KEYWORDS.iter().any(|key| schema.contains_key(*key)) { + vec![JsonType::Number] + } else { + JSON_TYPES.to_vec() + } +} + +fn single_xtml_type(schema: &Value, root: &Value) -> Option<&'static str> { + let types = schema_types(schema, root, &HashSet::new()); + (types.len() == 1).then(|| types[0].xtml_name()) +} + fn bounded_string_regex(schema: &Map) -> Option { let max_len = schema.get("maxLength")?.as_u64()?; if max_len > 4096 { @@ -194,6 +412,146 @@ fn arguments_block(parameters: Option<&Value>) -> Format { star(one_of(tags)) } +fn nonempty_argument_format(key: &str, xtml_type: &str) -> Format { + Format::Tag(TagFormat { + begin: format!( + "{OPEN}argument key=\"{}\" type=\"{xtml_type}\"{SEP}", + escape_attr(key) + ), + // K3 string values are raw rather than JSON quoted. Requiring one + // atom prevents the empty required arguments that pass the generic K3 + // grammar while still stopping before an XTML close marker. + content: Box::new(Format::Regex(RegexFormat { + pattern: format!("{STRING_ATOM}+"), + })), + end: ARGUMENT_CLOSE.to_string(), + }) +} + +fn auto_arguments_block(tool: &ToolDefinition) -> Format { + let Some(parameters) = tool.parameters.as_ref().and_then(Value::as_object) else { + return Format::AnyText(AnyTextFormat { excludes: vec![] }); + }; + let root = Value::Object(parameters.clone()); + let Some(properties) = parameters.get("properties").and_then(Value::as_object) else { + return Format::AnyText(AnyTextFormat { excludes: vec![] }); + }; + let required = match parameters.get("required") { + None => Vec::new(), + Some(Value::Array(required)) if required.iter().all(|value| value.as_str().is_some()) => { + required + .iter() + .filter_map(Value::as_str) + .collect::>() + } + Some(_) => return Format::AnyText(AnyTextFormat { excludes: vec![] }), + }; + + let mut elements = Vec::new(); + if required.is_empty() { + let alternatives: Vec<_> = properties + .iter() + .filter_map(|(key, schema)| { + matches!(schema, Value::Bool(_) | Value::Object(_)) + .then(|| single_xtml_type(schema, &root)) + .flatten() + .map(|xtml_type| nonempty_argument_format(key, xtml_type)) + }) + .collect(); + if alternatives.is_empty() { + return Format::AnyText(AnyTextFormat { excludes: vec![] }); + } + elements.push(one_of(alternatives)); + } else { + for key in required { + let Some(schema) = properties.get(key) else { + return Format::AnyText(AnyTextFormat { excludes: vec![] }); + }; + if !matches!(schema, Value::Bool(_) | Value::Object(_)) { + return Format::AnyText(AnyTextFormat { excludes: vec![] }); + } + let Some(xtml_type) = single_xtml_type(schema, &root) else { + return Format::AnyText(AnyTextFormat { excludes: vec![] }); + }; + elements.push(nonempty_argument_format(key, xtml_type)); + } + } + + // After required arguments are present, keep optional/additional argument + // handling deliberately loose, matching SGLang's K3 auto grammar. + elements.push(Format::AnyText(AnyTextFormat { + excludes: vec![CALL_OPEN.to_string(), CALL_CLOSE.to_string()], + })); + Format::Sequence(SequenceFormat { elements }) +} + +fn auto_call_tag(tool: &ToolDefinition) -> TagFormat { + TagFormat { + begin: format!("{OPEN}call tool=\"{}\" index=\"", escape_attr(&tool.name)), + content: Box::new(Format::Sequence(SequenceFormat { + elements: vec![ + Format::Regex(RegexFormat { + pattern: "[1-9][0-9]*".to_string(), + }), + Format::ConstString(ConstStringFormat { + value: format!("\"{SEP}"), + }), + auto_arguments_block(tool), + ], + })), + end: CALL_CLOSE.to_string(), + } +} + +fn build_auto_structural_tag( + tools: Vec<&ToolDefinition>, + ctx: &ToolCallFormatBuildContext<'_>, +) -> StructuralTag { + let call_tags: Vec<_> = tools.into_iter().map(auto_call_tag).collect(); + let parallel_tool_calls = !ctx.stop_after_first(); + let calls = if parallel_tool_calls { + Format::TagsWithSeparator(TagsWithSeparatorFormat { + tags: call_tags, + separator: String::new(), + at_least_one: true, + stop_after_first: false, + }) + } else { + one_of(call_tags.into_iter().map(Format::Tag).collect()) + }; + let tools_tag = TagFormat { + begin: TOOLS_OPEN.to_string(), + content: Box::new(calls), + end: TOOLS_CLOSE.to_string(), + }; + let suffix = Format::TriggeredTags(TriggeredTagsFormat { + triggers: vec![TOOLS_OPEN.to_string()], + tags: vec![tools_tag], + at_least_one: false, + stop_after_first: !parallel_tool_calls, + excludes: vec![ + THINK_OPEN.to_string(), + THINK_CLOSE.to_string(), + CALL_OPEN.to_string(), + ], + }); + let format = if ctx.starts_in_reasoning { + Format::Sequence(SequenceFormat { + elements: vec![ + Format::Tag(TagFormat { + begin: String::new(), + content: Box::new(Format::AnyText(AnyTextFormat { excludes: vec![] })), + end: THINK_CLOSE.to_string(), + }), + suffix, + ], + }) + } else { + suffix + }; + StructuralTag { format } +} + fn call_tag(tool: &ToolDefinition, strict_schema: bool) -> TagFormat { // Match vLLM's K3 behavior: use the declared schema unless the caller // explicitly sets strict=false. Global strict mode overrides that opt-out. @@ -228,6 +586,9 @@ pub(crate) fn build_kimi_k3( if tools.is_empty() { return Ok(None); } + if matches!(ctx.tool_choice, crate::tool_calling::ToolChoice::Auto) { + return Ok(Some(build_auto_structural_tag(tools, ctx))); + } // Moonshot's named-tool contract returns the selected call with no // assistant content. Keeping the response channel itself is required by @@ -369,29 +730,20 @@ mod tests { } #[test] - fn non_named_choices_keep_the_existing_response_body() { + fn required_choice_keeps_the_existing_response_body() { let tools = tools(); + let choice = ToolChoice::Required; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let response_body = &value["format"]["elements"][1]["content"]; - for choice in [ToolChoice::Auto, ToolChoice::Required] { - let value = - serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) - .unwrap(); - let response_body = &value["format"]["elements"][1]["content"]; - - assert_eq!( - response_body["type"], "any_text", - "{choice:?} must retain response text" - ); - assert_eq!( - response_body["excludes"], - json!([]), - "{choice:?} must retain the existing unrestricted response body" - ); - } + assert_eq!(response_body["type"], "any_text"); + assert_eq!(response_body["excludes"], json!([])); } #[test] - fn named_choice_is_mandatory_and_auto_is_optional() { + fn named_choice_is_mandatory_and_auto_uses_an_optional_trigger() { let tools = tools(); let named = ToolChoice::Named("get_weather".to_string()); let named_value = @@ -402,7 +754,93 @@ mod tests { let auto = ToolChoice::Auto; let auto_value = serde_json::to_value(build_kimi_k3(&context(&auto, &tools)).unwrap().unwrap()).unwrap(); - assert_eq!(auto_value["format"]["elements"][2]["type"], "optional"); + assert_eq!(auto_value["format"]["type"], "triggered_tags"); + assert_eq!(auto_value["format"]["at_least_one"], false); + assert_eq!(auto_value["format"]["triggers"], json!([TOOLS_OPEN])); + assert_eq!( + auto_value["format"]["excludes"], + json!([THINK_OPEN, THINK_CLOSE, CALL_OPEN]) + ); + } + + #[test] + fn auto_requires_nonempty_required_arguments_then_allows_optional_content() { + let tools = tools(); + let choice = ToolChoice::Auto; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let calls = &value["format"]["tags"][0]["content"]; + assert_eq!(calls["type"], "tags_with_separator"); + let call = &calls["tags"][0]; + assert_eq!(call["begin"], "<|open|>call tool=\"get_weather\" index=\""); + assert_eq!(call["content"]["elements"][0]["pattern"], "[1-9][0-9]*"); + + let arguments = &call["content"]["elements"][2]; + assert_eq!(arguments["type"], "sequence"); + assert_eq!( + arguments["elements"][0]["begin"], + "<|open|>argument key=\"city\" type=\"string\"<|sep|>" + ); + assert_eq!( + arguments["elements"][0]["content"]["pattern"], + format!("{STRING_ATOM}+") + ); + assert_eq!(arguments["elements"][1]["type"], "any_text"); + assert_eq!( + arguments["elements"][1]["excludes"], + json!([CALL_OPEN, CALL_CLOSE]) + ); + } + + #[test] + fn auto_without_required_properties_still_requires_one_nonempty_argument() { + let tools = vec![ToolDefinition { + name: "run_command".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": { + "command": {"type": "string"}, + "timeout": {"type": "integer"} + } + })), + strict: None, + }]; + let choice = ToolChoice::Auto; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let arguments = &value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2]; + + assert_eq!(arguments["type"], "sequence"); + assert_eq!(arguments["elements"][0]["type"], "or"); + assert_eq!( + arguments["elements"][0]["elements"] + .as_array() + .unwrap() + .len(), + 2 + ); + assert_eq!(arguments["elements"][1]["type"], "any_text"); + } + + #[test] + fn auto_thinking_closes_reasoning_before_the_triggered_tools_suffix() { + let tools = tools(); + let choice = ToolChoice::Auto; + let ctx = ToolCallFormatBuildContext { + tool_choice: &choice, + tools: &tools, + parallel_tool_calls: None, + schema_mode: StructuralTagSchemaMode::Auto, + starts_in_reasoning: true, + }; + let value = serde_json::to_value(build_kimi_k3(&ctx).unwrap().unwrap()).unwrap(); + + assert_eq!(value["format"]["type"], "sequence"); + assert_eq!(value["format"]["elements"][0]["type"], "tag"); + assert_eq!(value["format"]["elements"][0]["end"], THINK_CLOSE); + assert_eq!(value["format"]["elements"][1]["type"], "triggered_tags"); } #[test] diff --git a/parsers/v1/src/tool_calling/structural_tag/triggered_tags.rs b/parsers/v1/src/tool_calling/structural_tag/triggered_tags.rs index 0d259365b..423e0d162 100644 --- a/parsers/v1/src/tool_calling/structural_tag/triggered_tags.rs +++ b/parsers/v1/src/tool_calling/structural_tag/triggered_tags.rs @@ -79,6 +79,7 @@ pub(crate) fn build_triggered_tags( tags, at_least_one, stop_after_first: ctx.stop_after_first(), + excludes: vec![], }), })) } From e8ca4caaa2bc10feb859d9d193d8ccfc5bdee71a Mon Sep 17 00:00:00 2001 From: Indrajit Bhosale Date: Tue, 25 Aug 2026 09:34:28 -0700 Subject: [PATCH 02/10] Fix review comments Signed-off-by: Indrajit Bhosale --- .../tool_calling/structural_tag/kimi_k3.rs | 244 +++++++++++++++--- 1 file changed, 210 insertions(+), 34 deletions(-) diff --git a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs index 3f79555a4..e43a400fa 100644 --- a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs +++ b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs @@ -36,6 +36,7 @@ const ARGUMENT_CLOSE: &str = "<|close|>argument<|sep|>"; const MESSAGE_CLOSE: &str = "<|close|>message<|sep|>"; const STRING_ATOM: &str = r"(?:[^<]|<[^|])"; +const MAX_PERMUTED_REQUIRED_ARGUMENTS: usize = 5; fn escape_attr(value: &str) -> String { value.replace('&', "&").replace('"', """) @@ -428,7 +429,61 @@ fn nonempty_argument_format(key: &str, xtml_type: &str) -> Format { }) } -fn auto_arguments_block(tool: &ToolDefinition) -> Format { +fn loose_auto_arguments() -> Format { + Format::AnyText(AnyTextFormat { + excludes: vec![CALL_OPEN.to_string(), CALL_CLOSE.to_string()], + }) +} + +fn unordered_required_arguments(arguments: &[Format]) -> Format { + fn visit( + arguments: &[Format], + used: &mut [bool], + current: &mut Vec, + alternatives: &mut Vec, + ) { + if current.len() == arguments.len() { + // JSON object properties are unordered. Loose gaps let optional or + // additional arguments appear before, between, or after the + // required arguments without weakening the required tags + // themselves. + let mut elements = Vec::with_capacity(current.len() * 2 + 1); + elements.push(loose_auto_arguments()); + for argument in current.iter().cloned() { + elements.push(argument); + elements.push(loose_auto_arguments()); + } + alternatives.push(Format::Sequence(SequenceFormat { elements })); + return; + } + + for index in 0..arguments.len() { + if used[index] { + continue; + } + used[index] = true; + current.push(arguments[index].clone()); + visit(arguments, used, current, alternatives); + current.pop(); + used[index] = false; + } + } + + let mut alternatives = Vec::new(); + visit( + arguments, + &mut vec![false; arguments.len()], + &mut Vec::with_capacity(arguments.len()), + &mut alternatives, + ); + one_of(alternatives) +} + +fn auto_arguments_block(tool: &ToolDefinition, strict_schema: bool) -> Format { + if !super::builder::kimi_uses_declared_tool_schema(tool, strict_schema) { + return arguments_block(None); + } + let Some(parameters) = tool.parameters.as_ref().and_then(Value::as_object) else { return Format::AnyText(AnyTextFormat { excludes: vec![] }); }; @@ -447,7 +502,6 @@ fn auto_arguments_block(tool: &ToolDefinition) -> Format { Some(_) => return Format::AnyText(AnyTextFormat { excludes: vec![] }), }; - let mut elements = Vec::new(); if required.is_empty() { let alternatives: Vec<_> = properties .iter() @@ -461,31 +515,35 @@ fn auto_arguments_block(tool: &ToolDefinition) -> Format { if alternatives.is_empty() { return Format::AnyText(AnyTextFormat { excludes: vec![] }); } - elements.push(one_of(alternatives)); - } else { - for key in required { - let Some(schema) = properties.get(key) else { - return Format::AnyText(AnyTextFormat { excludes: vec![] }); - }; - if !matches!(schema, Value::Bool(_) | Value::Object(_)) { - return Format::AnyText(AnyTextFormat { excludes: vec![] }); - } - let Some(xtml_type) = single_xtml_type(schema, &root) else { - return Format::AnyText(AnyTextFormat { excludes: vec![] }); - }; - elements.push(nonempty_argument_format(key, xtml_type)); - } + return optional(Format::Sequence(SequenceFormat { + elements: vec![one_of(alternatives), loose_auto_arguments()], + })); } - // After required arguments are present, keep optional/additional argument - // handling deliberately loose, matching SGLang's K3 auto grammar. - elements.push(Format::AnyText(AnyTextFormat { - excludes: vec![CALL_OPEN.to_string(), CALL_CLOSE.to_string()], - })); - Format::Sequence(SequenceFormat { elements }) + // Enumerating every order grows factorially. Preserve schema-valid output + // for unusually wide tools by falling back to permissive content instead + // of imposing the schema's non-semantic `required` array order. + if required.len() > MAX_PERMUTED_REQUIRED_ARGUMENTS { + return loose_auto_arguments(); + } + + let mut arguments = Vec::with_capacity(required.len()); + for key in required { + let Some(schema) = properties.get(key) else { + return Format::AnyText(AnyTextFormat { excludes: vec![] }); + }; + if !matches!(schema, Value::Bool(_) | Value::Object(_)) { + return Format::AnyText(AnyTextFormat { excludes: vec![] }); + } + let Some(xtml_type) = single_xtml_type(schema, &root) else { + return Format::AnyText(AnyTextFormat { excludes: vec![] }); + }; + arguments.push(nonempty_argument_format(key, xtml_type)); + } + unordered_required_arguments(&arguments) } -fn auto_call_tag(tool: &ToolDefinition) -> TagFormat { +fn auto_call_tag(tool: &ToolDefinition, strict_schema: bool) -> TagFormat { TagFormat { begin: format!("{OPEN}call tool=\"{}\" index=\"", escape_attr(&tool.name)), content: Box::new(Format::Sequence(SequenceFormat { @@ -496,7 +554,7 @@ fn auto_call_tag(tool: &ToolDefinition) -> TagFormat { Format::ConstString(ConstStringFormat { value: format!("\"{SEP}"), }), - auto_arguments_block(tool), + auto_arguments_block(tool, strict_schema), ], })), end: CALL_CLOSE.to_string(), @@ -507,7 +565,11 @@ fn build_auto_structural_tag( tools: Vec<&ToolDefinition>, ctx: &ToolCallFormatBuildContext<'_>, ) -> StructuralTag { - let call_tags: Vec<_> = tools.into_iter().map(auto_call_tag).collect(); + let strict_schema = ctx.strict_schema(); + let call_tags: Vec<_> = tools + .into_iter() + .map(|tool| auto_call_tag(tool, strict_schema)) + .collect(); let parallel_tool_calls = !ctx.stop_after_first(); let calls = if parallel_tool_calls { Format::TagsWithSeparator(TagsWithSeparatorFormat { @@ -778,23 +840,99 @@ mod tests { let arguments = &call["content"]["elements"][2]; assert_eq!(arguments["type"], "sequence"); + assert_eq!(arguments["elements"][0]["type"], "any_text"); assert_eq!( - arguments["elements"][0]["begin"], + arguments["elements"][1]["begin"], "<|open|>argument key=\"city\" type=\"string\"<|sep|>" ); assert_eq!( - arguments["elements"][0]["content"]["pattern"], + arguments["elements"][1]["content"]["pattern"], format!("{STRING_ATOM}+") ); - assert_eq!(arguments["elements"][1]["type"], "any_text"); + assert_eq!(arguments["elements"][2]["type"], "any_text"); assert_eq!( - arguments["elements"][1]["excludes"], + arguments["elements"][2]["excludes"], json!([CALL_OPEN, CALL_CLOSE]) ); } #[test] - fn auto_without_required_properties_still_requires_one_nonempty_argument() { + fn auto_required_arguments_accept_every_property_order() { + let tools = vec![ToolDefinition { + name: "get_weather".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": { + "city": {"type": "string"}, + "days": {"type": "integer"}, + "units": {"type": "string"} + }, + "required": ["city", "days"] + })), + strict: None, + }]; + let choice = ToolChoice::Auto; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let arguments = &value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2]; + let orders = arguments["elements"].as_array().unwrap(); + + assert_eq!(arguments["type"], "or"); + assert_eq!(orders.len(), 2); + assert_eq!( + orders[0]["elements"][1]["begin"], + "<|open|>argument key=\"city\" type=\"string\"<|sep|>" + ); + assert_eq!( + orders[0]["elements"][3]["begin"], + "<|open|>argument key=\"days\" type=\"number\"<|sep|>" + ); + assert_eq!( + orders[1]["elements"][1]["begin"], + "<|open|>argument key=\"days\" type=\"number\"<|sep|>" + ); + assert_eq!( + orders[1]["elements"][3]["begin"], + "<|open|>argument key=\"city\" type=\"string\"<|sep|>" + ); + for order in orders { + assert_eq!(order["elements"][0]["type"], "any_text"); + assert_eq!(order["elements"][2]["type"], "any_text"); + assert_eq!(order["elements"][4]["type"], "any_text"); + } + } + + #[test] + fn auto_wide_required_schema_falls_back_instead_of_fixing_argument_order() { + let tools = vec![ToolDefinition { + name: "wide_tool".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": { + "a": {"type": "string"}, + "b": {"type": "string"}, + "c": {"type": "string"}, + "d": {"type": "string"}, + "e": {"type": "string"}, + "f": {"type": "string"} + }, + "required": ["a", "b", "c", "d", "e", "f"] + })), + strict: None, + }]; + let choice = ToolChoice::Auto; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let arguments = &value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2]; + + assert_eq!(arguments["type"], "any_text"); + assert_eq!(arguments["excludes"], json!([CALL_OPEN, CALL_CLOSE])); + } + + #[test] + fn auto_without_required_properties_allows_an_empty_argument_body() { let tools = vec![ToolDefinition { name: "run_command".to_string(), parameters: Some(json!({ @@ -812,16 +950,54 @@ mod tests { .unwrap(); let arguments = &value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2]; - assert_eq!(arguments["type"], "sequence"); - assert_eq!(arguments["elements"][0]["type"], "or"); + assert_eq!(arguments["type"], "optional"); + assert_eq!(arguments["content"]["type"], "sequence"); + assert_eq!(arguments["content"]["elements"][0]["type"], "or"); assert_eq!( - arguments["elements"][0]["elements"] + arguments["content"]["elements"][0]["elements"] .as_array() .unwrap() .len(), 2 ); - assert_eq!(arguments["elements"][1]["type"], "any_text"); + assert_eq!(arguments["content"]["elements"][1]["type"], "any_text"); + } + + #[test] + fn auto_explicit_non_strict_tool_uses_permissive_arguments() { + let tools = vec![ToolDefinition { + name: "get_weather".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"] + })), + strict: Some(false), + }]; + let choice = ToolChoice::Auto; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let arguments = &value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2]; + + assert_eq!(arguments["type"], "star"); + assert_eq!(arguments["content"]["begin"], "<|open|>argument "); + assert!(!arguments.to_string().contains("key=\\\"city\\\"")); + + let strict_ctx = ToolCallFormatBuildContext { + tool_choice: &choice, + tools: &tools, + parallel_tool_calls: None, + schema_mode: StructuralTagSchemaMode::Strict, + starts_in_reasoning: false, + }; + let strict_value = + serde_json::to_value(build_kimi_k3(&strict_ctx).unwrap().unwrap()).unwrap(); + let strict_arguments = + &strict_value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2]; + + assert_eq!(strict_arguments["type"], "sequence"); + assert!(strict_arguments.to_string().contains("key=\\\"city\\\"")); } #[test] From 913e344369f8a2018f1694c7486de4a433676416 Mon Sep 17 00:00:00 2001 From: Indrajit Bhosale Date: Wed, 2 Sep 2026 22:33:05 -0700 Subject: [PATCH 03/10] Resolve comments Signed-off-by: Indrajit Bhosale --- .../tool_calling/structural_tag/builder.rs | 8 +- .../tool_calling/structural_tag/kimi_k3.rs | 168 ++++++++---------- .../src/tool_calling/structural_tag/tests.rs | 26 --- 3 files changed, 77 insertions(+), 125 deletions(-) diff --git a/parsers/v1/src/tool_calling/structural_tag/builder.rs b/parsers/v1/src/tool_calling/structural_tag/builder.rs index ff3a7476d..82b7701f3 100644 --- a/parsers/v1/src/tool_calling/structural_tag/builder.rs +++ b/parsers/v1/src/tool_calling/structural_tag/builder.rs @@ -21,11 +21,13 @@ use super::triggered_tags::{self, TriggeredTagsConfig}; #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)] #[serde(rename_all = "snake_case")] pub enum StructuralTagSchemaMode { - /// Real schema only for tools with `strict: true`; others get an - /// unconstrained schema (`true` in xgrammar). + /// Use the declared schema unless the tool explicitly sets `strict: false`. + /// An omitted `strict` field is schema-enforced; explicit `false` uses an + /// unconstrained argument schema (`true` in xgrammar). #[default] Auto, - /// Real parameter schema for all tools regardless of `strict` flag. + /// Use the declared parameter schema for all tools, overriding + /// `strict: false`. Strict, } diff --git a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs index c50a64b99..2b5e9fd25 100644 --- a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs +++ b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs @@ -36,7 +36,6 @@ const ARGUMENT_CLOSE: &str = "<|close|>argument<|sep|>"; const MESSAGE_CLOSE: &str = "<|close|>message<|sep|>"; const STRING_ATOM: &str = r"(?:[^<]|<[^|])"; -const MAX_PERMUTED_REQUIRED_ARGUMENTS: usize = 5; fn escape_attr(value: &str) -> String { value.replace('&', "&").replace('"', """) @@ -461,54 +460,24 @@ fn auto_argument_format( })) } -fn loose_auto_arguments() -> Format { - Format::AnyText(AnyTextFormat { - excludes: vec![CALL_OPEN.to_string(), CALL_CLOSE.to_string()], - }) +fn loose_auto_arguments(required_keys: &[&str]) -> Format { + let mut excludes = vec![CALL_OPEN.to_string(), CALL_CLOSE.to_string()]; + excludes.extend( + required_keys + .iter() + .map(|key| format!("{OPEN}argument key=\"{}\"", escape_attr(key))), + ); + Format::AnyText(AnyTextFormat { excludes }) } -fn unordered_required_arguments(arguments: &[Format]) -> Format { - fn visit( - arguments: &[Format], - used: &mut [bool], - current: &mut Vec, - alternatives: &mut Vec, - ) { - if current.len() == arguments.len() { - // JSON object properties are unordered. Loose gaps let optional or - // additional arguments appear before, between, or after the - // required arguments without weakening the required tags - // themselves. - let mut elements = Vec::with_capacity(current.len() * 2 + 1); - elements.push(loose_auto_arguments()); - for argument in current.iter().cloned() { - elements.push(argument); - elements.push(loose_auto_arguments()); - } - alternatives.push(Format::Sequence(SequenceFormat { elements })); - return; - } - - for index in 0..arguments.len() { - if used[index] { - continue; - } - used[index] = true; - current.push(arguments[index].clone()); - visit(arguments, used, current, alternatives); - current.pop(); - used[index] = false; - } - } - - let mut alternatives = Vec::new(); - visit( - arguments, - &mut vec![false; arguments.len()], - &mut Vec::with_capacity(arguments.len()), - &mut alternatives, - ); - one_of(alternatives) +fn canonical_required_arguments(mut arguments: Vec, required_keys: &[&str]) -> Format { + // Guided decoding only needs one schema-valid order. Keep the schema's + // required-array order, then allow optional/additional arguments without + // allowing a required key to overwrite its constrained value. + arguments.push(loose_auto_arguments(required_keys)); + Format::Sequence(SequenceFormat { + elements: arguments, + }) } fn auto_arguments_block(tool: &ToolDefinition, strict_schema: bool) -> Format { @@ -548,20 +517,13 @@ fn auto_arguments_block(tool: &ToolDefinition, strict_schema: bool) -> Format { return Format::AnyText(AnyTextFormat { excludes: vec![] }); } return optional(Format::Sequence(SequenceFormat { - elements: vec![one_of(alternatives), loose_auto_arguments()], + elements: vec![one_of(alternatives), loose_auto_arguments(&[])], })); } - // Enumerating every order grows factorially. Preserve schema-valid output - // for unusually wide tools by falling back to permissive content instead - // of imposing the schema's non-semantic `required` array order. - if required.len() > MAX_PERMUTED_REQUIRED_ARGUMENTS { - return loose_auto_arguments(); - } - let mut arguments = Vec::with_capacity(required.len()); - for key in required { - let Some(schema) = properties.get(key) else { + for key in &required { + let Some(schema) = properties.get(*key) else { return Format::AnyText(AnyTextFormat { excludes: vec![] }); }; if !matches!(schema, Value::Bool(_) | Value::Object(_)) { @@ -572,7 +534,7 @@ fn auto_arguments_block(tool: &ToolDefinition, strict_schema: bool) -> Format { }; arguments.push(argument); } - unordered_required_arguments(&arguments) + canonical_required_arguments(arguments, &required) } fn auto_call_tag(tool: &ToolDefinition, strict_schema: bool) -> TagFormat { @@ -872,24 +834,23 @@ mod tests { let arguments = &call["content"]["elements"][2]; assert_eq!(arguments["type"], "sequence"); - assert_eq!(arguments["elements"][0]["type"], "any_text"); assert_eq!( - arguments["elements"][1]["begin"], + arguments["elements"][0]["begin"], "<|open|>argument key=\"city\" type=\"string\"<|sep|>" ); assert_eq!( - arguments["elements"][1]["content"]["pattern"], + arguments["elements"][0]["content"]["pattern"], format!("{STRING_ATOM}+") ); - assert_eq!(arguments["elements"][2]["type"], "any_text"); + assert_eq!(arguments["elements"][1]["type"], "any_text"); assert_eq!( - arguments["elements"][2]["excludes"], - json!([CALL_OPEN, CALL_CLOSE]) + arguments["elements"][1]["excludes"], + json!([CALL_OPEN, CALL_CLOSE, "<|open|>argument key=\"city\""]) ); } #[test] - fn auto_required_arguments_accept_every_property_order() { + fn auto_required_arguments_use_canonical_schema_order() { let tools = vec![ToolDefinition { name: "get_weather".to_string(), parameters: Some(json!({ @@ -908,36 +869,30 @@ mod tests { serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) .unwrap(); let arguments = &value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2]; - let orders = arguments["elements"].as_array().unwrap(); + let elements = arguments["elements"].as_array().unwrap(); - assert_eq!(arguments["type"], "or"); - assert_eq!(orders.len(), 2); + assert_eq!(arguments["type"], "sequence"); + assert_eq!(elements.len(), 3); assert_eq!( - orders[0]["elements"][1]["begin"], + elements[0]["begin"], "<|open|>argument key=\"city\" type=\"string\"<|sep|>" ); assert_eq!( - orders[0]["elements"][3]["begin"], + elements[1]["begin"], "<|open|>argument key=\"days\" type=\"number\"<|sep|>" ); - assert_eq!(orders[0]["elements"][3]["content"]["type"], "json_schema"); + assert_eq!(elements[1]["content"]["type"], "json_schema"); + assert_eq!(elements[1]["content"]["json_schema"]["type"], "integer"); + assert_eq!(elements[2]["type"], "any_text"); assert_eq!( - orders[0]["elements"][3]["content"]["json_schema"]["type"], - "integer" - ); - assert_eq!( - orders[1]["elements"][1]["begin"], - "<|open|>argument key=\"days\" type=\"number\"<|sep|>" + elements[2]["excludes"], + json!([ + CALL_OPEN, + CALL_CLOSE, + "<|open|>argument key=\"city\"", + "<|open|>argument key=\"days\"" + ]) ); - assert_eq!( - orders[1]["elements"][3]["begin"], - "<|open|>argument key=\"city\" type=\"string\"<|sep|>" - ); - for order in orders { - assert_eq!(order["elements"][0]["type"], "any_text"); - assert_eq!(order["elements"][2]["type"], "any_text"); - assert_eq!(order["elements"][4]["type"], "any_text"); - } } #[test] @@ -965,18 +920,18 @@ mod tests { serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) .unwrap(); let arguments = &value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2]; - let first_order = &arguments["elements"][0]["elements"]; + let elements = &arguments["elements"]; - for index in [1, 3, 5, 7] { - assert_eq!(first_order[index]["content"]["type"], "json_schema"); + for index in 0..4 { + assert_eq!(elements[index]["content"]["type"], "json_schema"); } - assert_eq!(first_order[1]["content"]["json_schema"]["minimum"], 1); + assert_eq!(elements[0]["content"]["json_schema"]["minimum"], 1); assert_eq!( - first_order[5]["content"]["json_schema"]["items"]["type"], + elements[2]["content"]["json_schema"]["items"]["type"], "string" ); assert_eq!( - first_order[7]["content"]["json_schema"]["required"], + elements[3]["content"]["json_schema"]["required"], json!(["source"]) ); } @@ -1002,7 +957,7 @@ mod tests { serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) .unwrap(); let argument = &value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2] - ["elements"][1]; + ["elements"][0]; assert_eq!(argument["content"]["type"], "json_schema"); assert_eq!( @@ -1016,7 +971,7 @@ mod tests { } #[test] - fn auto_wide_required_schema_falls_back_instead_of_fixing_argument_order() { + fn auto_wide_required_schema_uses_canonical_order() { let tools = vec![ToolDefinition { name: "wide_tool".to_string(), parameters: Some(json!({ @@ -1039,8 +994,29 @@ mod tests { .unwrap(); let arguments = &value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2]; - assert_eq!(arguments["type"], "any_text"); - assert_eq!(arguments["excludes"], json!([CALL_OPEN, CALL_CLOSE])); + assert_eq!(arguments["type"], "sequence"); + assert_eq!(arguments["elements"].as_array().unwrap().len(), 7); + assert_eq!( + arguments["elements"][0]["begin"], + "<|open|>argument key=\"a\" type=\"string\"<|sep|>" + ); + assert_eq!( + arguments["elements"][5]["begin"], + "<|open|>argument key=\"f\" type=\"string\"<|sep|>" + ); + assert_eq!( + arguments["elements"][6]["excludes"], + json!([ + CALL_OPEN, + CALL_CLOSE, + "<|open|>argument key=\"a\"", + "<|open|>argument key=\"b\"", + "<|open|>argument key=\"c\"", + "<|open|>argument key=\"d\"", + "<|open|>argument key=\"e\"", + "<|open|>argument key=\"f\"" + ]) + ); } #[test] diff --git a/parsers/v1/src/tool_calling/structural_tag/tests.rs b/parsers/v1/src/tool_calling/structural_tag/tests.rs index 914e8c4a1..b5030f240 100644 --- a/parsers/v1/src/tool_calling/structural_tag/tests.rs +++ b/parsers/v1/src/tool_calling/structural_tag/tests.rs @@ -650,29 +650,3 @@ fn parser_map_structural_tag_smoke() { assert_eq!(tag["type"], "structural_tag", "'{name}'"); } } - -#[test] -fn structural_tag_builder_parser_set_is_stable() { - use std::collections::BTreeSet; - - let actual: BTreeSet<_> = crate::tool_calling::parsers::get_tool_parser_map() - .iter() - .filter_map(|(name, config)| config.structural_tag_builder.is_some().then_some(*name)) - .collect(); - let expected = BTreeSet::from([ - "deepseek-v4", - "deepseek_v3_2", - "deepseek_v4", - "deepseekv4", - "hermes", - "inkling", - "kimi-k3", - "kimi_k2", - "kimi_k3", - "nemotron_nano", - "qwen25", - "qwen3_coder", - ]); - - assert_eq!(actual, expected); -} From 72a9a41efa567729cc3f807920ceee2b28850e5d Mon Sep 17 00:00:00 2001 From: Indrajit Bhosale Date: Thu, 3 Sep 2026 18:02:44 -0700 Subject: [PATCH 04/10] Fix review comments Signed-off-by: Indrajit Bhosale --- .../tool_calling/structural_tag/kimi_k3.rs | 153 ++++++------------ 1 file changed, 50 insertions(+), 103 deletions(-) diff --git a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs index 2b5e9fd25..1785f5cc3 100644 --- a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs +++ b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs @@ -460,21 +460,37 @@ fn auto_argument_format( })) } -fn loose_auto_arguments(required_keys: &[&str]) -> Format { - let mut excludes = vec![CALL_OPEN.to_string(), CALL_CLOSE.to_string()]; - excludes.extend( - required_keys - .iter() - .map(|key| format!("{OPEN}argument key=\"{}\"", escape_attr(key))), - ); - Format::AnyText(AnyTextFormat { excludes }) +fn optional_auto_arguments( + properties: &Map, + required_keys: &[&str], + root: &Value, + root_defs: &Map, +) -> Option { + let arguments = properties + .iter() + .filter(|(key, _)| !required_keys.contains(&key.as_str())) + .filter_map(|(key, schema)| { + matches!(schema, Value::Bool(_) | Value::Object(_)) + .then(|| auto_argument_format(key, schema, root, root_defs, false)) + .flatten() + }) + .collect::>(); + + (!arguments.is_empty()).then(|| star(one_of(arguments))) } -fn canonical_required_arguments(mut arguments: Vec, required_keys: &[&str]) -> Format { +fn canonical_required_arguments( + mut arguments: Vec, + optional_arguments: Option, +) -> Format { // Guided decoding only needs one schema-valid order. Keep the schema's - // required-array order, then allow optional/additional arguments without - // allowing a required key to overwrite its constrained value. - arguments.push(loose_auto_arguments(required_keys)); + // required-array order, then allow only complete, declared optional + // argument tags. Arbitrary text would be accepted by the grammar but + // rejected by the K3 parser, while a permissive argument tag could repeat + // a required key and overwrite its constrained value. + if let Some(optional_arguments) = optional_arguments { + arguments.push(optional_arguments); + } Format::Sequence(SequenceFormat { elements: arguments, }) @@ -486,11 +502,11 @@ fn auto_arguments_block(tool: &ToolDefinition, strict_schema: bool) -> Format { } let Some(parameters) = tool.parameters.as_ref().and_then(Value::as_object) else { - return Format::AnyText(AnyTextFormat { excludes: vec![] }); + return arguments_block(None); }; let root = Value::Object(parameters.clone()); let Some(properties) = parameters.get("properties").and_then(Value::as_object) else { - return Format::AnyText(AnyTextFormat { excludes: vec![] }); + return arguments_block(None); }; let root_defs = root_definitions(parameters); let required = match parameters.get("required") { @@ -501,40 +517,32 @@ fn auto_arguments_block(tool: &ToolDefinition, strict_schema: bool) -> Format { .filter_map(Value::as_str) .collect::>() } - Some(_) => return Format::AnyText(AnyTextFormat { excludes: vec![] }), + Some(_) => return arguments_block(None), }; + let optional_arguments = optional_auto_arguments(properties, &required, &root, &root_defs); if required.is_empty() { - let alternatives: Vec<_> = properties - .iter() - .filter_map(|(key, schema)| { - matches!(schema, Value::Bool(_) | Value::Object(_)) - .then(|| auto_argument_format(key, schema, &root, &root_defs, false)) - .flatten() + return optional_arguments.unwrap_or_else(|| { + Format::ConstString(ConstStringFormat { + value: String::new(), }) - .collect(); - if alternatives.is_empty() { - return Format::AnyText(AnyTextFormat { excludes: vec![] }); - } - return optional(Format::Sequence(SequenceFormat { - elements: vec![one_of(alternatives), loose_auto_arguments(&[])], - })); + }); } let mut arguments = Vec::with_capacity(required.len()); for key in &required { let Some(schema) = properties.get(*key) else { - return Format::AnyText(AnyTextFormat { excludes: vec![] }); + return arguments_block(None); }; if !matches!(schema, Value::Bool(_) | Value::Object(_)) { - return Format::AnyText(AnyTextFormat { excludes: vec![] }); + return arguments_block(None); } let Some(argument) = auto_argument_format(key, schema, &root, &root_defs, true) else { - return Format::AnyText(AnyTextFormat { excludes: vec![] }); + return arguments_block(None); }; arguments.push(argument); } - canonical_required_arguments(arguments, &required) + canonical_required_arguments(arguments, optional_arguments) } fn auto_call_tag(tool: &ToolDefinition, strict_schema: bool) -> TagFormat { @@ -842,10 +850,10 @@ mod tests { arguments["elements"][0]["content"]["pattern"], format!("{STRING_ATOM}+") ); - assert_eq!(arguments["elements"][1]["type"], "any_text"); + assert_eq!(arguments["elements"][1]["type"], "star"); assert_eq!( - arguments["elements"][1]["excludes"], - json!([CALL_OPEN, CALL_CLOSE, "<|open|>argument key=\"city\""]) + arguments["elements"][1]["content"]["begin"], + "<|open|>argument key=\"days\" type=\"number\"<|sep|>" ); } @@ -883,15 +891,10 @@ mod tests { ); assert_eq!(elements[1]["content"]["type"], "json_schema"); assert_eq!(elements[1]["content"]["json_schema"]["type"], "integer"); - assert_eq!(elements[2]["type"], "any_text"); + assert_eq!(elements[2]["type"], "star"); assert_eq!( - elements[2]["excludes"], - json!([ - CALL_OPEN, - CALL_CLOSE, - "<|open|>argument key=\"city\"", - "<|open|>argument key=\"days\"" - ]) + elements[2]["content"]["begin"], + "<|open|>argument key=\"units\" type=\"string\"<|sep|>" ); } @@ -970,55 +973,6 @@ mod tests { ); } - #[test] - fn auto_wide_required_schema_uses_canonical_order() { - let tools = vec![ToolDefinition { - name: "wide_tool".to_string(), - parameters: Some(json!({ - "type": "object", - "properties": { - "a": {"type": "string"}, - "b": {"type": "string"}, - "c": {"type": "string"}, - "d": {"type": "string"}, - "e": {"type": "string"}, - "f": {"type": "string"} - }, - "required": ["a", "b", "c", "d", "e", "f"] - })), - strict: None, - }]; - let choice = ToolChoice::Auto; - let value = - serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) - .unwrap(); - let arguments = &value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2]; - - assert_eq!(arguments["type"], "sequence"); - assert_eq!(arguments["elements"].as_array().unwrap().len(), 7); - assert_eq!( - arguments["elements"][0]["begin"], - "<|open|>argument key=\"a\" type=\"string\"<|sep|>" - ); - assert_eq!( - arguments["elements"][5]["begin"], - "<|open|>argument key=\"f\" type=\"string\"<|sep|>" - ); - assert_eq!( - arguments["elements"][6]["excludes"], - json!([ - CALL_OPEN, - CALL_CLOSE, - "<|open|>argument key=\"a\"", - "<|open|>argument key=\"b\"", - "<|open|>argument key=\"c\"", - "<|open|>argument key=\"d\"", - "<|open|>argument key=\"e\"", - "<|open|>argument key=\"f\"" - ]) - ); - } - #[test] fn auto_without_required_properties_allows_an_empty_argument_body() { let tools = vec![ToolDefinition { @@ -1038,26 +992,19 @@ mod tests { .unwrap(); let arguments = &value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2]; - assert_eq!(arguments["type"], "optional"); - assert_eq!(arguments["content"]["type"], "sequence"); - assert_eq!(arguments["content"]["elements"][0]["type"], "or"); + assert_eq!(arguments["type"], "star"); + assert_eq!(arguments["content"]["type"], "or"); assert_eq!( - arguments["content"]["elements"][0]["elements"] - .as_array() - .unwrap() - .len(), + arguments["content"]["elements"].as_array().unwrap().len(), 2 ); - let alternatives = arguments["content"]["elements"][0]["elements"] - .as_array() - .unwrap(); + let alternatives = arguments["content"]["elements"].as_array().unwrap(); assert_eq!( alternatives[0]["content"]["pattern"], format!("{STRING_ATOM}*") ); assert_eq!(alternatives[1]["content"]["type"], "json_schema"); assert_eq!(alternatives[1]["content"]["json_schema"]["type"], "integer"); - assert_eq!(arguments["content"]["elements"][1]["type"], "any_text"); } #[test] From 992559bfe276f668f1083535bffc95dd1616be89 Mon Sep 17 00:00:00 2001 From: Indrajit Bhosale Date: Tue, 22 Sep 2026 14:38:40 -0700 Subject: [PATCH 05/10] Fix comments Signed-off-by: Indrajit Bhosale --- .../src/tool_calling/structural_tag/dsml.rs | 1 - .../src/tool_calling/structural_tag/format.rs | 4 - .../tool_calling/structural_tag/kimi_k2.rs | 1 - .../tool_calling/structural_tag/kimi_k3.rs | 479 +++++++++++++----- .../structural_tag/triggered_tags.rs | 1 - 5 files changed, 359 insertions(+), 127 deletions(-) diff --git a/parsers/v1/src/tool_calling/structural_tag/dsml.rs b/parsers/v1/src/tool_calling/structural_tag/dsml.rs index 2f1f99a17..71439cc5c 100644 --- a/parsers/v1/src/tool_calling/structural_tag/dsml.rs +++ b/parsers/v1/src/tool_calling/structural_tag/dsml.rs @@ -97,7 +97,6 @@ pub(crate) fn build_dsml_tool_calls( tags: vec![block_tag], at_least_one: outer_at_least_one, stop_after_first: ctx.stop_after_first(), - excludes: vec![], }), })) } diff --git a/parsers/v1/src/tool_calling/structural_tag/format.rs b/parsers/v1/src/tool_calling/structural_tag/format.rs index 85e6c1e5f..036120fa8 100644 --- a/parsers/v1/src/tool_calling/structural_tag/format.rs +++ b/parsers/v1/src/tool_calling/structural_tag/format.rs @@ -88,8 +88,6 @@ pub struct TriggeredTagsFormat { pub tags: Vec, pub at_least_one: bool, pub stop_after_first: bool, - #[serde(skip_serializing_if = "Vec::is_empty")] - pub excludes: Vec, } /// `tags_with_separator`: a tag sequence with fixed separators and no free text. @@ -229,7 +227,6 @@ mod tests { tags: vec![], at_least_one: false, stop_after_first: false, - excludes: vec![], }; let value = serde_json::to_value(&format).unwrap(); @@ -243,7 +240,6 @@ mod tests { tags: vec![], at_least_one: true, stop_after_first: false, - excludes: vec![], }; let value = serde_json::to_value(&format).unwrap(); diff --git a/parsers/v1/src/tool_calling/structural_tag/kimi_k2.rs b/parsers/v1/src/tool_calling/structural_tag/kimi_k2.rs index e927b5edd..7a582a010 100644 --- a/parsers/v1/src/tool_calling/structural_tag/kimi_k2.rs +++ b/parsers/v1/src/tool_calling/structural_tag/kimi_k2.rs @@ -101,7 +101,6 @@ pub(crate) fn build_kimi_k2( }], at_least_one: outer_at_least_one, stop_after_first: ctx.stop_after_first(), - excludes: vec![], }) } else { Format::Sequence(SequenceFormat { diff --git a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs index 1785f5cc3..b40fad786 100644 --- a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs +++ b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs @@ -17,7 +17,7 @@ use super::builder::{ToolCallFormatBuildContext, resolve_tools_to_include}; use super::format::{ AnyTextFormat, ConstStringFormat, Format, JsonSchemaFormat, JsonSchemaStyle, OptionalFormat, OrFormat, RegexFormat, SequenceFormat, StarFormat, StructuralTag, TagFormat, - TagsWithSeparatorFormat, TriggeredTagsFormat, + TagsWithSeparatorFormat, }; use crate::tool_calling::ToolDefinition; @@ -271,7 +271,20 @@ fn schema_types(schema: &Value, root: &Value, seen_refs: &HashSet) -> Ve fn single_xtml_type(schema: &Value, root: &Value) -> Option<&'static str> { let types = schema_types(schema, root, &HashSet::new()); - (types.len() == 1).then(|| types[0].xtml_name()) + let xtml_types = types + .into_iter() + .map(JsonType::xtml_name) + .collect::>(); + (xtml_types.len() == 1).then(|| *xtml_types.iter().next().expect("one XTML type")) +} + +fn xtml_types(schema: &Value, root: &Value) -> Vec<&'static str> { + let mut seen = HashSet::new(); + schema_types(schema, root, &HashSet::new()) + .into_iter() + .map(JsonType::xtml_name) + .filter(|xtml_type| seen.insert(*xtml_type)) + .collect() } fn bounded_string_regex(schema: &Map) -> Option { @@ -287,6 +300,78 @@ fn bounded_string_regex(schema: &Map) -> Option { Some(format!("{STRING_ATOM}{{{min_len},{max_len}}}")) } +fn string_content_format(schema: &Value) -> Format { + let Some(schema) = schema.as_object() else { + return Format::AnyText(AnyTextFormat { + excludes: vec![CLOSE.to_string()], + }); + }; + + let enum_values = schema + .get("enum") + .and_then(Value::as_array) + .cloned() + .or_else(|| { + schema + .get("const") + .and_then(Value::as_str) + .map(|value| vec![Value::String(value.to_string())]) + }); + if let Some(values) = enum_values.filter(|values| { + !values.is_empty() + && values.len() <= 256 + && values + .iter() + .all(|value| value.as_str().is_some_and(|string| !string.contains("<|"))) + }) { + return one_of( + values + .iter() + .filter_map(Value::as_str) + .map(|value| { + Format::ConstString(ConstStringFormat { + value: value.to_string(), + }) + }) + .collect(), + ); + } + + if let Some(pattern) = schema.get("pattern").and_then(Value::as_str) + && !schema.contains_key("minLength") + && !schema.contains_key("maxLength") + { + let anchored_start = pattern.starts_with('^'); + let anchored_end = pattern.ends_with('$') && !pattern.ends_with("\\$"); + let pattern = pattern + .strip_prefix('^') + .unwrap_or(pattern) + .strip_suffix('$') + .unwrap_or_else(|| pattern.strip_prefix('^').unwrap_or(pattern)); + let prefix = if anchored_start { + String::new() + } else { + format!("{STRING_ATOM}*") + }; + let suffix = if anchored_end { + String::new() + } else { + format!("{STRING_ATOM}*") + }; + return Format::Regex(RegexFormat { + pattern: format!("{prefix}(?:{pattern}){suffix}"), + }); + } + + if let Some(pattern) = bounded_string_regex(schema) { + return Format::Regex(RegexFormat { pattern }); + } + + Format::AnyText(AnyTextFormat { + excludes: vec![CLOSE.to_string()], + }) +} + fn argument_tag( key: &str, schema: &Value, @@ -309,41 +394,7 @@ fn argument_tag( ); let content = if json_type == "string" { - let enum_values = schema_object - .get("enum") - .and_then(Value::as_array) - .cloned() - .or_else(|| { - schema_object - .get("const") - .and_then(Value::as_str) - .map(|value| vec![Value::String(value.to_string())]) - }); - if let Some(values) = enum_values.filter(|values| { - !values.is_empty() - && values.len() <= 256 - && values - .iter() - .all(|value| value.as_str().is_some_and(|string| !string.contains("<|"))) - }) { - one_of( - values - .iter() - .filter_map(Value::as_str) - .map(|value| { - Format::ConstString(ConstStringFormat { - value: value.to_string(), - }) - }) - .collect(), - ) - } else if let Some(pattern) = bounded_string_regex(schema_object) { - Format::Regex(RegexFormat { pattern }) - } else { - Format::AnyText(AnyTextFormat { - excludes: vec![CLOSE.to_string()], - }) - } + string_content_format(schema) } else { let mut embedded = schema_object.clone(); if let Some(root_defs) = root_defs { @@ -393,6 +444,75 @@ fn root_definitions(parameters: &Map) -> Map { .collect() } +fn resolve_local_refs(schema: &Value, root: &Value, seen_refs: &HashSet) -> Option { + match schema { + Value::Array(values) => values + .iter() + .map(|value| resolve_local_refs(value, root, seen_refs)) + .collect::>>() + .map(Value::Array), + Value::Object(object) => { + if let Some(reference) = object.get("$ref").and_then(Value::as_str) { + if seen_refs.contains(reference) { + return None; + } + let target = resolve_local_ref(reference, root)?; + let mut nested_seen = seen_refs.clone(); + nested_seen.insert(reference.to_string()); + let resolved = resolve_local_refs(target, root, &nested_seen)?; + let siblings = object + .iter() + .filter(|(key, _)| key.as_str() != "$ref") + .map(|(key, value)| { + resolve_local_refs(value, root, seen_refs).map(|value| (key.clone(), value)) + }) + .collect::>>()?; + if siblings.is_empty() { + return Some(resolved); + } + return Some(serde_json::json!({ + "allOf": [resolved, Value::Object(siblings)] + })); + } + + object + .iter() + .map(|(key, value)| { + resolve_local_refs(value, root, seen_refs).map(|value| (key.clone(), value)) + }) + .collect::>>() + .map(Value::Object) + } + _ => Some(schema.clone()), + } +} + +fn schema_for_xtml_type(schema: &Value, root: &Value, xtml_type: &str) -> Option { + let resolved = resolve_local_refs(schema, root, &HashSet::new())?; + if single_xtml_type(&resolved, &resolved) == Some(xtml_type) { + return Some(resolved); + } + + let json_types = schema_types(schema, root, &HashSet::new()) + .into_iter() + .filter(|json_type| json_type.xtml_name() == xtml_type) + .map(|json_type| match json_type { + JsonType::String => "string", + JsonType::Number => "number", + JsonType::Integer => "integer", + JsonType::Boolean => "boolean", + JsonType::Array => "array", + JsonType::Object => "object", + JsonType::Null => "null", + }) + .collect::>(); + match json_types.as_slice() { + [] => None, + [json_type] => Some(serde_json::json!({"type": json_type})), + _ => Some(serde_json::json!({"type": json_types})), + } +} + fn arguments_block(parameters: Option<&Value>) -> Format { let Some(parameters) = parameters.and_then(Value::as_object) else { return star(Format::Tag(permissive_argument_tag())); @@ -416,62 +536,43 @@ fn arguments_block(parameters: Option<&Value>) -> Format { star(one_of(tags)) } -fn auto_argument_format( - key: &str, - schema: &Value, - root: &Value, - root_defs: &Map, - require_nonempty_string: bool, -) -> Option { - let xtml_type = single_xtml_type(schema, root)?; - let content = if xtml_type == "string" { - // K3 string values are raw rather than JSON quoted. Required string - // arguments need at least one atom; optional strings may be empty when - // their schema permits it. - Format::Regex(RegexFormat { - pattern: format!( - "{STRING_ATOM}{}", - if require_nonempty_string { "+" } else { "*" } - ), - }) - } else { - // Non-string K3 argument bodies are JSON. Preserve the complete - // property schema instead of using the string-only regex, and carry - // root definitions so local references remain resolvable. - let mut embedded = schema.as_object()?.clone(); - for (definition_key, definitions) in root_defs { - embedded - .entry(definition_key.clone()) - .or_insert_with(|| definitions.clone()); - } - Format::JsonSchema(JsonSchemaFormat { - json_schema: Value::Object(embedded), - style: JsonSchemaStyle::Json, +fn auto_argument_format(key: &str, schema: &Value, root: &Value) -> Option { + let alternatives = xtml_types(schema, root) + .into_iter() + .filter_map(|xtml_type| { + let content = if xtml_type == "string" { + string_content_format(schema) + } else { + Format::JsonSchema(JsonSchemaFormat { + json_schema: schema_for_xtml_type(schema, root, xtml_type)?, + style: JsonSchemaStyle::Json, + }) + }; + Some(Format::Tag(TagFormat { + begin: format!( + "{OPEN}argument key=\"{}\" type=\"{xtml_type}\"{SEP}", + escape_attr(key) + ), + content: Box::new(content), + end: ARGUMENT_CLOSE.to_string(), + })) }) - }; + .collect::>(); - Some(Format::Tag(TagFormat { - begin: format!( - "{OPEN}argument key=\"{}\" type=\"{xtml_type}\"{SEP}", - escape_attr(key) - ), - content: Box::new(content), - end: ARGUMENT_CLOSE.to_string(), - })) + (!alternatives.is_empty()).then(|| one_of(alternatives)) } fn optional_auto_arguments( properties: &Map, required_keys: &[&str], root: &Value, - root_defs: &Map, ) -> Option { let arguments = properties .iter() .filter(|(key, _)| !required_keys.contains(&key.as_str())) .filter_map(|(key, schema)| { matches!(schema, Value::Bool(_) | Value::Object(_)) - .then(|| auto_argument_format(key, schema, root, root_defs, false)) + .then(|| auto_argument_format(key, schema, root)) .flatten() }) .collect::>(); @@ -501,14 +602,15 @@ fn auto_arguments_block(tool: &ToolDefinition, strict_schema: bool) -> Format { return arguments_block(None); } - let Some(parameters) = tool.parameters.as_ref().and_then(Value::as_object) else { + let Some(root) = tool.parameters.as_ref() else { + return arguments_block(None); + }; + let Some(parameters) = root.as_object() else { return arguments_block(None); }; - let root = Value::Object(parameters.clone()); let Some(properties) = parameters.get("properties").and_then(Value::as_object) else { return arguments_block(None); }; - let root_defs = root_definitions(parameters); let required = match parameters.get("required") { None => Vec::new(), Some(Value::Array(required)) if required.iter().all(|value| value.as_str().is_some()) => { @@ -519,7 +621,7 @@ fn auto_arguments_block(tool: &ToolDefinition, strict_schema: bool) -> Format { } Some(_) => return arguments_block(None), }; - let optional_arguments = optional_auto_arguments(properties, &required, &root, &root_defs); + let optional_arguments = optional_auto_arguments(properties, &required, root); if required.is_empty() { return optional_arguments.unwrap_or_else(|| { @@ -537,7 +639,7 @@ fn auto_arguments_block(tool: &ToolDefinition, strict_schema: bool) -> Format { if !matches!(schema, Value::Bool(_) | Value::Object(_)) { return arguments_block(None); } - let Some(argument) = auto_argument_format(key, schema, &root, &root_defs, true) else { + let Some(argument) = auto_argument_format(key, schema, root) else { return arguments_block(None); }; arguments.push(argument); @@ -588,15 +690,20 @@ fn build_auto_structural_tag( content: Box::new(calls), end: TOOLS_CLOSE.to_string(), }; - let suffix = Format::TriggeredTags(TriggeredTagsFormat { - triggers: vec![TOOLS_OPEN.to_string()], - tags: vec![tools_tag], - at_least_one: false, - stop_after_first: !parallel_tool_calls, - excludes: vec![ - THINK_OPEN.to_string(), - THINK_CLOSE.to_string(), - CALL_OPEN.to_string(), + // Express the optional tools suffix with existing public format nodes. + // Adding `excludes` to the public `TriggeredTagsFormat` struct would break + // downstream struct literals in the published parser crate. + let suffix = Format::Sequence(SequenceFormat { + elements: vec![ + Format::AnyText(AnyTextFormat { + excludes: vec![ + TOOLS_OPEN.to_string(), + THINK_OPEN.to_string(), + THINK_CLOSE.to_string(), + CALL_OPEN.to_string(), + ], + }), + optional(Format::Tag(tools_tag)), ], }); let format = if ctx.starts_in_reasoning { @@ -807,7 +914,7 @@ mod tests { } #[test] - fn named_choice_is_mandatory_and_auto_uses_an_optional_trigger() { + fn named_choice_is_mandatory_and_auto_uses_an_optional_tools_suffix() { let tools = tools(); let named = ToolChoice::Named("get_weather".to_string()); let named_value = @@ -818,23 +925,27 @@ mod tests { let auto = ToolChoice::Auto; let auto_value = serde_json::to_value(build_kimi_k3(&context(&auto, &tools)).unwrap().unwrap()).unwrap(); - assert_eq!(auto_value["format"]["type"], "triggered_tags"); - assert_eq!(auto_value["format"]["at_least_one"], false); - assert_eq!(auto_value["format"]["triggers"], json!([TOOLS_OPEN])); + assert_eq!(auto_value["format"]["type"], "sequence"); + assert_eq!(auto_value["format"]["elements"][0]["type"], "any_text"); + assert_eq!( + auto_value["format"]["elements"][0]["excludes"], + json!([TOOLS_OPEN, THINK_OPEN, THINK_CLOSE, CALL_OPEN]) + ); + assert_eq!(auto_value["format"]["elements"][1]["type"], "optional"); assert_eq!( - auto_value["format"]["excludes"], - json!([THINK_OPEN, THINK_CLOSE, CALL_OPEN]) + auto_value["format"]["elements"][1]["content"]["begin"], + TOOLS_OPEN ); } #[test] - fn auto_requires_nonempty_required_arguments_then_allows_optional_content() { + fn auto_requires_declared_arguments_then_allows_optional_content() { let tools = tools(); let choice = ToolChoice::Auto; let value = serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) .unwrap(); - let calls = &value["format"]["tags"][0]["content"]; + let calls = &value["format"]["elements"][1]["content"]["content"]; assert_eq!(calls["type"], "tags_with_separator"); let call = &calls["tags"][0]; assert_eq!(call["begin"], "<|open|>call tool=\"get_weather\" index=\""); @@ -846,9 +957,10 @@ mod tests { arguments["elements"][0]["begin"], "<|open|>argument key=\"city\" type=\"string\"<|sep|>" ); + assert_eq!(arguments["elements"][0]["content"]["type"], "any_text"); assert_eq!( - arguments["elements"][0]["content"]["pattern"], - format!("{STRING_ATOM}+") + arguments["elements"][0]["content"]["excludes"], + json!([CLOSE]) ); assert_eq!(arguments["elements"][1]["type"], "star"); assert_eq!( @@ -876,7 +988,8 @@ mod tests { let value = serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) .unwrap(); - let arguments = &value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2]; + let arguments = &value["format"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]; let elements = arguments["elements"].as_array().unwrap(); assert_eq!(arguments["type"], "sequence"); @@ -922,7 +1035,8 @@ mod tests { let value = serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) .unwrap(); - let arguments = &value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2]; + let arguments = &value["format"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]; let elements = &arguments["elements"]; for index in 0..4 { @@ -959,17 +1073,142 @@ mod tests { let value = serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) .unwrap(); - let argument = &value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2] - ["elements"][0]; + let argument = &value["format"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]["elements"][0]; assert_eq!(argument["content"]["type"], "json_schema"); + assert_eq!(argument["content"]["json_schema"]["type"], "integer"); + assert_eq!(argument["content"]["json_schema"]["minimum"], 1); + assert!(argument["content"]["json_schema"].get("$ref").is_none()); + } + + #[test] + fn auto_required_ref_into_properties_is_resolved_from_the_original_root() { + let tools = vec![ToolDefinition { + name: "lookup".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": { + "other": {"type": "integer", "minimum": 1}, + "value": {"$ref": "#/properties/other"} + }, + "required": ["value"] + })), + strict: None, + }]; + let choice = ToolChoice::Auto; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let argument = &value["format"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]["elements"][0]; + + assert_eq!(argument["content"]["json_schema"]["type"], "integer"); + assert_eq!(argument["content"]["json_schema"]["minimum"], 1); + assert!(argument["content"]["json_schema"].get("$ref").is_none()); + } + + #[test] + fn auto_string_arguments_preserve_schema_constraints_and_allow_empty_values() { + let tools = vec![ToolDefinition { + name: "strings".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": { + "mode": {"type": "string", "enum": ["fast", "safe"]}, + "bounded": {"type": "string", "minLength": 0, "maxLength": 8}, + "prefixed": {"type": "string", "pattern": "^item-[0-9]+$"}, + "empty_ok": {"type": "string", "minLength": 0, "maxLength": 0} + }, + "required": ["mode", "bounded", "prefixed", "empty_ok"] + })), + strict: None, + }]; + let choice = ToolChoice::Auto; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let arguments = &value["format"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]["elements"]; + + assert_eq!(arguments[0]["content"]["type"], "or"); + assert_eq!(arguments[0]["content"]["elements"][0]["value"], "fast"); + assert_eq!( + arguments[1]["content"]["pattern"], + format!("{STRING_ATOM}{{0,8}}") + ); + assert_eq!(arguments[2]["content"]["pattern"], "(?:item-[0-9]+)"); + assert_eq!( + arguments[3]["content"]["pattern"], + format!("{STRING_ATOM}{{0,0}}") + ); + } + + #[test] + fn auto_numeric_union_uses_one_number_tag_with_the_declared_schema() { + let tools = vec![ToolDefinition { + name: "measure".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": { + "value": {"anyOf": [{"type": "integer"}, {"type": "number"}]} + }, + "required": ["value"] + })), + strict: None, + }]; + let choice = ToolChoice::Auto; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let argument = &value["format"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]["elements"][0]; + assert_eq!( - argument["content"]["json_schema"]["$ref"], - "#/$defs/identifier" + argument["begin"], + "<|open|>argument key=\"value\" type=\"number\"<|sep|>" ); + assert_eq!(argument["content"]["type"], "json_schema"); assert_eq!( - argument["content"]["json_schema"]["$defs"]["identifier"]["type"], - "integer" + argument["content"]["json_schema"]["anyOf"] + .as_array() + .unwrap() + .len(), + 2 + ); + } + + #[test] + fn auto_optional_union_keeps_each_representable_xtml_type() { + let tools = vec![ToolDefinition { + name: "lookup".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": { + "value": {"type": ["string", "null"]} + } + })), + strict: None, + }]; + let choice = ToolChoice::Auto; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let alternatives = &value["format"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]["content"]["elements"]; + + assert_eq!(alternatives.as_array().unwrap().len(), 2); + assert!( + alternatives[0]["begin"] + .as_str() + .unwrap() + .contains("type=\"string\"") + ); + assert!( + alternatives[1]["begin"] + .as_str() + .unwrap() + .contains("type=\"null\"") ); } @@ -990,7 +1229,8 @@ mod tests { let value = serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) .unwrap(); - let arguments = &value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2]; + let arguments = &value["format"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]; assert_eq!(arguments["type"], "star"); assert_eq!(arguments["content"]["type"], "or"); @@ -999,10 +1239,8 @@ mod tests { 2 ); let alternatives = arguments["content"]["elements"].as_array().unwrap(); - assert_eq!( - alternatives[0]["content"]["pattern"], - format!("{STRING_ATOM}*") - ); + assert_eq!(alternatives[0]["content"]["type"], "any_text"); + assert_eq!(alternatives[0]["content"]["excludes"], json!([CLOSE])); assert_eq!(alternatives[1]["content"]["type"], "json_schema"); assert_eq!(alternatives[1]["content"]["json_schema"]["type"], "integer"); } @@ -1022,7 +1260,8 @@ mod tests { let value = serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) .unwrap(); - let arguments = &value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2]; + let arguments = &value["format"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]; assert_eq!(arguments["type"], "star"); assert_eq!(arguments["content"]["begin"], "<|open|>argument "); @@ -1037,8 +1276,8 @@ mod tests { }; let strict_value = serde_json::to_value(build_kimi_k3(&strict_ctx).unwrap().unwrap()).unwrap(); - let strict_arguments = - &strict_value["format"]["tags"][0]["content"]["tags"][0]["content"]["elements"][2]; + let strict_arguments = &strict_value["format"]["elements"][1]["content"]["content"]["tags"] + [0]["content"]["elements"][2]; assert_eq!(strict_arguments["type"], "sequence"); assert!(strict_arguments.to_string().contains("key=\\\"city\\\"")); @@ -1060,7 +1299,7 @@ mod tests { assert_eq!(value["format"]["type"], "sequence"); assert_eq!(value["format"]["elements"][0]["type"], "tag"); assert_eq!(value["format"]["elements"][0]["end"], THINK_CLOSE); - assert_eq!(value["format"]["elements"][1]["type"], "triggered_tags"); + assert_eq!(value["format"]["elements"][1]["type"], "sequence"); } #[test] diff --git a/parsers/v1/src/tool_calling/structural_tag/triggered_tags.rs b/parsers/v1/src/tool_calling/structural_tag/triggered_tags.rs index 423e0d162..0d259365b 100644 --- a/parsers/v1/src/tool_calling/structural_tag/triggered_tags.rs +++ b/parsers/v1/src/tool_calling/structural_tag/triggered_tags.rs @@ -79,7 +79,6 @@ pub(crate) fn build_triggered_tags( tags, at_least_one, stop_after_first: ctx.stop_after_first(), - excludes: vec![], }), })) } From 2cbdb547ca11cc65bb467c9f7c4e25ce9d5ad39a Mon Sep 17 00:00:00 2001 From: Indrajit Bhosale Date: Tue, 22 Sep 2026 15:29:40 -0700 Subject: [PATCH 06/10] Fix comments Signed-off-by: Indrajit Bhosale --- .../tool_calling/structural_tag/kimi_k3.rs | 61 ++++++++++++++----- 1 file changed, 47 insertions(+), 14 deletions(-) diff --git a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs index b40fad786..b3b88d2b8 100644 --- a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs +++ b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs @@ -287,17 +287,23 @@ fn xtml_types(schema: &Value, root: &Value) -> Vec<&'static str> { .collect() } -fn bounded_string_regex(schema: &Map) -> Option { - let max_len = schema.get("maxLength")?.as_u64()?; - if max_len > 4096 { +fn string_length_regex(schema: &Map) -> Option { + let min_len = match schema.get("minLength") { + Some(value) => value.as_u64()?, + None => 0, + }; + let max_len = match schema.get("maxLength") { + Some(value) => Some(value.as_u64()?), + None => None, + }; + if min_len > 4096 || max_len.is_some_and(|max_len| max_len > 4096 || min_len > max_len) { return None; } - let min_len = schema - .get("minLength") - .and_then(Value::as_u64) - .filter(|min| *min <= max_len) - .unwrap_or(0); - Some(format!("{STRING_ATOM}{{{min_len},{max_len}}}")) + match max_len { + Some(max_len) => Some(format!("{STRING_ATOM}{{{min_len},{max_len}}}")), + None if schema.contains_key("minLength") => Some(format!("{STRING_ATOM}{{{min_len},}}")), + None => None, + } } fn string_content_format(schema: &Value) -> Format { @@ -337,10 +343,11 @@ fn string_content_format(schema: &Value) -> Format { ); } - if let Some(pattern) = schema.get("pattern").and_then(Value::as_str) - && !schema.contains_key("minLength") - && !schema.contains_key("maxLength") - { + if let Some(pattern) = schema.get("pattern").and_then(Value::as_str) { + // XGrammar 0.2.3 regexes do not support lookahead, so a general + // intersection of `pattern` and length bounds cannot be expressed. + // Match XGrammar's JSON-schema policy: `pattern` takes precedence when + // both are present instead of silently discarding the pattern. let anchored_start = pattern.starts_with('^'); let anchored_end = pattern.ends_with('$') && !pattern.ends_with("\\$"); let pattern = pattern @@ -363,7 +370,7 @@ fn string_content_format(schema: &Value) -> Format { }); } - if let Some(pattern) = bounded_string_regex(schema) { + if let Some(pattern) = string_length_regex(schema) { return Format::Regex(RegexFormat { pattern }); } @@ -1108,6 +1115,32 @@ mod tests { assert!(argument["content"]["json_schema"].get("$ref").is_none()); } + #[test] + fn string_content_enforces_min_length_without_a_maximum() { + let value = serde_json::to_value(string_content_format(&json!({ + "type": "string", + "minLength": 2 + }))) + .unwrap(); + + assert_eq!(value["type"], "regex"); + assert_eq!(value["pattern"], format!("{STRING_ATOM}{{2,}}")); + } + + #[test] + fn string_content_uses_xgrammar_pattern_precedence_over_length_bounds() { + let value = serde_json::to_value(string_content_format(&json!({ + "type": "string", + "pattern": "^item-[0-9]+$", + "minLength": 6, + "maxLength": 12 + }))) + .unwrap(); + + assert_eq!(value["type"], "regex"); + assert_eq!(value["pattern"], "(?:item-[0-9]+)"); + } + #[test] fn auto_string_arguments_preserve_schema_constraints_and_allow_empty_values() { let tools = vec![ToolDefinition { From 829d5d02880360218b4ffcc929d44cdda39257da Mon Sep 17 00:00:00 2001 From: Indrajit Bhosale Date: Wed, 23 Sep 2026 16:09:37 -0700 Subject: [PATCH 07/10] Fix comments and Refactore some parts of K3 Signed-off-by: Indrajit Bhosale --- .../tool_calling/structural_tag/kimi_k3.rs | 1001 ++++++++++++++--- 1 file changed, 826 insertions(+), 175 deletions(-) diff --git a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs index b3b88d2b8..1258b6a72 100644 --- a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs +++ b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs @@ -36,6 +36,7 @@ const ARGUMENT_CLOSE: &str = "<|close|>argument<|sep|>"; const MESSAGE_CLOSE: &str = "<|close|>message<|sep|>"; const STRING_ATOM: &str = r"(?:[^<]|<[^|])"; +const CALL_INDEX_PATTERN: &str = "[1-9][0-9]*"; fn escape_attr(value: &str) -> String { value.replace('&', "&").replace('"', """) @@ -82,6 +83,14 @@ const JSON_TYPES: [JsonType; 7] = [ JsonType::Null, ]; +fn is_integral_number(number: &serde_json::Number) -> bool { + number.is_i64() + || number.is_u64() + || number + .as_f64() + .is_some_and(|value| value.is_finite() && value.fract() == 0.0) +} + impl JsonType { fn from_name(value: &str) -> Option { match value { @@ -100,7 +109,7 @@ impl JsonType { match value { Value::Null => Self::Null, Value::Bool(_) => Self::Boolean, - Value::Number(number) if number.is_i64() || number.is_u64() => Self::Integer, + Value::Number(number) if is_integral_number(number) => Self::Integer, Value::Number(_) => Self::Number, Value::String(_) => Self::String, Value::Array(_) => Self::Array, @@ -118,9 +127,36 @@ impl JsonType { Self::Null => "null", } } + + fn json_name(self) -> &'static str { + match self { + Self::String => "string", + Self::Number => "number", + Self::Integer => "integer", + Self::Boolean => "boolean", + Self::Array => "array", + Self::Object => "object", + Self::Null => "null", + } + } + + fn accepts(self, value: &Value) -> bool { + match self { + Self::String => value.is_string(), + Self::Number => value.is_number(), + Self::Integer => value.as_number().is_some_and(is_integral_number), + Self::Boolean => value.is_boolean(), + Self::Array => value.is_array(), + Self::Object => value.is_object(), + Self::Null => value.is_null(), + } + } } fn resolve_local_ref<'a>(reference: &str, root: &'a Value) -> Option<&'a Value> { + if reference == "#" { + return matches!(root, Value::Bool(_) | Value::Object(_)).then_some(root); + } let mut value = root; for raw_part in reference.strip_prefix("#/")?.split('/') { let part = raw_part.replace("~1", "/").replace("~0", "~"); @@ -269,15 +305,6 @@ fn schema_types(schema: &Value, root: &Value, seen_refs: &HashSet) -> Ve } } -fn single_xtml_type(schema: &Value, root: &Value) -> Option<&'static str> { - let types = schema_types(schema, root, &HashSet::new()); - let xtml_types = types - .into_iter() - .map(JsonType::xtml_name) - .collect::>(); - (xtml_types.len() == 1).then(|| *xtml_types.iter().next().expect("one XTML type")) -} - fn xtml_types(schema: &Value, root: &Value) -> Vec<&'static str> { let mut seen = HashSet::new(); schema_types(schema, root, &HashSet::new()) @@ -306,12 +333,55 @@ fn string_length_regex(schema: &Map) -> Option { } } +fn has_unescaped_trailing_dollar(pattern: &str) -> bool { + let Some(prefix) = pattern.strip_suffix('$') else { + return false; + }; + prefix + .as_bytes() + .iter() + .rev() + .take_while(|byte| **byte == b'\\') + .count() + % 2 + == 0 +} + +fn merge_compatible_all_of(schema: &Map) -> Option> { + let options = schema.get("allOf")?.as_array()?; + let mut merged = schema.clone(); + merged.remove("allOf"); + + for option in options { + let option = match option { + Value::Bool(true) => continue, + Value::Object(option) => { + merge_compatible_all_of(option).unwrap_or_else(|| option.clone()) + } + _ => return None, + }; + for (key, value) in option { + if merged.get(&key).is_some_and(|existing| existing != &value) { + return None; + } + merged.insert(key, value); + } + } + Some(merged) +} + fn string_content_format(schema: &Value) -> Format { + let root = schema; let Some(schema) = schema.as_object() else { return Format::AnyText(AnyTextFormat { excludes: vec![CLOSE.to_string()], }); }; + // XTML string arguments contain raw text rather than a JSON string, so they + // cannot use JsonSchemaFormat. Flatten compatible allOf branches before + // translating their string constraints to raw-text formats. + let merged_schema = merge_compatible_all_of(schema); + let schema = merged_schema.as_ref().unwrap_or(schema); let enum_values = schema .get("enum") @@ -323,24 +393,42 @@ fn string_content_format(schema: &Value) -> Format { .and_then(Value::as_str) .map(|value| vec![Value::String(value.to_string())]) }); - if let Some(values) = enum_values.filter(|values| { - !values.is_empty() + if let Some(values) = enum_values { + let pattern = schema.get("pattern").and_then(Value::as_str); + let compiled_pattern = pattern.map(regex::Regex::new); + if !values.is_empty() && values.len() <= 256 - && values - .iter() - .all(|value| value.as_str().is_some_and(|string| !string.contains("<|"))) - }) { - return one_of( - values + && values.iter().all(|value| value.as_str().is_some()) + && compiled_pattern.as_ref().is_none_or(Result::is_ok) + { + let min_len = schema.get("minLength").and_then(Value::as_u64); + let max_len = schema.get("maxLength").and_then(Value::as_u64); + let const_value = schema.get("const").and_then(Value::as_str); + let values = values .iter() .filter_map(Value::as_str) - .map(|value| { - Format::ConstString(ConstStringFormat { - value: value.to_string(), - }) + .filter(|value| !value.contains("<|")) + .filter(|value| { + let len = value.chars().count() as u64; + min_len.is_none_or(|min| len >= min) + && max_len.is_none_or(|max| len <= max) + && const_value.is_none_or(|constant| *value == constant) + && compiled_pattern.as_ref().is_none_or(|pattern| { + pattern.as_ref().is_ok_and(|re| re.is_match(value)) + }) }) - .collect(), - ); + .collect::>(); + return one_of( + values + .iter() + .map(|value| { + Format::ConstString(ConstStringFormat { + value: (*value).to_string(), + }) + }) + .collect(), + ); + } } if let Some(pattern) = schema.get("pattern").and_then(Value::as_str) { @@ -349,12 +437,15 @@ fn string_content_format(schema: &Value) -> Format { // Match XGrammar's JSON-schema policy: `pattern` takes precedence when // both are present instead of silently discarding the pattern. let anchored_start = pattern.starts_with('^'); - let anchored_end = pattern.ends_with('$') && !pattern.ends_with("\\$"); - let pattern = pattern - .strip_prefix('^') - .unwrap_or(pattern) - .strip_suffix('$') - .unwrap_or_else(|| pattern.strip_prefix('^').unwrap_or(pattern)); + let anchored_end = has_unescaped_trailing_dollar(pattern); + let pattern = pattern.strip_prefix('^').unwrap_or(pattern); + let pattern = if anchored_end { + pattern + .strip_suffix('$') + .expect("anchored pattern ends with a dollar") + } else { + pattern + }; let prefix = if anchored_start { String::new() } else { @@ -374,52 +465,35 @@ fn string_content_format(schema: &Value) -> Format { return Format::Regex(RegexFormat { pattern }); } + for keyword in ["anyOf", "oneOf"] { + let Some(options) = schema.get(keyword).and_then(Value::as_array) else { + continue; + }; + let formats = options + .iter() + .filter(|option| { + schema_types(option, root, &HashSet::new()).contains(&JsonType::String) + }) + .map(string_content_format) + .collect::>(); + if !formats.is_empty() { + return one_of(formats); + } + } + Format::AnyText(AnyTextFormat { excludes: vec![CLOSE.to_string()], }) } -fn argument_tag( - key: &str, - schema: &Value, - root_defs: Option<&Map>, -) -> Option { - let schema_object = schema.as_object()?; - let (json_type, xtml_type) = match schema_object.get("type").and_then(Value::as_str)? { - "string" => ("string", "string"), - "integer" => ("integer", "number"), - "number" => ("number", "number"), - "boolean" => ("boolean", "boolean"), - "null" => ("null", "null"), - "object" => ("object", "object"), - "array" => ("array", "array"), - _ => return None, - }; - let begin = format!( - "{OPEN}argument key=\"{}\" type=\"{xtml_type}\"{SEP}", - escape_attr(key) - ); - - let content = if json_type == "string" { - string_content_format(schema) +fn resolved_root_parameters(root: &Value) -> Option { + let resolved = resolve_local_refs(root, root, &HashSet::new())?; + let object = resolved.as_object()?; + if object.contains_key("allOf") { + merge_compatible_all_of(object).map(Value::Object) } else { - let mut embedded = schema_object.clone(); - if let Some(root_defs) = root_defs { - for (key, value) in root_defs { - embedded.entry(key.clone()).or_insert_with(|| value.clone()); - } - } - Format::JsonSchema(JsonSchemaFormat { - json_schema: Value::Object(embedded), - style: JsonSchemaStyle::Json, - }) - }; - - Some(TagFormat { - begin, - content: Box::new(content), - end: ARGUMENT_CLOSE.to_string(), - }) + Some(resolved) + } } fn permissive_argument_tag() -> TagFormat { @@ -439,16 +513,199 @@ fn permissive_argument_tag() -> TagFormat { } } -fn root_definitions(parameters: &Map) -> Map { - ["$defs", "definitions"] - .into_iter() - .filter_map(|key| { - parameters - .get(key) - .and_then(Value::as_object) - .map(|value| (key.to_string(), Value::Object(value.clone()))) +fn resolve_schema_map(value: &Value, root: &Value, seen_refs: &HashSet) -> Option { + let Some(values) = value.as_object() else { + return Some(value.clone()); + }; + values + .iter() + .map(|(key, value)| { + resolve_local_refs(value, root, seen_refs).map(|value| (key.clone(), value)) }) - .collect() + .collect::>>() + .map(Value::Object) +} + +// JSON Schema documents mix subschemas with arbitrary JSON data. These JSON Schema +// 2020-12 keywords, plus draft-07 compatibility forms, identify the positions where +// local `$ref` values have schema semantics and are therefore safe to resolve. + +// Each property of these keyword values is a subschema. +const SCHEMA_MAP_KEYWORDS: &[&str] = &[ + "$defs", + "definitions", + "dependentSchemas", + "dependencies", + "patternProperties", + "properties", +]; + +// The keyword value itself is a subschema. +const SCHEMA_VALUE_KEYWORDS: &[&str] = &[ + "additionalItems", + "additionalProperties", + "contains", + "contentSchema", + "else", + "if", + "items", + "not", + "propertyNames", + "then", + "unevaluatedItems", + "unevaluatedProperties", +]; + +// Each array element in these keyword values is a subschema. +const SCHEMA_ARRAY_KEYWORDS: &[&str] = &["allOf", "anyOf", "oneOf", "prefixItems"]; + +fn local_ref_with_prefix(reference: &str, prefix: &str) -> Option { + if reference == "#" { + Some(prefix.to_string()) + } else { + reference + .strip_prefix("#/") + .map(|suffix| format!("{prefix}/{suffix}")) + } +} + +fn rewrite_local_schema_refs(schema: &mut Value, prefix: &str) { + let Some(object) = schema.as_object_mut() else { + return; + }; + + if let Some(reference) = object.get_mut("$ref") + && let Some(rewritten) = reference + .as_str() + .and_then(|reference| local_ref_with_prefix(reference, prefix)) + { + *reference = Value::String(rewritten); + } + for key in SCHEMA_MAP_KEYWORDS { + if let Some(values) = object.get_mut(*key).and_then(Value::as_object_mut) { + for value in values.values_mut() { + rewrite_local_schema_refs(value, prefix); + } + } + } + for key in SCHEMA_VALUE_KEYWORDS { + if let Some(value) = object.get_mut(*key) { + rewrite_local_schema_refs(value, prefix); + } + } + for key in SCHEMA_ARRAY_KEYWORDS { + if let Some(values) = object.get_mut(*key).and_then(Value::as_array_mut) { + for value in values { + rewrite_local_schema_refs(value, prefix); + } + } + } +} + +fn has_unresolved_local_refs(schema: &Value, root: &Value) -> bool { + let Some(object) = schema.as_object() else { + return false; + }; + + if object + .get("$ref") + .and_then(Value::as_str) + .is_some_and(|reference| { + (reference == "#" || reference.starts_with("#/")) + && resolve_local_ref(reference, root).is_none() + }) + { + return true; + } + SCHEMA_MAP_KEYWORDS.iter().any(|key| { + object + .get(*key) + .and_then(Value::as_object) + .is_some_and(|values| { + values + .values() + .any(|value| has_unresolved_local_refs(value, root)) + }) + }) || SCHEMA_VALUE_KEYWORDS.iter().any(|key| { + object + .get(*key) + .is_some_and(|value| has_unresolved_local_refs(value, root)) + }) || SCHEMA_ARRAY_KEYWORDS.iter().any(|key| { + object + .get(*key) + .and_then(Value::as_array) + .is_some_and(|values| { + values + .iter() + .any(|value| has_unresolved_local_refs(value, root)) + }) + }) +} + +fn preserve_recursive_root_refs(mut schema: Value, root: &Value) -> Value { + if !has_unresolved_local_refs(&schema, &schema) { + return schema; + } + + let Some(schema_object) = schema.as_object_mut() else { + return schema; + }; + let definitions = schema_object + .entry("$defs") + .or_insert_with(|| Value::Object(Map::new())); + let Some(definitions) = definitions.as_object_mut() else { + return schema; + }; + let mut definition_name = "__dynamo_root".to_string(); + let mut suffix = 2; + while definitions.contains_key(&definition_name) { + definition_name = format!("__dynamo_root_{suffix}"); + suffix += 1; + } + let escaped_name = definition_name.replace('~', "~0").replace('/', "~1"); + let prefix = format!("#/$defs/{escaped_name}"); + + // A cycle can leave a reference to a location outside the extracted + // argument schema (for example, `#/properties/node`). Preserve the original + // root under a private definition and rebase only schema-position refs to it. + rewrite_local_schema_refs(&mut schema, &prefix); + let mut embedded_root = root.clone(); + rewrite_local_schema_refs(&mut embedded_root, &prefix); + schema + .as_object_mut() + .and_then(|schema| schema.get_mut("$defs")) + .and_then(Value::as_object_mut) + .expect("definitions were initialized above") + .insert(definition_name, embedded_root); + schema +} + +fn resolve_schema_keyword( + key: &str, + value: &Value, + root: &Value, + seen_refs: &HashSet, +) -> Option { + if SCHEMA_MAP_KEYWORDS.contains(&key) { + return resolve_schema_map(value, root, seen_refs); + } + if SCHEMA_VALUE_KEYWORDS.contains(&key) { + return resolve_local_refs(value, root, seen_refs); + } + if SCHEMA_ARRAY_KEYWORDS.contains(&key) { + let Some(values) = value.as_array() else { + return Some(value.clone()); + }; + return values + .iter() + .map(|value| resolve_local_refs(value, root, seen_refs)) + .collect::>>() + .map(Value::Array); + } + + // Unknown keywords and values under `const`, `enum`, `default`, and `examples` + // may contain arbitrary JSON data, where a `$ref` object is not a schema reference. + Some(value.clone()) } fn resolve_local_refs(schema: &Value, root: &Value, seen_refs: &HashSet) -> Option { @@ -461,7 +718,7 @@ fn resolve_local_refs(schema: &Value, root: &Value, seen_refs: &HashSet) Value::Object(object) => { if let Some(reference) = object.get("$ref").and_then(Value::as_str) { if seen_refs.contains(reference) { - return None; + return Some(schema.clone()); } let target = resolve_local_ref(reference, root)?; let mut nested_seen = seen_refs.clone(); @@ -471,7 +728,8 @@ fn resolve_local_refs(schema: &Value, root: &Value, seen_refs: &HashSet) .iter() .filter(|(key, _)| key.as_str() != "$ref") .map(|(key, value)| { - resolve_local_refs(value, root, seen_refs).map(|value| (key.clone(), value)) + resolve_schema_keyword(key, value, root, seen_refs) + .map(|value| (key.clone(), value)) }) .collect::>>()?; if siblings.is_empty() { @@ -485,7 +743,8 @@ fn resolve_local_refs(schema: &Value, root: &Value, seen_refs: &HashSet) object .iter() .map(|(key, value)| { - resolve_local_refs(value, root, seen_refs).map(|value| (key.clone(), value)) + resolve_schema_keyword(key, value, root, seen_refs) + .map(|value| (key.clone(), value)) }) .collect::>>() .map(Value::Object) @@ -494,61 +753,108 @@ fn resolve_local_refs(schema: &Value, root: &Value, seen_refs: &HashSet) } } -fn schema_for_xtml_type(schema: &Value, root: &Value, xtml_type: &str) -> Option { - let resolved = resolve_local_refs(schema, root, &HashSet::new())?; - if single_xtml_type(&resolved, &resolved) == Some(xtml_type) { - return Some(resolved); - } - - let json_types = schema_types(schema, root, &HashSet::new()) +fn intersect_json_types(candidates: &[JsonType], requested: &[JsonType]) -> Vec { + let candidate_has_number = candidates.contains(&JsonType::Number); + let requested_has_number = requested.contains(&JsonType::Number); + let accepts_integer = (candidate_has_number || candidates.contains(&JsonType::Integer)) + && (requested_has_number || requested.contains(&JsonType::Integer)); + JSON_TYPES .into_iter() - .filter(|json_type| json_type.xtml_name() == xtml_type) - .map(|json_type| match json_type { - JsonType::String => "string", - JsonType::Number => "number", - JsonType::Integer => "integer", - JsonType::Boolean => "boolean", - JsonType::Array => "array", - JsonType::Object => "object", - JsonType::Null => "null", + .filter(|json_type| match json_type { + JsonType::Number => candidate_has_number && requested_has_number, + JsonType::Integer => accepts_integer && !(candidate_has_number && requested_has_number), + _ => candidates.contains(json_type) && requested.contains(json_type), }) - .collect::>(); - match json_types.as_slice() { - [] => None, - [json_type] => Some(serde_json::json!({"type": json_type})), - _ => Some(serde_json::json!({"type": json_types})), + .collect() +} + +fn json_type_constraint(json_types: &[JsonType]) -> Value { + match json_types { + [json_type] => Value::String(json_type.json_name().to_string()), + _ => Value::Array( + json_types + .iter() + .map(|json_type| Value::String(json_type.json_name().to_string())) + .collect(), + ), } } -fn arguments_block(parameters: Option<&Value>) -> Format { - let Some(parameters) = parameters.and_then(Value::as_object) else { - return star(Format::Tag(permissive_argument_tag())); - }; - let Some(properties) = parameters.get("properties").and_then(Value::as_object) else { - return star(Format::Tag(permissive_argument_tag())); +fn narrow_resolved_schema( + schema: Value, + root: &Value, + requested_types: &[JsonType], +) -> Option> { + let candidates = schema_types(&schema, root, &HashSet::new()); + let json_types = intersect_json_types(&candidates, requested_types); + if json_types.is_empty() { + return None; + } + + let mut narrowed = match schema { + Value::Bool(true) => Map::new(), + Value::Object(object) => object, + _ => return None, }; - if properties.is_empty() { - return star(Format::Tag(permissive_argument_tag())); + for keyword in ["anyOf", "oneOf"] { + let Some(options) = narrowed.get(keyword).and_then(Value::as_array) else { + continue; + }; + let options = options + .iter() + .filter_map(|option| { + narrow_resolved_schema(option.clone(), root, &json_types).map(Value::Object) + }) + .collect::>(); + if options.is_empty() { + return None; + } + narrowed.insert(keyword.to_string(), Value::Array(options)); + } + if let Some(values) = narrowed.get_mut("enum").and_then(Value::as_array_mut) { + values.retain(|value| json_types.iter().any(|json_type| json_type.accepts(value))); + if values.is_empty() { + return None; + } } + if narrowed + .get("const") + .is_some_and(|value| !json_types.iter().any(|json_type| json_type.accepts(value))) + { + return None; + } + narrowed.insert("type".to_string(), json_type_constraint(&json_types)); + Some(narrowed) +} - let root_defs = root_definitions(parameters); - let tags = properties - .iter() - .map(|(key, schema)| { - Format::Tag( - argument_tag(key, schema, Some(&root_defs)).unwrap_or_else(permissive_argument_tag), - ) - }) - .collect(); - star(one_of(tags)) +fn schema_for_xtml_type(schema: &Value, root: &Value, xtml_type: &str) -> Option { + let resolved = resolve_local_refs(schema, root, &HashSet::new())?; + let mut json_types = schema_types(&resolved, &resolved, &HashSet::new()) + .into_iter() + .filter(|json_type| json_type.xtml_name() == xtml_type) + .collect::>(); + if json_types.contains(&JsonType::Number) { + json_types.retain(|json_type| *json_type != JsonType::Integer); + } + let mut narrowed = narrow_resolved_schema(resolved, root, &json_types)?; + if let Some(root) = root.as_object() { + for key in ["$defs", "definitions"] { + if let Some(value) = root.get(key) { + narrowed + .entry(key.to_string()) + .or_insert_with(|| value.clone()); + } + } + } + Some(preserve_recursive_root_refs(Value::Object(narrowed), root)) } -fn auto_argument_format(key: &str, schema: &Value, root: &Value) -> Option { +fn argument_format(key: &str, schema: &Value, root: &Value) -> Option { let alternatives = xtml_types(schema, root) .into_iter() .filter_map(|xtml_type| { let content = if xtml_type == "string" { - string_content_format(schema) + string_content_format(&schema_for_xtml_type(schema, root, xtml_type)?) } else { Format::JsonSchema(JsonSchemaFormat { json_schema: schema_for_xtml_type(schema, root, xtml_type)?, @@ -569,7 +875,7 @@ fn auto_argument_format(key: &str, schema: &Value, root: &Value) -> Option, required_keys: &[&str], root: &Value, @@ -579,7 +885,7 @@ fn optional_auto_arguments( .filter(|(key, _)| !required_keys.contains(&key.as_str())) .filter_map(|(key, schema)| { matches!(schema, Value::Bool(_) | Value::Object(_)) - .then(|| auto_argument_format(key, schema, root)) + .then(|| argument_format(key, schema, root)) .flatten() }) .collect::>(); @@ -604,19 +910,23 @@ fn canonical_required_arguments( }) } -fn auto_arguments_block(tool: &ToolDefinition, strict_schema: bool) -> Format { +fn arguments_block(tool: &ToolDefinition, strict_schema: bool) -> Format { if !super::builder::uses_declared_tool_schema(tool, strict_schema) { - return arguments_block(None); + return star(Format::Tag(permissive_argument_tag())); } let Some(root) = tool.parameters.as_ref() else { - return arguments_block(None); + return star(Format::Tag(permissive_argument_tag())); }; - let Some(parameters) = root.as_object() else { - return arguments_block(None); + let resolved_root = root + .as_object() + .filter(|object| object.contains_key("$ref") || object.contains_key("allOf")) + .and_then(|_| resolved_root_parameters(root)); + let Some(parameters) = resolved_root.as_ref().unwrap_or(root).as_object() else { + return star(Format::Tag(permissive_argument_tag())); }; let Some(properties) = parameters.get("properties").and_then(Value::as_object) else { - return arguments_block(None); + return star(Format::Tag(permissive_argument_tag())); }; let required = match parameters.get("required") { None => Vec::new(), @@ -626,9 +936,9 @@ fn auto_arguments_block(tool: &ToolDefinition, strict_schema: bool) -> Format { .filter_map(Value::as_str) .collect::>() } - Some(_) => return arguments_block(None), + Some(_) => return star(Format::Tag(permissive_argument_tag())), }; - let optional_arguments = optional_auto_arguments(properties, &required, root); + let optional_arguments = optional_arguments(properties, &required, root); if required.is_empty() { return optional_arguments.unwrap_or_else(|| { @@ -641,37 +951,19 @@ fn auto_arguments_block(tool: &ToolDefinition, strict_schema: bool) -> Format { let mut arguments = Vec::with_capacity(required.len()); for key in &required { let Some(schema) = properties.get(*key) else { - return arguments_block(None); + return star(Format::Tag(permissive_argument_tag())); }; if !matches!(schema, Value::Bool(_) | Value::Object(_)) { - return arguments_block(None); + return star(Format::Tag(permissive_argument_tag())); } - let Some(argument) = auto_argument_format(key, schema, root) else { - return arguments_block(None); + let Some(argument) = argument_format(key, schema, root) else { + return star(Format::Tag(permissive_argument_tag())); }; arguments.push(argument); } canonical_required_arguments(arguments, optional_arguments) } -fn auto_call_tag(tool: &ToolDefinition, strict_schema: bool) -> TagFormat { - TagFormat { - begin: format!("{OPEN}call tool=\"{}\" index=\"", escape_attr(&tool.name)), - content: Box::new(Format::Sequence(SequenceFormat { - elements: vec![ - Format::Regex(RegexFormat { - pattern: "[1-9][0-9]*".to_string(), - }), - Format::ConstString(ConstStringFormat { - value: format!("\"{SEP}"), - }), - auto_arguments_block(tool, strict_schema), - ], - })), - end: CALL_CLOSE.to_string(), - } -} - fn build_auto_structural_tag( tools: Vec<&ToolDefinition>, ctx: &ToolCallFormatBuildContext<'_>, @@ -679,7 +971,7 @@ fn build_auto_structural_tag( let strict_schema = ctx.strict_schema(); let call_tags: Vec<_> = tools .into_iter() - .map(|tool| auto_call_tag(tool, strict_schema)) + .map(|tool| call_tag(tool, strict_schema)) .collect(); let parallel_tool_calls = !ctx.stop_after_first(); let calls = if parallel_tool_calls { @@ -731,25 +1023,18 @@ fn build_auto_structural_tag( } fn call_tag(tool: &ToolDefinition, strict_schema: bool) -> TagFormat { - // Match vLLM's K3 behavior: use the declared schema unless the caller - // explicitly sets strict=false. Global strict mode overrides that opt-out. - let parameters = if super::builder::uses_declared_tool_schema(tool, strict_schema) { - tool.parameters.as_ref() - } else { - None - }; let begin = format!("{OPEN}call tool=\"{}\" index=\"", escape_attr(&tool.name)); TagFormat { begin, content: Box::new(Format::Sequence(SequenceFormat { elements: vec![ Format::Regex(RegexFormat { - pattern: "[0-9]+".to_string(), + pattern: CALL_INDEX_PATTERN.to_string(), }), Format::ConstString(ConstStringFormat { value: format!("\"{SEP}"), }), - arguments_block(parameters), + arguments_block(tool, strict_schema), ], })), end: CALL_CLOSE.to_string(), @@ -1141,6 +1426,21 @@ mod tests { assert_eq!(value["pattern"], "(?:item-[0-9]+)"); } + #[test] + fn string_content_preserves_an_escaped_trailing_dollar() { + assert!(!has_unescaped_trailing_dollar(r"^price\$")); + assert!(has_unescaped_trailing_dollar(r"^path\\$")); + + let value = serde_json::to_value(string_content_format(&json!({ + "type": "string", + "pattern": r"^price\$" + }))) + .unwrap(); + + assert_eq!(value["type"], "regex"); + assert_eq!(value["pattern"], format!(r"(?:price\$){STRING_ATOM}*")); + } + #[test] fn auto_string_arguments_preserve_schema_constraints_and_allow_empty_values() { let tools = vec![ToolDefinition { @@ -1245,6 +1545,355 @@ mod tests { ); } + #[test] + fn auto_recursive_optional_object_keeps_the_argument_and_definitions() { + let tools = vec![ToolDefinition { + name: "walk".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": { + "child": {"$ref": "#/$defs/node"} + }, + "$defs": { + "node": { + "type": "object", + "properties": { + "next": {"$ref": "#/$defs/node"} + } + } + } + })), + strict: None, + }]; + let choice = ToolChoice::Auto; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let arguments = &value["format"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]; + + assert_eq!(arguments["type"], "star"); + assert_eq!( + arguments["content"]["begin"], + "<|open|>argument key=\"child\" type=\"object\"<|sep|>" + ); + let schema = &arguments["content"]["content"]["json_schema"]; + assert_eq!(schema["type"], "object"); + assert_eq!(schema["properties"]["next"]["$ref"], "#/$defs/node"); + assert_eq!(schema["$defs"]["node"]["type"], "object"); + } + + #[test] + fn recursive_property_reference_remains_resolvable_after_extraction() { + let root = json!({ + "type": "object", + "properties": { + "node": { + "type": "object", + "properties": { + "next": {"$ref": "#/properties/node"}, + "metadata": { + "type": "object", + "const": {"$ref": "literal"} + } + } + } + } + }); + let schema = schema_for_xtml_type(&root["properties"]["node"], &root, "object").unwrap(); + let reference = schema["properties"]["next"]["properties"]["next"]["$ref"] + .as_str() + .unwrap(); + + assert!(resolve_local_ref(reference, &schema).is_some()); + assert!(!has_unresolved_local_refs(&schema, &schema)); + assert!(reference.starts_with("#/$defs/__dynamo_root")); + assert_eq!( + schema["properties"]["metadata"]["const"], + json!({"$ref": "literal"}) + ); + } + + #[test] + fn root_reference_resolves_to_the_original_schema() { + let root = json!({"type": "object"}); + + assert_eq!(resolve_local_ref("#", &root), Some(&root)); + } + + #[test] + fn auto_union_narrowing_preserves_value_constraints() { + let tools = vec![ToolDefinition { + name: "choose".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": { + "value": { + "type": ["integer", "null"], + "enum": [1, null], + "minimum": 1 + } + }, + "required": ["value"] + })), + strict: None, + }]; + let choice = ToolChoice::Auto; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let alternatives = &value["format"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]["elements"][0]["elements"]; + let number = alternatives + .as_array() + .unwrap() + .iter() + .find(|alternative| { + alternative["begin"] + .as_str() + .is_some_and(|begin| begin.contains("type=\"number\"")) + }) + .expect("number alternative"); + let schema = &number["content"]["json_schema"]; + + assert_eq!(schema["type"], "integer"); + assert_eq!(schema["enum"], json!([1])); + assert_eq!(schema["minimum"], 1); + let null = alternatives + .as_array() + .unwrap() + .iter() + .find(|alternative| { + alternative["begin"] + .as_str() + .is_some_and(|begin| begin.contains("type=\"null\"")) + }) + .expect("null alternative"); + assert_eq!(null["content"]["json_schema"]["enum"], json!([null])); + } + + #[test] + fn integer_union_keeps_integral_float_enum_values() { + let schema = json!({ + "type": ["integer", "null"], + "enum": [1.0, null] + }); + let narrowed = schema_for_xtml_type(&schema, &schema, "number").unwrap(); + + assert_eq!(narrowed["type"], "integer"); + assert_eq!(narrowed["enum"], json!([1.0])); + } + + #[test] + fn union_composition_removes_other_xtml_types() { + let schema = json!({ + "anyOf": [ + {"type": "integer", "enum": [1]}, + {"type": "null"} + ] + }); + let narrowed = schema_for_xtml_type(&schema, &schema, "number").unwrap(); + let options = narrowed["anyOf"].as_array().unwrap(); + + assert_eq!(narrowed["type"], "integer"); + assert_eq!(options.len(), 1); + assert_eq!(options[0]["type"], "integer"); + assert_eq!(options[0]["enum"], json!([1])); + } + + #[test] + fn auto_string_type_union_preserves_applicable_enum_values() { + let tools = vec![ToolDefinition { + name: "choose".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": { + "value": { + "type": ["string", "null"], + "enum": ["safe", null] + } + }, + "required": ["value"] + })), + strict: None, + }]; + let choice = ToolChoice::Auto; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let alternatives = &value["format"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]["elements"][0]["elements"]; + let string = alternatives + .as_array() + .unwrap() + .iter() + .find(|alternative| { + alternative["begin"] + .as_str() + .is_some_and(|begin| begin.contains("type=\"string\"")) + }) + .expect("string alternative"); + + assert_eq!(string["content"]["type"], "const_string"); + assert_eq!(string["content"]["value"], "safe"); + } + + #[test] + fn auto_string_references_and_unions_preserve_enum_constraints() { + let tools = vec![ToolDefinition { + name: "choose".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": { + "direct": {"$ref": "#/$defs/safe"}, + "union": { + "anyOf": [ + {"type": "string", "enum": ["safe"]}, + {"type": "null"} + ] + } + }, + "required": ["direct", "union"], + "$defs": { + "safe": {"type": "string", "enum": ["safe"]} + } + })), + strict: None, + }]; + let choice = ToolChoice::Auto; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let arguments = &value["format"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]["elements"]; + + assert_eq!(arguments[0]["content"]["type"], "const_string"); + assert_eq!(arguments[0]["content"]["value"], "safe"); + let union_alternatives = arguments[1]["elements"].as_array().unwrap(); + let string = union_alternatives + .iter() + .find(|alternative| { + alternative["begin"] + .as_str() + .is_some_and(|begin| begin.contains("type=\"string\"")) + }) + .expect("string alternative"); + assert_eq!(string["content"]["type"], "const_string"); + assert_eq!(string["content"]["value"], "safe"); + } + + #[test] + fn string_all_of_preserves_compatible_enum_constraints() { + let schema = json!({ + "allOf": [ + {"type": "string"}, + {"enum": ["safe"]} + ] + }); + let narrowed = schema_for_xtml_type(&schema, &schema, "string").unwrap(); + let format = serde_json::to_value(string_content_format(&narrowed)).unwrap(); + + assert_eq!(format["type"], "const_string"); + assert_eq!(format["value"], "safe"); + } + + #[test] + fn string_enum_intersects_length_and_pattern_constraints() { + let schema = json!({ + "allOf": [ + {"type": "string", "enum": ["x", "code-42", "other"]}, + {"minLength": 3, "pattern": "^code-[0-9]+$"} + ] + }); + let narrowed = schema_for_xtml_type(&schema, &schema, "string").unwrap(); + let format = serde_json::to_value(string_content_format(&narrowed)).unwrap(); + + assert_eq!(format["type"], "const_string"); + assert_eq!(format["value"], "code-42"); + + let tools = vec![ToolDefinition { + name: "select".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": {"value": schema}, + "required": ["value"] + })), + strict: None, + }]; + let choice = ToolChoice::Auto; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let argument = &value["format"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]["elements"][0]; + assert_eq!(argument["content"]["type"], "const_string"); + assert_eq!(argument["content"]["value"], "code-42"); + } + + #[test] + fn root_referenced_parameters_keep_required_arguments() { + let tools = vec![ToolDefinition { + name: "weather".to_string(), + parameters: Some(json!({ + "$ref": "#/$defs/Args", + "$defs": { + "Args": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + "additionalProperties": false + } + } + })), + strict: None, + }]; + let choice = ToolChoice::Auto; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let arguments = &value["format"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]; + + assert_eq!(arguments["type"], "sequence"); + assert_eq!( + arguments["elements"][0]["begin"], + "<|open|>argument key=\"city\" type=\"string\"<|sep|>" + ); + } + + #[test] + fn auto_object_const_treats_ref_as_literal_data() { + let tools = vec![ToolDefinition { + name: "literal".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": { + "value": { + "type": "object", + "const": {"$ref": "literal"} + } + }, + "required": ["value"] + })), + strict: None, + }]; + let choice = ToolChoice::Auto; + let value = + serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + let argument = &value["format"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]["elements"][0]; + + assert_eq!( + argument["begin"], + "<|open|>argument key=\"value\" type=\"object\"<|sep|>" + ); + assert_eq!( + argument["content"]["json_schema"]["const"], + json!({"$ref": "literal"}) + ); + } + #[test] fn auto_without_required_properties_allows_an_empty_argument_body() { let tools = vec![ToolDefinition { @@ -1336,23 +1985,25 @@ mod tests { } #[test] - fn named_choice_keeps_declared_argument_schema() { + fn named_choice_enforces_required_and_optional_argument_schemas() { let tools = tools(); let choice = ToolChoice::Named("get_weather".to_string()); let value = serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) .unwrap(); let call = &value["format"]["elements"][2]["content"]["tags"][0]; - let argument_alternatives = &call["content"]["elements"][2]["content"]["elements"]; - assert!( - argument_alternatives - .to_string() - .contains("key=\\\"city\\\"") + let arguments = &call["content"]["elements"][2]; + + assert_eq!(call["content"]["elements"][0]["pattern"], "[1-9][0-9]*"); + assert_eq!(arguments["type"], "sequence"); + assert_eq!( + arguments["elements"][0]["begin"], + "<|open|>argument key=\"city\" type=\"string\"<|sep|>" ); - assert!( - argument_alternatives - .to_string() - .contains("key=\\\"days\\\"") + assert_eq!(arguments["elements"][1]["type"], "star"); + assert_eq!( + arguments["elements"][1]["content"]["begin"], + "<|open|>argument key=\"days\" type=\"number\"<|sep|>" ); } From d483f3977afd282688ad64b70cd8055166e3c31e Mon Sep 17 00:00:00 2001 From: Indrajit Bhosale Date: Wed, 23 Sep 2026 17:16:49 -0700 Subject: [PATCH 08/10] Fix comments Signed-off-by: Indrajit Bhosale --- .../tool_calling/structural_tag/kimi_k3.rs | 224 ++++++++++++------ 1 file changed, 156 insertions(+), 68 deletions(-) diff --git a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs index 1258b6a72..e5f5d722b 100644 --- a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs +++ b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs @@ -361,21 +361,36 @@ fn merge_compatible_all_of(schema: &Map) -> Option return None, }; for (key, value) in option { - if merged.get(&key).is_some_and(|existing| existing != &value) { - return None; + match merged.get_mut(&key) { + Some(existing) if *existing == value => {} + Some(existing) if key == "enum" => { + // The raw XTML string path can enumerate this intersection. + // Leave other enum types on the existing fallback path. + let allowed = value.as_array()?; + let candidates = existing.as_array_mut()?; + if !allowed.iter().all(Value::is_string) + || !candidates.iter().all(Value::is_string) + { + return None; + } + candidates.retain(|item| allowed.contains(item)); + } + Some(_) => return None, + None => { + merged.insert(key, value); + } } - merged.insert(key, value); } } Some(merged) } -fn string_content_format(schema: &Value) -> Format { +fn string_content_format(schema: &Value) -> Option { let root = schema; let Some(schema) = schema.as_object() else { - return Format::AnyText(AnyTextFormat { + return Some(Format::AnyText(AnyTextFormat { excludes: vec![CLOSE.to_string()], - }); + })); }; // XTML string arguments contain raw text rather than a JSON string, so they // cannot use JsonSchemaFormat. Flatten compatible allOf branches before @@ -383,6 +398,31 @@ fn string_content_format(schema: &Value) -> Format { let merged_schema = merge_compatible_all_of(schema); let schema = merged_schema.as_ref().unwrap_or(schema); + // Apply outer constraints to each union branch before choosing a raw XTML + // string format. A top-level length regex alone would discard an inner enum. + for keyword in ["anyOf", "oneOf"] { + let Some(options) = schema.get(keyword).and_then(Value::as_array) else { + continue; + }; + let common = Value::Object( + schema + .iter() + .filter(|(key, _)| key.as_str() != keyword) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + ); + let formats = options + .iter() + .filter(|option| { + schema_types(option, root, &HashSet::new()).contains(&JsonType::String) + }) + .filter_map(|option| { + string_content_format(&serde_json::json!({"allOf": [common.clone(), option]})) + }) + .collect::>(); + return (!formats.is_empty()).then(|| one_of(formats)); + } + let enum_values = schema .get("enum") .and_then(Value::as_array) @@ -394,17 +434,19 @@ fn string_content_format(schema: &Value) -> Format { .map(|value| vec![Value::String(value.to_string())]) }); if let Some(values) = enum_values { + if values.is_empty() { + return None; + } let pattern = schema.get("pattern").and_then(Value::as_str); let compiled_pattern = pattern.map(regex::Regex::new); - if !values.is_empty() - && values.len() <= 256 + if values.len() <= 256 && values.iter().all(|value| value.as_str().is_some()) && compiled_pattern.as_ref().is_none_or(Result::is_ok) { let min_len = schema.get("minLength").and_then(Value::as_u64); let max_len = schema.get("maxLength").and_then(Value::as_u64); let const_value = schema.get("const").and_then(Value::as_str); - let values = values + let formats = values .iter() .filter_map(Value::as_str) .filter(|value| !value.contains("<|")) @@ -417,17 +459,13 @@ fn string_content_format(schema: &Value) -> Format { pattern.as_ref().is_ok_and(|re| re.is_match(value)) }) }) - .collect::>(); - return one_of( - values - .iter() - .map(|value| { - Format::ConstString(ConstStringFormat { - value: (*value).to_string(), - }) + .map(|value| { + Format::ConstString(ConstStringFormat { + value: value.to_string(), }) - .collect(), - ); + }) + .collect::>(); + return (!formats.is_empty()).then(|| one_of(formats)); } } @@ -456,34 +494,18 @@ fn string_content_format(schema: &Value) -> Format { } else { format!("{STRING_ATOM}*") }; - return Format::Regex(RegexFormat { + return Some(Format::Regex(RegexFormat { pattern: format!("{prefix}(?:{pattern}){suffix}"), - }); + })); } if let Some(pattern) = string_length_regex(schema) { - return Format::Regex(RegexFormat { pattern }); + return Some(Format::Regex(RegexFormat { pattern })); } - for keyword in ["anyOf", "oneOf"] { - let Some(options) = schema.get(keyword).and_then(Value::as_array) else { - continue; - }; - let formats = options - .iter() - .filter(|option| { - schema_types(option, root, &HashSet::new()).contains(&JsonType::String) - }) - .map(string_content_format) - .collect::>(); - if !formats.is_empty() { - return one_of(formats); - } - } - - Format::AnyText(AnyTextFormat { + Some(Format::AnyText(AnyTextFormat { excludes: vec![CLOSE.to_string()], - }) + })) } fn resolved_root_parameters(root: &Value) -> Option { @@ -602,7 +624,7 @@ fn rewrite_local_schema_refs(schema: &mut Value, prefix: &str) { } } -fn has_unresolved_local_refs(schema: &Value, root: &Value) -> bool { +fn has_local_schema_refs(schema: &Value) -> bool { let Some(object) = schema.as_object() else { return false; }; @@ -610,10 +632,7 @@ fn has_unresolved_local_refs(schema: &Value, root: &Value) -> bool { if object .get("$ref") .and_then(Value::as_str) - .is_some_and(|reference| { - (reference == "#" || reference.starts_with("#/")) - && resolve_local_ref(reference, root).is_none() - }) + .is_some_and(|reference| reference == "#" || reference.starts_with("#/")) { return true; } @@ -621,29 +640,20 @@ fn has_unresolved_local_refs(schema: &Value, root: &Value) -> bool { object .get(*key) .and_then(Value::as_object) - .is_some_and(|values| { - values - .values() - .any(|value| has_unresolved_local_refs(value, root)) - }) - }) || SCHEMA_VALUE_KEYWORDS.iter().any(|key| { - object - .get(*key) - .is_some_and(|value| has_unresolved_local_refs(value, root)) - }) || SCHEMA_ARRAY_KEYWORDS.iter().any(|key| { - object - .get(*key) - .and_then(Value::as_array) - .is_some_and(|values| { - values - .iter() - .any(|value| has_unresolved_local_refs(value, root)) - }) - }) + .is_some_and(|values| values.values().any(has_local_schema_refs)) + }) || SCHEMA_VALUE_KEYWORDS + .iter() + .any(|key| object.get(*key).is_some_and(has_local_schema_refs)) + || SCHEMA_ARRAY_KEYWORDS.iter().any(|key| { + object + .get(*key) + .and_then(Value::as_array) + .is_some_and(|values| values.iter().any(has_local_schema_refs)) + }) } fn preserve_recursive_root_refs(mut schema: Value, root: &Value) -> Value { - if !has_unresolved_local_refs(&schema, &schema) { + if !has_local_schema_refs(&schema) { return schema; } @@ -854,7 +864,7 @@ fn argument_format(key: &str, schema: &Value, root: &Value) -> Option { .into_iter() .filter_map(|xtml_type| { let content = if xtml_type == "string" { - string_content_format(&schema_for_xtml_type(schema, root, xtml_type)?) + string_content_format(&schema_for_xtml_type(schema, root, xtml_type)?)? } else { Format::JsonSchema(JsonSchemaFormat { json_schema: schema_for_xtml_type(schema, root, xtml_type)?, @@ -1441,6 +1451,52 @@ mod tests { assert_eq!(value["pattern"], format!(r"(?:price\$){STRING_ATOM}*")); } + #[test] + fn string_all_of_intersects_distinct_enums() { + let format = string_content_format(&json!({ + "allOf": [ + {"type": "string", "enum": ["safe", "other"]}, + {"enum": ["safe", "third"]} + ] + })); + let value = serde_json::to_value(format).unwrap(); + assert_eq!(value["type"], "const_string"); + assert_eq!(value["value"], "safe"); + } + + #[test] + fn string_outer_length_keeps_any_of_enum_restrictions() { + let format = string_content_format(&json!({ + "type": "string", + "minLength": 1, + "anyOf": [{"enum": ["safe"]}, {"enum": ["other"]}] + })); + let value = serde_json::to_value(format).unwrap(); + assert_eq!(value["type"], "or"); + let values = value["elements"].as_array().unwrap(); + assert_eq!(values.len(), 2); + assert!(values.iter().any(|item| item["value"] == "safe")); + assert!(values.iter().any(|item| item["value"] == "other")); + } + + #[test] + fn impossible_optional_string_enum_does_not_emit_empty_or() { + let root = json!({ + "type": "object", + "properties": {"value": {"type": "string", "enum": ["a"], "minLength": 2}} + }); + assert!(argument_format("value", &root["properties"]["value"], &root).is_none()); + let tools = vec![ToolDefinition { + name: "optional".to_string(), + parameters: Some(root), + strict: None, + }]; + let choice = ToolChoice::Auto; + let tag = serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) + .unwrap(); + assert!(!tag.to_string().contains("\"elements\":[]")); + } + #[test] fn auto_string_arguments_preserve_schema_constraints_and_allow_empty_values() { let tools = vec![ToolDefinition { @@ -1579,10 +1635,42 @@ mod tests { ); let schema = &arguments["content"]["content"]["json_schema"]; assert_eq!(schema["type"], "object"); - assert_eq!(schema["properties"]["next"]["$ref"], "#/$defs/node"); + assert_eq!( + schema["properties"]["next"]["$ref"], + "#/$defs/__dynamo_root/$defs/node" + ); + assert!( + resolve_local_ref( + schema["properties"]["next"]["$ref"].as_str().unwrap(), + schema + ) + .is_some() + ); assert_eq!(schema["$defs"]["node"]["type"], "object"); } + #[test] + fn recursive_root_hash_retains_original_parameter_scope() { + let root = json!({ + "type": "object", + "properties": { + "value": { + "type": "object", + "properties": {"parent": {"$ref": "#"}}, + "additionalProperties": false + } + }, + "additionalProperties": false + }); + let schema = schema_for_xtml_type(&root["properties"]["value"], &root, "object") + .expect("object argument schema"); + assert_eq!( + schema["properties"]["parent"]["properties"]["value"]["properties"]["parent"]["$ref"], + "#/$defs/__dynamo_root" + ); + assert_eq!(schema["$defs"]["__dynamo_root"]["type"], "object"); + } + #[test] fn recursive_property_reference_remains_resolvable_after_extraction() { let root = json!({ @@ -1606,7 +1694,7 @@ mod tests { .unwrap(); assert!(resolve_local_ref(reference, &schema).is_some()); - assert!(!has_unresolved_local_refs(&schema, &schema)); + assert!(has_local_schema_refs(&schema)); assert!(reference.starts_with("#/$defs/__dynamo_root")); assert_eq!( schema["properties"]["metadata"]["const"], From 5687838d972eae5414588f0438d0441171553fcf Mon Sep 17 00:00:00 2001 From: Ryan McCormick Date: Tue, 29 Sep 2026 12:17:05 -0700 Subject: [PATCH 09/10] fix(parsers): preserve K3 schema intersections and errors Signed-off-by: Ryan McCormick --- .../tool_calling/structural_tag/kimi_k3.rs | 353 ++++++++++++------ 1 file changed, 231 insertions(+), 122 deletions(-) diff --git a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs index e5f5d722b..fc82baaa4 100644 --- a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs +++ b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs @@ -11,6 +11,7 @@ use std::collections::HashSet; +use anyhow::{Context, bail}; use serde_json::{Map, Value}; use super::builder::{ToolCallFormatBuildContext, resolve_tools_to_include}; @@ -347,8 +348,13 @@ fn has_unescaped_trailing_dollar(pattern: &str) -> bool { == 0 } -fn merge_compatible_all_of(schema: &Map) -> Option> { - let options = schema.get("allOf")?.as_array()?; +fn merge_compatible_all_of( + schema: &Map, +) -> anyhow::Result>> { + let Some(options) = schema.get("allOf") else { + return Ok(None); + }; + let options = options.as_array().context("allOf must be an array")?; let mut merged = schema.clone(); merged.remove("allOf"); @@ -356,47 +362,66 @@ fn merge_compatible_all_of(schema: &Map) -> Option continue, Value::Object(option) => { - merge_compatible_all_of(option).unwrap_or_else(|| option.clone()) + merge_compatible_all_of(option)?.unwrap_or_else(|| option.clone()) } - _ => return None, + _ => bail!("cannot translate this allOf branch to XTML"), }; for (key, value) in option { match merged.get_mut(&key) { Some(existing) if *existing == value => {} Some(existing) if key == "enum" => { // The raw XTML string path can enumerate this intersection. - // Leave other enum types on the existing fallback path. - let allowed = value.as_array()?; - let candidates = existing.as_array_mut()?; + let allowed = value.as_array().context("enum must be an array")?; + let candidates = existing.as_array_mut().context("enum must be an array")?; if !allowed.iter().all(Value::is_string) || !candidates.iter().all(Value::is_string) { - return None; + bail!("cannot intersect non-string enums in an XTML allOf"); } candidates.retain(|item| allowed.contains(item)); } - Some(_) => return None, + Some(existing) if key == "minLength" || key == "maxLength" => { + let left = existing + .as_u64() + .context("string length must be a nonnegative integer")?; + let right = value + .as_u64() + .context("string length must be a nonnegative integer")?; + *existing = Value::from(if key == "minLength" { + left.max(right) + } else { + left.min(right) + }); + } + Some(_) => bail!("cannot intersect allOf constraints for {key}"), None => { merged.insert(key, value); } } } } - Some(merged) + Ok(Some(merged)) } -fn string_content_format(schema: &Value) -> Option { +fn string_content_format(schema: &Value) -> anyhow::Result> { let root = schema; let Some(schema) = schema.as_object() else { - return Some(Format::AnyText(AnyTextFormat { + return Ok(Some(Format::AnyText(AnyTextFormat { excludes: vec![CLOSE.to_string()], - })); + }))); }; // XTML string arguments contain raw text rather than a JSON string, so they // cannot use JsonSchemaFormat. Flatten compatible allOf branches before // translating their string constraints to raw-text formats. - let merged_schema = merge_compatible_all_of(schema); + let merged_schema = merge_compatible_all_of(schema)?; let schema = merged_schema.as_ref().unwrap_or(schema); + if let (Some(min), Some(max)) = ( + schema.get("minLength").and_then(Value::as_u64), + schema.get("maxLength").and_then(Value::as_u64), + ) && min > max + { + return Ok(None); + } // Apply outer constraints to each union branch before choosing a raw XTML // string format. A top-level length regex alone would discard an inner enum. @@ -411,16 +436,16 @@ fn string_content_format(schema: &Value) -> Option { .map(|(key, value)| (key.clone(), value.clone())) .collect(), ); - let formats = options - .iter() - .filter(|option| { - schema_types(option, root, &HashSet::new()).contains(&JsonType::String) - }) - .filter_map(|option| { - string_content_format(&serde_json::json!({"allOf": [common.clone(), option]})) - }) - .collect::>(); - return (!formats.is_empty()).then(|| one_of(formats)); + let mut formats = Vec::new(); + for option in options { + if schema_types(option, root, &HashSet::new()).contains(&JsonType::String) + && let Some(format) = + string_content_format(&serde_json::json!({"allOf": [common.clone(), option]}))? + { + formats.push(format); + } + } + return Ok((!formats.is_empty()).then(|| one_of(formats))); } let enum_values = schema @@ -435,7 +460,7 @@ fn string_content_format(schema: &Value) -> Option { }); if let Some(values) = enum_values { if values.is_empty() { - return None; + return Ok(None); } let pattern = schema.get("pattern").and_then(Value::as_str); let compiled_pattern = pattern.map(regex::Regex::new); @@ -465,7 +490,7 @@ fn string_content_format(schema: &Value) -> Option { }) }) .collect::>(); - return (!formats.is_empty()).then(|| one_of(formats)); + return Ok((!formats.is_empty()).then(|| one_of(formats))); } } @@ -494,28 +519,31 @@ fn string_content_format(schema: &Value) -> Option { } else { format!("{STRING_ATOM}*") }; - return Some(Format::Regex(RegexFormat { + return Ok(Some(Format::Regex(RegexFormat { pattern: format!("{prefix}(?:{pattern}){suffix}"), - })); + }))); } if let Some(pattern) = string_length_regex(schema) { - return Some(Format::Regex(RegexFormat { pattern })); + return Ok(Some(Format::Regex(RegexFormat { pattern }))); } - Some(Format::AnyText(AnyTextFormat { + Ok(Some(Format::AnyText(AnyTextFormat { excludes: vec![CLOSE.to_string()], - })) + }))) } -fn resolved_root_parameters(root: &Value) -> Option { - let resolved = resolve_local_refs(root, root, &HashSet::new())?; - let object = resolved.as_object()?; - if object.contains_key("allOf") { - merge_compatible_all_of(object).map(Value::Object) - } else { - Some(resolved) - } +fn resolved_root_parameters(root: &Value) -> anyhow::Result> { + let Some(resolved) = resolve_local_refs(root, root, &HashSet::new()) else { + bail!("cannot resolve tool parameter references"); + }; + let Some(object) = resolved.as_object() else { + return Ok(None); + }; + Ok(Some(match merge_compatible_all_of(object)? { + Some(merged) => Value::Object(merged), + None => resolved, + })) } fn permissive_argument_tag() -> TagFormat { @@ -859,48 +887,52 @@ fn schema_for_xtml_type(schema: &Value, root: &Value, xtml_type: &str) -> Option Some(preserve_recursive_root_refs(Value::Object(narrowed), root)) } -fn argument_format(key: &str, schema: &Value, root: &Value) -> Option { - let alternatives = xtml_types(schema, root) - .into_iter() - .filter_map(|xtml_type| { - let content = if xtml_type == "string" { - string_content_format(&schema_for_xtml_type(schema, root, xtml_type)?)? - } else { - Format::JsonSchema(JsonSchemaFormat { - json_schema: schema_for_xtml_type(schema, root, xtml_type)?, - style: JsonSchemaStyle::Json, - }) +fn argument_format(key: &str, schema: &Value, root: &Value) -> anyhow::Result> { + let mut alternatives = Vec::new(); + for xtml_type in xtml_types(schema, root) { + let Some(schema) = schema_for_xtml_type(schema, root, xtml_type) else { + continue; + }; + let content = if xtml_type == "string" { + let Some(content) = string_content_format(&schema) + .with_context(|| format!("cannot constrain argument {key}"))? + else { + continue; }; - Some(Format::Tag(TagFormat { - begin: format!( - "{OPEN}argument key=\"{}\" type=\"{xtml_type}\"{SEP}", - escape_attr(key) - ), - content: Box::new(content), - end: ARGUMENT_CLOSE.to_string(), - })) - }) - .collect::>(); - - (!alternatives.is_empty()).then(|| one_of(alternatives)) + content + } else { + Format::JsonSchema(JsonSchemaFormat { + json_schema: schema, + style: JsonSchemaStyle::Json, + }) + }; + alternatives.push(Format::Tag(TagFormat { + begin: format!( + "{OPEN}argument key=\"{}\" type=\"{xtml_type}\"{SEP}", + escape_attr(key) + ), + content: Box::new(content), + end: ARGUMENT_CLOSE.to_string(), + })); + } + Ok((!alternatives.is_empty()).then(|| one_of(alternatives))) } fn optional_arguments( properties: &Map, required_keys: &[&str], root: &Value, -) -> Option { - let arguments = properties - .iter() - .filter(|(key, _)| !required_keys.contains(&key.as_str())) - .filter_map(|(key, schema)| { - matches!(schema, Value::Bool(_) | Value::Object(_)) - .then(|| argument_format(key, schema, root)) - .flatten() - }) - .collect::>(); - - (!arguments.is_empty()).then(|| star(one_of(arguments))) +) -> anyhow::Result> { + let mut arguments = Vec::new(); + for (key, schema) in properties { + if !required_keys.contains(&key.as_str()) + && matches!(schema, Value::Bool(_) | Value::Object(_)) + && let Some(argument) = argument_format(key, schema, root)? + { + arguments.push(argument); + } + } + Ok((!arguments.is_empty()).then(|| star(one_of(arguments)))) } fn canonical_required_arguments( @@ -920,23 +952,25 @@ fn canonical_required_arguments( }) } -fn arguments_block(tool: &ToolDefinition, strict_schema: bool) -> Format { +fn arguments_block(tool: &ToolDefinition, strict_schema: bool) -> anyhow::Result { if !super::builder::uses_declared_tool_schema(tool, strict_schema) { - return star(Format::Tag(permissive_argument_tag())); + return Ok(star(Format::Tag(permissive_argument_tag()))); } let Some(root) = tool.parameters.as_ref() else { - return star(Format::Tag(permissive_argument_tag())); + return Ok(star(Format::Tag(permissive_argument_tag()))); }; let resolved_root = root .as_object() .filter(|object| object.contains_key("$ref") || object.contains_key("allOf")) - .and_then(|_| resolved_root_parameters(root)); + .map(|_| resolved_root_parameters(root)) + .transpose()? + .flatten(); let Some(parameters) = resolved_root.as_ref().unwrap_or(root).as_object() else { - return star(Format::Tag(permissive_argument_tag())); + return Ok(star(Format::Tag(permissive_argument_tag()))); }; let Some(properties) = parameters.get("properties").and_then(Value::as_object) else { - return star(Format::Tag(permissive_argument_tag())); + return Ok(star(Format::Tag(permissive_argument_tag()))); }; let required = match parameters.get("required") { None => Vec::new(), @@ -946,43 +980,42 @@ fn arguments_block(tool: &ToolDefinition, strict_schema: bool) -> Format { .filter_map(Value::as_str) .collect::>() } - Some(_) => return star(Format::Tag(permissive_argument_tag())), + Some(_) => return Ok(star(Format::Tag(permissive_argument_tag()))), }; - let optional_arguments = optional_arguments(properties, &required, root); + let optional_arguments = optional_arguments(properties, &required, root)?; if required.is_empty() { - return optional_arguments.unwrap_or_else(|| { + return Ok(optional_arguments.unwrap_or_else(|| { Format::ConstString(ConstStringFormat { value: String::new(), }) - }); + })); } let mut arguments = Vec::with_capacity(required.len()); for key in &required { let Some(schema) = properties.get(*key) else { - return star(Format::Tag(permissive_argument_tag())); + return Ok(star(Format::Tag(permissive_argument_tag()))); }; if !matches!(schema, Value::Bool(_) | Value::Object(_)) { - return star(Format::Tag(permissive_argument_tag())); + return Ok(star(Format::Tag(permissive_argument_tag()))); } - let Some(argument) = argument_format(key, schema, root) else { - return star(Format::Tag(permissive_argument_tag())); - }; + let argument = argument_format(key, schema, root)? + .with_context(|| format!("required argument {key} has no representable values"))?; arguments.push(argument); } - canonical_required_arguments(arguments, optional_arguments) + Ok(canonical_required_arguments(arguments, optional_arguments)) } fn build_auto_structural_tag( tools: Vec<&ToolDefinition>, ctx: &ToolCallFormatBuildContext<'_>, -) -> StructuralTag { +) -> anyhow::Result { let strict_schema = ctx.strict_schema(); let call_tags: Vec<_> = tools .into_iter() .map(|tool| call_tag(tool, strict_schema)) - .collect(); + .collect::>()?; let parallel_tool_calls = !ctx.stop_after_first(); let calls = if parallel_tool_calls { Format::TagsWithSeparator(TagsWithSeparatorFormat { @@ -1029,12 +1062,12 @@ fn build_auto_structural_tag( } else { suffix }; - StructuralTag { format } + Ok(StructuralTag { format }) } -fn call_tag(tool: &ToolDefinition, strict_schema: bool) -> TagFormat { +fn call_tag(tool: &ToolDefinition, strict_schema: bool) -> anyhow::Result { let begin = format!("{OPEN}call tool=\"{}\" index=\"", escape_attr(&tool.name)); - TagFormat { + Ok(TagFormat { begin, content: Box::new(Format::Sequence(SequenceFormat { elements: vec![ @@ -1044,11 +1077,12 @@ fn call_tag(tool: &ToolDefinition, strict_schema: bool) -> TagFormat { Format::ConstString(ConstStringFormat { value: format!("\"{SEP}"), }), - arguments_block(tool, strict_schema), + arguments_block(tool, strict_schema) + .with_context(|| format!("cannot constrain tool {}", tool.name))?, ], })), end: CALL_CLOSE.to_string(), - } + }) } /// Build the format-style xgrammar tag for K3's response + tools channels. @@ -1060,7 +1094,7 @@ pub(crate) fn build_kimi_k3( return Ok(None); } if matches!(ctx.tool_choice, crate::tool_calling::ToolChoice::Auto) { - return Ok(Some(build_auto_structural_tag(tools, ctx))); + return Ok(Some(build_auto_structural_tag(tools, ctx)?)); } // Moonshot's named-tool contract returns the selected call with no @@ -1090,7 +1124,7 @@ pub(crate) fn build_kimi_k3( tags: tools .into_iter() .map(|tool| call_tag(tool, ctx.strict_schema())) - .collect(), + .collect::>()?, separator: String::new(), at_least_one: true, stop_after_first: ctx.stop_after_first(), @@ -1122,7 +1156,9 @@ mod tests { use serde_json::json; use super::*; - use crate::tool_calling::structural_tag::builder::StructuralTagSchemaMode; + use crate::tool_calling::structural_tag::builder::{ + StructuralTagBuilder, StructuralTagSchemaMode, + }; use crate::tool_calling::{ToolChoice, ToolDefinition}; fn tools() -> Vec { @@ -1412,10 +1448,13 @@ mod tests { #[test] fn string_content_enforces_min_length_without_a_maximum() { - let value = serde_json::to_value(string_content_format(&json!({ - "type": "string", - "minLength": 2 - }))) + let value = serde_json::to_value( + string_content_format(&json!({ + "type": "string", + "minLength": 2 + })) + .unwrap(), + ) .unwrap(); assert_eq!(value["type"], "regex"); @@ -1424,12 +1463,15 @@ mod tests { #[test] fn string_content_uses_xgrammar_pattern_precedence_over_length_bounds() { - let value = serde_json::to_value(string_content_format(&json!({ - "type": "string", - "pattern": "^item-[0-9]+$", - "minLength": 6, - "maxLength": 12 - }))) + let value = serde_json::to_value( + string_content_format(&json!({ + "type": "string", + "pattern": "^item-[0-9]+$", + "minLength": 6, + "maxLength": 12 + })) + .unwrap(), + ) .unwrap(); assert_eq!(value["type"], "regex"); @@ -1441,10 +1483,13 @@ mod tests { assert!(!has_unescaped_trailing_dollar(r"^price\$")); assert!(has_unescaped_trailing_dollar(r"^path\\$")); - let value = serde_json::to_value(string_content_format(&json!({ - "type": "string", - "pattern": r"^price\$" - }))) + let value = serde_json::to_value( + string_content_format(&json!({ + "type": "string", + "pattern": r"^price\$" + })) + .unwrap(), + ) .unwrap(); assert_eq!(value["type"], "regex"); @@ -1458,19 +1503,25 @@ mod tests { {"type": "string", "enum": ["safe", "other"]}, {"enum": ["safe", "third"]} ] - })); + })) + .unwrap(); let value = serde_json::to_value(format).unwrap(); assert_eq!(value["type"], "const_string"); assert_eq!(value["value"], "safe"); } #[test] - fn string_outer_length_keeps_any_of_enum_restrictions() { + fn string_outer_bounds_intersect_any_of_enum_restrictions() { let format = string_content_format(&json!({ "type": "string", - "minLength": 1, - "anyOf": [{"enum": ["safe"]}, {"enum": ["other"]}] - })); + "minLength": 2, + "maxLength": 6, + "anyOf": [ + {"minLength": 3, "maxLength": 4, "enum": ["a", "safe", "longer"]}, + {"enum": ["other"]} + ] + })) + .unwrap(); let value = serde_json::to_value(format).unwrap(); assert_eq!(value["type"], "or"); let values = value["elements"].as_array().unwrap(); @@ -1479,13 +1530,71 @@ mod tests { assert!(values.iter().any(|item| item["value"] == "other")); } + #[test] + fn impossible_required_argument_errors_for_every_tool_choice() { + let mut tools = vec![ToolDefinition { + name: "test".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": {"value": {"type": "string", "enum": ["a"], "minLength": 2}}, + "required": ["value"] + })), + strict: None, + }]; + for choice in [ + ToolChoice::Auto, + ToolChoice::Required, + ToolChoice::Named("test".into()), + ] { + let error = StructuralTagBuilder::KimiK3 + .build_tool_call_format(&context(&choice, &tools)) + .unwrap_err(); + let message = format!("{error:#}"); + assert!( + message.contains("test") && message.contains("value"), + "{message}" + ); + } + tools[0].strict = Some(false); + assert!( + StructuralTagBuilder::KimiK3 + .build_tool_call_format(&context(&ToolChoice::Auto, &tools)) + .unwrap() + .is_some() + ); + } + + #[test] + fn unsupported_optional_intersection_returns_an_error() { + let tools = vec![ToolDefinition { + name: "test".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": {"value": {"allOf": [ + {"type": "string", "pattern": "^s"}, + {"pattern": "e$"} + ]}} + })), + strict: None, + }]; + assert!( + StructuralTagBuilder::KimiK3 + .build_tool_call_format(&context(&ToolChoice::Auto, &tools)) + .is_err() + ); + } + #[test] fn impossible_optional_string_enum_does_not_emit_empty_or() { let root = json!({ "type": "object", "properties": {"value": {"type": "string", "enum": ["a"], "minLength": 2}} }); - assert!(argument_format("value", &root["properties"]["value"], &root).is_none()); + assert!( + argument_format("value", &root["properties"]["value"], &root) + .unwrap() + .is_none() + ); let tools = vec![ToolDefinition { name: "optional".to_string(), parameters: Some(root), @@ -1879,7 +1988,7 @@ mod tests { ] }); let narrowed = schema_for_xtml_type(&schema, &schema, "string").unwrap(); - let format = serde_json::to_value(string_content_format(&narrowed)).unwrap(); + let format = serde_json::to_value(string_content_format(&narrowed).unwrap()).unwrap(); assert_eq!(format["type"], "const_string"); assert_eq!(format["value"], "safe"); @@ -1894,7 +2003,7 @@ mod tests { ] }); let narrowed = schema_for_xtml_type(&schema, &schema, "string").unwrap(); - let format = serde_json::to_value(string_content_format(&narrowed)).unwrap(); + let format = serde_json::to_value(string_content_format(&narrowed).unwrap()).unwrap(); assert_eq!(format["type"], "const_string"); assert_eq!(format["value"], "code-42"); From 75e882f2444ee5d13d699ccf4553bec1245599bf Mon Sep 17 00:00:00 2001 From: Ryan McCormick Date: Tue, 29 Sep 2026 15:13:49 -0700 Subject: [PATCH 10/10] fix(parsers): correct K3 reasoning and unresolved string schemas Signed-off-by: Ryan McCormick --- .../tool_calling/structural_tag/kimi_k3.rs | 145 +++++++++++++----- 1 file changed, 110 insertions(+), 35 deletions(-) diff --git a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs index 5f23adc3c..e60da5433 100644 --- a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs +++ b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs @@ -933,19 +933,33 @@ fn narrow_resolved_schema( Some(narrowed) } -fn schema_for_xtml_type(schema: &Value, root: &Value, xtml_type: &str) -> Option { +fn schema_for_xtml_type( + schema: &Value, + root: &Value, + xtml_type: &str, +) -> anyhow::Result> { let Some(resolved) = resolve_local_refs(schema, root, &HashSet::new()) else { - // Keep an explicitly typed schema intact when scoped references or + // Raw strings cannot delegate unresolved references to JsonSchemaFormat. + if xtml_type == "string" { + bail!("cannot resolve references for raw XTML string"); + } + // Keep an explicitly typed JSON schema intact when scoped references or // bounded expansion prevent safe inlining. - let original = schema + let Some(original) = schema .as_object() - .filter(|schema| !schema.contains_key("$ref"))?; - let declared_type = original + .filter(|schema| !schema.contains_key("$ref")) + else { + return Ok(None); + }; + let Some(declared_type) = original .get("type") .and_then(Value::as_str) - .and_then(JsonType::from_name)?; + .and_then(JsonType::from_name) + else { + return Ok(None); + }; if declared_type.xtml_name() != xtml_type { - return None; + return Ok(None); } let mut original = original.clone(); if let Some(root) = root.as_object() { @@ -957,7 +971,7 @@ fn schema_for_xtml_type(schema: &Value, root: &Value, xtml_type: &str) -> Option } } } - return Some(Value::Object(original)); + return Ok(Some(Value::Object(original))); }; let mut json_types = schema_types(&resolved, &resolved, &HashSet::new()) .into_iter() @@ -966,7 +980,9 @@ fn schema_for_xtml_type(schema: &Value, root: &Value, xtml_type: &str) -> Option if json_types.contains(&JsonType::Number) { json_types.retain(|json_type| *json_type != JsonType::Integer); } - let mut narrowed = narrow_resolved_schema(resolved, root, &json_types)?; + let Some(mut narrowed) = narrow_resolved_schema(resolved, root, &json_types) else { + return Ok(None); + }; if let Some(root) = root.as_object() { for key in ["$defs", "definitions"] { if let Some(value) = root.get(key) { @@ -976,13 +992,18 @@ fn schema_for_xtml_type(schema: &Value, root: &Value, xtml_type: &str) -> Option } } } - Some(preserve_recursive_root_refs(Value::Object(narrowed), root)) + Ok(Some(preserve_recursive_root_refs( + Value::Object(narrowed), + root, + ))) } fn argument_format(key: &str, schema: &Value, root: &Value) -> anyhow::Result> { let mut alternatives = Vec::new(); for xtml_type in xtml_types(schema, root) { - let Some(schema) = schema_for_xtml_type(schema, root, xtml_type) else { + let Some(schema) = schema_for_xtml_type(schema, root, xtml_type) + .with_context(|| format!("cannot constrain argument {key}"))? + else { continue; }; let content = if xtml_type == "string" { @@ -1140,21 +1161,7 @@ fn build_auto_structural_tag( optional(Format::Tag(tools_tag)), ], }); - let format = if ctx.starts_in_reasoning { - Format::Sequence(SequenceFormat { - elements: vec![ - Format::Tag(TagFormat { - begin: String::new(), - content: Box::new(Format::AnyText(AnyTextFormat { excludes: vec![] })), - end: THINK_CLOSE.to_string(), - }), - suffix, - ], - }) - } else { - suffix - }; - Ok(StructuralTag { format }) + Ok(StructuralTag { format: suffix }) } fn call_tag(tool: &ToolDefinition, strict_schema: bool) -> anyhow::Result { @@ -1867,6 +1874,7 @@ mod tests { "additionalProperties": false }); let schema = schema_for_xtml_type(&root["properties"]["value"], &root, "object") + .unwrap() .expect("object argument schema"); assert_eq!( schema["properties"]["parent"]["properties"]["value"]["properties"]["parent"]["$ref"], @@ -1892,7 +1900,9 @@ mod tests { } } }); - let schema = schema_for_xtml_type(&root["properties"]["node"], &root, "object").unwrap(); + let schema = schema_for_xtml_type(&root["properties"]["node"], &root, "object") + .unwrap() + .unwrap(); let reference = schema["properties"]["next"]["properties"]["next"]["$ref"] .as_str() .unwrap(); @@ -1956,7 +1966,9 @@ mod tests { "id": {"$ref": "#/$defs/caf%C3%A9%20name"} } }); - let schema = schema_for_xtml_type(&root["properties"]["id"], &root, "number").unwrap(); + let schema = schema_for_xtml_type(&root["properties"]["id"], &root, "number") + .unwrap() + .unwrap(); assert_eq!(schema["type"], "integer"); assert_eq!(schema["minimum"], 1); @@ -2041,13 +2053,61 @@ mod tests { } } }); - let schema = schema_for_xtml_type(&root["properties"]["data"], &root, "object").unwrap(); + let schema = schema_for_xtml_type(&root["properties"]["data"], &root, "object") + .unwrap() + .unwrap(); assert_eq!(schema["type"], "object"); assert_eq!(schema["properties"]["item"]["$ref"], "#/$defs/step_0"); assert_eq!(schema["$defs"], root["$defs"]); } + #[test] + fn unresolved_string_references_return_errors() { + let mut definitions = Map::new(); + for index in 0..MAX_SCHEMA_REF_DEPTH { + definitions.insert( + format!("step_{index}"), + json!({"$ref": format!("#/$defs/step_{}", index + 1)}), + ); + } + definitions.insert( + format!("step_{MAX_SCHEMA_REF_DEPTH}"), + json!({"type": "string", "enum": ["safe"]}), + ); + let schemas = [ + json!({"type": "string", "allOf": [{"$ref": "#/$defs/step_0"}]}), + json!({ + "type": "string", + "$id": "urn:example:scoped", + "$defs": {"allowed": {"type": "string", "enum": ["safe"]}}, + "allOf": [{"$ref": "#/$defs/allowed"}] + }), + ]; + for schema in schemas { + for required in [json!([]), json!(["value"])] { + let tools = vec![ToolDefinition { + name: "test".into(), + parameters: Some(json!({ + "type": "object", + "$defs": definitions, + "properties": {"value": schema}, + "required": required + })), + strict: None, + }]; + let error = StructuralTagBuilder::KimiK3 + .build_tool_call_format(&context(&ToolChoice::Auto, &tools)) + .unwrap_err(); + let message = format!("{error:#}"); + assert!( + message.contains("value") && message.contains("string"), + "{message}" + ); + } + } + } + #[test] fn auto_union_narrowing_preserves_value_constraints() { let tools = vec![ToolDefinition { @@ -2105,7 +2165,9 @@ mod tests { "type": ["integer", "null"], "enum": [1.0, null] }); - let narrowed = schema_for_xtml_type(&schema, &schema, "number").unwrap(); + let narrowed = schema_for_xtml_type(&schema, &schema, "number") + .unwrap() + .unwrap(); assert_eq!(narrowed["type"], "integer"); assert_eq!(narrowed["enum"], json!([1.0])); @@ -2119,7 +2181,9 @@ mod tests { {"type": "null"} ] }); - let narrowed = schema_for_xtml_type(&schema, &schema, "number").unwrap(); + let narrowed = schema_for_xtml_type(&schema, &schema, "number") + .unwrap() + .unwrap(); let options = narrowed["anyOf"].as_array().unwrap(); assert_eq!(narrowed["type"], "integer"); @@ -2217,7 +2281,9 @@ mod tests { {"enum": ["safe"]} ] }); - let narrowed = schema_for_xtml_type(&schema, &schema, "string").unwrap(); + let narrowed = schema_for_xtml_type(&schema, &schema, "string") + .unwrap() + .unwrap(); let format = serde_json::to_value(string_content_format(&narrowed).unwrap()).unwrap(); assert_eq!(format["type"], "const_string"); @@ -2232,7 +2298,9 @@ mod tests { {"minLength": 3, "pattern": "^code-[0-9]+$"} ] }); - let narrowed = schema_for_xtml_type(&schema, &schema, "string").unwrap(); + let narrowed = schema_for_xtml_type(&schema, &schema, "string") + .unwrap() + .unwrap(); let format = serde_json::to_value(string_content_format(&narrowed).unwrap()).unwrap(); assert_eq!(format["type"], "const_string"); @@ -2393,7 +2461,7 @@ mod tests { } #[test] - fn auto_thinking_closes_reasoning_before_the_triggered_tools_suffix() { + fn auto_thinking_has_one_reasoning_prefix() { let tools = tools(); let choice = ToolChoice::Auto; let ctx = ToolCallFormatBuildContext { @@ -2403,12 +2471,19 @@ mod tests { schema_mode: StructuralTagSchemaMode::Auto, starts_in_reasoning: true, }; - let value = serde_json::to_value(build_kimi_k3(&ctx).unwrap().unwrap()).unwrap(); + let value = StructuralTagBuilder::KimiK3 + .build_tool_call_format(&ctx) + .unwrap() + .unwrap(); assert_eq!(value["format"]["type"], "sequence"); assert_eq!(value["format"]["elements"][0]["type"], "tag"); assert_eq!(value["format"]["elements"][0]["end"], THINK_CLOSE); assert_eq!(value["format"]["elements"][1]["type"], "sequence"); + assert_eq!( + value["format"]["elements"][1]["elements"][0]["type"], + "any_text" + ); } #[test]