diff --git a/rust/src/parser/src/unified/kimi_k3.rs b/rust/src/parser/src/unified/kimi_k3.rs index a721b29f3f..932bab656a 100644 --- a/rust/src/parser/src/unified/kimi_k3.rs +++ b/rust/src/parser/src/unified/kimi_k3.rs @@ -56,6 +56,7 @@ use crate::unified::parsing_failed; use crate::utils::{MarkerScanState, parse_buffered_event, safe_text_len_mul, take_until_marker}; const OPEN: &str = "<|open|>"; +const CLOSE: &str = "<|close|>"; const SEP: &str = "<|sep|>"; const END_OF_MSG: &str = "<|end_of_msg|>"; 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 d3515ad2dc..0a86a274e1 100644 --- a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs +++ b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs @@ -12,8 +12,8 @@ use xgrammar_structural_tag::format::{Format, JsonSchemaFormat, StructuralTag, T use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam, function_parameters}; use super::{ - ARG_CLOSE, CALL_CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN, RESPONSE_CLOSE, - RESPONSE_OPEN, SEP, THINK_CLOSE, TOOLS_CLOSE, TOOLS_OPEN, + ARG_CLOSE, CALL_CLOSE, CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN, + RESPONSE_CLOSE, RESPONSE_OPEN, SEP, THINK_CLOSE, TOOLS_CLOSE, TOOLS_OPEN, }; pub(super) static KIMI_K3_STRUCTURAL_TAG_BUILDER: KimiK3StructuralTagBuilder = @@ -64,7 +64,9 @@ fn response_prefix(reasoning: bool) -> Vec { } elements.push(Format::tag( "", - Format::any_text_excluding(&[RESPONSE_CLOSE, TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]), + // Keep marker-looking prefixes out of response text so the grammar + // must resolve them through a valid channel boundary. + Format::any_text_excluding(&[OPEN, CLOSE, END_OF_MSG]), RESPONSE_CLOSE, )); elements @@ -384,7 +386,7 @@ 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":["<|close|>response<|sep|>","<|open|>tools<|sep|>","<|close|>message<|sep|>","<|end_of_msg|>"]},"end":"<|close|>response<|sep|>"},{"type":"tag","begin":"<|open|>tools<|sep|>","content":{"type":"tags_with_separator","tags":[{"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},"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},"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##"{"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|>"]},"end":"<|close|>response<|sep|>"},{"type":"tag","begin":"<|open|>tools<|sep|>","content":{"type":"tags_with_separator","tags":[{"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},"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},"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()); } #[test] diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py index b84dfad092..c1e2e09437 100644 --- a/tests/tool_parsers/test_structural_tag_registry.py +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -177,6 +177,7 @@ _K3_TOOLS_CLOSE = "<|close|>tools<|sep|>" _K3_CALL_CLOSE = "<|close|>call<|sep|>" _K3_ARG_CLOSE = "<|close|>argument<|sep|>" _K3_MESSAGE_CLOSE = "<|close|>message<|sep|>" +_K3_END_OF_MSG = "<|end_of_msg|>" def _k3_tools_by_name() -> list[ChatCompletionToolsParam]: @@ -357,6 +358,20 @@ def test_kimi_k3_auto_strict_allows_response_only(sample_tools_strict): assert _is_grammar_accept_string(grammar, _k3_response("Just answering.")) +@pytest.mark.parametrize( + "body", + [ + _K3_RESPONSE_OPEN + "answer<|close|>message", + _K3_RESPONSE_OPEN + "answer<|open|>tools", + _K3_RESPONSE_OPEN + "answer" + _K3_END_OF_MSG, + ], +) +def test_kimi_k3_response_body_rejects_reserved_marker_prefixes(body: str): + assert not _is_grammar_accept_string( + _k3_grammar("required"), body, require_termination=False + ) + + @pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS)) def test_get_model_structural_tag_supports_named_tool_choice( model: str, diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py index 1f9a651ee4..7710082b3b 100644 --- a/vllm/tool_parsers/structural_tag_registry.py +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -376,6 +376,7 @@ _K3_ARG_CLOSE = f"{_K3_CLOSE}argument{_K3_SEP}" # the end-of-message token. It is generated (not part of the prompt prefix), so # the tag must permit it or the FSM would mask the model's natural terminator. _K3_MESSAGE_CLOSE = f"{_K3_CLOSE}message{_K3_SEP}" +_K3_END_OF_MSG = "<|end_of_msg|>" # JSON-schema type -> K3 XTML ``type=`` attribute value. Mirrors # ``encoding_k3._xtml_type`` (integer collapses onto number). @@ -556,12 +557,17 @@ def _k3_response_prefix() -> list[Any]: ``response`` is generated in thinking mode (prefix ends at ``<|open|>think<|sep|>``) but is part of the generation prefix in - non-thinking mode, so its open marker is optional. The body is bounded by - the response close marker. + non-thinking mode, so its open marker is optional. Excluding the reserved + marker heads commits ``<|close|>`` to the response-close branch immediately + and prevents a nested channel from starting inside response text. """ return [ OptionalFormat(content=ConstStringFormat(value=_K3_RESPONSE_OPEN)), - TagFormat(begin="", content=AnyTextFormat(), end=_K3_RESPONSE_CLOSE), + TagFormat( + begin="", + content=AnyTextFormat(excludes=[_K3_OPEN, _K3_CLOSE, _K3_END_OF_MSG]), + end=_K3_RESPONSE_CLOSE, + ), ]