[Frontend] Use XGrammar schema constraints for DeepSeek V4.1 (#56408)
Signed-off-by: Ubospica <ubospica@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com>
Yixin Dong committed
Sep 15, 2026 at 04:36 UTC
2c2cbd84fdf199901f3f4e6b0e708ee0664aadd6
2 files changed
+90
-4
tests/tool_parsers/test_structural_tag_registry.py
+77
-3
@@ -5,7 +5,7 @@ from types import SimpleNamespace
5
from unittest.mock import MagicMock
6
7
import pytest
8
-from xgrammar import Grammar, StructuralTag
8
+from xgrammar import Grammar, StructuralTag, builtin_structural_tag
9
from xgrammar.testing import _is_grammar_accept_string
10
11
from vllm.entrypoints.openai.chat_completion.protocol import (
@@ -81,10 +81,14 @@ def test_supported_structural_tag_models_include_vllm_builtins():
81
82
83
@pytest.mark.parametrize("choice", ["auto", "required", "get_weather"])
84
-def test_deepseek_v41_format_constraints_ignore_parameter_schema(
84
+def test_deepseek_v41_older_xgrammar_keeps_format_constraints(
85
choice,
86
sample_tools_strict,
87
+ monkeypatch,
88
):
89
+ monkeypatch.setattr(
90
+ builtin_structural_tag, "get_deepseek_v4_1_structural_tag", None, raising=False
91
+ )
92
tool_choice = (
93
ChatCompletionNamedToolChoiceParam(
94
function=ChatCompletionNamedFunction(name=choice)
@@ -128,7 +132,11 @@ def test_deepseek_v41_named_choice_emits_one_call(sample_tools):
132
reasoning=False,
133
)
134
grammar = Grammar.from_structural_tag(tag)
131
- call = '<｜DSML｜ invoke name="get_weather">\n</｜DSML｜ invoke>\n'
135
+ call = (
136
+ '<｜DSML｜ invoke name="get_weather">\n'
137
+ '<｜DSML｜ parameter name="city" string="true">Paris</｜DSML｜ parameter>\n'
138
+ "</｜DSML｜ invoke>\n"
139
+ )
140
assert _is_grammar_accept_string(
141
grammar, "\n\n<｜DSML｜ calls>\n" + call + "</｜DSML｜ calls>"
142
)
@@ -141,6 +149,72 @@ def test_deepseek_v41_named_choice_emits_one_call(sample_tools):
149
)
150
151
152
+@pytest.mark.skipif(
153
+ not hasattr(builtin_structural_tag, "get_deepseek_v4_1_structural_tag"),
154
+ reason="Requires XGrammar's DeepSeek V4.1 builtin",
155
+)
156
+@pytest.mark.parametrize("choice", ["auto", "required", "get_weather"])
157
+@pytest.mark.parametrize("reasoning", [False, True])
158
+def test_deepseek_v41_constrains_parameters_after_reasoning(
159
+ sample_tools_strict, choice, reasoning
160
+):
161
+ """Strict calls must enforce the schema without requiring another think close."""
162
+ tool_choice = (
163
+ ChatCompletionNamedToolChoiceParam(
164
+ function=ChatCompletionNamedFunction(name=choice)
165
+ )
166
+ if choice == "get_weather"
167
+ else choice
168
+ )
169
+ sample_tools_strict[0].function.parameters["additionalProperties"] = False
170
+ tag = get_model_structural_tag(
171
+ "deepseek_v41", sample_tools_strict, tool_choice, reasoning=reasoning
172
+ )
173
+ grammar = Grammar.from_structural_tag(tag)
174
+ begin = '\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name="get_weather">\n'
175
+ end = "</｜DSML｜ invoke>\n</｜DSML｜ calls>"
176
+ parameter = (
177
+ '<｜DSML｜ parameter name="city" string="true">Paris</｜DSML｜ parameter>\n'
178
+ )
179
+ assert _is_grammar_accept_string(grammar, begin + parameter + end)
180
+ for invalid in (
181
+ "", # missing required city
182
+ parameter * 2,
183
+ parameter.replace('name="city"', 'name="unknown"'),
184
+ parameter.replace('string="true">Paris', 'string="false">42'),
185
+ parameter.replace("｜ parameter", "｜parameter"),
186
+ '{"city":"Paris"}', # the encoder emits DSML parameters, not a JSON body
187
+ ):
188
+ assert not _is_grammar_accept_string(grammar, begin + invalid + end)
189
+ assert not _is_grammar_accept_string(
190
+ grammar, "reason</think>" + begin + parameter + end
191
+ )
192
+ assert _is_grammar_accept_string(grammar, "Hello") == (choice == "auto")
193
+
194
+
195
+@pytest.mark.skipif(
196
+ not hasattr(builtin_structural_tag, "get_deepseek_v4_1_structural_tag"),
197
+ reason="Requires XGrammar's DeepSeek V4.1 builtin",
198
+)
199
+def test_deepseek_v41_non_strict_parallel_calls_keep_typed_dsml(sample_tools):
200
+ sample_tools[0].function.strict = False
201
+ tag = get_model_structural_tag(
202
+ "deepseek_v41", sample_tools, "required", reasoning=False
203
+ )
204
+ grammar = Grammar.from_structural_tag(tag)
205
+ call = (
206
+ '<｜DSML｜ invoke name="get_weather">\n'
207
+ '<｜DSML｜ parameter name="extra" string="false">'
208
+ '{"nested":[true,null,1.5]}</｜DSML｜ parameter>\n'
209
+ "</｜DSML｜ invoke>\n"
210
+ )
211
+ output = "\n\n<｜DSML｜ calls>\n" + call * 2 + "</｜DSML｜ calls>"
212
+ assert _is_grammar_accept_string(grammar, output)
213
+ assert not _is_grammar_accept_string(
214
+ grammar, output.replace('{"nested":[true,null,1.5]}', "invalid")
215
+ )
216
+
217
+
218
@pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS))
219
def test_get_model_structural_tag_supports_all_xgrammar_builtins(
220
model: str,
vllm/tool_parsers/structural_tag_registry.py
+13
-1
@@ -10,6 +10,7 @@ from openai.types.responses.tool import Tool as ResponsesTool
10
from openai.types.responses.tool_choice_allowed import ToolChoiceAllowed
11
from openai.types.responses.tool_choice_function import ToolChoiceFunction
12
from xgrammar import StructuralTag, normalize_tool_choice
13
+from xgrammar import builtin_structural_tag as xgrammar_builtin_structural_tag
14
from xgrammar import get_model_structural_tag as get_xgrammar_model_structural_tag
15
from xgrammar.openai_tool_call_schema import (
16
BuiltinToolParam,
@@ -242,9 +243,20 @@ def get_deepseek_v41_structural_tag(
243
token_suffix: str = "",
244
) -> StructuralTag:
245
# Serving enables this visible-text grammar after the reasoning boundary.
246
+ builder = getattr(
247
+ xgrammar_builtin_structural_tag, "get_deepseek_v4_1_structural_tag", None
248
+ )
249
+ if builder is not None:
250
+ return builder(
251
+ tools=tools,
252
+ builtin_tools=builtin_tools,
253
+ tool_choice=tool_choice,
254
+ reasoning="disabled",
255
+ )
256
+
257
del builtin_tools, reasoning, token_suffix
258
247
- # TODO: lower parameter schemas into DSML constraints. This builder constrains
259
+ # Compatibility with xgrammar releases without the V4.1 builtin. This constrains
260
# tool names and DSML/value syntax only. Parameter names, presence, uniqueness
261
# and value schemas remain unconstrained, including for strict=true. The request
262
# layer still uses strict to decide whether auto tool choice activates a grammar.