[Rust Frontend] Build the Kimi K3 full-output grammar from its parser mode

rust/src/chat/tests/grammar_replay/Kimi-K3.txt

@@ -1,8 +1,29 @@
-# auto (final_output_only)
+# auto (from_token_zero)
 sequence
-  optional `<|open|>response<|sep|>`
-  tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>`
-  optional
+  tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>`
+  or
+    sequence
+      tag `<|open|>response<|sep|>` 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
+                  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
+                  sequence
+                    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|>`
     tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>`
       tags_with_separator `` at_least_one
         tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>`
@@ -24,42 +45,44 @@ sequence
               tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>`
   optional `<|close|>message<|sep|>`
 
-# required (final_output_only)
+# required (from_token_zero)
 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
-            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
-            sequence
-              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|>`
+  tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>`
+  sequence
+    optional tag `<|open|>response<|sep|>` 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
+              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
+              sequence
+                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)
+# named (from_token_zero)
 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
-            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 `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>`
+  sequence
+    optional tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>`
+    tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>`
+      tags_with_separator `` at_least_one stop_after_first
+        tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>`
+          sequence
+            /[1-9][0-9]*/
+            `"<|sep|>`
+            or
+              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|>`

rust/src/parser/src/output_grammar.rs

@@ -12,7 +12,8 @@ use xgrammar_structural_tag::builders::{StructuralTagBuilder, StructuralTagOptio
 use xgrammar_structural_tag::format::Format;
 use xgrammar_structural_tag::tool::ToolChoiceValue;
 use xgrammar_structural_tag::{
-    FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag,
+    FunctionDefinition, FunctionToolParam, NormalizedToolChoice, ToolChoice, ToolParam,
+    build_structural_tag,
 };
 
 use crate::tool::Tool;
@@ -218,6 +219,22 @@ fn tool_params(tools: &[Tool], strict_level: ToolStrictLevel) -> Vec<ToolParam>
         .collect()
 }
 
+/// Normalize the request tools the way `build_structural_tag` does before
+/// calling a crate builder, for model-owned grammar builders. Returns `None`
+/// when the request asks for no tool grammar.
+pub(crate) fn normalize_tool_choice(
+    ctx: &OutputGrammarContext<'_>,
+) -> Result<Option<NormalizedToolChoice>> {
+    if !tool_grammar_applies(ctx) {
+        return Ok(None);
+    }
+
+    let tools = tool_params(ctx.tools, ctx.tool_strict_level);
+    let normalized =
+        xgrammar_structural_tag::normalize_tool_choice(&tools, ctx.tool_choice.clone())?;
+    Ok(Some(normalized))
+}
+
 /// Whether the request asks for a tool grammar at all.
 fn tool_grammar_applies(ctx: &OutputGrammarContext<'_>) -> bool {
     if ctx.tools.is_empty() {

rust/src/parser/src/unified/kimi_k3.rs

@@ -46,9 +46,7 @@
 //! now only applies to `argument` and `json` blocks inside a call body, which
 //! are still parsed from text.
 
-mod structural_tag;
-
-pub use structural_tag::KimiK3StructuralTagBuilder;
+mod grammar;
 
 use serde_json::{Map, Value};
 use vllm_tokenizer::{DecodedText, DynTokenizer, Tokenizer};
@@ -58,11 +56,8 @@ use winnow::error::{ContextError, ErrMode, ModalResult, StrContext};
 use winnow::prelude::*;
 use winnow::token::{literal, rest, take_till, take_until, take_while};
 
-use self::structural_tag::KIMI_K3_STRUCTURAL_TAG_BUILDER;
 use super::{Result, UnifiedParser, UnifiedParserOutput, special_token};
-use crate::output_grammar::{
-    self, BuiltOutputGrammar, OutputGrammarContext, visible_format_from_builder,
-};
+use crate::output_grammar::{self, BuiltOutputGrammar, OutputGrammarContext};
 use crate::tool::{Tool, ToolCallDelta};
 use crate::unified::parsing_failed;
 use crate::utils::{
@@ -347,10 +342,7 @@ impl UnifiedParser for KimiK3UnifiedParser {
         &self,
         ctx: &OutputGrammarContext<'_>,
     ) -> output_grammar::Result<Option<BuiltOutputGrammar>> {
-        Ok(
-            visible_format_from_builder(Some(&KIMI_K3_STRUCTURAL_TAG_BUILDER), ctx)?
-                .map(BuiltOutputGrammar::final_output_only),
-        )
+        grammar::build_output_grammar(&self.mode, ctx)
     }
 
     fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> {