[Bugfix] Hoist $defs/definitions in Cohere parser tool schema composition (#49602)
Signed-off-by: Václav Čadek <vaclavcadek@gmail.com> Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Václav Čadek committed
Sep 29, 2026 at 09:11 UTC
6cbbea872141851421dcad5fca009aa6cec93dc6
2 files changed
+109
-16
tests/parser/cohere/test_structural_tags.py
+96
@@ -10,6 +10,7 @@ from types import SimpleNamespace
10
11
import pytest
12
from pydantic import ValidationError
13
+from xgrammar.testing import _is_grammar_accept_string
14
15
from vllm.entrypoints.generate.base.protocol import (
16
JsonSchemaResponseFormat,
@@ -20,9 +21,11 @@ from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionReque
21
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
22
from vllm.parser.cohere_command import (
23
CohereCommandParser,
24
+ CohereNormalizedTool,
25
_has_effective_tools,
26
_response_format_type,
27
_schema_dict_from_structured_outputs,
28
+ collect_tool_schema,
29
convert_schema_to_structural_tags,
30
)
31
from vllm.sampling_params import StructuredOutputsParams
@@ -41,6 +44,48 @@ GET_WEATHER_TOOL = {
44
},
45
},
46
}
47
+NESTED_DEFS_PARAMETERS = {
48
+ "$defs": {
49
+ "Location": {
50
+ "type": "object",
51
+ "properties": {
52
+ "city": {"type": "string"},
53
+ "country": {"type": "string"},
54
+ },
55
+ "required": ["city"],
56
+ }
57
+ },
58
+ "type": "object",
59
+ "properties": {"location": {"$ref": "#/$defs/Location"}},
60
+ "required": ["location"],
61
+}
62
+NESTED_DEFS_TOOL = {
63
+ "type": "function",
64
+ "function": {"name": "get_weather", "parameters": NESTED_DEFS_PARAMETERS},
65
+}
66
+RECURSIVE_DEFS_PARAMETERS = {
67
+ "$defs": {
68
+ "Finding": {
69
+ "type": "object",
70
+ "properties": {
71
+ "note": {"type": "string"},
72
+ "related": {"type": "array", "items": {"$ref": "#/$defs/Finding"}},
73
+ },
74
+ }
75
+ },
76
+ "type": "object",
77
+ "properties": {"finding": {"$ref": "#/$defs/Finding"}},
78
+}
79
+DRAFT07_DEFINITIONS_PARAMETERS = {
80
+ "definitions": {
81
+ "Location": {
82
+ "type": "object",
83
+ "properties": {"city": {"type": "string"}},
84
+ }
85
+ },
86
+ "type": "object",
87
+ "properties": {"location": {"$ref": "#/definitions/Location"}},
88
+}
89
VALID_STRUCTURAL_TAG = {
90
"type": "structural_tag",
91
"format": {
@@ -360,3 +405,54 @@ class TestAdjustRequestTools:
405
types = _content_types(o.structured_outputs.structural_tag)
406
assert "grammar" in types
407
assert "json_schema" in types
408
+
409
+
410
+class TestCollectToolSchemaDefs:
411
+ def test_nested_defs_tool_adjust_request(self, parser) -> None:
412
+ o = parser.adjust_request(
413
+ _make_chat_request(tools=[NESTED_DEFS_TOOL], tool_choice="auto"),
414
+ )
415
+ assert "grammar" in _content_types(o.structured_outputs.structural_tag)
416
+
417
+ @staticmethod
418
+ def _tool_calls_json(name: str, parameters: dict) -> str:
419
+ return json.dumps(
420
+ [{"tool_call_id": "0", "tool_name": name, "parameters": parameters}]
421
+ )
422
+
423
+ def test_recursive_defs(self) -> None:
424
+ grammar = collect_tool_schema(
425
+ [CohereNormalizedTool(name="report", parameters=RECURSIVE_DEFS_PARAMETERS)]
426
+ )
427
+ nested = {"note": "a", "related": [{"note": "b", "related": []}]}
428
+ assert _is_grammar_accept_string(
429
+ grammar, self._tool_calls_json("report", {"finding": nested})
430
+ )
431
+ bad = {"note": 1, "related": []}
432
+ assert not _is_grammar_accept_string(
433
+ grammar, self._tool_calls_json("report", {"finding": bad})
434
+ )
435
+
436
+ def test_draft07_definitions(self) -> None:
437
+ grammar = collect_tool_schema(
438
+ [
439
+ CohereNormalizedTool(
440
+ name="get_weather", parameters=DRAFT07_DEFINITIONS_PARAMETERS
441
+ )
442
+ ]
443
+ )
444
+ assert _is_grammar_accept_string(
445
+ grammar,
446
+ self._tool_calls_json("get_weather", {"location": {"city": "Prague"}}),
447
+ )
448
+ assert not _is_grammar_accept_string(
449
+ grammar,
450
+ self._tool_calls_json("get_weather", {"location": {"city": 1}}),
451
+ )
452
+
453
+ def test_input_parameters_not_mutated(self) -> None:
454
+ params = json.loads(json.dumps(NESTED_DEFS_PARAMETERS))
455
+ collect_tool_schema(
456
+ [CohereNormalizedTool(name="get_weather", parameters=params)]
457
+ )
458
+ assert params == NESTED_DEFS_PARAMETERS
vllm/parser/cohere_command.py
+13
-16
@@ -127,22 +127,19 @@ def collect_tool_schema(tool_schema: list[CohereNormalizedTool]) -> str:
127
tool_dictionary: dict[str, str] = {}
128
for tool in tool_schema:
129
tool_name = tool["name"]
130
- tool_parameters = json.dumps(tool["parameters"])
131
- json_schema = f"""{{
132
- "type": "object",
133
- "properties": {{
134
- "tool_call_id": {{
135
- "type": "string",
136
- "pattern": "^[0-9]+$"
137
- }},
138
- "tool_name": {{
139
- "type": "string",
140
- "const": "{tool_name}"
141
- }},
142
- "parameters": {tool_parameters}
143
- }}
144
- }}"""
145
- tool_grammar = str(xgr.Grammar.from_json_schema(json_schema))
130
+ tool_parameters = dict(tool["parameters"])
131
+ json_schema: dict[str, Any] = {
132
+ "type": "object",
133
+ "properties": {
134
+ "tool_call_id": {"type": "string", "pattern": "^[0-9]+$"},
135
+ "tool_name": {"type": "string", "const": tool_name},
136
+ "parameters": tool_parameters,
137
+ },
138
+ }
139
+ for defs_key in ("$defs", "definitions"):
140
+ if defs_key in tool_parameters:
141
+ json_schema[defs_key] = tool_parameters.pop(defs_key)
142
+ tool_grammar = str(xgr.Grammar.from_json_schema(json.dumps(json_schema)))
143
for match in re.findall(r"\b(\w+)\s*::=", tool_grammar):
144
tool_grammar = re.sub(
145
rf"\b{re.escape(match)}\b", tool_name + match, tool_grammar