diff --git a/parsers/v1/src/tool_calling/structural_tag/builder.rs b/parsers/v1/src/tool_calling/structural_tag/builder.rs index 4c7c71f0b..eb3493d00 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, } @@ -82,33 +84,27 @@ pub(crate) fn resolve_tools_to_include<'a>( } } -/// Resolve one tool's argument schema using Dynamo's generic schema policy. +/// Whether a tool's declared parameter schema should constrain its arguments. /// -/// In [`StructuralTagSchemaMode::Auto`], only `strict: true` selects the -/// declared schema. Model-native builders may intentionally follow different -/// upstream semantics and should use a model-specific policy helper. +/// [`StructuralTagSchemaMode::Auto`] treats an omitted `strict` field as +/// schema-enforced. Only an explicit `strict: false` opts out. Global strict +/// mode overrides that per-tool opt-out. +pub(crate) fn uses_declared_tool_schema(tool: &ToolDefinition, strict_schema: bool) -> bool { + strict_schema || tool.strict != Some(false) +} + +/// Resolve one tool's argument schema using Dynamo's generic schema policy. pub(crate) fn resolve_tool_schema(tool: &ToolDefinition, strict_schema: bool) -> Value { // xgrammar uses `true` for syntactically valid but schema-unconstrained JSON. let default_schema = json!(true); - let use_tool_schema = strict_schema || tool.strict.unwrap_or(false); - if use_tool_schema { + if uses_declared_tool_schema(tool, strict_schema) { tool.parameters.clone().unwrap_or(default_schema) } else { default_schema } } -/// Whether Kimi's vLLM/xgrammar-compatible format should use the declared -/// parameter schema. -/// -/// Kimi treats only an explicit `strict: false` as an opt-out. An omitted -/// `strict` value therefore uses the declared schema, unlike Dynamo's generic -/// [`StructuralTagSchemaMode::Auto`] policy. -pub(crate) fn kimi_uses_declared_tool_schema(tool: &ToolDefinition, strict_schema: bool) -> bool { - strict_schema || tool.strict != Some(false) -} - /// Builder for model-family-specific tool-call structural tags. #[derive(Debug, Clone)] #[non_exhaustive] 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..7a582a010 100644 --- a/parsers/v1/src/tool_calling/structural_tag/kimi_k2.rs +++ b/parsers/v1/src/tool_calling/structural_tag/kimi_k2.rs @@ -15,7 +15,7 @@ use serde_json::{Value, json}; use super::builder::{ - ToolCallFormatBuildContext, kimi_uses_declared_tool_schema, resolve_tools_to_include, + ToolCallFormatBuildContext, resolve_tools_to_include, uses_declared_tool_schema, }; use super::format::{ ConstStringFormat, Format, JsonSchemaFormat, JsonSchemaStyle, RegexFormat, SequenceFormat, @@ -33,7 +33,7 @@ fn tool_schema(tool: &ToolDefinition, strict_schema: bool) -> Value { // Match vLLM/xgrammar: use the declared parameters unless the request // explicitly opts out with strict=false. Global strict mode overrides the // opt-out. Xgrammar uses `true` for unconstrained but valid JSON. - if kimi_uses_declared_tool_schema(tool, strict_schema) { + if uses_declared_tool_schema(tool, strict_schema) { tool.parameters.clone().unwrap_or_else(|| json!(true)) } else { json!(true) 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 19574582b..e60da5433 100644 --- a/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs +++ b/parsers/v1/src/tool_calling/structural_tag/kimi_k3.rs @@ -9,6 +9,9 @@ //! format so named and required tool choices can be constrained without //! changing what the K3 parser expects. +use std::collections::HashSet; + +use anyhow::{Context, bail}; use serde_json::{Map, Value}; use super::builder::{ToolCallFormatBuildContext, resolve_tools_to_include}; @@ -24,13 +27,19 @@ 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|>"; const STRING_ATOM: &str = r"(?:[^<]|<[^|])"; +const CALL_INDEX_PATTERN: &str = "[1-9][0-9]*"; +const MAX_SCHEMA_REF_DEPTH: usize = 16; +const MAX_SCHEMA_NODES: usize = 4096; fn escape_attr(value: &str) -> String { value.replace('&', "&").replace('"', """) @@ -56,94 +65,525 @@ fn one_of(elements: Vec) -> Format { } } -fn bounded_string_regex(schema: &Map) -> Option { - let max_len = schema.get("maxLength")?.as_u64()?; - if max_len > 4096 { - return None; +#[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, +]; + +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 { + "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 is_integral_number(number) => 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 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(), + } } - 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}}}")) } -fn argument_tag( - key: &str, +fn resolve_local_ref<'a>(reference: &str, root: &'a Value) -> Option<&'a Value> { + // URI percent-decoding precedes JSON Pointer's ~0/~1 unescaping. + let pointer = reference.strip_prefix('#')?; + let decoded; + let pointer = if pointer.contains('%') { + let mut bytes = pointer.bytes(); + let mut result = Vec::with_capacity(pointer.len()); + while let Some(byte) = bytes.next() { + result.push(if byte == b'%' { + let high = char::from(bytes.next()?).to_digit(16)?; + let low = char::from(bytes.next()?).to_digit(16)?; + ((high << 4) | low) as u8 + } else { + byte + }); + } + decoded = String::from_utf8(result).ok()?; + &decoded + } else { + pointer + }; + let value = root.pointer(pointer)?; + matches!(value, Value::Bool(_) | Value::Object(_)).then_some(value) +} + +fn schema_types(schema: &Value, root: &Value, seen_refs: &HashSet) -> Vec { + let mut remaining = MAX_SCHEMA_NODES; + schema_types_inner(schema, root, seen_refs, 0, &mut remaining) +} + +fn has_unsupported_ref_scope(schema: &Map) -> bool { + ["$id", "$dynamicRef", "$recursiveRef"] + .iter() + .any(|key| schema.contains_key(*key)) +} + +fn schema_types_inner( 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, + root: &Value, + seen_refs: &HashSet, + depth: usize, + remaining: &mut usize, +) -> Vec { + let Some(next) = (*remaining).checked_sub(1) else { + return JSON_TYPES.to_vec(); }; - let begin = format!( - "{OPEN}argument key=\"{}\" type=\"{xtml_type}\"{SEP}", - escape_attr(key) - ); - - 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 }) + *remaining = next; + if depth >= MAX_SCHEMA_REF_DEPTH { + return JSON_TYPES.to_vec(); + } + if let Some(allowed) = schema.as_bool() { + return if allowed { + JSON_TYPES.to_vec() } else { - Format::AnyText(AnyTextFormat { - excludes: vec![CLOSE.to_string()], + Vec::new() + }; + } + let Some(schema) = schema.as_object() else { + return JSON_TYPES.to_vec(); + }; + + if !has_unsupported_ref_scope(schema) + && 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_inner(target, root, &nested_seen, depth + 1, remaining); + } + + 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_inner(option, root, seen_refs, depth, remaining)) + .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_inner(option, root, seen_refs, depth, remaining) + .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 { - 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()); + JSON_TYPES.to_vec() + } +} + +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 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; + } + 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 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, +) -> 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"); + + 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()) + } + _ => 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. + 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) + { + bail!("cannot intersect non-string enums in an XTML allOf"); + } + candidates.retain(|item| allowed.contains(item)); + } + 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); + } } } - Format::JsonSchema(JsonSchemaFormat { - json_schema: Value::Object(embedded), - style: JsonSchemaStyle::Json, - }) + } + Ok(Some(merged)) +} + +fn string_content_format(schema: &Value) -> anyhow::Result> { + let root = schema; + let Some(schema) = schema.as_object() else { + 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 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); + } - Some(TagFormat { - begin, - content: Box::new(content), - end: ARGUMENT_CLOSE.to_string(), - }) + // 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 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 + .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 { + if values.is_empty() { + return Ok(None); + } + let pattern = schema.get("pattern").and_then(Value::as_str); + let compiled_pattern = pattern.map(regex::Regex::new); + 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 formats = values + .iter() + .filter_map(Value::as_str) + .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)) + }) + }) + .map(|value| { + Format::ConstString(ConstStringFormat { + value: value.to_string(), + }) + }) + .collect::>(); + return Ok((!formats.is_empty()).then(|| one_of(formats))); + } + } + + 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 = 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 { + format!("{STRING_ATOM}*") + }; + let suffix = if anchored_end { + String::new() + } else { + format!("{STRING_ATOM}*") + }; + return Ok(Some(Format::Regex(RegexFormat { + pattern: format!("{prefix}(?:{pattern}){suffix}"), + }))); + } + + if let Some(pattern) = string_length_regex(schema) { + return Ok(Some(Format::Regex(RegexFormat { pattern }))); + } + + Ok(Some(Format::AnyText(AnyTextFormat { + excludes: vec![CLOSE.to_string()], + }))) +} + +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 { @@ -163,153 +603,682 @@ fn permissive_argument_tag() -> TagFormat { } } -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 resolve_schema_map( + value: &Value, + root: &Value, + seen_refs: &HashSet, + depth: usize, + remaining: &mut usize, +) -> Option { + let Some(values) = value.as_object() else { + return Some(value.clone()); }; - if properties.is_empty() { - return star(Format::Tag(permissive_argument_tag())); - } - - let root_defs: 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()))) - }) - .collect(); - let tags = properties + values .iter() - .map(|(key, schema)| { - Format::Tag( - argument_tag(key, schema, Some(&root_defs)).unwrap_or_else(permissive_argument_tag), - ) + .map(|(key, value)| { + resolve_local_refs_inner(value, root, seen_refs, depth, remaining) + .map(|value| (key.clone(), value)) }) - .collect(); - star(one_of(tags)) + .collect::>>() + .map(Value::Object) } -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::kimi_uses_declared_tool_schema(tool, strict_schema) { - tool.parameters.as_ref() +// 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 { - 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(), - }), - Format::ConstString(ConstStringFormat { - value: format!("\"{SEP}"), - }), - arguments_block(parameters), - ], - })), - end: CALL_CLOSE.to_string(), + reference + .strip_prefix("#/") + .map(|suffix| format!("{prefix}/{suffix}")) } } -/// Build the format-style xgrammar tag for K3's response + tools channels. -pub(crate) fn build_kimi_k3( - ctx: &ToolCallFormatBuildContext<'_>, -) -> anyhow::Result> { - let (tools, at_least_one) = resolve_tools_to_include(ctx)?; - if tools.is_empty() { - return Ok(None); +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); + } + } } +} - // Moonshot's named-tool contract returns the selected call with no - // assistant content. Leaving the response body as `any_text` lets the model - // put a second, generic `...` representation there - // before emitting the structurally constrained XTML call. Restrict only - // named choice; auto/required may legitimately include response text. - let response_content = if matches!(ctx.tool_choice, crate::tool_calling::ToolChoice::Named(_)) { - Format::ConstString(ConstStringFormat { - value: String::new(), - }) - } else { - // Reserve XTML controls for channel transitions so a direct tools - // channel cannot be swallowed as response text (which masks EOS). - Format::AnyText(AnyTextFormat { - excludes: vec![OPEN.to_string(), CLOSE.to_string()], - }) +fn has_local_schema_refs(schema: &Value) -> bool { + let Some(object) = schema.as_object() else { + return false; }; - // Native output may skip the response channel or go straight from its - // body to tools, so both response markers are optional. - let response = vec![ - optional(Format::ConstString(ConstStringFormat { - value: RESPONSE_OPEN.to_string(), - })), - response_content, - optional(Format::ConstString(ConstStringFormat { - value: RESPONSE_CLOSE.to_string(), - })), - ]; - let calls = Format::TagsWithSeparator(TagsWithSeparatorFormat { - tags: tools - .into_iter() - .map(|tool| call_tag(tool, ctx.strict_schema())) - .collect(), - separator: String::new(), - at_least_one: true, - stop_after_first: ctx.stop_after_first(), - }); - let tools_channel = Format::Tag(TagFormat { - begin: TOOLS_OPEN.to_string(), - content: Box::new(calls), - end: TOOLS_CLOSE.to_string(), - }); - let tools_part = if at_least_one { - tools_channel - } else { - optional(tools_channel) + if object + .get("$ref") + .and_then(Value::as_str) + .is_some_and(|reference| reference == "#" || reference.starts_with("#/")) + { + return true; + } + SCHEMA_MAP_KEYWORDS.iter().any(|key| { + object + .get(*key) + .and_then(Value::as_object) + .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_local_schema_refs(&schema) { + return schema; + } + + let Some(schema_object) = schema.as_object_mut() else { + return schema; }; - let mut elements = response; - elements.push(tools_part); - elements.push(optional(Format::ConstString(ConstStringFormat { - value: MESSAGE_CLOSE.to_string(), - }))); + 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}"); - Ok(Some(StructuralTag { - format: Format::Sequence(SequenceFormat { elements }), - })) + // 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 } -#[cfg(test)] -mod tests { - use serde_json::json; +fn resolve_schema_keyword( + key: &str, + value: &Value, + root: &Value, + seen_refs: &HashSet, + depth: usize, + remaining: &mut usize, +) -> Option { + if SCHEMA_MAP_KEYWORDS.contains(&key) { + return resolve_schema_map(value, root, seen_refs, depth, remaining); + } + if SCHEMA_VALUE_KEYWORDS.contains(&key) { + return resolve_local_refs_inner(value, root, seen_refs, depth, remaining); + } + 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_inner(value, root, seen_refs, depth, remaining)) + .collect::>>() + .map(Value::Array); + } - use super::*; - use crate::tool_calling::structural_tag::builder::StructuralTagSchemaMode; - use crate::tool_calling::{ToolChoice, ToolDefinition}; + // 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 tools() -> Vec { - vec![ - ToolDefinition { - name: "get_weather".to_string(), - parameters: Some(json!({ - "type": "object", - "properties": { - "city": {"type": "string"}, - "days": {"type": "integer"} - }, - "required": ["city"] - })), - strict: None, - }, - ToolDefinition { +fn resolve_local_refs(schema: &Value, root: &Value, seen_refs: &HashSet) -> Option { + let mut remaining = MAX_SCHEMA_NODES; + resolve_local_refs_inner(schema, root, seen_refs, 0, &mut remaining) +} + +fn resolve_local_refs_inner( + schema: &Value, + root: &Value, + seen_refs: &HashSet, + depth: usize, + remaining: &mut usize, +) -> Option { + *remaining = (*remaining).checked_sub(1)?; + if depth >= MAX_SCHEMA_REF_DEPTH { + return None; + } + match schema { + Value::Array(values) => values + .iter() + .map(|value| resolve_local_refs_inner(value, root, seen_refs, depth, remaining)) + .collect::>>() + .map(Value::Array), + Value::Object(object) => { + if has_unsupported_ref_scope(object) { + return None; + } + if let Some(reference) = object.get("$ref").and_then(Value::as_str) { + if seen_refs.contains(reference) { + return Some(schema.clone()); + } + 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_inner(target, root, &nested_seen, depth + 1, remaining)?; + let siblings = object + .iter() + .filter(|(key, _)| key.as_str() != "$ref") + .map(|(key, value)| { + resolve_schema_keyword(key, value, root, seen_refs, depth, remaining) + .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_schema_keyword(key, value, root, seen_refs, depth, remaining) + .map(|value| (key.clone(), value)) + }) + .collect::>>() + .map(Value::Object) + } + _ => Some(schema.clone()), + } +} + +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| 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() +} + +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 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, + }; + 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) +} + +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 { + // 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 Some(original) = schema + .as_object() + .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) + else { + return Ok(None); + }; + if declared_type.xtml_name() != xtml_type { + return Ok(None); + } + let mut original = original.clone(); + if let Some(root) = root.as_object() { + for key in ["$defs", "definitions"] { + if let Some(value) = root.get(key) { + original + .entry(key.to_string()) + .or_insert_with(|| value.clone()); + } + } + } + return Ok(Some(Value::Object(original))); + }; + 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 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) { + narrowed + .entry(key.to_string()) + .or_insert_with(|| value.clone()); + } + } + } + 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) + .with_context(|| format!("cannot constrain argument {key}"))? + else { + continue; + }; + let content = if xtml_type == "string" { + let Some(content) = string_content_format(&schema) + .with_context(|| format!("cannot constrain argument {key}"))? + else { + continue; + }; + 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, +) -> 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( + mut arguments: Vec, + optional_arguments: Option, +) -> Format { + // Guided decoding only needs one schema-valid order. Keep the schema's + // 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, + }) +} + +fn arguments_block(tool: &ToolDefinition, strict_schema: bool) -> anyhow::Result { + if !super::builder::uses_declared_tool_schema(tool, strict_schema) { + return Ok(star(Format::Tag(permissive_argument_tag()))); + } + + let Some(root) = tool.parameters.as_ref() else { + return Ok(star(Format::Tag(permissive_argument_tag()))); + }; + let resolved_root = root + .as_object() + .filter(|object| object.contains_key("$ref") || object.contains_key("allOf")) + .map(|_| resolved_root_parameters(root)) + .transpose()? + .flatten(); + let Some(parameters) = resolved_root.as_ref().unwrap_or(root).as_object() else { + return Ok(star(Format::Tag(permissive_argument_tag()))); + }; + let Some(properties) = parameters.get("properties").and_then(Value::as_object) else { + return Ok(star(Format::Tag(permissive_argument_tag()))); + }; + 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 Ok(star(Format::Tag(permissive_argument_tag()))), + }; + let optional_arguments = optional_arguments(properties, &required, root)?; + + if required.is_empty() { + 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 Ok(star(Format::Tag(permissive_argument_tag()))); + }; + if !matches!(schema, Value::Bool(_) | Value::Object(_)) { + return Ok(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); + } + Ok(canonical_required_arguments(arguments, optional_arguments)) +} + +fn build_auto_structural_tag( + tools: Vec<&ToolDefinition>, + ctx: &ToolCallFormatBuildContext<'_>, +) -> anyhow::Result { + let strict_schema = ctx.strict_schema(); + let call_tags: Vec<_> = tools + .into_iter() + .map(|tool| call_tag(tool, strict_schema)) + .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(), + }; + // 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)), + ], + }); + Ok(StructuralTag { format: suffix }) +} + +fn call_tag(tool: &ToolDefinition, strict_schema: bool) -> anyhow::Result { + let begin = format!("{OPEN}call tool=\"{}\" index=\"", escape_attr(&tool.name)); + Ok(TagFormat { + begin, + content: Box::new(Format::Sequence(SequenceFormat { + elements: vec![ + Format::Regex(RegexFormat { + pattern: CALL_INDEX_PATTERN.to_string(), + }), + Format::ConstString(ConstStringFormat { + value: format!("\"{SEP}"), + }), + 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. +pub(crate) fn build_kimi_k3( + ctx: &ToolCallFormatBuildContext<'_>, +) -> anyhow::Result> { + let (tools, at_least_one) = resolve_tools_to_include(ctx)?; + 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. Leaving the response body as `any_text` lets the model + // put a second, generic `...` representation there + // before emitting the structurally constrained XTML call. Restrict only + // named choice; auto/required may legitimately include response text. + let response_content = if matches!(ctx.tool_choice, crate::tool_calling::ToolChoice::Named(_)) { + Format::ConstString(ConstStringFormat { + value: String::new(), + }) + } else { + // Reserve XTML controls for channel transitions so a direct tools + // channel cannot be swallowed as response text (which masks EOS). + Format::AnyText(AnyTextFormat { + excludes: vec![OPEN.to_string(), CLOSE.to_string()], + }) + }; + // Native output may skip the response channel or go straight from its + // body to tools, so both response markers are optional. + let response = vec![ + optional(Format::ConstString(ConstStringFormat { + value: RESPONSE_OPEN.to_string(), + })), + response_content, + optional(Format::ConstString(ConstStringFormat { + value: RESPONSE_CLOSE.to_string(), + })), + ]; + let calls = Format::TagsWithSeparator(TagsWithSeparatorFormat { + tags: tools + .into_iter() + .map(|tool| call_tag(tool, ctx.strict_schema())) + .collect::>()?, + separator: String::new(), + at_least_one: true, + stop_after_first: ctx.stop_after_first(), + }); + let tools_channel = Format::Tag(TagFormat { + begin: TOOLS_OPEN.to_string(), + content: Box::new(calls), + end: TOOLS_CLOSE.to_string(), + }); + + let tools_part = if at_least_one { + tools_channel + } else { + optional(tools_channel) + }; + let mut elements = response; + elements.push(tools_part); + elements.push(optional(Format::ConstString(ConstStringFormat { + value: MESSAGE_CLOSE.to_string(), + }))); + + Ok(Some(StructuralTag { + format: Format::Sequence(SequenceFormat { elements }), + })) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + use crate::tool_calling::structural_tag::builder::{ + StructuralTagBuilder, StructuralTagSchemaMode, + }; + use crate::tool_calling::{ToolChoice, ToolDefinition}; + + fn tools() -> Vec { + vec![ + ToolDefinition { + name: "get_weather".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": { + "city": {"type": "string"}, + "days": {"type": "integer"} + }, + "required": ["city"] + })), + strict: None, + }, + ToolDefinition { name: "run_command".to_string(), parameters: Some(json!({ "type": "object", @@ -373,29 +1342,19 @@ mod tests { } #[test] - fn response_text_cannot_swallow_xtml_channels() { + fn required_choice_response_text_reserves_xtml_controls() { let tools = tools(); - - 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]; - - assert_eq!( - response_body["type"], "any_text", - "{choice:?} must retain response text" - ); - assert_eq!( - response_body["excludes"], - json!([OPEN, CLOSE]), - "{choice:?} must reserve control markers for channel transitions" - ); - } + 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]; + assert_eq!(response_body["type"], "any_text"); + assert_eq!(response_body["excludes"], json!([OPEN, CLOSE])); } #[test] - fn named_choice_is_mandatory_and_auto_is_optional() { + 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 = @@ -406,27 +1365,1147 @@ 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"][3]["type"], "optional"); + 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"]["elements"][1]["content"]["begin"], + TOOLS_OPEN + ); } #[test] - fn named_choice_keeps_declared_argument_schema() { + fn auto_requires_declared_arguments_then_allows_optional_content() { let tools = tools(); - let choice = ToolChoice::Named("get_weather".to_string()); + let choice = ToolChoice::Auto; let value = serde_json::to_value(build_kimi_k3(&context(&choice, &tools)).unwrap().unwrap()) .unwrap(); - let call = &value["format"]["elements"][3]["content"]["tags"][0]; - let argument_alternatives = &call["content"]["elements"][2]["content"]["elements"]; - assert!( - argument_alternatives - .to_string() - .contains("key=\\\"city\\\"") + 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=\""); + 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!( - argument_alternatives - .to_string() - .contains("key=\\\"days\\\"") + assert_eq!(arguments["elements"][0]["content"]["type"], "any_text"); + assert_eq!( + arguments["elements"][0]["content"]["excludes"], + json!([CLOSE]) + ); + assert_eq!(arguments["elements"][1]["type"], "star"); + assert_eq!( + arguments["elements"][1]["content"]["begin"], + "<|open|>argument key=\"days\" type=\"number\"<|sep|>" + ); + } + + #[test] + fn auto_required_arguments_use_canonical_schema_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"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]; + let elements = arguments["elements"].as_array().unwrap(); + + assert_eq!(arguments["type"], "sequence"); + assert_eq!(elements.len(), 3); + assert_eq!( + elements[0]["begin"], + "<|open|>argument key=\"city\" type=\"string\"<|sep|>" + ); + assert_eq!( + elements[1]["begin"], + "<|open|>argument key=\"days\" type=\"number\"<|sep|>" + ); + assert_eq!(elements[1]["content"]["type"], "json_schema"); + assert_eq!(elements[1]["content"]["json_schema"]["type"], "integer"); + assert_eq!(elements[2]["type"], "star"); + assert_eq!( + elements[2]["content"]["begin"], + "<|open|>argument key=\"units\" type=\"string\"<|sep|>" + ); + } + + #[test] + fn auto_required_non_string_arguments_use_their_json_schemas() { + let tools = vec![ToolDefinition { + name: "typed_tool".to_string(), + parameters: Some(json!({ + "type": "object", + "properties": { + "count": {"type": "integer", "minimum": 1}, + "enabled": {"type": "boolean"}, + "items": {"type": "array", "items": {"type": "string"}}, + "metadata": { + "type": "object", + "properties": {"source": {"type": "string"}}, + "required": ["source"] + } + }, + "required": ["count", "enabled", "items", "metadata"] + })), + 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]; + let elements = &arguments["elements"]; + + for index in 0..4 { + assert_eq!(elements[index]["content"]["type"], "json_schema"); + } + assert_eq!(elements[0]["content"]["json_schema"]["minimum"], 1); + assert_eq!( + elements[2]["content"]["json_schema"]["items"]["type"], + "string" + ); + assert_eq!( + elements[3]["content"]["json_schema"]["required"], + json!(["source"]) + ); + } + + #[test] + fn auto_required_ref_keeps_root_definitions_in_typed_content() { + let tools = vec![ToolDefinition { + name: "lookup".to_string(), + parameters: Some(json!({ + "type": "object", + "$defs": { + "identifier": {"type": "integer", "minimum": 1} + }, + "properties": { + "id": {"$ref": "#/$defs/identifier"} + }, + "required": ["id"] + })), + 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"], "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 string_content_enforces_min_length_without_a_maximum() { + let value = serde_json::to_value( + string_content_format(&json!({ + "type": "string", + "minLength": 2 + })) + .unwrap(), + ) + .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(), + ) + .unwrap(); + + assert_eq!(value["type"], "regex"); + 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(), + ) + .unwrap(); + + assert_eq!(value["type"], "regex"); + 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"]} + ] + })) + .unwrap(); + let value = serde_json::to_value(format).unwrap(); + assert_eq!(value["type"], "const_string"); + assert_eq!(value["value"], "safe"); + } + + #[test] + fn string_outer_bounds_intersect_any_of_enum_restrictions() { + let format = string_content_format(&json!({ + "type": "string", + "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(); + assert_eq!(values.len(), 2); + assert!(values.iter().any(|item| item["value"] == "safe")); + 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) + .unwrap() + .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 { + 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["begin"], + "<|open|>argument key=\"value\" type=\"number\"<|sep|>" + ); + assert_eq!(argument["content"]["type"], "json_schema"); + assert_eq!( + 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\"") + ); + } + + #[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/__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") + .unwrap() + .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!({ + "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() + .unwrap(); + let reference = schema["properties"]["next"]["properties"]["next"]["$ref"] + .as_str() + .unwrap(); + + assert!(resolve_local_ref(reference, &schema).is_some()); + assert!(has_local_schema_refs(&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 array_index_reference_keeps_the_argument_constraint() { + let root = json!({ + "type": "object", + "properties": { + "operations": { + "items": { + "anyOf": [ + {"type": "null"}, + {"properties": {"kind": {"type": "integer", "enum": [7]}}} + ] + } + }, + "marker": {"$ref": "#/properties/operations/items/anyOf/1/properties/kind"} + } + }); + let format = serde_json::to_value( + argument_format("marker", &root["properties"]["marker"], &root) + .unwrap() + .unwrap(), + ) + .unwrap(); + + assert!( + format["begin"] + .as_str() + .unwrap() + .contains("type=\"number\"") + ); + assert_eq!(format["content"]["json_schema"]["enum"], json!([7])); + } + + #[test] + fn percent_encoded_reference_keeps_the_argument_constraint() { + let root = json!({ + "$defs": { + "café name": {"type": "integer", "minimum": 1}, + "a/b~c": {"type": "boolean"} + }, + "properties": { + "id": {"$ref": "#/$defs/caf%C3%A9%20name"} + } + }); + let schema = schema_for_xtml_type(&root["properties"]["id"], &root, "number") + .unwrap() + .unwrap(); + + assert_eq!(schema["type"], "integer"); + assert_eq!(schema["minimum"], 1); + assert_eq!( + resolve_local_ref("#/$defs/a%7E1b%7E0c", &root), + Some(&root["$defs"]["a/b~c"]) + ); + for malformed in ["#/$defs/%FF", "#/$defs/%2", "#/$defs/%GG"] { + assert!(resolve_local_ref(malformed, &root).is_none()); + } + } + + #[test] + fn scoped_reference_keeps_the_original_typed_schema() { + let root = json!({ + "type": "object", + "properties": { + "data": { + "type": "object", + "$id": "nested", + "$defs": {"item": {"type": "integer"}}, + "properties": {"item": {"$ref": "#/$defs/item"}} + } + } + }); + let format = serde_json::to_value( + argument_format("data", &root["properties"]["data"], &root) + .unwrap() + .unwrap(), + ) + .unwrap(); + + assert!( + format["begin"] + .as_str() + .unwrap() + .contains("type=\"object\"") + ); + assert_eq!(format["content"]["json_schema"], root["properties"]["data"]); + } + + #[test] + fn expansion_limit_keeps_the_original_typed_schema() { + let mut properties = Map::new(); + for index in 0..MAX_SCHEMA_NODES { + properties.insert(format!("item_{index}"), json!({"type": "integer"})); + } + let schema = json!({"type": "object", "properties": properties}); + let format = + serde_json::to_value(argument_format("data", &schema, &schema).unwrap().unwrap()) + .unwrap(); + + assert!( + format["begin"] + .as_str() + .unwrap() + .contains("type=\"object\"") + ); + assert_eq!(format["content"]["json_schema"], schema); + } + + #[test] + fn reference_depth_limit_keeps_the_original_typed_schema() { + 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": "integer"}), + ); + let root = json!({ + "$defs": definitions, + "type": "object", + "properties": { + "data": { + "type": "object", + "properties": {"item": {"$ref": "#/$defs/step_0"}} + } + } + }); + 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 { + 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() + .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() + .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() + .unwrap(); + let format = serde_json::to_value(string_content_format(&narrowed).unwrap()).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() + .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"); + + 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 { + 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"]["elements"][1]["content"]["content"]["tags"][0]["content"] + ["elements"][2]; + + assert_eq!(arguments["type"], "star"); + assert_eq!(arguments["content"]["type"], "or"); + assert_eq!( + arguments["content"]["elements"].as_array().unwrap().len(), + 2 + ); + let alternatives = arguments["content"]["elements"].as_array().unwrap(); + 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"); + } + + #[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"]["elements"][1]["content"]["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"]["elements"][1]["content"]["content"]["tags"] + [0]["content"]["elements"][2]; + + assert_eq!(strict_arguments["type"], "sequence"); + assert!(strict_arguments.to_string().contains("key=\\\"city\\\"")); + } + + #[test] + fn auto_thinking_has_one_reasoning_prefix() { + 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 = 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] + 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"][3]["content"]["tags"][0]; + 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_eq!(arguments["elements"][1]["type"], "star"); + assert_eq!( + arguments["elements"][1]["content"]["begin"], + "<|open|>argument key=\"days\" type=\"number\"<|sep|>" ); } diff --git a/parsers/v1/src/tool_calling/structural_tag/tests.rs b/parsers/v1/src/tool_calling/structural_tag/tests.rs index acfcadc3d..b5030f240 100644 --- a/parsers/v1/src/tool_calling/structural_tag/tests.rs +++ b/parsers/v1/src/tool_calling/structural_tag/tests.rs @@ -6,7 +6,7 @@ use serde_json::{Value, json}; use super::TOOL_NAME_PLACEHOLDER; use super::builder::{ StructuralTagBuilder, StructuralTagSchemaMode, ToolCallFormatBuildContext, - kimi_uses_declared_tool_schema, + uses_declared_tool_schema, }; use super::dsml::DsmlToolCallsConfig; use super::format::JsonSchemaStyle; @@ -244,7 +244,7 @@ fn parallel_true_does_not_stop_after_first() { } #[test] -fn non_strict_tools_get_unconstrained_schema() { +fn omitted_strict_uses_declared_schema() { let tools = sample_tools(); let parsed = build_unwrap( &builder(), @@ -255,13 +255,14 @@ fn non_strict_tools_get_unconstrained_schema() { ); let tags = parsed["format"]["tags"].as_array().unwrap(); - assert_eq!(tags[0]["content"]["json_schema"], json!(true)); + assert!(tags[0]["content"]["json_schema"]["properties"]["a"].is_object()); + assert!(tags[1]["content"]["json_schema"]["properties"]["location"].is_object()); } #[test] -fn strict_tool_uses_own_schema() { +fn explicit_strict_false_uses_unconstrained_schema() { let mut tools = sample_tools(); - tools[0].strict = Some(true); + tools[0].strict = Some(false); let parsed = build_unwrap( &builder(), &ToolChoice::Required, @@ -271,13 +272,14 @@ fn strict_tool_uses_own_schema() { ); let tags = parsed["format"]["tags"].as_array().unwrap(); - assert!(tags[0]["content"]["json_schema"]["properties"]["a"].is_object()); - assert_eq!(tags[1]["content"]["json_schema"], json!(true)); + assert_eq!(tags[0]["content"]["json_schema"], json!(true)); + assert!(tags[1]["content"]["json_schema"]["properties"]["location"].is_object()); } #[test] fn strict_schema_mode_uses_real_schema_for_all() { - let tools = sample_tools(); + let mut tools = sample_tools(); + tools[0].strict = Some(false); let parsed = build_unwrap( &builder(), &ToolChoice::Required, @@ -292,18 +294,18 @@ fn strict_schema_mode_uses_real_schema_for_all() { } #[test] -fn kimi_schema_policy_matches_vllm_xgrammar() { +fn schema_policy_treats_only_explicit_false_as_opt_out() { let mut tool = sample_tools().remove(0); tool.strict = None; - assert!(kimi_uses_declared_tool_schema(&tool, false)); + assert!(uses_declared_tool_schema(&tool, false)); tool.strict = Some(true); - assert!(kimi_uses_declared_tool_schema(&tool, false)); + assert!(uses_declared_tool_schema(&tool, false)); tool.strict = Some(false); - assert!(!kimi_uses_declared_tool_schema(&tool, false)); - assert!(kimi_uses_declared_tool_schema(&tool, true)); + assert!(!uses_declared_tool_schema(&tool, false)); + assert!(uses_declared_tool_schema(&tool, true)); } #[test]