diff --git a/rust/src/chat/Cargo.toml b/rust/src/chat/Cargo.toml index d145fc237a..fdaa5ad114 100644 --- a/rust/src/chat/Cargo.toml +++ b/rust/src/chat/Cargo.toml @@ -52,6 +52,7 @@ tempfile.workspace = true tokio.workspace = true uuid.workspace = true vllm-engine-core-client = { workspace = true, features = ["test-util"] } +vllm-parser = { workspace = true, features = ["test-util"] } vllm-tokenizer = { workspace = true, features = ["test-utils"] } zeromq.workspace = true diff --git a/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.txt b/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.txt new file mode 100644 index 0000000000..5cf40f3218 --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.txt @@ -0,0 +1,26 @@ +# auto (from_token_zero) +sequence + tag `` text `` + triggered_tags [`<｜DSML｜function_calls>`] excluding [``, ``] + tag `<｜DSML｜function_calls>\n` .. `` + tags_with_separator `` at_least_one + tag `<｜DSML｜invoke name="get_weather">\n` deepseek_xml({ location: string }) `\n` + tag `<｜DSML｜invoke name="add">\n` deepseek_xml({ y: number, x: number, items: string[] }) `\n` + +# required (from_token_zero) +sequence + tag `` text `` + sequence + `\n\n<｜DSML｜function_calls>\n` + tags_with_separator `` at_least_one + tag `<｜DSML｜invoke name="get_weather">\n` deepseek_xml({ location: string }) `\n` + tag `<｜DSML｜invoke name="add">\n` deepseek_xml({ y: number, x: number, items: string[] }) `\n` + `` + +# named (from_token_zero) +sequence + tag `` text `` + sequence + `\n\n<｜DSML｜function_calls>\n` + tag `<｜DSML｜invoke name="get_weather">\n` deepseek_xml({ location: string }) `\n` + `` diff --git a/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.txt b/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.txt new file mode 100644 index 0000000000..5e216e8ffc --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.txt @@ -0,0 +1,26 @@ +# auto (from_token_zero) +sequence + tag `` text `` + triggered_tags [`<｜DSML｜tool_calls>`] excluding [``, ``] + tag `<｜DSML｜tool_calls>\n` .. `` + tags_with_separator `` at_least_one + tag `<｜DSML｜invoke name="get_weather">\n` deepseek_xml({ location: string }) `\n` + tag `<｜DSML｜invoke name="add">\n` deepseek_xml({ y: number, x: number, items: string[] }) `\n` + +# required (from_token_zero) +sequence + tag `` text `` + sequence + `\n\n<｜DSML｜tool_calls>\n` + tags_with_separator `` at_least_one + tag `<｜DSML｜invoke name="get_weather">\n` deepseek_xml({ location: string }) `\n` + tag `<｜DSML｜invoke name="add">\n` deepseek_xml({ y: number, x: number, items: string[] }) `\n` + `` + +# named (from_token_zero) +sequence + tag `` text `` + sequence + `\n\n<｜DSML｜tool_calls>\n` + tag `<｜DSML｜invoke name="get_weather">\n` deepseek_xml({ location: string }) `\n` + `` diff --git a/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.txt b/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.txt new file mode 100644 index 0000000000..8c92c3a5a3 --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.txt @@ -0,0 +1,26 @@ +# auto (from_token_zero) +sequence + tag `` text `` + triggered_tags [`<｜DSML｜ calls>`] excluding [``, ``] + tag `<｜DSML｜ calls>\n` .. `` + tags_with_separator `` at_least_one + tag `<｜DSML｜ invoke name="get_weather">\n` deepseek_v4_1_xml({ location: string }) `\n` + tag `<｜DSML｜ invoke name="add">\n` deepseek_v4_1_xml({ y: number, x: number, items: string[] }) `\n` + +# required (from_token_zero) +sequence + tag `` text `` + sequence + `\n\n<｜DSML｜ calls>\n` + tags_with_separator `` at_least_one + tag `<｜DSML｜ invoke name="get_weather">\n` deepseek_v4_1_xml({ location: string }) `\n` + tag `<｜DSML｜ invoke name="add">\n` deepseek_v4_1_xml({ y: number, x: number, items: string[] }) `\n` + `` + +# named (from_token_zero) +sequence + tag `` text `` + sequence + `\n\n<｜DSML｜ calls>\n` + tag `<｜DSML｜ invoke name="get_weather">\n` deepseek_v4_1_xml({ location: string }) `\n` + `` diff --git a/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.txt b/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.txt new file mode 100644 index 0000000000..58b4bef2da --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.txt @@ -0,0 +1,18 @@ +# auto (from_token_zero) +sequence + tag `` text `` + triggered_tags [``] excluding [``, ``, ``, ``, ``, ``, ``] + tag `get_weather` glm_xml({ location: string }) `` + tag `add` glm_xml({ y: number, x: number, items: string[] }) `` + +# required (from_token_zero) +sequence + tag `` text `` + triggered_tags [``] excluding [``, ``, ``, ``, ``, ``, ``] at_least_one + tag `get_weather` glm_xml({ location: string }) `` + tag `add` glm_xml({ y: number, x: number, items: string[] }) `` + +# named (from_token_zero) +sequence + tag `` text `` + tag `get_weather` glm_xml({ location: string }) `` diff --git a/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.txt b/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.txt new file mode 100644 index 0000000000..58b4bef2da --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.txt @@ -0,0 +1,18 @@ +# auto (from_token_zero) +sequence + tag `` text `` + triggered_tags [``] excluding [``, ``, ``, ``, ``, ``, ``] + tag `get_weather` glm_xml({ location: string }) `` + tag `add` glm_xml({ y: number, x: number, items: string[] }) `` + +# required (from_token_zero) +sequence + tag `` text `` + triggered_tags [``] excluding [``, ``, ``, ``, ``, ``, ``] at_least_one + tag `get_weather` glm_xml({ location: string }) `` + tag `add` glm_xml({ y: number, x: number, items: string[] }) `` + +# named (from_token_zero) +sequence + tag `` text `` + tag `get_weather` glm_xml({ location: string }) `` diff --git a/rust/src/chat/tests/grammar_replay/Hy3.txt b/rust/src/chat/tests/grammar_replay/Hy3.txt new file mode 100644 index 0000000000..dcd6bb08ec --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/Hy3.txt @@ -0,0 +1,112 @@ +# auto (from_token_zero) +sequence + tag `` text `` + triggered_tags [``] + tag `\n` .. `\n` + tags_with_separator `\n` at_least_one + tag `get_weather\n` .. `` + sequence + sequence + `` + `location` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + tag `add\n` .. `` + sequence + sequence + `` + `y` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + sequence + `` + `x` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + sequence + `` + `items` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + +# required (from_token_zero) +sequence + tag `` text `` + sequence + `\n` + tags_with_separator `\n` at_least_one + tag `get_weather\n` .. `` + sequence + sequence + `` + `location` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + tag `add\n` .. `` + sequence + sequence + `` + `y` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + sequence + `` + `x` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + sequence + `` + `items` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + `\n` + +# named (from_token_zero) +sequence + tag `` text `` + sequence + `\n` + tag `get_weather\n` .. `` + sequence + sequence + `` + `location` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + `\n` diff --git a/rust/src/chat/tests/grammar_replay/Kimi-K2.5.txt b/rust/src/chat/tests/grammar_replay/Kimi-K2.5.txt new file mode 100644 index 0000000000..553c1ce6f4 --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/Kimi-K2.5.txt @@ -0,0 +1,46 @@ +# auto (from_token_zero) +sequence + tag `` text `` + triggered_tags [`<|tool_calls_section_begin|>`] excluding [``, ``, `<|tool_call_begin|>`] + tag `<|tool_calls_section_begin|>` .. `<|tool_calls_section_end|>` + tags_with_separator `` at_least_one + tag `<|tool_call_begin|>functions.get_weather:` .. `<|tool_call_end|>` + sequence + /\d+/ + `<|tool_call_argument_begin|>` + json({ location: string }) + tag `<|tool_call_begin|>functions.add:` .. `<|tool_call_end|>` + sequence + /\d+/ + `<|tool_call_argument_begin|>` + json({ y: number, x: number, items: string[] }) + +# required (from_token_zero) +sequence + tag `` text `` + sequence + `<|tool_calls_section_begin|>` + tags_with_separator `` at_least_one + tag `<|tool_call_begin|>functions.get_weather:` .. `<|tool_call_end|>` + sequence + /\d+/ + `<|tool_call_argument_begin|>` + json({ location: string }) + tag `<|tool_call_begin|>functions.add:` .. `<|tool_call_end|>` + sequence + /\d+/ + `<|tool_call_argument_begin|>` + json({ y: number, x: number, items: string[] }) + `<|tool_calls_section_end|>` + +# named (from_token_zero) +sequence + tag `` text `` + sequence + `<|tool_calls_section_begin|>` + tag `<|tool_call_begin|>functions.get_weather:` .. `<|tool_call_end|>` + sequence + /\d+/ + `<|tool_call_argument_begin|>` + json({ location: string }) + `<|tool_calls_section_end|>` diff --git a/rust/src/chat/tests/grammar_replay/Kimi-K3.txt b/rust/src/chat/tests/grammar_replay/Kimi-K3.txt new file mode 100644 index 0000000000..9b2eca8b2a --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/Kimi-K3.txt @@ -0,0 +1,67 @@ +# auto (final_output_only) +sequence + optional `<|open|>response<|sep|>` + tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + optional + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + plus tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` + tag `<|open|>call tool="add" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + plus + or + tag `<|open|>argument key="y" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="x" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="items" type="array"<|sep|>` json(string[]) `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>` + optional `<|close|>message<|sep|>` + +# required (final_output_only) +sequence + optional `<|open|>response<|sep|>` + tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + plus tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` + tag `<|open|>call tool="add" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + plus + or + tag `<|open|>argument key="y" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="x" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="items" type="array"<|sep|>` json(string[]) `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>` + optional `<|close|>message<|sep|>` + +# named (final_output_only) +sequence + optional `<|open|>response<|sep|>` + tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + plus tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` + optional `<|close|>message<|sep|>` diff --git a/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.txt b/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.txt new file mode 100644 index 0000000000..1d12560eb0 --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.txt @@ -0,0 +1,18 @@ +# auto (from_token_zero) +sequence + optional tag `` text `` + triggered_tags [``] excluding [``, ``, ``, `` qwen_xml({ location: string }) `` + tag `` qwen_xml({ y: number, x: number, items: string[] }) `` + +# required (from_token_zero) +sequence + optional tag `` text `` + triggered_tags [``] excluding [``, ``, ``, `` qwen_xml({ location: string }) `` + tag `` qwen_xml({ y: number, x: number, items: string[] }) `` + +# named (from_token_zero) +sequence + optional tag `` text `` + tag `` qwen_xml({ location: string }) `` diff --git a/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.txt b/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.txt new file mode 100644 index 0000000000..431238009a --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.txt @@ -0,0 +1,32 @@ +# auto (from_token_zero) +sequence + sequence + tag `` text `` + `\n\n` + triggered_tags [``] excluding [``, ``] + tag `\n` .. `` + tags_with_separator `` at_least_one + tag `\n` minimax_xml({ location: string }) `\n` + tag `\n` minimax_xml({ y: number, x: number, items: string[] }) `\n` + +# required (from_token_zero) +sequence + sequence + tag `` text `` + `\n\n` + sequence + `\n\n` + tags_with_separator `` at_least_one + tag `\n` minimax_xml({ location: string }) `\n` + tag `\n` minimax_xml({ y: number, x: number, items: string[] }) `\n` + `` + +# named (from_token_zero) +sequence + sequence + tag `` text `` + `\n\n` + sequence + `\n\n` + tag `\n` minimax_xml({ location: string }) `\n` + `` diff --git a/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.txt b/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.txt new file mode 100644 index 0000000000..fe38f5e9ff --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.txt @@ -0,0 +1,27 @@ +# auto (from_token_zero) +sequence + optional + sequence + tag `` text `` + `\n\n` + triggered_tags [``] excluding [``, ``] + tag `\n{"name": "get_weather", "arguments": ` json({ location: string }) `}\n` + tag `\n{"name": "add", "arguments": ` json({ y: number, x: number, items: string[] }) `}\n` + +# required (from_token_zero) +sequence + optional + sequence + tag `` text `` + `\n\n` + triggered_tags [``] excluding [``, ``] at_least_one + tag `\n{"name": "get_weather", "arguments": ` json({ location: string }) `}\n` + tag `\n{"name": "add", "arguments": ` json({ y: number, x: number, items: string[] }) `}\n` + +# named (from_token_zero) +sequence + optional + sequence + tag `` text `` + `\n\n` + tag `\n{"name": "get_weather", "arguments": ` json({ location: string }) `}\n` diff --git a/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.txt b/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.txt new file mode 100644 index 0000000000..b057780649 --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.txt @@ -0,0 +1,24 @@ +# auto (from_token_zero) +sequence + sequence + tag `` text `` + `\n\n` + triggered_tags [`\n`, ``] + tag `\n\n` qwen_xml({ location: string }) `\n\n` + tag `\n\n` qwen_xml({ y: number, x: number, items: string[] }) `\n\n` + +# required (from_token_zero) +sequence + sequence + tag `` text `` + `\n\n` + triggered_tags [`\n`, ``] at_least_one + tag `\n\n` qwen_xml({ location: string }) `\n\n` + tag `\n\n` qwen_xml({ y: number, x: number, items: string[] }) `\n\n` + +# named (from_token_zero) +sequence + sequence + tag `` text `` + `\n\n` + tag `\n\n` qwen_xml({ location: string }) `\n\n` diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index 6925ae1b02..de813e3045 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -1150,6 +1150,7 @@ fn check_grammar_replay_file( ) -> Result<()> { let model_name = case.model_id.rsplit_once('/').map_or(case.model_id, |(_, name)| name); let path = grammar_replay_dir().join(format!("{model_name}.json")); + let outline_path = path.with_extension("txt"); let results = results .into_iter() .filter_map(|(variant, result)| { @@ -1157,11 +1158,13 @@ fn check_grammar_replay_file( }) .collect::>(); if results.is_empty() { - ensure!( - !path.exists(), - "{} builds no output grammar, but grammar replay file {path:?} exists", - case.model_id - ); + for path in [&path, &outline_path] { + ensure!( + !path.exists(), + "{} builds no output grammar, but grammar replay file {path:?} exists", + case.model_id + ); + } return Ok(()); } @@ -1205,6 +1208,25 @@ fn check_grammar_replay_file( ); } + // A readable outline of the same grammars, for review only. + let mut outline = String::new(); + for (variant, case) in &cases { + if !outline.is_empty() { + outline.push('\n'); + } + let variant = serde_json::to_value(variant)?; + let coverage = serde_json::to_value(case.grammar.coverage)?; + outline.push_str(&format!( + "# {} ({})\n", + variant.as_str().context("variant name")?, + coverage.as_str().context("coverage name")?, + )); + outline.push_str(&vllm_parser::output_grammar::test_utils::outline( + &case.grammar.format, + )); + } + expect_test::expect_file![outline_path].assert_eq(&outline); + let file = GrammarReplayFile { model_id: case.model_id, vocab_size: backends.text_backend.model_vocab_size(), diff --git a/rust/src/parser/src/output_grammar.rs b/rust/src/parser/src/output_grammar.rs index db56f97708..29f91a4c25 100644 --- a/rust/src/parser/src/output_grammar.rs +++ b/rust/src/parser/src/output_grammar.rs @@ -17,6 +17,9 @@ use xgrammar_structural_tag::{ use crate::tool::Tool; +#[cfg(any(test, feature = "test-util"))] +pub mod test_utils; + /// Result alias for output grammar construction. pub type Result = std::result::Result; diff --git a/rust/src/parser/src/output_grammar/test_utils.rs b/rust/src/parser/src/output_grammar/test_utils.rs new file mode 100644 index 0000000000..c745152556 --- /dev/null +++ b/rust/src/parser/src/output_grammar/test_utils.rs @@ -0,0 +1,676 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Readable renderings of structural-tag formats for test snapshots. + +use std::fmt::Write; + +use serde_json::{Map, Value}; +use xgrammar_structural_tag::format::{EndBoundary, Format, TagBoundary, TagFormat, TokenValue}; + +/// Render `format` as an indented outline, one node per line. +/// +/// Each node starts with its format type, and composite nodes list their +/// children one level deeper. Literal text is quoted in backticks with `\`, +/// `` ` ``, and control characters escaped, so markers read as written. Other +/// leaves are compact: `/regex/`, `text`, `json(SCHEMA)` or `STYLE(SCHEMA)` +/// for styled schemas, and `token ID` or ``token `TEXT` ``. Lists such as +/// exclusions and triggers are bracketed, and schemas read as types, see +/// [`schema`]. A single leaf child of `optional`, `star`, `plus`, `repeat`, +/// and dispatch rules stays on its parent's line, as does a tag's leaf +/// content between its boundaries. Options that keep their defaults are +/// omitted. +pub fn outline(format: &Format) -> String { + let mut out = String::new(); + write(&mut out, 0, format_node(format)); + out +} + +/// A child line of an outline node. +enum Child<'a> { + Format(&'a Format), + Tag(&'a TagFormat), + Rule(String, &'a Format), + Text(&'a str), +} + +/// An outline node: its header line and children. +struct Node<'a> { + header: String, + children: Vec>, + /// Whether a single leaf child joins the header line. + inline_leaf: bool, +} + +impl<'a> Node<'a> { + fn leaf(header: String) -> Self { + Self { + header, + children: vec![], + inline_leaf: false, + } + } + + fn branch(header: impl Into, children: Vec>) -> Self { + Self { + header: header.into(), + children, + inline_leaf: false, + } + } + + fn wrapper(header: impl Into, child: Child<'a>) -> Self { + Self { + header: header.into(), + children: vec![child], + inline_leaf: true, + } + } +} + +fn write(out: &mut String, depth: usize, node: Node<'_>) { + let Node { + mut header, + children, + inline_leaf, + } = node; + let mut children = children.into_iter().map(self::node).collect::>(); + if inline_leaf + && let [child] = &children[..] + && child.children.is_empty() + { + let _ = write!(header, " {}", child.header); + children.clear(); + } + let _ = writeln!(out, "{}{header}", " ".repeat(depth)); + for child in children { + write(out, depth + 1, child); + } +} + +fn node(child: Child<'_>) -> Node<'_> { + match child { + Child::Format(format) => format_node(format), + Child::Tag(tag) => tag_node(tag), + Child::Rule(pattern, format) => { + Node::wrapper(format!("on {pattern}"), Child::Format(format)) + } + Child::Text(text) => Node::leaf(text.to_string()), + } +} + +fn format_node(format: &Format) -> Node<'_> { + match format { + Format::ConstString(format) => Node::leaf(quote(&format.value)), + Format::Regex(format) => Node::leaf(format!("/{}/", format.pattern)), + Format::AnyText(format) => { + let mut header = "text".to_string(); + push_excluding(&mut header, format.excludes.iter().map(|text| quote(text))); + push_limit(&mut header, "max_tokens", format.max_tokens); + push_limit(&mut header, "max_chars", format.max_chars); + Node::leaf(header) + } + Format::JsonSchema(format) => { + let style = serde_json::to_value(format.style) + .ok() + .and_then(|style| style.as_str().map(str::to_string)) + .unwrap_or_else(|| format!("{:?}", format.style)); + let mut header = format!("{style}({})", schema(&format.json_schema)); + if format.any_order { + header.push_str(" any_order"); + } + if let Some(max) = format.max_whitespace_cnt { + let _ = write!(header, " max_whitespace={max}"); + } + push_excluding(&mut header, format.excludes.iter().map(|text| quote(text))); + Node::leaf(header) + } + Format::Grammar(format) => { + Node::branch("grammar", format.grammar.lines().map(Child::Text).collect()) + } + Format::Token(format) => Node::leaf(format!("token {}", token(&format.token))), + Format::ExcludeToken(format) => { + let mut header = "token".to_string(); + push_excluding(&mut header, format.exclude_tokens.iter().map(token)); + Node::leaf(header) + } + Format::AnyTokens(format) => { + let mut header = "tokens".to_string(); + push_excluding(&mut header, format.exclude_tokens.iter().map(token)); + push_limit(&mut header, "max_tokens", format.max_tokens); + Node::leaf(header) + } + Format::Sequence(format) => Node::branch( + "sequence", + format.elements.iter().map(Child::Format).collect(), + ), + Format::Or(format) => { + Node::branch("or", format.elements.iter().map(Child::Format).collect()) + } + Format::Optional(format) => Node::wrapper("optional", Child::Format(&format.content)), + Format::Plus(format) => Node::wrapper("plus", Child::Format(&format.content)), + Format::Star(format) => Node::wrapper("star", Child::Format(&format.content)), + Format::Repeat(format) => { + let max = if format.max < 0 { + String::new() + } else { + format.max.to_string() + }; + Node::wrapper( + format!("repeat {}..{max}", format.min), + Child::Format(&format.content), + ) + } + Format::Tag(tag) => tag_node(tag), + Format::TriggeredTags(format) => { + let triggers = format.triggers.iter().map(|trigger| quote(trigger)); + let mut header = format!( + "triggered_tags [{}]", + triggers.collect::>().join(", ") + ); + push_excluding(&mut header, format.excludes.iter().map(|text| quote(text))); + push_flags(&mut header, format.at_least_one, format.stop_after_first); + Node::branch(header, format.tags.iter().map(Child::Tag).collect()) + } + Format::TokenTriggeredTags(format) => { + let triggers = format.trigger_tokens.iter().map(token); + let mut header = format!( + "token_triggered_tags [{}]", + triggers.collect::>().join(", ") + ); + push_excluding(&mut header, format.exclude_tokens.iter().map(token)); + push_flags(&mut header, format.at_least_one, format.stop_after_first); + Node::branch(header, format.tags.iter().map(Child::Tag).collect()) + } + Format::TagsWithSeparator(format) => { + let mut header = format!("tags_with_separator {}", quote(&format.separator)); + push_flags(&mut header, format.at_least_one, format.stop_after_first); + Node::branch(header, format.tags.iter().map(Child::Tag).collect()) + } + Format::Dispatch(format) => { + let mut header = "dispatch".to_string(); + if !format.r#loop { + header.push_str(" once"); + } + push_excluding(&mut header, format.excludes.iter().map(|text| quote(text))); + Node::branch( + header, + format + .rules + .iter() + .map(|(pattern, format)| Child::Rule(quote(pattern), format)) + .collect(), + ) + } + Format::TokenDispatch(format) => { + let mut header = "token_dispatch".to_string(); + if !format.r#loop { + header.push_str(" once"); + } + push_excluding(&mut header, format.exclude_tokens.iter().map(token)); + Node::branch( + header, + format + .rules + .iter() + .map(|(trigger, format)| Child::Rule(token(trigger), format)) + .collect(), + ) + } + } +} + +/// `tag BEGIN CONTENT END` for leaf content, or `tag BEGIN .. END` above the +/// content's own node. +fn tag_node(tag: &TagFormat) -> Node<'_> { + let begin = match &tag.begin { + TagBoundary::Text(text) => quote(text), + TagBoundary::Token(boundary) => format!("token {}", token(&boundary.token)), + }; + let end = match &tag.end { + EndBoundary::Text(text) => quote(text), + EndBoundary::Texts(texts) => { + texts.iter().map(|text| quote(text)).collect::>().join(" | ") + } + EndBoundary::Token(boundary) => format!("token {}", token(&boundary.token)), + }; + let content = format_node(&tag.content); + if content.children.is_empty() { + Node::leaf(format!("tag {begin} {} {end}", content.header)) + } else { + Node::branch( + format!("tag {begin} .. {end}"), + vec![Child::Format(&tag.content)], + ) + } +} + +/// Schema keywords that do not constrain the generated text. +const ANNOTATIONS: &[&str] = &[ + "$comment", + "$id", + "$schema", + "default", + "deprecated", + "description", + "examples", + "readOnly", + "title", + "writeOnly", +]; + +/// Render a JSON schema as a TypeScript-like type. +/// +/// Objects are `{ key: T, optional?: T }`, with `...` for any additional +/// property and `...: T` for additional properties under `T`. Arrays are +/// `T[]`. Unions, type arrays, and enums join their options with `|`, +/// `oneOf` and `allOf` keep their names, local references show the +/// definition name, a `type` next to a reference intersects it as +/// `NAME & T`, and `true` / `false` are `any` / `never`. Every other +/// keyword follows the type as `(key=VALUE, ...)`, and definitions follow the +/// root as `where NAME = T, ...`. Annotations such as `description` are +/// dropped. +pub fn schema(schema: &Value) -> String { + typed(schema).text +} + +/// A rendered type, and whether it is a top-level union or intersection that +/// an operand position must parenthesize. +struct Typed { + text: String, + compound: bool, +} + +impl Typed { + fn plain(text: impl Into) -> Self { + Self { + text: text.into(), + compound: false, + } + } + + fn compound(text: String) -> Self { + Self { + text, + compound: true, + } + } + + fn union(options: Vec) -> Self { + Self { + compound: options.len() > 1, + text: options.join(" | "), + } + } + + /// The type as an operand of `&` or of a suffix such as `[]` or `(...)`. + fn grouped(self) -> String { + if self.compound { + format!("({})", self.text) + } else { + self.text + } + } +} + +fn typed(schema: &Value) -> Typed { + let schema = match schema { + Value::Bool(true) => return Typed::plain("any"), + Value::Bool(false) => return Typed::plain("never"), + Value::Object(schema) => schema, + schema => return Typed::plain(schema.to_string()), + }; + let mut rendered = SchemaType::new(schema).render(); + let definitions = ["$defs", "definitions"] + .into_iter() + .filter_map(|key| schema.get(key)?.as_object()) + .flatten() + .map(|(name, definition)| format!("{name} = {}", self::schema(definition))) + .collect::>(); + if !definitions.is_empty() { + rendered = Typed::plain(format!( + "{} where {}", + rendered.text, + definitions.join(", ") + )); + } + rendered +} + +/// One schema object being rendered, with the keywords its type consumed. +struct SchemaType<'a> { + schema: &'a Map, + consumed: Vec<&'static str>, +} + +impl<'a> SchemaType<'a> { + fn new(schema: &'a Map) -> Self { + Self { + schema, + consumed: vec!["$defs", "definitions"], + } + } + + fn take(&mut self, key: &'static str) -> Option<&'a Value> { + let value = self.schema.get(key)?; + self.consumed.push(key); + Some(value) + } + + fn render(mut self) -> Typed { + let mut rendered = self.body(); + let rest = self + .schema + .iter() + .filter(|(key, _)| { + !self.consumed.contains(&key.as_str()) && !ANNOTATIONS.contains(&key.as_str()) + }) + .map(|(key, value)| format!("{key}={value}")) + .collect::>(); + if !rest.is_empty() { + rendered = Typed::plain(format!("{}({})", rendered.grouped(), rest.join(", "))); + } + rendered + } + + fn body(&mut self) -> Typed { + if let Some(Value::String(reference)) = self.take("$ref") { + let name = ["#/$defs/", "#/definitions/"] + .into_iter() + .find_map(|prefix| reference.strip_prefix(prefix)) + .unwrap_or(reference); + // Keywords next to `$ref` apply together with it. + return match self.take("type") { + Some(type_name) => { + Typed::compound(format!("{name} & {}", type_names(type_name).grouped())) + } + None => Typed::plain(name), + }; + } + if let Some(value) = self.take("const") { + self.take_redundant_type([value]); + return Typed::plain(value.to_string()); + } + if let Some(Value::Array(values)) = self.take("enum") { + self.take_redundant_type(values); + return Typed::union(values.iter().map(Value::to_string).collect()); + } + if let Some(Value::Array(options)) = self.take("anyOf") { + return Typed::union(options.iter().map(schema).collect()); + } + for combinator in ["oneOf", "allOf"] { + if let Some(Value::Array(options)) = self.take(combinator) { + let options = options.iter().map(schema).collect::>(); + return Typed::plain(format!("{combinator}({})", options.join(", "))); + } + } + let is = |name: &str| self.schema.get("type").is_none_or(|type_name| type_name == name); + if is("object") + && (self.schema.contains_key("properties") + || self.schema.contains_key("additionalProperties")) + { + self.take("type"); + return Typed::plain(self.object()); + } + if is("array") && self.schema.contains_key("items") { + self.take("type"); + let items = self.take("items").map_or_else(|| Typed::plain("any"), typed); + return Typed::plain(format!("{}[]", items.grouped())); + } + match self.take("type") { + Some(type_name) => type_names(type_name), + None => Typed::plain("any"), + } + } + + /// Consume `type` when every literal already has that type. + fn take_redundant_type<'v>(&mut self, values: impl IntoIterator) { + let Some(Value::String(type_name)) = self.schema.get("type") else { + return; + }; + let matches = |value: &Value| match value { + Value::String(_) => type_name == "string", + Value::Number(number) => { + type_name == "number" || (type_name == "integer" && !number.is_f64()) + } + Value::Bool(_) => type_name == "boolean", + Value::Null => type_name == "null", + Value::Object(_) => type_name == "object", + Value::Array(_) => type_name == "array", + }; + if values.into_iter().all(matches) { + self.take("type"); + } + } + + fn object(&mut self) -> String { + let required = self + .take("required") + .and_then(Value::as_array) + .map(|required| required.iter().filter_map(Value::as_str).collect::>()) + .unwrap_or_default(); + let properties = self.take("properties").and_then(Value::as_object); + let mut fields = properties + .into_iter() + .flatten() + .map(|(key, value)| { + let optional = if required.contains(&key.as_str()) { + "" + } else { + "?" + }; + format!("{key}{optional}: {}", schema(value)) + }) + .collect::>(); + // Required keys without a declared schema. + fields.extend( + required + .iter() + .filter(|key| properties.is_none_or(|properties| !properties.contains_key(**key))) + .map(|key| format!("{key}: any")), + ); + match self.take("additionalProperties") { + Some(Value::Bool(false)) | None => {} + Some(Value::Bool(true)) => fields.push("...".to_string()), + Some(additional) => fields.push(format!("...: {}", schema(additional))), + } + if fields.is_empty() { + "{}".to_string() + } else { + format!("{{ {} }}", fields.join(", ")) + } + } +} + +/// A `type` keyword's type names. +fn type_names(type_name: &Value) -> Typed { + match type_name { + Value::String(name) if name == "array" => Typed::plain("any[]"), + Value::String(name) => Typed::plain(name.clone()), + Value::Array(names) => Typed::union( + names + .iter() + .map(|name| name.as_str().map_or_else(|| name.to_string(), str::to_string)) + .collect(), + ), + type_name => Typed::plain(type_name.to_string()), + } +} + +fn quote(text: &str) -> String { + let mut quoted = String::with_capacity(text.len() + 2); + quoted.push('`'); + for char in text.chars() { + match char { + '\\' => quoted.push_str("\\\\"), + '`' => quoted.push_str("\\`"), + '\n' => quoted.push_str("\\n"), + '\r' => quoted.push_str("\\r"), + '\t' => quoted.push_str("\\t"), + char if char.is_control() => { + let _ = write!(quoted, "\\u{{{:x}}}", char as u32); + } + char => quoted.push(char), + } + } + quoted.push('`'); + quoted +} + +fn token(token: &TokenValue) -> String { + match token { + TokenValue::Id(id) => id.to_string(), + TokenValue::Text(text) => quote(text), + } +} + +fn push_excluding(header: &mut String, excludes: impl Iterator) { + let excludes = excludes.collect::>(); + if !excludes.is_empty() { + let _ = write!(header, " excluding [{}]", excludes.join(", ")); + } +} + +fn push_limit(header: &mut String, name: &str, limit: Option) { + if let Some(limit) = limit { + let _ = write!(header, " {name}={limit}"); + } +} + +fn push_flags(header: &mut String, at_least_one: bool, stop_after_first: bool) { + if at_least_one { + header.push_str(" at_least_one"); + } + if stop_after_first { + header.push_str(" stop_after_first"); + } +} + +#[cfg(test)] +mod tests { + use expect_test::expect; + use serde_json::json; + use xgrammar_structural_tag::format::{ + AnyTokensFormat, DispatchFormat, ExcludeTokenFormat, GrammarFormat, JsonSchemaFormat, + JsonSchemaStyle, TokenBoundary, TokenDispatchFormat, TokenFormat, TokenTriggeredTagsFormat, + }; + + use super::*; + + #[test] + fn outline_renders_every_format_type() { + let format = Format::sequence(vec![ + Format::Token(TokenFormat { + token: TokenValue::Id(7), + }), + Format::ExcludeToken(ExcludeTokenFormat { + exclude_tokens: vec![TokenValue::Id(1), TokenValue::Text("".to_string())], + }), + Format::AnyTokens(AnyTokensFormat { + exclude_tokens: vec![TokenValue::Id(2)], + max_tokens: Some(16), + }), + Format::Grammar(GrammarFormat { + grammar: "root ::= a\na ::= \"x\"".to_string(), + }), + Format::repeat(Format::regex("[0-9]"), 1, -1), + Format::JsonSchema( + JsonSchemaFormat::new(json!({ "type": "string" })) + .with_style(JsonSchemaStyle::KimiK3Xml) + .with_any_order(true) + .with_max_whitespace_cnt(Some(2)) + .with_excludes(&[""]), + ), + Format::TokenTriggeredTags(TokenTriggeredTagsFormat { + trigger_tokens: vec![TokenValue::Id(9)], + tags: vec![TagFormat::new( + TokenBoundary::new(9), + Format::optional(Format::const_string("`\t")), + vec!["", ""], + )], + exclude_tokens: vec![], + at_least_one: false, + stop_after_first: true, + }), + Format::Dispatch(DispatchFormat { + rules: vec![( + "".to_string(), + Format::sequence(vec![Format::any_text(), Format::const_string("")]), + )], + r#loop: false, + excludes: vec!["".to_string()], + }), + Format::TokenDispatch(TokenDispatchFormat { + rules: vec![( + TokenValue::Id(3), + Format::JsonSchema(JsonSchemaFormat::new(json!(true))), + )], + r#loop: true, + exclude_tokens: vec![], + }), + ]); + + expect![[r#" + sequence + token 7 + token excluding [1, ``] + tokens excluding [2] max_tokens=16 + grammar + root ::= a + a ::= "x" + repeat 1.. /[0-9]/ + kimi_k3_xml(string) any_order max_whitespace=2 excluding [``] + token_triggered_tags [9] stop_after_first + tag token 9 .. `` | `` + optional `\`\t` + dispatch once excluding [``] + on `` + sequence + text + `` + token_dispatch + on 3 json(any) + "#]] + .assert_eq(&outline(&format)); + } + + #[test] + fn schema_renders_common_shapes_as_types() { + let rendered = schema(&json!({ + "$defs": { "place": { "type": "object", "properties": { "city": { "type": "string" } } } }, + "type": "object", + "description": "ignored", + "properties": { + "location": { "type": "string" }, + "id": { "type": ["integer", "null"] }, + "unit": { "type": "string", "enum": ["celsius", "fahrenheit"] }, + "tags": { "type": "array", "items": { "anyOf": [{ "type": "string" }, { "type": "number" }] } }, + "place": { "$ref": "#/$defs/place" }, + "days": { "type": "integer", "minimum": 1 }, + "free": { "type": "object", "additionalProperties": true } + }, + "required": ["location", "id"], + "additionalProperties": false + })); + + expect![[r#"{ location: string, id: integer | null, unit?: "celsius" | "fahrenheit", tags?: (string | number)[], place?: place, days?: integer(minimum=1), free?: { ... } } where place = { city?: string }"#]].assert_eq(&rendered); + } + + #[test] + fn schema_keeps_other_keywords_after_the_type() { + let rendered = schema(&json!({ + "type": "object", + "properties": { + "code": { "type": "string", "pattern": "[A-Z]{3}" }, + "ids": { "type": "array", "items": { "type": ["integer", "null"] }, "maxItems": 3 }, + "place": { "$ref": "#/$defs/place", "type": "object" }, + "either": { "not": { "type": "null" } } + }, + "required": ["code", "tenant"], + "additionalProperties": { "type": "boolean" }, + "minProperties": 1 + })); + + expect![[r#"{ code: string(pattern="[A-Z]{3}"), ids?: (integer | null)[](maxItems=3), place?: place & object, either?: any(not={"type":"null"}), tenant: any, ...: boolean }(minProperties=1)"#]].assert_eq(&rendered); + } +} diff --git a/rust/src/parser/src/tool/hy/structural_tag.rs b/rust/src/parser/src/tool/hy/structural_tag.rs index e883acbafd..ed2dc6cec3 100644 --- a/rust/src/parser/src/tool/hy/structural_tag.rs +++ b/rust/src/parser/src/tool/hy/structural_tag.rs @@ -170,6 +170,7 @@ mod tests { }; use super::HyStructuralTagBuilder; + use crate::output_grammar::test_utils::outline; use crate::tool::{HyDialect, HyToolMarkers}; fn tool(name: &str, parameters: Value) -> ToolParam { @@ -219,7 +220,33 @@ mod tests { ToolChoice::required(), ); - expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"\n"},{"type":"tags_with_separator","tags":[{"type":"tag","begin":"get_weather\n","content":{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"city"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""],"max_tokens":null,"max_chars":null},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]},{"type":"star","content":{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"days"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""],"max_tokens":null,"max_chars":null},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]}}]},"end":""}],"separator":"\n","at_least_one":true,"stop_after_first":false},{"type":"const_string","value":"\n"}]}}"#]].assert_eq(&tag.to_json_string().unwrap()); + expect![[r#" + sequence + `\n` + tags_with_separator `\n` at_least_one + tag `get_weather\n` .. `` + sequence + sequence + `` + `city` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + star + sequence + `` + `days` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + `\n` + "#]].assert_eq(&outline(&tag.format)); } #[test] @@ -238,7 +265,21 @@ mod tests { ToolChoice::required(), ); - expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"tags_with_separator","tags":[{"type":"tag","begin":"get_weather","content":{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"city"},{"type":"const_string","value":""},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","",""],"max_tokens":null,"max_chars":null},{"type":"const_string","value":""}]}]},"end":""}],"separator":"","at_least_one":true,"stop_after_first":false},{"type":"const_string","value":""}]}}"#]].assert_eq(&tag.to_json_string().unwrap()); + expect![[r#" + sequence + `` + tags_with_separator `` at_least_one + tag `get_weather` .. `` + sequence + sequence + `` + `city` + `` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``] + `` + `` + "#]].assert_eq(&outline(&tag.format)); } #[test] @@ -257,14 +298,34 @@ mod tests { )], ToolChoice::required(), ); - let value = serde_json::to_value(tag).unwrap(); - - expect![[r#"{"type":"sequence","elements":[{"type":"star","content":{"type":"or","elements":[{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"query"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""],"max_tokens":null,"max_chars":null},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]},{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"limit"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""],"max_tokens":null,"max_chars":null},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]}]}}]}"#]].assert_eq( - &serde_json::to_string( - &value["format"]["elements"][1]["tags"][0]["content"], - ) - .unwrap(), - ); + expect![[r#" + sequence + `\n` + tags_with_separator `\n` at_least_one + tag `lookup\n` .. `` + sequence + star + or + sequence + `` + `query` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + sequence + `` + `limit` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + `\n` + "#]].assert_eq(&outline(&tag.format)); } #[test] @@ -273,11 +334,24 @@ mod tests { tool("search", json!({ "type": "object" })), tool("lookup", json!({ "type": "object" })), ]; - let auto = build("", &tools, ToolChoice::auto()).to_json_string().unwrap(); - let forced = build("", &tools, ToolChoice::function("lookup")).to_json_string().unwrap(); + let auto = build("", &tools, ToolChoice::auto()); + let forced = build("", &tools, ToolChoice::function("lookup")); - expect![[r#"{"type":"structural_tag","format":{"type":"triggered_tags","triggers":[""],"tags":[{"type":"tag","begin":"\n","content":{"type":"tags_with_separator","tags":[{"type":"tag","begin":"search\n","content":{"type":"any_text","excludes":[],"max_tokens":null,"max_chars":null},"end":""},{"type":"tag","begin":"lookup\n","content":{"type":"any_text","excludes":[],"max_tokens":null,"max_chars":null},"end":""}],"separator":"\n","at_least_one":true,"stop_after_first":false},"end":"\n"}],"at_least_one":false,"stop_after_first":false,"excludes":[]}}"#]].assert_eq(&auto); - expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"\n"},{"type":"tag","begin":"lookup\n","content":{"type":"any_text","excludes":[],"max_tokens":null,"max_chars":null},"end":""},{"type":"const_string","value":"\n"}]}}"#]].assert_eq(&forced); + expect![[r#" + triggered_tags [``] + tag `\n` .. `\n` + tags_with_separator `\n` at_least_one + tag `search\n` text `` + tag `lookup\n` text `` + "#]] + .assert_eq(&outline(&auto.format)); + expect![[r#" + sequence + `\n` + tag `lookup\n` text `` + `\n` + "#]] + .assert_eq(&outline(&forced.format)); } #[test] @@ -291,11 +365,13 @@ mod tests { )], ToolChoice::function("get_current_date"), ); - let value = serde_json::to_value(tag).unwrap(); - - expect![[r#"{"type":"tag","begin":"get_current_date","content":{"type":"const_string","value":""},"end":""}"#]].assert_eq( - &serde_json::to_string(&value["format"]["elements"][1]).unwrap(), - ); + expect![[r#" + sequence + `` + tag `get_current_date` `` `` + `` + "#]] + .assert_eq(&outline(&tag.format)); } #[test] @@ -312,13 +388,31 @@ mod tests { )], ToolChoice::required(), ); - let value = serde_json::to_value(tag).unwrap(); - - expect![[r#"{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"query"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""],"max_tokens":null,"max_chars":null},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]},{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"tenant"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""],"max_tokens":null,"max_chars":null},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]}]}"#]].assert_eq( - &serde_json::to_string( - &value["format"]["elements"][1]["tags"][0]["content"], - ) - .unwrap(), - ); + expect![[r#" + sequence + `\n` + tags_with_separator `\n` at_least_one + tag `search\n` .. `` + sequence + sequence + `` + `query` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + sequence + `` + `tenant` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + `\n` + "#]].assert_eq(&outline(&tag.format)); } } diff --git a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs index f542f34e67..a80be7bb1e 100644 --- a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs +++ b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs @@ -364,6 +364,7 @@ mod tests { }; use super::KimiK3StructuralTagBuilder; + use crate::output_grammar::test_utils::outline; fn tool(name: &str, parameters: serde_json::Value) -> ToolParam { ToolParam::Function(FunctionToolParam::new( @@ -395,7 +396,27 @@ mod tests { ) .unwrap(); - expect![[r##"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"optional","content":{"type":"const_string","value":"<|open|>response<|sep|>"}},{"type":"tag","begin":"","content":{"type":"any_text","excludes":["<|open|>","<|close|>","<|end_of_msg|>"],"max_tokens":null,"max_chars":null},"end":"<|close|>response<|sep|>"},{"type":"tag","begin":"<|open|>tools<|sep|>","content":{"type":"tags_with_separator","tags":[{"type":"tag","begin":"<|open|>call tool=\"get_weather\" index=\"","content":{"type":"sequence","elements":[{"type":"regex","pattern":"[1-9][0-9]*"},{"type":"const_string","value":"\"<|sep|>"},{"type":"or","elements":[{"type":"plus","content":{"type":"or","elements":[{"type":"tag","begin":"<|open|>argument key=\"unit\" type=\"string\"<|sep|>","content":{"type":"or","elements":[{"type":"const_string","value":"celsius"},{"type":"const_string","value":"fahrenheit"}]},"end":"<|close|>argument<|sep|>"},{"type":"tag","begin":"<|open|>argument key=\"place\" type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$ref":"#/$defs/place","type":"object","$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}}},"style":"json","any_order":false,"max_whitespace_cnt":null,"excludes":[]},"end":"<|close|>argument<|sep|>"}]}},{"type":"tag","begin":"<|open|>json type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}},"type":"object","properties":{"unit":{"type":"string","enum":["celsius","fahrenheit"]},"place":{"$ref":"#/$defs/place","type":"object"}},"required":["place"]},"style":"json","any_order":false,"max_whitespace_cnt":null,"excludes":[]},"end":"<|close|>json<|sep|>"}]}]},"end":"<|close|>call<|sep|>"}],"separator":"","at_least_one":true,"stop_after_first":false},"end":"<|close|>tools<|sep|>"},{"type":"optional","content":{"type":"const_string","value":"<|close|>message<|sep|>"}}]}}"##]].assert_eq(&tag.to_json_string().unwrap()); + expect![[r#" + sequence + optional `<|open|>response<|sep|>` + tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + plus + or + tag `<|open|>argument key="unit" type="string"<|sep|>` .. `<|close|>argument<|sep|>` + or + `celsius` + `fahrenheit` + tag `<|open|>argument key="place" type="object"<|sep|>` json(place & object where place = { city?: string }) `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ unit?: "celsius" | "fahrenheit", place: place & object } where place = { city?: string }) `<|close|>json<|sep|>` + optional `<|close|>message<|sep|>` + "#]].assert_eq(&outline(&tag.format)); } #[test]