@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b2d3a3b9ad

[Rust Frontend] Split argument grammars into schema resolution and per-model rendering (#59395)

Signed-off-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Bugen Zhao committed Oct 1, 2026 at 13:43 UTC b2d3a3b9ad7c58e96a825a33112c571610219460
5 files changed +1050 -427
rust/src/chat/tests/grammar_replay/Kimi-K3.json
+125 -140
@@ -56,21 +56,18 @@
56 "type": "or",
57 "elements": [
58 {
59 - "type": "plus",
59 + "type": "tag",
60 + "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>",
61 "content": {
61 - "type": "tag",
62 - "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>",
63 - "content": {
64 - "type": "any_text",
65 - "excludes": [
66 - "<|close|>argument<|sep|>",
67 - "<|close|>call<|sep|>"
68 - ],
69 - "max_tokens": null,
70 - "max_chars": null
71 - },
72 - "end": "<|close|>argument<|sep|>"
73 - }
62 + "type": "any_text",
63 + "excludes": [
64 + "<|close|>argument<|sep|>",
65 + "<|close|>call<|sep|>"
66 + ],
67 + "max_tokens": null,
68 + "max_chars": null
69 + },
70 + "end": "<|close|>argument<|sep|>"
71 },
72 {
73 "type": "tag",
@@ -119,60 +116,57 @@
116 "type": "or",
117 "elements": [
118 {
122 - "type": "plus",
123 - "content": {
124 - "type": "or",
125 - "elements": [
126 - {
127 - "type": "tag",
128 - "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>",
129 - "content": {
130 - "type": "json_schema",
131 - "json_schema": {
132 - "type": "number"
133 - },
134 - "style": "json",
135 - "any_order": false,
136 - "max_whitespace_cnt": null,
137 - "excludes": []
119 + "type": "sequence",
120 + "elements": [
121 + {
122 + "type": "tag",
123 + "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>",
124 + "content": {
125 + "type": "json_schema",
126 + "json_schema": {
127 + "type": "number"
128 },
139 - "end": "<|close|>argument<|sep|>"
129 + "style": "json",
130 + "any_order": false,
131 + "max_whitespace_cnt": null,
132 + "excludes": []
133 },
141 - {
142 - "type": "tag",
143 - "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>",
144 - "content": {
145 - "type": "json_schema",
146 - "json_schema": {
147 - "type": "number"
148 - },
149 - "style": "json",
150 - "any_order": false,
151 - "max_whitespace_cnt": null,
152 - "excludes": []
134 + "end": "<|close|>argument<|sep|>"
135 + },
136 + {
137 + "type": "tag",
138 + "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>",
139 + "content": {
140 + "type": "json_schema",
141 + "json_schema": {
142 + "type": "number"
143 },
154 - "end": "<|close|>argument<|sep|>"
144 + "style": "json",
145 + "any_order": false,
146 + "max_whitespace_cnt": null,
147 + "excludes": []
148 },
156 - {
157 - "type": "tag",
158 - "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>",
159 - "content": {
160 - "type": "json_schema",
161 - "json_schema": {
162 - "type": "array",
163 - "items": {
164 - "type": "string"
165 - }
166 - },
167 - "style": "json",
168 - "any_order": false,
169 - "max_whitespace_cnt": null,
170 - "excludes": []
149 + "end": "<|close|>argument<|sep|>"
150 + },
151 + {
152 + "type": "tag",
153 + "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>",
154 + "content": {
155 + "type": "json_schema",
156 + "json_schema": {
157 + "type": "array",
158 + "items": {
159 + "type": "string"
160 + }
161 },
172 - "end": "<|close|>argument<|sep|>"
173 - }
174 - ]
175 - }
162 + "style": "json",
163 + "any_order": false,
164 + "max_whitespace_cnt": null,
165 + "excludes": []
166 + },
167 + "end": "<|close|>argument<|sep|>"
168 + }
169 + ]
170 },
171 {
172 "type": "tag",
@@ -287,21 +281,18 @@
281 "type": "or",
282 "elements": [
283 {
290 - "type": "plus",
284 + "type": "tag",
285 + "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>",
286 "content": {
292 - "type": "tag",
293 - "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>",
294 - "content": {
295 - "type": "any_text",
296 - "excludes": [
297 - "<|close|>argument<|sep|>",
298 - "<|close|>call<|sep|>"
299 - ],
300 - "max_tokens": null,
301 - "max_chars": null
302 - },
303 - "end": "<|close|>argument<|sep|>"
304 - }
287 + "type": "any_text",
288 + "excludes": [
289 + "<|close|>argument<|sep|>",
290 + "<|close|>call<|sep|>"
291 + ],
292 + "max_tokens": null,
293 + "max_chars": null
294 + },
295 + "end": "<|close|>argument<|sep|>"
296 },
297 {
298 "type": "tag",
@@ -350,60 +341,57 @@
341 "type": "or",
342 "elements": [
343 {
353 - "type": "plus",
354 - "content": {
355 - "type": "or",
356 - "elements": [
357 - {
358 - "type": "tag",
359 - "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>",
360 - "content": {
361 - "type": "json_schema",
362 - "json_schema": {
363 - "type": "number"
364 - },
365 - "style": "json",
366 - "any_order": false,
367 - "max_whitespace_cnt": null,
368 - "excludes": []
344 + "type": "sequence",
345 + "elements": [
346 + {
347 + "type": "tag",
348 + "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>",
349 + "content": {
350 + "type": "json_schema",
351 + "json_schema": {
352 + "type": "number"
353 },
370 - "end": "<|close|>argument<|sep|>"
354 + "style": "json",
355 + "any_order": false,
356 + "max_whitespace_cnt": null,
357 + "excludes": []
358 },
372 - {
373 - "type": "tag",
374 - "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>",
375 - "content": {
376 - "type": "json_schema",
377 - "json_schema": {
378 - "type": "number"
379 - },
380 - "style": "json",
381 - "any_order": false,
382 - "max_whitespace_cnt": null,
383 - "excludes": []
359 + "end": "<|close|>argument<|sep|>"
360 + },
361 + {
362 + "type": "tag",
363 + "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>",
364 + "content": {
365 + "type": "json_schema",
366 + "json_schema": {
367 + "type": "number"
368 },
385 - "end": "<|close|>argument<|sep|>"
369 + "style": "json",
370 + "any_order": false,
371 + "max_whitespace_cnt": null,
372 + "excludes": []
373 },
387 - {
388 - "type": "tag",
389 - "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>",
390 - "content": {
391 - "type": "json_schema",
392 - "json_schema": {
393 - "type": "array",
394 - "items": {
395 - "type": "string"
396 - }
397 - },
398 - "style": "json",
399 - "any_order": false,
400 - "max_whitespace_cnt": null,
401 - "excludes": []
374 + "end": "<|close|>argument<|sep|>"
375 + },
376 + {
377 + "type": "tag",
378 + "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>",
379 + "content": {
380 + "type": "json_schema",
381 + "json_schema": {
382 + "type": "array",
383 + "items": {
384 + "type": "string"
385 + }
386 },
403 - "end": "<|close|>argument<|sep|>"
404 - }
405 - ]
406 - }
387 + "style": "json",
388 + "any_order": false,
389 + "max_whitespace_cnt": null,
390 + "excludes": []
391 + },
392 + "end": "<|close|>argument<|sep|>"
393 + }
394 + ]
395 },
396 {
397 "type": "tag",
@@ -517,21 +505,18 @@
505 "type": "or",
506 "elements": [
507 {
520 - "type": "plus",
508 + "type": "tag",
509 + "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>",
510 "content": {
522 - "type": "tag",
523 - "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>",
524 - "content": {
525 - "type": "any_text",
526 - "excludes": [
527 - "<|close|>argument<|sep|>",
528 - "<|close|>call<|sep|>"
529 - ],
530 - "max_tokens": null,
531 - "max_chars": null
532 - },
533 - "end": "<|close|>argument<|sep|>"
534 - }
511 + "type": "any_text",
512 + "excludes": [
513 + "<|close|>argument<|sep|>",
514 + "<|close|>call<|sep|>"
515 + ],
516 + "max_tokens": null,
517 + "max_chars": null
518 + },
519 + "end": "<|close|>argument<|sep|>"
520 },
521 {
522 "type": "tag",
rust/src/chat/tests/grammar_replay/Kimi-K3.txt
+11 -13
@@ -10,18 +10,17 @@ sequence
10 /[1-9][0-9]*/
11 `"<|sep|>`
12 or
13 - plus tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>`
13 + tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>`
14 tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>`
15 tag `<|open|>call tool="add" index="` .. `<|close|>call<|sep|>`
16 sequence
17 /[1-9][0-9]*/
18 `"<|sep|>`
19 or
20 - plus
21 - or
22 - tag `<|open|>argument key="y" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>`
23 - tag `<|open|>argument key="x" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>`
24 - tag `<|open|>argument key="items" type="array"<|sep|>` json(string[]) `<|close|>argument<|sep|>`
20 + sequence
21 + tag `<|open|>argument key="y" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>`
22 + tag `<|open|>argument key="x" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>`
23 + tag `<|open|>argument key="items" type="array"<|sep|>` json(string[]) `<|close|>argument<|sep|>`
24 tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>`
25 optional `<|close|>message<|sep|>`
26
@@ -36,18 +35,17 @@ sequence
35 /[1-9][0-9]*/
36 `"<|sep|>`
37 or
39 - plus tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>`
38 + tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>`
39 tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>`
40 tag `<|open|>call tool="add" index="` .. `<|close|>call<|sep|>`
41 sequence
42 /[1-9][0-9]*/
43 `"<|sep|>`
44 or
46 - plus
47 - or
48 - tag `<|open|>argument key="y" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>`
49 - tag `<|open|>argument key="x" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>`
50 - tag `<|open|>argument key="items" type="array"<|sep|>` json(string[]) `<|close|>argument<|sep|>`
45 + sequence
46 + tag `<|open|>argument key="y" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>`
47 + tag `<|open|>argument key="x" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>`
48 + tag `<|open|>argument key="items" type="array"<|sep|>` json(string[]) `<|close|>argument<|sep|>`
49 tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>`
50 optional `<|close|>message<|sep|>`
51
@@ -62,6 +60,6 @@ sequence
60 /[1-9][0-9]*/
61 `"<|sep|>`
62 or
65 - plus tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>`
63 + tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>`
64 tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>`
65 optional `<|close|>message<|sep|>`
rust/src/parser/src/output_grammar.rs
+1
@@ -17,6 +17,7 @@ use xgrammar_structural_tag::{
17
18 use crate::tool::Tool;
19
20 +pub mod arguments;
21 #[cfg(any(test, feature = "test-util"))]
22 pub mod test_utils;
23
rust/src/parser/src/output_grammar/arguments.rs new
+845
@@ -0,0 +1,845 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +//! Tool-call argument grammars for keyed-parameter protocols.
5 +//!
6 +//! Many tool-call protocols render the root arguments object as a sequence of
7 +//! keyed parameters. [`arguments`] owns what follows from the JSON schema
8 +//! alone: which parameters exist, their order and presence, and every option
9 +//! a parameter value can take, with `$ref`s resolved and unions, type arrays,
10 +//! and enums split per option. An [`ArgumentSyntax`], implemented next to each
11 +//! parser, decides how one parameter and its value options are rendered.
12 +//! Values nested below a parameter stay with XGrammar's JSON schema converter
13 +//! through [`ValueOption::json`].
14 +//!
15 +//! This is for protocols whose argument rendering XGrammar's per-model schema
16 +//! styles do not cover, or cover without matching the parser. Parameter lists
17 +//! follow those styles' fixed-order property semantics, with these deliberate
18 +//! differences:
19 +//!
20 +//! - Unions, type arrays, and mixed enums split into one option per
21 +//! alternative, so a syntax can keep a value's type label consistent with
22 +//! the value.
23 +//! - Additional-property keys are not required to differ from declared keys.
24 +//! - String `format` and length constraints are not enforced. `pattern` is.
25 +//! - `minProperties` and `maxProperties` are not enforced.
26 +//! - `allOf` with more than one schema accepts any value.
27 +
28 +use std::collections::HashSet;
29 +
30 +use serde_json::{Map, Value, json};
31 +use xgrammar_structural_tag::format::{Format, JsonSchemaFormat};
32 +
33 +/// Bound on `$ref` and combinator nesting within one value, which keeps a long
34 +/// chain of references from overflowing the stack.
35 +const MAX_SCHEMA_DEPTH: usize = 32;
36 +
37 +/// The schema accepting any value.
38 +static ANY_SCHEMA: Value = Value::Bool(true);
39 +
40 +/// How one protocol renders the parameters of a call.
41 +pub trait ArgumentSyntax {
42 + /// Text around and between parameters, or `None` when parameters are
43 + /// adjacent.
44 + fn separator(&self) -> Option<Format>;
45 +
46 + /// The grammar of one parameter whose value takes one of `options`, or
47 + /// `None` when none of them can be rendered.
48 + fn parameter(&self, key: ParameterKey<'_>, options: &[ValueOption<'_>]) -> Option<Format>;
49 +}
50 +
51 +/// The key of one parameter.
52 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53 +pub enum ParameterKey<'a> {
54 + /// A property declared by the schema.
55 + Declared(&'a str),
56 + /// Any key allowed by `additionalProperties` or an unconstrained schema.
57 + Free,
58 +}
59 +
60 +impl ParameterKey<'_> {
61 + /// Pattern of free keys, as in XGrammar's XML styles.
62 + pub const FREE_PATTERN: &'static str = "[a-zA-Z_][a-zA-Z0-9_]*";
63 +
64 + /// `prefix KEY suffix content end`: a tag for a declared key, or a
65 + /// sequence matching [`Self::FREE_PATTERN`] for a free key.
66 + pub fn tag(&self, prefix: &str, suffix: &str, content: Format, end: &str) -> Format {
67 + match self {
68 + Self::Declared(key) => Format::tag(format!("{prefix}{key}{suffix}"), content, end),
69 + Self::Free => Format::sequence(vec![
70 + Format::const_string(prefix),
71 + Format::regex(Self::FREE_PATTERN),
72 + Format::const_string(suffix),
73 + content,
74 + Format::const_string(end),
75 + ]),
76 + }
77 + }
78 +}
79 +
80 +/// JSON type of a value option.
81 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
82 +pub enum JsonType {
83 + String,
84 + Number,
85 + Boolean,
86 + Null,
87 + Object,
88 + Array,
89 +}
90 +
91 +impl JsonType {
92 + const ALL: [Self; 6] = [
93 + Self::String,
94 + Self::Number,
95 + Self::Boolean,
96 + Self::Null,
97 + Self::Object,
98 + Self::Array,
99 + ];
100 +
101 + /// The type's JSON name. Integers are numbers.
102 + pub fn name(self) -> &'static str {
103 + match self {
104 + Self::String => "string",
105 + Self::Number => "number",
106 + Self::Boolean => "boolean",
107 + Self::Null => "null",
108 + Self::Object => "object",
109 + Self::Array => "array",
110 + }
111 + }
112 +
113 + fn of_schema_type(name: &str) -> Option<Self> {
114 + Some(match name {
115 + "string" => Self::String,
116 + "integer" | "number" => Self::Number,
117 + "boolean" => Self::Boolean,
118 + "null" => Self::Null,
119 + "object" => Self::Object,
120 + "array" => Self::Array,
121 + _ => return None,
122 + })
123 + }
124 +
125 + fn of_value(value: &Value) -> Self {
126 + match value {
127 + Value::String(_) => Self::String,
128 + Value::Number(_) => Self::Number,
129 + Value::Bool(_) => Self::Boolean,
130 + Value::Null => Self::Null,
131 + Value::Object(_) => Self::Object,
132 + Value::Array(_) => Self::Array,
133 + }
134 + }
135 +
136 + /// Schema accepting every JSON value of this type.
137 + fn any_schema(self) -> Value {
138 + match self {
139 + Self::Object => json!({ "type": "object", "additionalProperties": true }),
140 + Self::Array => json!({ "type": "array", "items": true }),
141 + _ => json!({ "type": self.name() }),
142 + }
143 + }
144 +}
145 +
146 +/// One option of a parameter value: a single JSON type under a narrowed
147 +/// schema, a constant, or any value of the type.
148 +pub struct ValueOption<'a> {
149 + /// The option's JSON type.
150 + pub ty: JsonType,
151 + source: OptionSource,
152 + cx: &'a ArgumentContext<'a>,
153 +}
154 +
155 +enum OptionSource {
156 + /// A schema narrowed to [`ValueOption::ty`].
157 + Schema(Map<String, Value>),
158 + /// A `const` or `enum` value.
159 + Literal(Value),
160 + /// Any value of [`ValueOption::ty`].
161 + Any,
162 +}
163 +
164 +impl ValueOption<'_> {
165 + /// The option as JSON text.
166 + pub fn json(&self) -> Format {
167 + match &self.source {
168 + OptionSource::Schema(schema) => self.cx.json(Value::Object(schema.clone())),
169 + OptionSource::Literal(value) => self.cx.json(json!({ "const": value })),
170 + OptionSource::Any => self.cx.json(self.ty.any_schema()),
171 + }
172 + }
173 +
174 + /// The option as raw text that ends before any of `terminators`, or `None`
175 + /// for a non-string option or a constant containing a terminator or an
176 + /// excluded substring.
177 + pub fn raw_string(&self, terminators: &[&str]) -> Option<Format> {
178 + if self.ty != JsonType::String {
179 + return None;
180 + }
181 + let excludes = self
182 + .cx
183 + .options
184 + .excludes
185 + .iter()
186 + .map(String::as_str)
187 + .chain(terminators.iter().copied())
188 + .collect::<Vec<_>>();
189 + let text = || Format::any_text_excluding(&excludes);
190 + Some(match &self.source {
191 + OptionSource::Literal(Value::String(value)) => {
192 + if excludes.iter().any(|exclude| value.contains(exclude)) {
193 + return None;
194 + }
195 + Format::const_string(value)
196 + }
197 + OptionSource::Literal(_) => unreachable!("literal options are typed by their value"),
198 + OptionSource::Schema(schema) => match schema.get("pattern").and_then(Value::as_str) {
199 + Some(pattern) => Format::regex(pattern),
200 + None => text(),
201 + },
202 + OptionSource::Any => text(),
203 + })
204 + }
205 +}
206 +
207 +/// Request options that also apply to nested JSON schema regions.
208 +#[derive(Debug, Clone, Default, PartialEq, Eq)]
209 +pub struct ArgumentOptions {
210 + /// Whether object properties may appear in any order, at the parameter
211 + /// level and in nested objects.
212 + pub any_order: bool,
213 + /// Maximum consecutive whitespace characters, or no limit when unset.
214 + pub max_whitespace_cnt: Option<i32>,
215 + /// Substrings forbidden in string values, nested JSON strings included.
216 + pub excludes: Vec<String>,
217 +}
218 +
219 +/// Reference targets expanded while collecting one value's options.
220 +#[derive(Default)]
221 +struct Expansion {
222 + expanded: HashSet<*const Value>,
223 + /// Whether a reference was skipped because its target already expanded.
224 + skipped: bool,
225 +}
226 +
227 +/// Request state shared by the value options of one call.
228 +struct ArgumentContext<'a> {
229 + options: &'a ArgumentOptions,
230 + /// The parameters schema, which local `$ref`s resolve against.
231 + root: &'a Value,
232 +}
233 +
234 +impl ArgumentContext<'_> {
235 + /// A JSON value under `schema`, which keeps the root's definitions so its
236 + /// local `$ref`s still resolve.
237 + fn json(&self, mut schema: Value) -> Format {
238 + if let (Some(schema), Some(root)) = (schema.as_object_mut(), self.root.as_object()) {
239 + for key in ["$defs", "definitions"] {
240 + if let Some(definitions) = root.get(key) {
241 + schema.entry(key).or_insert_with(|| definitions.clone());
242 + }
243 + }
244 + }
245 + Format::JsonSchema(JsonSchemaFormat {
246 + excludes: self.options.excludes.clone(),
247 + ..JsonSchemaFormat::new(schema)
248 + .with_any_order(self.options.any_order)
249 + .with_max_whitespace_cnt(self.options.max_whitespace_cnt)
250 + })
251 + }
252 +
253 + fn resolve_ref(&self, reference: &str) -> Option<&Value> {
254 + self.root.pointer(reference.strip_prefix('#')?)
255 + }
256 +
257 + /// Follow `$ref`s and single-schema `allOf`s at the root.
258 + fn resolve<'s>(&'s self, schema: &'s Value, depth: usize) -> &'s Value {
259 + if depth > MAX_SCHEMA_DEPTH {
260 + return &ANY_SCHEMA;
261 + }
262 + if let Some(reference) = schema.get("$ref").and_then(Value::as_str) {
263 + return match self.resolve_ref(reference) {
264 + Some(target) => self.resolve(target, depth + 1),
265 + None => &ANY_SCHEMA,
266 + };
267 + }
268 + match schema.get("allOf").and_then(Value::as_array).map(Vec::as_slice) {
269 + Some([schema]) => self.resolve(schema, depth + 1),
270 + _ => schema,
271 + }
272 + }
273 +
274 + /// Every option of a value under `schema`.
275 + ///
276 + /// Each reference target expands at most once. A union gains no values
277 + /// from a repeated alternative, and a definition shared by several
278 + /// alternatives would otherwise expand once per path to it, exponentially
279 + /// in the nesting. A schema whose references only lead back into it
280 + /// accepts any value.
281 + fn value_options(&self, schema: &Value) -> Vec<ValueOption<'_>> {
282 + let mut expansion = Expansion::default();
283 + let options = self.value_options_at(schema, 0, &mut expansion);
284 + if options.is_empty() && expansion.skipped {
285 + return self.any_value();
286 + }
287 + options
288 + }
289 +
290 + fn value_options_at(
291 + &self,
292 + schema: &Value,
293 + depth: usize,
294 + expansion: &mut Expansion,
295 + ) -> Vec<ValueOption<'_>> {
296 + let schema = match schema {
297 + _ if depth > MAX_SCHEMA_DEPTH => return self.any_value(),
298 + Value::Bool(false) => return vec![],
299 + Value::Object(schema) => schema,
300 + _ => return self.any_value(),
301 + };
302 + if let Some(reference) = schema.get("$ref").and_then(Value::as_str) {
303 + let Some(target) = self.resolve_ref(reference) else {
304 + return self.any_value();
305 + };
306 + if !expansion.expanded.insert(std::ptr::from_ref(target)) {
307 + expansion.skipped = true;
308 + return vec![];
309 + }
310 + return self.value_options_at(target, depth + 1, expansion);
311 + }
312 + if let Some(value) = schema.get("const") {
313 + return vec![self.literal(value)];
314 + }
315 + if let Some(values) = schema.get("enum").and_then(Value::as_array) {
316 + return values.iter().map(|value| self.literal(value)).collect();
317 + }
318 + if let Some(options) =
319 + schema.get("anyOf").or_else(|| schema.get("oneOf")).and_then(Value::as_array)
320 + {
321 + return options
322 + .iter()
323 + .flat_map(|option| self.value_options_at(option, depth + 1, expansion))
324 + .collect();
325 + }
326 + if let Some(schemas) = schema.get("allOf").and_then(Value::as_array) {
327 + return match schemas.as_slice() {
328 + [schema] => self.value_options_at(schema, depth + 1, expansion),
329 + _ => self.any_value(),
330 + };
331 + }
332 + match schema.get("type") {
333 + Some(Value::String(name)) => self.typed(schema.clone(), name).into_iter().collect(),
334 + Some(Value::Array(names)) => names
335 + .iter()
336 + .filter_map(Value::as_str)
337 + .filter_map(|name| {
338 + let mut schema = schema.clone();
339 + schema.insert("type".to_string(), Value::String(name.to_string()));
340 + self.typed(schema, name)
341 + })
342 + .collect(),
343 + _ => self.any_value(),
344 + }
345 + }
346 +
347 + fn typed(&self, schema: Map<String, Value>, type_name: &str) -> Option<ValueOption<'_>> {
348 + Some(ValueOption {
349 + ty: JsonType::of_schema_type(type_name)?,
350 + source: OptionSource::Schema(schema),
351 + cx: self,
352 + })
353 + }
354 +
355 + fn literal(&self, value: &Value) -> ValueOption<'_> {
356 + ValueOption {
357 + ty: JsonType::of_value(value),
358 + source: OptionSource::Literal(value.clone()),
359 + cx: self,
360 + }
361 + }
362 +
363 + fn any_value(&self) -> Vec<ValueOption<'_>> {
364 + JsonType::ALL
365 + .into_iter()
366 + .map(|ty| ValueOption {
367 + ty,
368 + source: OptionSource::Any,
369 + cx: self,
370 + })
371 + .collect()
372 + }
373 +}
374 +
375 +/// Build the grammar of one call's arguments under `schema`, the tool's
376 +/// parameters schema, rendering each parameter with `syntax`.
377 +pub fn arguments(schema: &Value, syntax: &dyn ArgumentSyntax, options: &ArgumentOptions) -> Format {
378 + let cx = ArgumentContext {
379 + options,
380 + root: schema,
381 + };
382 + let parameter =
383 + |key: ParameterKey<'_>, schema: &Value| syntax.parameter(key, &cx.value_options(schema));
384 +
385 + let schema = cx.resolve(schema, 0);
386 + let (properties, additional) = match schema {
387 + Value::Object(schema) => {
388 + let additional = match schema.get("additionalProperties") {
389 + Some(Value::Bool(false)) => None,
390 + Some(additional) => parameter(ParameterKey::Free, additional),
391 + // Strict schemas allow no undeclared properties, and a schema
392 + // without declared properties allows any.
393 + None if schema.contains_key("properties") => None,
394 + None => parameter(ParameterKey::Free, &ANY_SCHEMA),
395 + };
396 + let required = schema
397 + .get("required")
398 + .and_then(Value::as_array)
399 + .map(|required| required.iter().filter_map(Value::as_str).collect::<Vec<_>>())
400 + .unwrap_or_default();
401 + let properties = schema
402 + .get("properties")
403 + .and_then(Value::as_object)
404 + .into_iter()
405 + .flatten()
406 + .filter_map(|(key, schema)| {
407 + let parameter = parameter(ParameterKey::Declared(key), schema)?;
408 + Some((parameter, required.contains(&key.as_str())))
409 + })
410 + .collect();
411 + (properties, additional)
412 + }
413 + Value::Bool(false) => (vec![], None),
414 + _ => (vec![], parameter(ParameterKey::Free, &ANY_SCHEMA)),
415 + };
416 + parameter_list(options, syntax.separator(), properties, additional)
417 +}
418 +
419 +/// Parameters in declared order, required ones mandatory, followed by any
420 +/// number of additional parameters; or, under `any_order`, any sequence of
421 +/// them with at least one when some are required.
422 +fn parameter_list(
423 + options: &ArgumentOptions,
424 + separator: Option<Format>,
425 + properties: Vec<(Format, bool)>,
426 + additional: Option<Format>,
427 +) -> Format {
428 + let separated = |parameter: Format| match &separator {
429 + Some(separator) => Format::sequence(vec![parameter, separator.clone()]),
430 + None => parameter,
431 + };
432 + let mut elements = Vec::from_iter(separator.clone());
433 + if options.any_order {
434 + let any_required = properties.iter().any(|(_, required)| *required);
435 + let parameters = properties
436 + .into_iter()
437 + .map(|(parameter, _)| parameter)
438 + .chain(additional)
439 + .collect::<Vec<_>>();
440 + if !parameters.is_empty() {
441 + let parameter = separated(one_of(parameters));
442 + elements.push(if any_required {
443 + Format::plus(parameter)
444 + } else {
445 + Format::star(parameter)
446 + });
447 + }
448 + } else {
449 + for (parameter, required) in properties {
450 + let parameter = separated(parameter);
451 + elements.push(if required {
452 + parameter
453 + } else {
454 + Format::optional(parameter)
455 + });
456 + }
457 + if let Some(additional) = additional {
458 + elements.push(Format::star(separated(additional)));
459 + }
460 + }
461 + match elements.len() {
462 + 0 => Format::const_string(""),
463 + 1 => elements.pop().expect("one element"),
464 + _ => Format::sequence(elements),
465 + }
466 +}
467 +
468 +/// One of `formats`, without an `or` around a single format.
469 +pub fn one_of(mut formats: Vec<Format>) -> Format {
470 + if formats.len() == 1 {
471 + formats.pop().expect("one format")
472 + } else {
473 + Format::or(formats)
474 + }
475 +}
476 +
477 +/// `options` grouped by `key`, in order of first appearance.
478 +pub fn group_by<'o, 'a, K: PartialEq>(
479 + options: &'o [ValueOption<'a>],
480 + key: impl Fn(&ValueOption<'a>) -> K,
481 +) -> Vec<(K, Vec<&'o ValueOption<'a>>)> {
482 + let mut groups: Vec<(K, Vec<&ValueOption<'a>>)> = Vec::new();
483 + for option in options {
484 + let key = key(option);
485 + match groups.iter_mut().find(|(existing, _)| *existing == key) {
486 + Some((_, group)) => group.push(option),
487 + None => groups.push((key, vec![option])),
488 + }
489 + }
490 + groups
491 +}
492 +
493 +#[cfg(test)]
494 +mod tests {
495 + use super::*;
496 + use crate::output_grammar::test_utils::outline;
497 + use expect_test::{Expect, expect};
498 + use serde_json::json;
499 +
500 + /// Qwen-like: `<parameter=KEY>VALUE</parameter>` on separate lines, with
501 + /// raw strings.
502 + struct Untyped;
503 +
504 + impl ArgumentSyntax for Untyped {
505 + fn separator(&self) -> Option<Format> {
506 + Some(Format::const_string("\n"))
507 + }
508 +
509 + fn parameter(&self, key: ParameterKey<'_>, options: &[ValueOption<'_>]) -> Option<Format> {
510 + let values = options
511 + .iter()
512 + .filter_map(|option| match option.ty {
513 + JsonType::String => option.raw_string(&["</parameter>"]),
514 + _ => Some(option.json()),
515 + })
516 + .collect::<Vec<_>>();
517 + (!values.is_empty())
518 + .then(|| key.tag("<parameter=", ">", one_of(values), "</parameter>"))
519 + }
520 + }
521 +
522 + /// K3-like: `<arg key="KEY" type="TYPE">VALUE</arg>` with adjacent
523 + /// parameters and one header per JSON type.
524 + struct Typed;
525 +
526 + impl ArgumentSyntax for Typed {
527 + fn separator(&self) -> Option<Format> {
528 + None
529 + }
530 +
531 + fn parameter(&self, key: ParameterKey<'_>, options: &[ValueOption<'_>]) -> Option<Format> {
532 + let tags = group_by(options, |option| option.ty)
533 + .into_iter()
534 + .filter_map(|(ty, options)| {
535 + let values = options
536 + .into_iter()
537 + .filter_map(|option| match ty {
538 + JsonType::String => option.raw_string(&["</arg>"]),
539 + _ => Some(option.json()),
540 + })
541 + .collect::<Vec<_>>();
542 + (!values.is_empty()).then(|| {
543 + let suffix = format!("\" type=\"{}\">", ty.name());
544 + key.tag("<arg key=\"", &suffix, one_of(values), "</arg>")
545 + })
546 + })
547 + .collect::<Vec<_>>();
548 + (!tags.is_empty()).then(|| one_of(tags))
549 + }
550 + }
551 +
552 + fn check(schema: Value, syntax: &dyn ArgumentSyntax, expected: Expect) {
553 + check_with(schema, syntax, &ArgumentOptions::default(), expected);
554 + }
555 +
556 + fn check_with(
557 + schema: Value,
558 + syntax: &dyn ArgumentSyntax,
559 + options: &ArgumentOptions,
560 + expected: Expect,
561 + ) {
562 + expected.assert_eq(&outline(&arguments(&schema, syntax, options)));
563 + }
564 +
565 + fn weather_schema() -> Value {
566 + json!({
567 + "type": "object",
568 + "properties": {
569 + "location": { "type": "string" },
570 + "unit": { "enum": ["celsius", "fahrenheit"] },
571 + "days": { "type": "integer", "minimum": 1 }
572 + },
573 + "required": ["location"]
574 + })
575 + }
576 +
577 + #[test]
578 + fn parameters_follow_declared_order() {
579 + check(
580 + weather_schema(),
581 + &Untyped,
582 + expect![[r#"
583 + sequence
584 + `\n`
585 + sequence
586 + tag `<parameter=location>` text excluding [`</parameter>`] `</parameter>`
587 + `\n`
588 + optional
589 + sequence
590 + tag `<parameter=unit>` .. `</parameter>`
591 + or
592 + `celsius`
593 + `fahrenheit`
594 + `\n`
595 + optional
596 + sequence
597 + tag `<parameter=days>` json(integer(minimum=1)) `</parameter>`
598 + `\n`
599 + "#]],
600 + );
601 + }
602 +
603 + #[test]
604 + fn any_order_repeats_any_declared_parameter() {
605 + check_with(
606 + weather_schema(),
607 + &Typed,
608 + &ArgumentOptions {
609 + any_order: true,
610 + ..Default::default()
611 + },
612 + expect![[r#"
613 + plus
614 + or
615 + tag `<arg key="location" type="string">` text excluding [`</arg>`] `</arg>`
616 + tag `<arg key="unit" type="string">` .. `</arg>`
617 + or
618 + `celsius`
619 + `fahrenheit`
620 + tag `<arg key="days" type="number">` json(integer(minimum=1)) any_order `</arg>`
621 + "#]],
622 + );
623 + }
624 +
625 + #[test]
626 + fn unions_and_mixed_enums_split_per_option() {
627 + check(
628 + json!({
629 + "type": "object",
630 + "properties": {
631 + "id": { "type": ["integer", "string", "null"] },
632 + "mode": { "enum": ["fast", 1] },
633 + "q": { "anyOf": [{ "type": "string", "pattern": "[a-z]+" }, { "type": "array", "items": { "type": "string" } }] }
634 + },
635 + "required": ["q"]
636 + }),
637 + &Typed,
638 + expect![[r#"
639 + sequence
640 + optional
641 + or
642 + tag `<arg key="id" type="number">` json(integer) `</arg>`
643 + tag `<arg key="id" type="string">` text excluding [`</arg>`] `</arg>`
644 + tag `<arg key="id" type="null">` json(null) `</arg>`
645 + optional
646 + or
647 + tag `<arg key="mode" type="string">` `fast` `</arg>`
648 + tag `<arg key="mode" type="number">` json(1) `</arg>`
649 + or
650 + tag `<arg key="q" type="string">` /[a-z]+/ `</arg>`
651 + tag `<arg key="q" type="array">` json(string[]) `</arg>`
652 + "#]],
653 + );
654 + }
655 +
656 + #[test]
657 + fn references_resolve_against_the_root_definitions() {
658 + let schema = json!({
659 + "$defs": {
660 + "place": { "type": "object", "properties": { "city": { "type": "string" } } },
661 + "name": { "type": "string" }
662 + },
663 + "type": "object",
664 + "properties": {
665 + "place": { "$ref": "#/$defs/place" },
666 + "name": { "$ref": "#/$defs/name" }
667 + },
668 + "required": ["place", "name"]
669 + });
670 + let arguments = arguments(&schema, &Typed, &ArgumentOptions::default());
671 + expect![[r#"
672 + sequence
673 + tag `<arg key="place" type="object">` json({ city?: string } where place = { city?: string }, name = string) `</arg>`
674 + tag `<arg key="name" type="string">` text excluding [`</arg>`] `</arg>`
675 + "#]].assert_eq(&outline(&arguments));
676 +
677 + // Nested JSON regions keep the definitions for their own references.
678 + let serialized = serde_json::to_value(&arguments).unwrap();
679 + assert_eq!(
680 + serialized["elements"][0]["content"]["json_schema"]["$defs"],
681 + schema["$defs"]
682 + );
683 + }
684 +
685 + #[test]
686 + fn shared_definitions_expand_once() {
687 + // Each level unions four references to the next, so expanding every
688 + // path would yield 4^12 options.
689 + let mut definitions = Map::new();
690 + for level in 0..12 {
691 + let next = json!({ "$ref": format!("#/$defs/level{}", level + 1) });
692 + definitions.insert(format!("level{level}"), json!({ "anyOf": vec![next; 4] }));
693 + }
694 + definitions.insert("level12".to_string(), json!({ "type": "string" }));
695 + check(
696 + json!({
697 + "$defs": definitions,
698 + "type": "object",
699 + "properties": { "p": { "$ref": "#/$defs/level0" } },
700 + "required": ["p"]
701 + }),
702 + &Untyped,
703 + expect![[r#"
704 + sequence
705 + `\n`
706 + sequence
707 + tag `<parameter=p>` text excluding [`</parameter>`] `</parameter>`
708 + `\n`
709 + "#]],
710 + );
711 + }
712 +
713 + #[test]
714 + fn references_that_only_cycle_accept_any_value() {
715 + check(
716 + json!({
717 + "$defs": {
718 + "a": { "anyOf": [{ "$ref": "#/$defs/b" }] },
719 + "b": { "$ref": "#/$defs/a" }
720 + },
721 + "type": "object",
722 + "properties": { "p": { "$ref": "#/$defs/a" } },
723 + "required": ["p"]
724 + }),
725 + &Untyped,
726 + expect![[r#"
727 + sequence
728 + `\n`
729 + sequence
730 + tag `<parameter=p>` .. `</parameter>`
731 + or
732 + text excluding [`</parameter>`]
733 + json(number where a = b, b = a)
734 + json(boolean where a = b, b = a)
735 + json(null where a = b, b = a)
736 + json({ ... } where a = b, b = a)
737 + json(any[] where a = b, b = a)
738 + `\n`
739 + "#]],
740 + );
741 + }
742 +
743 + #[test]
744 + fn unconstrained_arguments_accept_any_parameter() {
745 + check(
746 + json!(true),
747 + &Typed,
748 + expect![[r#"
749 + star
750 + or
751 + sequence
752 + `<arg key="`
753 + /[a-zA-Z_][a-zA-Z0-9_]*/
754 + `" type="string">`
755 + text excluding [`</arg>`]
756 + `</arg>`
757 + sequence
758 + `<arg key="`
759 + /[a-zA-Z_][a-zA-Z0-9_]*/
760 + `" type="number">`
761 + json(number)
762 + `</arg>`
763 + sequence
764 + `<arg key="`
765 + /[a-zA-Z_][a-zA-Z0-9_]*/
766 + `" type="boolean">`
767 + json(boolean)
768 + `</arg>`
769 + sequence
770 + `<arg key="`
771 + /[a-zA-Z_][a-zA-Z0-9_]*/
772 + `" type="null">`
773 + json(null)
774 + `</arg>`
775 + sequence
776 + `<arg key="`
777 + /[a-zA-Z_][a-zA-Z0-9_]*/
778 + `" type="object">`
779 + json({ ... })
780 + `</arg>`
781 + sequence
782 + `<arg key="`
783 + /[a-zA-Z_][a-zA-Z0-9_]*/
784 + `" type="array">`
785 + json(any[])
786 + `</arg>`
787 + "#]],
788 + );
789 + check(
790 + json!(false),
791 + &Untyped,
792 + expect![[r#"
793 + `\n`
794 + "#]],
795 + );
796 + }
797 +
798 + #[test]
799 + fn additional_properties_follow_declared_ones() {
800 + check(
801 + json!({
802 + "type": "object",
803 + "properties": { "path": { "type": "string" } },
804 + "required": ["path"],
805 + "additionalProperties": { "type": "boolean" }
806 + }),
807 + &Untyped,
808 + expect![[r#"
809 + sequence
810 + `\n`
811 + sequence
812 + tag `<parameter=path>` text excluding [`</parameter>`] `</parameter>`
813 + `\n`
814 + star
815 + sequence
816 + sequence
817 + `<parameter=`
818 + /[a-zA-Z_][a-zA-Z0-9_]*/
819 + `>`
820 + json(boolean)
821 + `</parameter>`
822 + `\n`
823 + "#]],
824 + );
825 + }
826 +
827 + #[test]
828 + fn raw_literals_exclude_terminators_and_request_exclusions() {
829 + check_with(
830 + json!({
831 + "type": "object",
832 + "properties": { "tag": { "enum": ["ok", "a</call>b", "c</arg>d"] } },
833 + "required": ["tag"]
834 + }),
835 + &Typed,
836 + &ArgumentOptions {
837 + excludes: vec!["</call>".to_string()],
838 + ..Default::default()
839 + },
840 + expect![[r#"
841 + tag `<arg key="tag" type="string">` `ok` `</arg>`
842 + "#]],
843 + );
844 + }
845 +}
rust/src/parser/src/unified/kimi_k3/structural_tag.rs
+68 -274
@@ -3,7 +3,7 @@
3
4 //! Structural-tag grammar for Kimi K3 XTML tool calls.
5
6 -use serde_json::{Map, Value};
6 +use serde_json::Value;
7 use xgrammar_structural_tag::Result;
8 use xgrammar_structural_tag::builders::{
9 ReasoningMode, StructuralTagBuilder, StructuralTagContext, StructuralTagOptions,
@@ -11,15 +11,67 @@ use xgrammar_structural_tag::builders::{
11 use xgrammar_structural_tag::format::{Format, JsonSchemaFormat, StructuralTag, TagFormat};
12 use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam, function_parameters};
13
14 +use crate::output_grammar::arguments::{
15 + self, ArgumentOptions, ArgumentSyntax, JsonType, ParameterKey, ValueOption, group_by, one_of,
16 +};
17 +
18 use super::{
15 - ARG_CLOSE, CALL_CLOSE, CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN,
19 + ARG_CLOSE, ARG_OPEN, CALL_CLOSE, CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN,
20 RESPONSE_CLOSE, RESPONSE_OPEN, SEP, THINK_CLOSE, THINK_OPEN, TOOLS_CLOSE, TOOLS_OPEN,
21 };
22
23 pub(super) static KIMI_K3_STRUCTURAL_TAG_BUILDER: KimiK3StructuralTagBuilder =
24 KimiK3StructuralTagBuilder;
25
22 -const XTML_TYPES: &[&str] = &["string", "number", "boolean", "null", "object", "array"];
26 +/// XTML arguments:
27 +/// `<|open|>argument key="KEY" type="TYPE"<|sep|>VALUE<|close|>argument<|sep|>`,
28 +/// adjacent to each other.
29 +///
30 +/// The parser decodes each value by its `type` attribute, so every value option
31 +/// sits under the attribute of its own JSON type. String values are raw text
32 +/// and every other value is JSON.
33 +struct XtmlArguments;
34 +
35 +impl ArgumentSyntax for XtmlArguments {
36 + fn separator(&self) -> Option<Format> {
37 + None
38 + }
39 +
40 + fn parameter(&self, key: ParameterKey<'_>, options: &[ValueOption<'_>]) -> Option<Format> {
41 + let escaped;
42 + let key = match key {
43 + ParameterKey::Declared(key) => {
44 + escaped = escape_attr_value(key);
45 + ParameterKey::Declared(&escaped)
46 + }
47 + ParameterKey::Free => ParameterKey::Free,
48 + };
49 + let arguments = group_by(options, |option| option.ty)
50 + .into_iter()
51 + .filter_map(|(ty, options)| {
52 + let values = options
53 + .into_iter()
54 + .filter_map(|option| match ty {
55 + // The parser ends the call at its close marker, so a
56 + // string value may contain neither close marker.
57 + JsonType::String => option.raw_string(&[ARG_CLOSE, CALL_CLOSE]),
58 + _ => Some(option.json()),
59 + })
60 + .collect::<Vec<_>>();
61 + (!values.is_empty()).then(|| {
62 + let suffix = format!("\" type=\"{}\"{SEP}", ty.name());
63 + key.tag(
64 + &format!("{ARG_OPEN} key=\""),
65 + &suffix,
66 + one_of(values),
67 + ARG_CLOSE,
68 + )
69 + })
70 + })
71 + .collect::<Vec<_>>();
72 + (!arguments.is_empty()).then(|| one_of(arguments))
73 + }
74 +}
75
76 /// Kimi K3 XTML structural-tag builder.
77 #[derive(Debug, Clone, Copy, Default)]
@@ -92,8 +144,17 @@ fn tools_channel(tools: &[FunctionToolParam], options: StructuralTagOptions) ->
144
145 fn call_tag(tool: &FunctionToolParam, options: StructuralTagOptions) -> TagFormat {
146 let parameters = function_parameters(&tool.function);
147 + let typed_arguments = arguments::arguments(
148 + &parameters,
149 + &XtmlArguments,
150 + &ArgumentOptions {
151 + any_order: options.any_order,
152 + max_whitespace_cnt: options.max_whitespace_cnt,
153 + excludes: vec![],
154 + },
155 + );
156 let call_body = Format::or(vec![
96 - typed_arguments(&parameters, options),
157 + typed_arguments,
158 raw_json_arguments(&parameters, options),
159 ]);
160
@@ -111,189 +172,6 @@ fn call_tag(tool: &FunctionToolParam, options: StructuralTagOptions) -> TagForma
172 )
173 }
174
114 -fn typed_arguments(parameters: &Value, options: StructuralTagOptions) -> Format {
115 - let Some(schema) = parameters.as_object() else {
116 - return if parameters == &Value::Bool(false) {
117 - Format::const_string("")
118 - } else {
119 - Format::star(permissive_argument())
120 - };
121 - };
122 - let Some(properties) = schema.get("properties").and_then(Value::as_object) else {
123 - return Format::star(permissive_argument());
124 - };
125 - if properties.is_empty() {
126 - return Format::star(permissive_argument());
127 - }
128 -
129 - let root_defs = root_definitions(schema);
130 - let arguments = properties
131 - .iter()
132 - .flat_map(|(key, schema)| argument_tags(key, schema, &root_defs, options))
133 - .map(Format::Tag)
134 - .collect::<Vec<_>>();
135 - let arguments = match arguments.as_slice() {
136 - [argument] => argument.clone(),
137 - _ => Format::or(arguments),
138 - };
139 - // Keep typed arguments order-agnostic and non-unique, but do not allow an
140 - // empty call when the root schema declares required properties.
141 - if schema
142 - .get("required")
143 - .and_then(Value::as_array)
144 - .is_some_and(|required| !required.is_empty())
145 - {
146 - Format::plus(arguments)
147 - } else {
148 - Format::star(arguments)
149 - }
150 -}
151 -
152 -fn argument_tags(
153 - key: &str,
154 - schema: &Value,
155 - root_defs: &Map<String, Value>,
156 - options: StructuralTagOptions,
157 -) -> Vec<TagFormat> {
158 - let types = schema_types(schema);
159 - types
160 - .into_iter()
161 - .map(|xtml_type| {
162 - let content = if xtml_type == "string" {
163 - string_argument_content(schema)
164 - } else {
165 - json_schema(
166 - attach_root_definitions(&narrow_schema_type(schema, xtml_type), root_defs),
167 - options,
168 - )
169 - };
170 - TagFormat::new(
171 - format!(
172 - "{OPEN}argument key=\"{}\" type=\"{xtml_type}\"{SEP}",
173 - escape_attr_value(key)
174 - ),
175 - content,
176 - ARG_CLOSE,
177 - )
178 - })
179 - .collect()
180 -}
181 -
182 -fn schema_types(schema: &Value) -> Vec<&'static str> {
183 - let Some(schema) = schema.as_object() else {
184 - return XTML_TYPES.to_vec();
185 - };
186 - let mut types = Vec::new();
187 - match schema.get("type") {
188 - Some(Value::String(value)) => push_schema_type(&mut types, value),
189 - Some(Value::Array(values)) => {
190 - for value in values.iter().filter_map(Value::as_str) {
191 - push_schema_type(&mut types, value);
192 - }
193 - }
194 - _ => {}
195 - }
196 - if types.is_empty()
197 - && let Some(value) = schema.get("const")
198 - {
199 - push_value_type(&mut types, value);
200 - }
201 - if types.is_empty()
202 - && let Some(values) = schema.get("enum").and_then(Value::as_array)
203 - {
204 - for value in values {
205 - push_value_type(&mut types, value);
206 - }
207 - }
208 - if types.is_empty() {
209 - XTML_TYPES.to_vec()
210 - } else {
211 - types
212 - }
213 -}
214 -
215 -fn push_value_type(types: &mut Vec<&'static str>, value: &Value) {
216 - let xtml_type = match value {
217 - Value::String(_) => "string",
218 - Value::Number(_) => "number",
219 - Value::Bool(_) => "boolean",
220 - Value::Null => "null",
221 - Value::Object(_) => "object",
222 - Value::Array(_) => "array",
223 - };
224 - if !types.contains(&xtml_type) {
225 - types.push(xtml_type);
226 - }
227 -}
228 -
229 -fn narrow_schema_type(schema: &Value, xtml_type: &str) -> Value {
230 - let Some(mut schema) = schema.as_object().cloned() else {
231 - return schema.clone();
232 - };
233 - let json_type = if xtml_type == "number" && explicitly_integer_only(&schema) {
234 - "integer"
235 - } else {
236 - xtml_type
237 - };
238 - schema.insert("type".to_string(), Value::String(json_type.to_string()));
239 - Value::Object(schema)
240 -}
241 -
242 -fn explicitly_integer_only(schema: &Map<String, Value>) -> bool {
243 - match schema.get("type") {
244 - Some(Value::String(value)) => value == "integer",
245 - Some(Value::Array(values)) => {
246 - let values = values.iter().filter_map(Value::as_str).collect::<Vec<_>>();
247 - values.contains(&"integer") && !values.contains(&"number")
248 - }
249 - _ => false,
250 - }
251 -}
252 -
253 -fn push_schema_type(types: &mut Vec<&'static str>, json_type: &str) {
254 - let xtml_type = match json_type {
255 - "string" => Some("string"),
256 - "integer" | "number" => Some("number"),
257 - "boolean" => Some("boolean"),
258 - "null" => Some("null"),
259 - "object" => Some("object"),
260 - "array" => Some("array"),
261 - _ => None,
262 - };
263 - if let Some(xtml_type) = xtml_type
264 - && !types.contains(&xtml_type)
265 - {
266 - types.push(xtml_type);
267 - }
268 -}
269 -
270 -fn string_argument_content(schema: &Value) -> Format {
271 - let Some(schema) = schema.as_object() else {
272 - return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]);
273 - };
274 - let values = schema
275 - .get("enum")
276 - .and_then(Value::as_array)
277 - .cloned()
278 - .or_else(|| schema.get("const").cloned().map(|value| vec![value]));
279 - let Some(values) = values else {
280 - return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]);
281 - };
282 - if values.is_empty()
283 - || values.len() > 256
284 - || values
285 - .iter()
286 - .any(|value| value.as_str().is_none_or(|value| value.contains("<|")))
287 - {
288 - return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]);
289 - }
290 - let values = values.iter().filter_map(Value::as_str).collect::<Vec<_>>();
291 - match values.as_slice() {
292 - [value] => Format::const_string(*value),
293 - _ => Format::or(values.into_iter().map(Format::const_string).collect()),
294 - }
295 -}
296 -
175 fn raw_json_arguments(parameters: &Value, options: StructuralTagOptions) -> Format {
176 Format::tag(
177 format!("{JSON_OPEN} type=\"object\"{SEP}"),
@@ -302,29 +180,6 @@ fn raw_json_arguments(parameters: &Value, options: StructuralTagOptions) -> Form
180 )
181 }
182
305 -fn permissive_argument() -> Format {
306 - let key = Format::regex(r#"(?:[^<\"&]|&(?:amp|quot);|<[^|])*"#);
307 - let alternatives = XTML_TYPES
308 - .iter()
309 - .map(|xtml_type| {
310 - Format::sequence(vec![
311 - key.clone(),
312 - Format::const_string(format!("\" type=\"{xtml_type}\"{SEP}")),
313 - if *xtml_type == "string" {
314 - Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE])
315 - } else {
316 - Format::json_schema(Value::Bool(true))
317 - },
318 - ])
319 - })
320 - .collect();
321 - Format::tag(
322 - format!("{OPEN}argument key=\""),
323 - Format::or(alternatives),
324 - ARG_CLOSE,
325 - )
326 -}
327 -
183 fn json_schema(schema: Value, options: StructuralTagOptions) -> Format {
184 Format::JsonSchema(
185 JsonSchemaFormat::new(schema)
@@ -333,23 +188,6 @@ fn json_schema(schema: Value, options: StructuralTagOptions) -> Format {
188 )
189 }
190
336 -fn root_definitions(schema: &Map<String, Value>) -> Map<String, Value> {
337 - ["$defs", "definitions"]
338 - .into_iter()
339 - .filter_map(|key| schema.get(key).map(|value| (key.to_string(), value.clone())))
340 - .collect()
341 -}
342 -
343 -fn attach_root_definitions(schema: &Value, root_defs: &Map<String, Value>) -> Value {
344 - let Some(mut schema) = schema.as_object().cloned() else {
345 - return schema.clone();
346 - };
347 - for (key, value) in root_defs {
348 - schema.entry(key.clone()).or_insert_with(|| value.clone());
349 - }
350 - Value::Object(schema)
351 -}
352 -
191 fn escape_attr_value(value: &str) -> String {
192 value.replace('&', "&amp;").replace('"', "&quot;")
193 }
@@ -407,52 +245,18 @@ mod tests {
245 /[1-9][0-9]*/
246 `"<|sep|>`
247 or
410 - plus
411 - or
248 + sequence
249 + optional
250 tag `<|open|>argument key="unit" type="string"<|sep|>` .. `<|close|>argument<|sep|>`
251 or
252 `celsius`
253 `fahrenheit`
416 - tag `<|open|>argument key="place" type="object"<|sep|>` json(place & object where place = { city?: string }) `<|close|>argument<|sep|>`
254 + tag `<|open|>argument key="place" type="object"<|sep|>` json({ city?: string } where place = { city?: string }) `<|close|>argument<|sep|>`
255 tag `<|open|>json type="object"<|sep|>` json({ unit?: "celsius" | "fahrenheit", place: place & object } where place = { city?: string }) `<|close|>json<|sep|>`
256 optional `<|close|>message<|sep|>`
257 "#]].assert_eq(&outline(&tag.format));
258 }
259
422 - #[test]
423 - fn typed_arguments_require_one_tag_only_for_nonempty_required() {
424 - let required = super::typed_arguments(
425 - &json!({
426 - "type": "object",
427 - "properties": { "query": { "type": "string" } },
428 - "required": ["query"]
429 - }),
430 - StructuralTagOptions::default(),
431 - );
432 - let optional = super::typed_arguments(
433 - &json!({
434 - "type": "object",
435 - "properties": { "query": { "type": "string" } }
436 - }),
437 - StructuralTagOptions::default(),
438 - );
439 - let empty_required = super::typed_arguments(
440 - &json!({
441 - "type": "object",
442 - "properties": { "query": { "type": "string" } },
443 - "required": []
444 - }),
445 - StructuralTagOptions::default(),
446 - );
447 -
448 - assert_eq!(serde_json::to_value(required).unwrap()["type"], "plus");
449 - assert_eq!(serde_json::to_value(optional).unwrap()["type"], "star");
450 - assert_eq!(
451 - serde_json::to_value(empty_required).unwrap()["type"],
452 - "star"
453 - );
454 - }
455 -
260 #[test]
261 fn reasoning_grammar_starts_inside_prefilled_think_channel() {
262 let tools = vec![tool("ping", json!({ "type": "object", "properties": {} }))];
@@ -522,14 +326,4 @@ mod tests {
326 assert!(tag.contains(r#"type=\"null\""#));
327 assert!(tag.contains(r#""json_schema":{"type":"null"}"#), "{tag}");
328 }
525 -
526 - #[test]
527 - fn unsafe_string_enum_falls_back_as_a_whole() {
528 - let format = super::string_argument_content(&json!({
529 - "type": "string",
530 - "enum": ["safe", "<|unsafe"]
531 - }));
532 -
533 - assert_eq!(serde_json::to_value(format).unwrap()["type"], "any_text");
534 - }
329 }