[Frontend] create unified Cohere parser (#56392)
Signed-off-by: walterbm <walter.beller.morales@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com>
Walter Beller-Morales committed
Sep 12, 2026 at 04:04 UTC
756794a9a7f08900c00fbfaa6d8332631503f528
21 files changed
+1917
-2137
docs/features/reasoning_outputs.md
+1
@@ -15,6 +15,7 @@ vLLM currently supports the following reasoning models:
15
| Model Series | Parser Name | Structured Output Support | Tool Calling |
16
| ------------ | ----------- | ---------------- | ----------- |
17
| [Cohere Command A Reasoning](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) | `cohere_command3` | `json`, `regex` | ✅ |
18
+| [Cohere Command A Plus](https://huggingface.co/CohereLabs/command-a-plus-05-2026-bf16) | `cohere_command4` | `json`, `regex` | ✅ |
19
| [DeepSeek R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d) | `deepseek_r1` | `json`, `regex` | ❌ |
20
| [Gemma 4 series](https://huggingface.co/google/gemma-4-26B-A4B-it) | `gemma4` | `json`, `regex` | ✅ |
21
| [DeepSeek-V3.1](https://huggingface.co/collections/deepseek-ai/deepseek-v31-68a491bed32bd77e7fca048f) | `deepseek_v3` | `json`, `regex` | ❌ |
docs/features/tool_calling.md
+5
-4
@@ -366,15 +366,16 @@ Flags:
366
* For non-reasoning: `--tool-call-parser hunyuan_a13b`
367
* For reasoning: `--tool-call-parser hunyuan_a13b --reasoning-parser hunyuan_a13b`
368
369
-### Cohere Command A Reasoning (`cohere_command3`)
369
+### Cohere Command (`cohere_command3` / `cohere_command4`)
370
371
Supported models:
372
373
-* [`CohereLabs/command-a-reasoning-08-2025`](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025)
373
+* [`CohereLabs/command-a-reasoning-08-2025`](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) (`cohere_command3`)
374
+* [`CohereLabs/command-a-plus-05-2026`](https://huggingface.co/CohereLabs/command-a-plus-05-2026-bf16) and [`CohereLabs/North-Mini-Code-1.0`](https://huggingface.co/CohereLabs/North-Mini-Code-1.0) (`cohere_command4`)
375
375
-Flags: `--tool-call-parser cohere_command3 --reasoning-parser cohere_command3`
376
+Flags: `--tool-call-parser cohere_command4 --reasoning-parser cohere_command4`.
377
377
-Note: the Cohere tool parser requires the `cohere_melody` package, which is not installed by default. Before using this parser please install the [cohere_melody](https://pypi.org/project/cohere-melody/) package.
378
+Note: the Cohere parser requires the `cohere_melody` package, which is not installed by default. Before using this parser please install the [cohere_melody](https://pypi.org/project/cohere-melody/) package.
379
380
### LongCat-Flash-Chat Models (`longcat`)
381
requirements/test/cpu.txt
+1
-1
@@ -127,7 +127,7 @@ cloudpickle==3.1.2
127
# via -r requirements/test/../common.txt
128
cohere==7.0.8
129
# via -r requirements/test/cuda.in
130
-cohere-melody==0.11.1
130
+cohere-melody==0.14.0
131
# via -r requirements/test/cuda.in
132
colorama==0.4.6
133
# via
requirements/test/cuda.txt
+1
-1
@@ -132,7 +132,7 @@ cloudpickle==3.1.2
132
# via -r requirements/test/../common.txt
133
cohere==7.0.8
134
# via -r requirements/test/cuda.in
135
-cohere-melody==0.11.1
135
+cohere-melody==0.14.0
136
# via -r requirements/test/cuda.in
137
colorama==0.4.6
138
# via
requirements/test/rocm.txt
+1
-1
@@ -134,7 +134,7 @@ cloudpickle==3.1.2
134
# tilelang
135
cohere==7.0.8
136
# via -r requirements/test/rocm.in
137
-cohere-melody==0.11.1
137
+cohere-melody==0.14.0
138
# via -r requirements/test/rocm.in
139
colorama==0.4.6
140
# via
tests/parser/cohere/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/parser/cohere/conftest.py
new
+18
@@ -0,0 +1,18 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
7
+
8
+from .utils import MockCohereTokenizer
9
+
10
+
11
+@pytest.fixture(scope="package")
12
+def tokenizer() -> MockCohereTokenizer:
13
+ return MockCohereTokenizer()
14
+
15
+
16
+@pytest.fixture
17
+def request_obj() -> ChatCompletionRequest:
18
+ return ChatCompletionRequest(messages=[], model="test-model")
tests/parser/cohere/test_citations.py
new
+172
@@ -0,0 +1,172 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Citation source resolution and the streamed citation wire shape."""
4
+
5
+from __future__ import annotations
6
+
7
+from types import SimpleNamespace
8
+from typing import Any
9
+
10
+from vllm.parser.cohere_command import (
11
+ _melody_citations_to_vllm,
12
+ _melody_sources_to_vllm,
13
+)
14
+
15
+from .utils import make_parser, stream_parser
16
+
17
+
18
+class TestMelodySourceResolution:
19
+ """Pin the parser-side source resolution (see
20
+ ``_melody_sources_to_vllm``). The parser receives melody's numeric
21
+ ``(tool_call_index, tool_result_indices)`` addressing and resolves
22
+ it against a ``position_to_source`` map handed in by
23
+ ``CohereServingChatV2._apply_cohere_template_kwargs`` -- the same
24
+ map that would otherwise live in the serving layer's resolver.
25
+ """
26
+
27
+ @staticmethod
28
+ def _fake_melody_source(bucket: int, indices: list[int]) -> Any:
29
+ return SimpleNamespace(tool_call_index=bucket, tool_result_indices=indices)
30
+
31
+ def test_multi_index_source_fans_out(self):
32
+ from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
33
+
34
+ position_map: dict[tuple[int, int], CitationSource] = {
35
+ (0, 0): CitationSource(type="document", id="d0", document={"id": "d0"}),
36
+ (0, 1): CitationSource(type="document", id="d1", document={"id": "d1"}),
37
+ }
38
+ raw = [self._fake_melody_source(0, [0, 1])]
39
+ out = _melody_sources_to_vllm(raw, position_map)
40
+ assert [s.id for s in out] == ["d0", "d1"]
41
+ # Verify type / payload were plumbed through, not just the id.
42
+ assert out[0].type == "document"
43
+ assert out[0].document == {"id": "d0"}
44
+
45
+ def test_unresolvable_position_skipped(self):
46
+ from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
47
+
48
+ position_map: dict[tuple[int, int], CitationSource] = {
49
+ (0, 0): CitationSource(type="document", id="d0"),
50
+ }
51
+ raw = [self._fake_melody_source(9, [0])]
52
+ assert _melody_sources_to_vllm(raw, position_map) == []
53
+
54
+ def test_missing_position_map_drops_all_sources(self):
55
+ # A parser instance without a position map (parser wired
56
+ # outside of ``CohereServingChatV2``) can't attribute anything,
57
+ # so every source is dropped. Callers downstream will see the
58
+ # citation with empty ``sources`` and drop it entirely.
59
+ raw = [self._fake_melody_source(0, [0])]
60
+ assert _melody_sources_to_vllm(raw, None) == []
61
+
62
+ def test_citations_pass_through_is_thinking_tag(self):
63
+ from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
64
+
65
+ position_map: dict[tuple[int, int], CitationSource] = {
66
+ (0, 0): CitationSource(type="document", id="d0"),
67
+ }
68
+ raw = [
69
+ SimpleNamespace(
70
+ start_index=0,
71
+ end_index=5,
72
+ text="hello",
73
+ is_thinking=True,
74
+ sources=[self._fake_melody_source(0, [0])],
75
+ )
76
+ ]
77
+ out = _melody_citations_to_vllm(raw, position_map)
78
+ assert out is not None
79
+ assert out[0].type == "THINKING_CONTENT"
80
+ assert out[0].sources[0].id == "d0"
81
+
82
+
83
+class TestParserStreamingEndToEnd:
84
+ """End-to-end streaming shape: raw model output text is tokenized
85
+ with the mock byte-level tokenizer and fed one token at a time
86
+ through :meth:`CohereCommand3ReasoningParser.extract_reasoning_streaming`.
87
+ We collect the ordered sequence of ``DeltaMessage``s the parser
88
+ hands back -- i.e. what the client would receive on the wire --
89
+ and pin the exact shape.
90
+
91
+ Combines the three flavors that reach ``delta``: a thinking block
92
+ (populates ``reasoning``), a text block (populates ``content``),
93
+ and a resolved citation (populates ``citations``) referencing
94
+ ``(start, end)`` offsets into the already-streamed content.
95
+ """
96
+
97
+ _MODEL_OUTPUT = (
98
+ "<|START_THINKING|>Let me check.<|END_THINKING|>"
99
+ "<|START_RESPONSE|>The capital is <co>Paris</co: 0:[1]>.<|END_RESPONSE|>"
100
+ )
101
+
102
+ def test_full_wire_stream_shape(self, tokenizer, request_obj):
103
+ from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
104
+ from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY
105
+
106
+ # Mimic what ``CohereServingChatV2._apply_cohere_template_kwargs``
107
+ # installs on a request whose tool_call_index=0 tool result at
108
+ # index 1 is the "France" document.
109
+ position_to_source = {
110
+ (0, 1): CitationSource(
111
+ type="document",
112
+ id="doc-paris",
113
+ document={"id": "doc-paris", "title": "France"},
114
+ ),
115
+ }
116
+
117
+ parser = make_parser(
118
+ tokenizer,
119
+ "cohere_command3",
120
+ chat_template_kwargs={POSITION_TO_SOURCE_KEY: position_to_source},
121
+ )
122
+ deltas = stream_parser(parser, request_obj, tokenizer, self._MODEL_OUTPUT)
123
+
124
+ # -- 1. Reasoning and content stream separately and reach
125
+ # the wire only as their real payload. Framing tokens
126
+ # (``<|START_THINKING|>`` etc.) are consumed by the parser's
127
+ # state machine and never produce a delta.
128
+ reasoning_stream = [d.reasoning for d in deltas if d.reasoning is not None]
129
+ content_stream = [d.content for d in deltas if d.content is not None]
130
+ citation_deltas = [d for d in deltas if getattr(d, "citations", None)]
131
+
132
+ assert "".join(reasoning_stream) == "Let me check."
133
+ assert "".join(content_stream) == "The capital is Paris."
134
+
135
+ # -- 2. Exactly one citation reaches the wire, resolved to the
136
+ # source from ``position_to_source`` (not the raw
137
+ # ``(bucket, idx)`` coordinates).
138
+ assert len(citation_deltas) == 1
139
+ (cite_delta,) = citation_deltas
140
+ assert cite_delta.citations is not None
141
+ assert len(cite_delta.citations) == 1
142
+ cite = cite_delta.citations[0]
143
+ assert cite.text == "Paris"
144
+ # Char offsets are into the accumulated ``content``: "The
145
+ # capital is " is 15 chars, "Paris" is 5.
146
+ assert cite.start == 15
147
+ assert cite.end == 20
148
+ assert len(cite.sources) == 1
149
+ assert cite.sources[0].id == "doc-paris"
150
+ assert cite.sources[0].type == "document"
151
+
152
+ # -- 3. The citation's ``(start, end)`` offsets point into
153
+ # bytes that were already emitted as content-deltas earlier
154
+ # in the stream -- i.e. the citation is anchored to real
155
+ # already-shipped content.
156
+ cite_delta_idx = deltas.index(cite_delta)
157
+ earlier_content = "".join(
158
+ d.content for d in deltas[:cite_delta_idx] if d.content is not None
159
+ )
160
+ assert earlier_content[cite.start : cite.end] == "Paris"
161
+
162
+ # Everything after the citation delta is the trailing period.
163
+ later_content = "".join(
164
+ d.content for d in deltas[cite_delta_idx + 1 :] if d.content is not None
165
+ )
166
+ assert later_content == "."
167
+
168
+ # -- 4. Reasoning and content are never mixed on the same
169
+ # delta: the parser flips modes on the ``<|END_THINKING|>``
170
+ # boundary and never emits a delta with both fields set.
171
+ for d in deltas:
172
+ assert not (d.reasoning is not None and d.content is not None)
tests/parser/cohere/test_reasoning.py
new
+203
@@ -0,0 +1,203 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Reasoning extraction, reasoning-end gating and parser selection for the
4
+unified Cohere Command parser."""
5
+
6
+from __future__ import annotations
7
+
8
+from dataclasses import dataclass
9
+
10
+import pytest
11
+
12
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
13
+from vllm.parser import ParserManager
14
+from vllm.reasoning.cohere_command_reasoning_parser import (
15
+ CohereCommand3ReasoningParser,
16
+ CohereCommand4ReasoningParser,
17
+)
18
+
19
+from .utils import make_parser, stream_parser
20
+
21
+
22
+@dataclass
23
+class ReasoningCase:
24
+ parser_name: str
25
+ model_output: str
26
+ expected_reasoning: str | None
27
+ expected_content: str | None
28
+
29
+
30
+REASONING_CASES = [
31
+ pytest.param(
32
+ ReasoningCase(
33
+ parser_name="cohere_command3",
34
+ model_output="""\
35
+<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
36
+[
37
+ {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
38
+]
39
+<|END_ACTION|>""",
40
+ expected_reasoning="i will call foo with query1",
41
+ expected_content=None,
42
+ ),
43
+ id="cmd3-single_tool_call",
44
+ ),
45
+ pytest.param(
46
+ ReasoningCase(
47
+ parser_name="cohere_command4",
48
+ model_output="""\
49
+<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
50
+[
51
+ {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
52
+]
53
+<|END_ACTION|>""",
54
+ expected_reasoning="i will call foo with query1",
55
+ expected_content=None,
56
+ ),
57
+ id="cmd4-single_tool_call",
58
+ ),
59
+ pytest.param(
60
+ ReasoningCase(
61
+ parser_name="cohere_command3",
62
+ model_output="""\
63
+<|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
64
+<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_RESPONSE|>""",
65
+ expected_reasoning="This is a rainbow emoji: 🌈",
66
+ expected_content="foo bar",
67
+ ),
68
+ id="cmd3-citations_with_emoji",
69
+ ),
70
+ pytest.param(
71
+ ReasoningCase(
72
+ parser_name="cohere_command4",
73
+ model_output="""\
74
+<|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
75
+<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_RESPONSE|>""",
76
+ expected_reasoning="This is a rainbow emoji: 🌈",
77
+ expected_content="foo bar",
78
+ ),
79
+ id="cmd4-citations_with_emoji",
80
+ ),
81
+]
82
+
83
+
84
+@pytest.mark.parametrize("case", REASONING_CASES)
85
+class TestExtractReasoning:
86
+ def test_nonstreaming(self, tokenizer, request_obj, case: ReasoningCase):
87
+ parser = make_parser(tokenizer, case.parser_name)
88
+ reasoning, content, tool_calls = parser.parse(
89
+ case.model_output, request_obj, enable_auto_tools=True
90
+ )
91
+ assert reasoning == case.expected_reasoning
92
+ assert content == case.expected_content
93
+ # No tools on the request, so any tool call is dropped, not surfaced.
94
+ assert not tool_calls
95
+
96
+ def test_streaming(self, tokenizer, request_obj, case: ReasoningCase):
97
+ parser = make_parser(tokenizer, case.parser_name)
98
+ deltas = stream_parser(parser, request_obj, tokenizer, case.model_output)
99
+
100
+ reasoning = "".join(d.reasoning for d in deltas if d.reasoning) or None
101
+ content = "".join(d.content for d in deltas if d.content) or None
102
+ assert reasoning == case.expected_reasoning
103
+ assert content == case.expected_content
104
+ assert all(not d.tool_calls for d in deltas)
105
+
106
+
107
+class TestFramingTokensStripped:
108
+ @pytest.mark.parametrize(
109
+ ("parser_name", "content_tags"),
110
+ [
111
+ pytest.param(
112
+ "cohere_command4", ("<|START_TEXT|>", "<|END_TEXT|>"), id="cmd4"
113
+ ),
114
+ pytest.param(
115
+ "cohere_command3", ("<|START_RESPONSE|>", "<|END_RESPONSE|>"), id="cmd3"
116
+ ),
117
+ ],
118
+ )
119
+ @pytest.mark.parametrize("chunk_size", [1, 4], ids=["per_token", "batched"])
120
+ @pytest.mark.parametrize("with_tools", [False, True], ids=["no_tools", "tools"])
121
+ def test_content_framing_tokens_stripped(
122
+ self, tokenizer, parser_name, content_tags, chunk_size, with_tools
123
+ ):
124
+ parser = make_parser(tokenizer, parser_name)
125
+ tools = [{"type": "function", "function": {"name": "foo"}}]
126
+ request = ChatCompletionRequest(
127
+ messages=[], model="test-model", tools=tools if with_tools else None
128
+ )
129
+ parser.adjust_request(request)
130
+
131
+ start_tag, end_tag = content_tags
132
+ generation = (
133
+ f"<|START_THINKING|>Think deeply. The user greets us.<|END_THINKING|>"
134
+ f"{start_tag}I'm doing well, thank you{end_tag}"
135
+ )
136
+ deltas = stream_parser(parser, request, tokenizer, generation, chunk_size)
137
+
138
+ assert "".join(d.reasoning or "" for d in deltas) == (
139
+ "Think deeply. The user greets us."
140
+ )
141
+ assert not any(d.tool_calls for d in deltas)
142
+ assert "".join(d.content or "" for d in deltas) == "I'm doing well, thank you"
143
+
144
+
145
+class TestIsReasoningEnd:
146
+ @pytest.mark.parametrize(
147
+ "parser_cls",
148
+ [CohereCommand3ReasoningParser, CohereCommand4ReasoningParser],
149
+ ids=["cmd3", "cmd4"],
150
+ )
151
+ def test_is_reasoning_end(self, tokenizer, parser_cls):
152
+ parser = parser_cls(tokenizer)
153
+ start_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>")
154
+ end_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>")
155
+ chatbot_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>")
156
+ content_ids = [99, 100]
157
+
158
+ # Generation-only tokens have no chatbot marker, so the whole sequence
159
+ # is considered.
160
+ assert parser.is_reasoning_end([end_id])
161
+ assert parser.is_reasoning_end([start_id, *content_ids, end_id])
162
+ assert not parser.is_reasoning_end([start_id, *content_ids])
163
+
164
+ # Full prompt/history tokens are scoped to the latest chatbot marker,
165
+ # so stray thinking tokens from the preamble or previous turns are ignored.
166
+ assert not parser.is_reasoning_end([start_id, end_id, chatbot_id, *content_ids])
167
+ assert parser.is_reasoning_end(
168
+ [start_id, end_id, chatbot_id, start_id, *content_ids, end_id]
169
+ )
170
+
171
+
172
+@pytest.mark.parametrize("parser_name", ["cohere_command3", "cohere_command4"])
173
+def test_count_reasoning_tokens(tokenizer, parser_name):
174
+ parser = make_parser(tokenizer, parser_name)
175
+ start = tokenizer.convert_tokens_to_ids("<|START_THINKING|>")
176
+ end = tokenizer.convert_tokens_to_ids("<|END_THINKING|>")
177
+
178
+ assert parser.count_reasoning_tokens([99, start, 11, 12, end, 100]) == 2
179
+ assert parser.count_reasoning_tokens([end, start, 1, start, 2, end, 3, end]) == 3
180
+ assert parser.count_reasoning_tokens([start, 1, 2]) == 2
181
+ assert parser.count_reasoning_tokens([1, 2, end]) == 0
182
+
183
+
184
+class TestParserSelection:
185
+ @pytest.mark.parametrize("with_tool_parser", [True, False])
186
+ def test_adjust_request_keeps_framing_tokens(self, tokenizer, with_tool_parser):
187
+ """Melody needs the special tokens, with or without a tool parser."""
188
+ cls = ParserManager.get_parser(
189
+ "cohere_command4" if with_tool_parser else None,
190
+ "cohere_command4",
191
+ enable_auto_tools=True,
192
+ )
193
+ request = ChatCompletionRequest(messages=[], model="m")
194
+ assert request.skip_special_tokens
195
+ cls(tokenizer).adjust_request(request)
196
+ assert request.skip_special_tokens is False
197
+
198
+ def test_non_cohere_selection_does_not_import_melody_parser(self, monkeypatch):
199
+ """``cohere_melody`` is optional; other parsers must not depend on it."""
200
+ import sys
201
+
202
+ monkeypatch.setitem(sys.modules, "vllm.parser.cohere_command", None)
203
+ assert ParserManager.get_parser("hermes", "deepseek_r1", True) is not None
tests/parser/cohere/test_structural_tags.py
new
+362
@@ -0,0 +1,362 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""``adjust_request``: folding tools and JSON schemas into Cohere structural tags."""
4
+
5
+from __future__ import annotations
6
+
7
+import json
8
+from collections import UserDict
9
+from types import SimpleNamespace
10
+
11
+import pytest
12
+from pydantic import ValidationError
13
+
14
+from vllm.entrypoints.generate.base.protocol import (
15
+ JsonSchemaResponseFormat,
16
+ ResponseFormat,
17
+ StructuralTagResponseFormat,
18
+)
19
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
20
+from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
21
+from vllm.parser.cohere_command import (
22
+ CohereCommandParser,
23
+ _has_effective_tools,
24
+ _response_format_type,
25
+ _schema_dict_from_structured_outputs,
26
+ convert_schema_to_structural_tags,
27
+)
28
+from vllm.sampling_params import StructuredOutputsParams
29
+
30
+from .utils import MockCohereTokenizer, make_parser
31
+
32
+SCHEMA_A = {"type": "object", "properties": {"a": {"type": "string"}}}
33
+SCHEMA_B = {"type": "object", "properties": {"b": {"type": "number"}}}
34
+GET_WEATHER_TOOL = {
35
+ "type": "function",
36
+ "function": {
37
+ "name": "get_weather",
38
+ "parameters": {
39
+ "type": "object",
40
+ "properties": {"city": {"type": "string"}},
41
+ },
42
+ },
43
+}
44
+VALID_STRUCTURAL_TAG = {
45
+ "type": "structural_tag",
46
+ "format": {
47
+ "type": "triggered_tags",
48
+ "tags": [
49
+ {
50
+ "begin": "<tool>",
51
+ "content": {"type": "any_text"},
52
+ "end": "</tool>",
53
+ }
54
+ ],
55
+ "triggers": ["<tool>"],
56
+ },
57
+}
58
+
59
+
60
+def _model_config(arch: str) -> SimpleNamespace:
61
+ return SimpleNamespace(
62
+ architecture=arch,
63
+ architectures=[arch],
64
+ hf_text_config=SimpleNamespace(architectures=[arch]),
65
+ )
66
+
67
+
68
+def _make_chat_request(**kwargs) -> ChatCompletionRequest:
69
+ data = {"messages": [{"role": "user", "content": "hi"}], "model": "m"}
70
+ data.update(kwargs)
71
+ return ChatCompletionRequest.model_validate(data)
72
+
73
+
74
+def _first_json_schema(tag_json: str) -> dict | None:
75
+ outer = json.loads(tag_json)
76
+ for t in (outer.get("format") or {}).get("tags") or []:
77
+ c = t.get("content") or {}
78
+ if c.get("type") == "json_schema":
79
+ js = c.get("json_schema")
80
+ return js if isinstance(js, dict) else None
81
+ return None
82
+
83
+
84
+def _content_types(tag_json: str) -> set[str]:
85
+ outer = json.loads(tag_json)
86
+ out: set[str] = set()
87
+ for t in (outer.get("format") or {}).get("tags") or []:
88
+ ty = (t.get("content") or {}).get("type")
89
+ if isinstance(ty, str):
90
+ out.add(ty)
91
+ return out
92
+
93
+
94
+@pytest.fixture(scope="module")
95
+def parser(tokenizer: MockCohereTokenizer) -> CohereCommandParser:
96
+ """Parser configured with a supported Cohere architecture."""
97
+ return make_parser(
98
+ tokenizer, "cohere_command4", model_config=_model_config("Cohere2ForCausalLM")
99
+ )
100
+
101
+
102
+@pytest.fixture(scope="module")
103
+def parser_no_model_config(tokenizer: MockCohereTokenizer) -> CohereCommandParser:
104
+ """Parser with no ``model_config`` (cannot resolve architecture)."""
105
+ return make_parser(tokenizer, "cohere_command4")
106
+
107
+
108
+@pytest.fixture(scope="module")
109
+def parser_unsupported_arch(tokenizer: MockCohereTokenizer) -> CohereCommandParser:
110
+ """Parser configured with an architecture that has no structural tag style."""
111
+ return make_parser(
112
+ tokenizer, "cohere_command4", model_config=_model_config("LlamaForCausalLM")
113
+ )
114
+
115
+
116
+class TestAdjustRequestPassthrough:
117
+ def test_structured_outputs_structural_tag_not_modified(self, parser) -> None:
118
+ tag = json.dumps(VALID_STRUCTURAL_TAG)
119
+ r = _make_chat_request(structured_outputs={"structural_tag": tag})
120
+ o = parser.adjust_request(r)
121
+ assert o.structured_outputs.structural_tag == tag
122
+
123
+ def test_response_format_structural_tag_short_circuit(self, parser) -> None:
124
+ # ``ChatCompletionRequest`` validates ``response_format`` as a union;
125
+ # bare ``{"type": "structural_tag"}`` is invalid (use pydantic model).
126
+ rf = StructuralTagResponseFormat(
127
+ type="structural_tag",
128
+ format=VALID_STRUCTURAL_TAG["format"],
129
+ )
130
+ r = _make_chat_request(response_format=rf)
131
+ o = parser.adjust_request(r)
132
+ assert _response_format_type(o.response_format) == "structural_tag"
133
+ assert o.structured_outputs is None
134
+
135
+
136
+class TestAdjustRequestNoOp:
137
+ def test_no_schema_no_tools(self, parser) -> None:
138
+ o = parser.adjust_request(_make_chat_request())
139
+ assert o.structured_outputs is None
140
+ assert o.response_format is None
141
+
142
+ def test_no_model_config(self, parser_no_model_config) -> None:
143
+ inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A)
144
+ r = _make_chat_request(
145
+ response_format=ResponseFormat(type="json_schema", json_schema=inner),
146
+ )
147
+ o = parser_no_model_config.adjust_request(r)
148
+ assert o.response_format is not None
149
+ assert o.structured_outputs is None
150
+
151
+
152
+class TestAdjustRequestUnsupportedArchitecture:
153
+ def test_json_schema_raises(self, parser_unsupported_arch) -> None:
154
+ inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A)
155
+ r = _make_chat_request(
156
+ response_format=ResponseFormat(type="json_schema", json_schema=inner),
157
+ )
158
+ with pytest.raises(ValueError, match="does not support"):
159
+ parser_unsupported_arch.adjust_request(r)
160
+
161
+
162
+class TestAdjustRequestFoldFromResponseFormat:
163
+ @pytest.mark.parametrize(
164
+ "response_format, expected_schema",
165
+ [
166
+ pytest.param(
167
+ ResponseFormat(
168
+ type="json_schema",
169
+ json_schema=JsonSchemaResponseFormat(
170
+ name="n", json_schema=SCHEMA_A
171
+ ),
172
+ ),
173
+ SCHEMA_A,
174
+ id="json_schema_pydantic",
175
+ ),
176
+ pytest.param(
177
+ {
178
+ "type": "json_schema",
179
+ "json_schema": {"name": "n", "schema": SCHEMA_A},
180
+ },
181
+ SCHEMA_A,
182
+ id="json_schema_dict",
183
+ ),
184
+ pytest.param(
185
+ {"type": "json_object"},
186
+ {"type": "object"},
187
+ id="json_object",
188
+ ),
189
+ ],
190
+ )
191
+ def test_response_format_cleared(
192
+ self, parser, response_format, expected_schema
193
+ ) -> None:
194
+ r = _make_chat_request(response_format=response_format)
195
+ o = parser.adjust_request(r)
196
+ assert o.response_format is None
197
+ assert (
198
+ _first_json_schema(o.structured_outputs.structural_tag) == expected_schema
199
+ )
200
+
201
+
202
+class TestHasEffectiveTools:
203
+ @pytest.mark.parametrize(
204
+ "tools, expected",
205
+ [
206
+ pytest.param(None, False, id="none"),
207
+ pytest.param([], False, id="empty_list"),
208
+ pytest.param(" ", False, id="blank_str"),
209
+ pytest.param(
210
+ [{"type": "function", "function": {"name": "f"}}],
211
+ True,
212
+ id="non_empty_list",
213
+ ),
214
+ pytest.param('{"x": 1}', True, id="non_empty_str"),
215
+ ],
216
+ )
217
+ def test_has_effective_tools(self, tools, expected) -> None:
218
+ assert _has_effective_tools(tools) is expected
219
+
220
+ def test_convert_schema_json_only_with_empty_tools_list(self) -> None:
221
+ tag = convert_schema_to_structural_tags(
222
+ schema=SCHEMA_B,
223
+ tools=[],
224
+ model_architecture="Cohere2ForCausalLM",
225
+ )
226
+ assert tag is not None
227
+ assert _first_json_schema(tag) == SCHEMA_B
228
+
229
+
230
+class TestAdjustRequestFoldFromStructuredOutputs:
231
+ @pytest.mark.parametrize(
232
+ "structured_outputs, expected_schema",
233
+ [
234
+ pytest.param({"json": SCHEMA_B}, SCHEMA_B, id="json_dict"),
235
+ pytest.param({"json": json.dumps(SCHEMA_B)}, SCHEMA_B, id="json_string"),
236
+ pytest.param(
237
+ {"json_object": True}, {"type": "object"}, id="json_object_flag"
238
+ ),
239
+ pytest.param(
240
+ StructuredOutputsParams(json=SCHEMA_B),
241
+ SCHEMA_B,
242
+ id="structured_outputs_dataclass",
243
+ ),
244
+ pytest.param(
245
+ {"json": {"name": "n", "schema": SCHEMA_A}},
246
+ SCHEMA_A,
247
+ id="openai_wrapper_dict_unwrapped",
248
+ ),
249
+ ],
250
+ )
251
+ def test_structured_outputs_folded(
252
+ self, parser, structured_outputs, expected_schema
253
+ ) -> None:
254
+ o = parser.adjust_request(
255
+ _make_chat_request(structured_outputs=structured_outputs),
256
+ )
257
+ assert (
258
+ _first_json_schema(o.structured_outputs.structural_tag) == expected_schema
259
+ )
260
+
261
+ def test_responses_request_default_empty_tools(self, parser) -> None:
262
+ """``ResponsesRequest.tools`` defaults to ``[]``, not ``None``."""
263
+ r = ResponsesRequest.model_validate(
264
+ {
265
+ "input": "hi",
266
+ "model": "m",
267
+ "structured_outputs": {"json": SCHEMA_B},
268
+ }
269
+ )
270
+ assert r.tools == []
271
+ o = parser.adjust_request(r)
272
+ assert _first_json_schema(o.structured_outputs.structural_tag) == SCHEMA_B
273
+
274
+ def test_json_userdict_mapping_unwrapped(self) -> None:
275
+ inner = {"type": "object", "properties": {"u": {"type": "number"}}}
276
+ so = StructuredOutputsParams(json=UserDict(inner))
277
+ assert _schema_dict_from_structured_outputs(so) == inner
278
+
279
+ @pytest.mark.parametrize(
280
+ "json_value, match",
281
+ [
282
+ pytest.param("{not json}", "valid JSON", id="invalid_json_string"),
283
+ pytest.param(
284
+ json.dumps(["a", "b"]), "JSON object", id="non_object_json_string"
285
+ ),
286
+ pytest.param(" ", "empty", id="empty_json_string"),
287
+ ],
288
+ )
289
+ def test_structured_outputs_json_string_raises(
290
+ self, parser, json_value, match
291
+ ) -> None:
292
+ with pytest.raises(ValueError, match=match):
293
+ parser.adjust_request(
294
+ _make_chat_request(structured_outputs={"json": json_value}),
295
+ )
296
+
297
+ @pytest.mark.parametrize(
298
+ "construct",
299
+ [
300
+ pytest.param(
301
+ lambda: _make_chat_request(structured_outputs={"json": [1, 2, 3]}),
302
+ id="chat_completion_request",
303
+ ),
304
+ pytest.param(
305
+ lambda: StructuredOutputsParams(json=[1, 2, 3]), # type: ignore[arg-type]
306
+ id="structured_outputs_params",
307
+ ),
308
+ ],
309
+ )
310
+ def test_json_wrong_type_raises(self, construct) -> None:
311
+ """Non-str / non-dict ``json`` fails at Pydantic validation."""
312
+ with pytest.raises(ValidationError):
313
+ construct()
314
+
315
+
316
+class TestAdjustRequestPrecedence:
317
+ def test_response_format_over_structured_outputs_json(self, parser) -> None:
318
+ s_rf = {"type": "object", "properties": {"rf": {"type": "string"}}}
319
+ s_so = {"type": "object", "properties": {"so": {"type": "number"}}}
320
+ inner = JsonSchemaResponseFormat(name="n", json_schema=s_rf)
321
+ r = _make_chat_request(
322
+ response_format=ResponseFormat(type="json_schema", json_schema=inner),
323
+ structured_outputs={"json": s_so},
324
+ )
325
+ o = parser.adjust_request(r)
326
+ assert _first_json_schema(o.structured_outputs.structural_tag) == s_rf
327
+
328
+
329
+class TestAdjustRequestTextPlusStructuredOutputs:
330
+ def test_text_response_format_preserved(self, parser) -> None:
331
+ sch = {"type": "object", "properties": {"k": {"type": "string"}}}
332
+ r = _make_chat_request(
333
+ response_format=ResponseFormat(type="text"),
334
+ structured_outputs={"json": sch},
335
+ )
336
+ o = parser.adjust_request(r)
337
+ assert o.response_format is not None
338
+ assert o.response_format.type == "text"
339
+ assert _first_json_schema(o.structured_outputs.structural_tag) == sch
340
+
341
+
342
+class TestAdjustRequestTools:
343
+ def test_tools_only_command_a_grammar(self, parser) -> None:
344
+ o = parser.adjust_request(
345
+ _make_chat_request(tools=[GET_WEATHER_TOOL], tool_choice="auto"),
346
+ )
347
+ assert "grammar" in _content_types(o.structured_outputs.structural_tag)
348
+
349
+ def test_tools_plus_json_schema_both_kinds(self, parser) -> None:
350
+ inner = JsonSchemaResponseFormat(
351
+ name="n",
352
+ json_schema={"type": "object", "properties": {"r": {"type": "string"}}},
353
+ )
354
+ r = _make_chat_request(
355
+ response_format=ResponseFormat(type="json_schema", json_schema=inner),
356
+ tools=[GET_WEATHER_TOOL],
357
+ tool_choice="auto",
358
+ )
359
+ o = parser.adjust_request(r)
360
+ types = _content_types(o.structured_outputs.structural_tag)
361
+ assert "grammar" in types
362
+ assert "json_schema" in types
tests/parser/cohere/test_tool_calls.py
new
+288
@@ -0,0 +1,288 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tool-call parsing for the unified Cohere Command parser.
4
+
5
+The ``cohere_command3`` / ``cohere_command4`` tool parsers are registry shims;
6
+``CohereCommandParser`` (resolved through ``ParserManager``) does the parsing.
7
+"""
8
+
9
+from __future__ import annotations
10
+
11
+import json
12
+from dataclasses import dataclass, field
13
+
14
+import pytest
15
+
16
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
17
+from vllm.parser.abstract_parser import Parser
18
+
19
+from .utils import MockCohereTokenizer, drive_parser, make_parser, token_deltas
20
+
21
+SPECIAL_TOKEN_MARKERS = (
22
+ "<|START_THINKING|>",
23
+ "<|END_THINKING|>",
24
+ "<|START_RESPONSE|>",
25
+ "<|END_RESPONSE|>",
26
+ "<|START_ACTION|>",
27
+ "<|END_ACTION|>",
28
+ "<|START_TEXT|>",
29
+ "<|END_TEXT|>",
30
+)
31
+
32
+DUMMY_TOOLS = [
33
+ {"type": "function", "function": {"name": name, "parameters": {"type": "object"}}}
34
+ for name in ("foo", "bar")
35
+]
36
+
37
+# (registry name, answer framing tag)
38
+VARIANTS = (("cohere_command3", "RESPONSE"), ("cohere_command4", "TEXT"))
39
+
40
+
41
+@dataclass
42
+class ExpectedToolCall:
43
+ id: str
44
+ name: str
45
+ arguments: dict
46
+
47
+
48
+@dataclass
49
+class ToolCallCase:
50
+ parser_name: str
51
+ model_output: str
52
+ expected_tool_calls: list[ExpectedToolCall] = field(default_factory=list)
53
+ expected_reasoning: str | None = None
54
+ expected_content: str | None = None
55
+
56
+
57
+SINGLE_TOOL_CALL = """\
58
+<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
59
+[
60
+ {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
61
+]
62
+<|END_ACTION|>"""
63
+MULTI_TOOL_CALL = """\
64
+<|START_THINKING|>first I think about foo<|END_THINKING|><|START_ACTION|>
65
+[
66
+ {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}},
67
+ {"tool_call_id": "1", "tool_name": "bar", "parameters": {"x": 42}}
68
+]
69
+<|END_ACTION|>"""
70
+CITATIONS_NO_TOOL_CALL = (
71
+ "<|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>\n"
72
+ "<|START_{tag}|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_{tag}|>"
73
+)
74
+
75
+# (case id, model output template, expectations)
76
+_CASE_TABLE = (
77
+ (
78
+ "single_tool_call",
79
+ SINGLE_TOOL_CALL,
80
+ dict(
81
+ expected_tool_calls=[ExpectedToolCall("0", "foo", {"query": "query1"})],
82
+ expected_reasoning="i will call foo with query1",
83
+ ),
84
+ ),
85
+ (
86
+ "citations_no_tool_calls",
87
+ CITATIONS_NO_TOOL_CALL,
88
+ dict(
89
+ expected_reasoning="This is a rainbow emoji: 🌈", expected_content="foo bar"
90
+ ),
91
+ ),
92
+ (
93
+ "multiple_tool_calls",
94
+ MULTI_TOOL_CALL,
95
+ dict(
96
+ expected_tool_calls=[
97
+ ExpectedToolCall("0", "foo", {"query": "query1"}),
98
+ ExpectedToolCall("1", "bar", {"x": 42}),
99
+ ],
100
+ expected_reasoning="first I think about foo",
101
+ ),
102
+ ),
103
+ (
104
+ "reasoning_only",
105
+ "<|START_THINKING|>just think, no response<|END_THINKING|>",
106
+ dict(expected_reasoning="just think, no response"),
107
+ ),
108
+)
109
+
110
+TOOL_CALL_CASES = [
111
+ pytest.param(
112
+ ToolCallCase(name, output.replace("{tag}", tag), **expected),
113
+ id=f"{name}-{case_id}",
114
+ )
115
+ for name, tag in VARIANTS
116
+ for case_id, output, expected in _CASE_TABLE
117
+]
118
+
119
+PARSER_NAMES = pytest.mark.parametrize("parser_name", [name for name, _ in VARIANTS])
120
+
121
+
122
+def _tool_parser(tokenizer: MockCohereTokenizer, parser_name: str) -> Parser:
123
+ return make_parser(tokenizer, parser_name, DUMMY_TOOLS)
124
+
125
+
126
+def _tool_request() -> ChatCompletionRequest:
127
+ return ChatCompletionRequest(
128
+ messages=[], model="test-model", tools=DUMMY_TOOLS, tool_choice="auto"
129
+ )
130
+
131
+
132
+@dataclass
133
+class StreamingResult:
134
+ tool_calls: dict[int, dict]
135
+ reasoning: str | None
136
+ content: str | None
137
+
138
+
139
+def run_streaming(
140
+ tokenizer: MockCohereTokenizer,
141
+ parser_name: str,
142
+ model_output: str,
143
+ chunk_size: int = 1,
144
+) -> StreamingResult:
145
+ deltas = drive_parser(
146
+ _tool_parser(tokenizer, parser_name),
147
+ _tool_request(),
148
+ token_deltas(tokenizer, model_output, chunk_size),
149
+ )
150
+ accumulated: dict[int, dict] = {}
151
+ for delta in deltas:
152
+ for tc in delta.tool_calls or []:
153
+ slot = accumulated.setdefault(
154
+ tc.index, {"id": "", "name": "", "arguments": ""}
155
+ )
156
+ if tc.id:
157
+ slot["id"] = tc.id
158
+ if tc.function and tc.function.name:
159
+ slot["name"] = tc.function.name
160
+ if tc.function and tc.function.arguments:
161
+ slot["arguments"] += tc.function.arguments
162
+ return StreamingResult(
163
+ tool_calls=accumulated,
164
+ reasoning="".join(d.reasoning for d in deltas if d.reasoning) or None,
165
+ content="".join(d.content for d in deltas if d.content) or None,
166
+ )
167
+
168
+
169
+def run_nonstreaming(
170
+ tokenizer: MockCohereTokenizer, parser_name: str, model_output: str
171
+) -> tuple[str | None, str | None, list]:
172
+ reasoning, content, tool_calls = _tool_parser(tokenizer, parser_name).parse(
173
+ model_output, _tool_request(), enable_auto_tools=True
174
+ )
175
+ return reasoning, content, tool_calls or []
176
+
177
+
178
+@pytest.mark.parametrize("case", TOOL_CALL_CASES)
179
+class TestExtractToolCalls:
180
+ def test_streaming(self, tokenizer: MockCohereTokenizer, case: ToolCallCase):
181
+ streamed = run_streaming(tokenizer, case.parser_name, case.model_output)
182
+ assert len(streamed.tool_calls) == len(case.expected_tool_calls)
183
+ for i, expected in enumerate(case.expected_tool_calls):
184
+ tc = streamed.tool_calls[i]
185
+ assert tc["id"] == expected.id
186
+ assert tc["name"] == expected.name
187
+ assert json.loads(tc["arguments"]) == expected.arguments
188
+
189
+ def test_streaming_reasoning(
190
+ self, tokenizer: MockCohereTokenizer, case: ToolCallCase
191
+ ):
192
+ streamed = run_streaming(tokenizer, case.parser_name, case.model_output)
193
+ assert streamed.reasoning == case.expected_reasoning
194
+
195
+ def test_streaming_content(
196
+ self, tokenizer: MockCohereTokenizer, case: ToolCallCase
197
+ ):
198
+ streamed = run_streaming(tokenizer, case.parser_name, case.model_output)
199
+ assert streamed.content == case.expected_content
200
+
201
+ def test_nonstreaming(self, tokenizer: MockCohereTokenizer, case: ToolCallCase):
202
+ reasoning, content, tool_calls = run_nonstreaming(
203
+ tokenizer, case.parser_name, case.model_output
204
+ )
205
+ assert reasoning == case.expected_reasoning
206
+ assert content == case.expected_content
207
+ assert len(tool_calls) == len(case.expected_tool_calls)
208
+ for actual, expected in zip(tool_calls, case.expected_tool_calls):
209
+ assert actual.name == expected.name
210
+ assert json.loads(actual.arguments) == expected.arguments
211
+
212
+ def test_streaming_nonstreaming_agree(
213
+ self, tokenizer: MockCohereTokenizer, case: ToolCallCase
214
+ ):
215
+ streamed = run_streaming(tokenizer, case.parser_name, case.model_output)
216
+ _, _, tool_calls = run_nonstreaming(
217
+ tokenizer, case.parser_name, case.model_output
218
+ )
219
+ assert len(streamed.tool_calls) == len(tool_calls)
220
+ for i, actual in enumerate(tool_calls):
221
+ assert streamed.tool_calls[i]["name"] == actual.name
222
+ assert json.loads(streamed.tool_calls[i]["arguments"]) == json.loads(
223
+ actual.arguments
224
+ )
225
+
226
+
227
+class TestSpeculativeDecodingMultiTokenDelta:
228
+ @PARSER_NAMES
229
+ @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6])
230
+ def test_no_special_token_leak_in_streaming_deltas(
231
+ self, tokenizer: MockCohereTokenizer, parser_name: str, chunk_size: int
232
+ ):
233
+ deltas = drive_parser(
234
+ _tool_parser(tokenizer, parser_name),
235
+ _tool_request(),
236
+ token_deltas(tokenizer, SINGLE_TOOL_CALL, chunk_size),
237
+ )
238
+ for delta in deltas:
239
+ fields: list[tuple[str, str | None]] = [
240
+ ("reasoning", delta.reasoning),
241
+ ("content", delta.content),
242
+ ]
243
+ for tc in delta.tool_calls or []:
244
+ if tc.function:
245
+ fields.append(("tool_call.name", tc.function.name))
246
+ fields.append(("tool_call.arguments", tc.function.arguments))
247
+ for marker in SPECIAL_TOKEN_MARKERS:
248
+ for field_name, value in fields:
249
+ assert value is None or marker not in value, (
250
+ f"special token {marker!r} leaked into {field_name} "
251
+ f"with chunk_size={chunk_size} delta={delta!r}"
252
+ )
253
+
254
+ @PARSER_NAMES
255
+ @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6])
256
+ def test_multi_token_chunks_still_produce_correct_tool_call(
257
+ self, tokenizer: MockCohereTokenizer, parser_name: str, chunk_size: int
258
+ ):
259
+ streamed = run_streaming(tokenizer, parser_name, SINGLE_TOOL_CALL, chunk_size)
260
+ assert len(streamed.tool_calls) == 1
261
+ tc = streamed.tool_calls[0]
262
+ assert tc["id"] == "0"
263
+ assert tc["name"] == "foo"
264
+ assert json.loads(tc["arguments"]) == {"query": "query1"}
265
+
266
+
267
+class TestStreamingDeltaShape:
268
+ @PARSER_NAMES
269
+ def test_reasoning_and_tool_calls_are_separate_deltas(
270
+ self, tokenizer: MockCohereTokenizer, parser_name: str
271
+ ):
272
+ deltas = drive_parser(
273
+ _tool_parser(tokenizer, parser_name),
274
+ _tool_request(),
275
+ token_deltas(tokenizer, SINGLE_TOOL_CALL),
276
+ )
277
+ for delta in deltas:
278
+ populated = [
279
+ delta.content is not None,
280
+ delta.reasoning is not None,
281
+ bool(delta.tool_calls),
282
+ ]
283
+ assert sum(populated) == 1, (
284
+ "A single streaming delta must carry exactly one of "
285
+ f"content/reasoning/tool_calls, got {delta!r}"
286
+ )
287
+ assert any(d.reasoning is not None for d in deltas)
288
+ assert any(d.tool_calls for d in deltas)
tests/parser/cohere/utils.py
new
+131
@@ -0,0 +1,131 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Shared helpers for the unified Cohere Command parser tests."""
4
+
5
+from __future__ import annotations
6
+
7
+import itertools
8
+
9
+import regex as re
10
+
11
+from vllm.entrypoints.generate.base.protocol import DeltaMessage
12
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
13
+from vllm.parser import ParserManager
14
+from vllm.parser.abstract_parser import Parser
15
+
16
+_SPECIAL_TOKEN_RE = re.compile(r"(<\|[A-Z_]+\|>)")
17
+REPLACEMENT_CHAR = "�"
18
+
19
+
20
+class MockCohereTokenizer:
21
+ """Byte-level stand-in for the Cohere tokenizer.
22
+
23
+ Text round-trips through UTF-8 bytes so a multi-byte character can split
24
+ across "tokens" (reproducing trailing U+FFFD buffering), while each
25
+ ``<|...|>`` marker is a single id so token-id reasoning-end gating behaves
26
+ as with the real tokenizer.
27
+ """
28
+
29
+ _SPECIAL_TOKEN_IDS = {
30
+ tok: 256 + i
31
+ for i, tok in enumerate(
32
+ (
33
+ "<|START_THINKING|>",
34
+ "<|END_THINKING|>",
35
+ "<|CHATBOT_TOKEN|>",
36
+ "<|START_RESPONSE|>",
37
+ "<|END_RESPONSE|>",
38
+ "<|START_TEXT|>",
39
+ "<|END_TEXT|>",
40
+ "<|START_ACTION|>",
41
+ "<|END_ACTION|>",
42
+ )
43
+ )
44
+ }
45
+ _ID_TO_SPECIAL_TOKEN = {v: k for k, v in _SPECIAL_TOKEN_IDS.items()}
46
+
47
+ def convert_tokens_to_ids(self, token: str) -> int:
48
+ return self._SPECIAL_TOKEN_IDS.get(token, 0)
49
+
50
+ def get_vocab(self) -> dict[str, int]:
51
+ return {}
52
+
53
+ def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
54
+ ids: list[int] = []
55
+ for part in _SPECIAL_TOKEN_RE.split(text):
56
+ if part in self._SPECIAL_TOKEN_IDS:
57
+ ids.append(self._SPECIAL_TOKEN_IDS[part])
58
+ else:
59
+ ids.extend(part.encode("utf-8"))
60
+ return ids
61
+
62
+ def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str:
63
+ out: list[str] = []
64
+ for special, run in itertools.groupby(
65
+ ids, self._ID_TO_SPECIAL_TOKEN.__contains__
66
+ ):
67
+ if not special:
68
+ out.append(bytes(run).decode("utf-8", errors="replace"))
69
+ elif not skip_special_tokens:
70
+ out.extend(self._ID_TO_SPECIAL_TOKEN[i] for i in run)
71
+ return "".join(out)
72
+
73
+
74
+def token_deltas(
75
+ tokenizer: MockCohereTokenizer, text: str, chunk_size: int = 1
76
+) -> list[tuple[str, list[int]]]:
77
+ """Split ``text`` into ``(delta_text, delta_token_ids)`` steps of
78
+ ``chunk_size`` tokens, buffering incomplete multi-byte sequences (trailing
79
+ U+FFFD) into the next step as real streaming does."""
80
+ ids = tokenizer.encode(text)
81
+ deltas: list[tuple[str, list[int]]] = []
82
+ prev = ""
83
+ pending: list[int] = []
84
+ for start in range(0, len(ids), chunk_size):
85
+ end = start + chunk_size
86
+ pending.extend(ids[start:end])
87
+ current = tokenizer.decode(ids[:end])
88
+ if current.endswith(REPLACEMENT_CHAR) and end < len(ids):
89
+ continue
90
+ deltas.append((current[len(prev) :], pending))
91
+ prev, pending = current, []
92
+ return deltas
93
+
94
+
95
+def make_parser(
96
+ tokenizer: MockCohereTokenizer,
97
+ name: str,
98
+ tools: list[dict] | None = None,
99
+ chat_template_kwargs: dict | None = None,
100
+ model_config=None,
101
+) -> Parser:
102
+ """Resolve the unified Cohere parser through ``ParserManager`` by registry
103
+ name (``cohere_command3`` / ``cohere_command4``)."""
104
+ cls = ParserManager.get_parser(name, name, enable_auto_tools=True)
105
+ assert cls is not None
106
+ return cls(
107
+ tokenizer,
108
+ tools,
109
+ model_config=model_config,
110
+ chat_template_kwargs=chat_template_kwargs or {},
111
+ )
112
+
113
+
114
+def drive_parser(
115
+ parser: Parser,
116
+ request: ChatCompletionRequest,
117
+ deltas: list[tuple[str, list[int]]],
118
+) -> list[DeltaMessage]:
119
+ """Feed ``deltas`` through ``parse_delta`` (flushing on the last one) and
120
+ return the non-``None`` messages in order."""
121
+ out: list[DeltaMessage] = []
122
+ for i, (text, ids) in enumerate(deltas):
123
+ delta = parser.parse_delta(text, ids, request, finished=i == len(deltas) - 1)
124
+ if delta is not None:
125
+ out.append(delta)
126
+ return out
127
+
128
+
129
+def stream_parser(parser, request, tokenizer, text: str, chunk_size: int = 1):
130
+ """Stream ``text`` through ``parser`` in ``chunk_size``-token deltas."""
131
+ return drive_parser(parser, request, token_deltas(tokenizer, text, chunk_size))
tests/reasoning/test_cohere_command_reasoning_parser.py
deleted
-796
@@ -1,796 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from __future__ import annotations
5
-
6
-import json
7
-from collections import UserDict
8
-from dataclasses import dataclass
9
-from types import SimpleNamespace
10
-from typing import Any
11
-
12
-import pytest
13
-from pydantic import ValidationError
14
-
15
-from vllm.entrypoints.generate.base.protocol import (
16
- JsonSchemaResponseFormat,
17
- ResponseFormat,
18
- StructuralTagResponseFormat,
19
-)
20
-from vllm.entrypoints.openai.chat_completion.protocol import (
21
- ChatCompletionRequest,
22
-)
23
-from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
24
-from vllm.reasoning.cohere_command_reasoning_parser import (
25
- CohereCommand3ReasoningParser,
26
- CohereCommand4ReasoningParser,
27
- _has_effective_tools,
28
- _melody_citations_to_vllm,
29
- _melody_sources_to_vllm,
30
- _response_format_type,
31
- _schema_dict_from_structured_outputs,
32
- convert_schema_to_structural_tags,
33
-)
34
-from vllm.sampling_params import StructuredOutputsParams
35
-
36
-
37
-@dataclass
38
-class ReasoningCase:
39
- parser_cls: Any
40
- model_output: str
41
- expected_reasoning: str | None
42
- expected_content: str | None
43
-
44
-
45
-REASONING_CASES = [
46
- pytest.param(
47
- ReasoningCase(
48
- parser_cls=CohereCommand3ReasoningParser,
49
- model_output="""\
50
-<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
51
-[
52
- {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
53
-]
54
-<|END_ACTION|>""",
55
- expected_reasoning="i will call foo with query1",
56
- expected_content="""\
57
-<|START_ACTION|>
58
-[
59
- {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
60
-]
61
-<|END_ACTION|>""",
62
- ),
63
- id="cmd3-single_tool_call",
64
- ),
65
- pytest.param(
66
- ReasoningCase(
67
- parser_cls=CohereCommand4ReasoningParser,
68
- model_output="""\
69
-<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
70
-[
71
- {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
72
-]
73
-<|END_ACTION|>""",
74
- expected_reasoning="i will call foo with query1",
75
- expected_content="""\
76
-<|START_ACTION|>
77
-[
78
- {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
79
-]
80
-<|END_ACTION|>""",
81
- ),
82
- id="cmd4-single_tool_call",
83
- ),
84
- pytest.param(
85
- ReasoningCase(
86
- parser_cls=CohereCommand3ReasoningParser,
87
- model_output="""\
88
-<|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
89
-<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_RESPONSE|>""",
90
- expected_reasoning="This is a rainbow emoji: 🌈",
91
- expected_content="foo bar",
92
- ),
93
- id="cmd3-citations_with_emoji",
94
- ),
95
- pytest.param(
96
- ReasoningCase(
97
- parser_cls=CohereCommand4ReasoningParser,
98
- model_output="""\
99
-<|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
100
-<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_RESPONSE|>""",
101
- expected_reasoning="This is a rainbow emoji: 🌈",
102
- expected_content="foo bar",
103
- ),
104
- id="cmd4-citations_with_emoji",
105
- ),
106
-]
107
-
108
-
109
-class MockCohereTokenizer:
110
- """Minimal byte-level stand-in for the Cohere tokenizer.
111
-
112
- ``encode``/``decode`` round-trip through UTF-8 bytes so splitting a
113
- multi-byte character (e.g. an emoji) across "tokens" reproduces the
114
- trailing U+FFFD buffering that real streaming exhibits. Cohere special
115
- tokens map to distinct synthetic ids; everything else shares a default id.
116
- ``adjust_request`` only needs the token ids, not real tokenization.
117
- """
118
-
119
- _SPECIAL_TOKEN_IDS = {
120
- "<|START_THINKING|>": -1,
121
- "<|END_THINKING|>": -2,
122
- "<|CHATBOT_TOKEN|>": -3,
123
- }
124
-
125
- def convert_tokens_to_ids(self, token: str) -> int:
126
- return self._SPECIAL_TOKEN_IDS.get(token, 0)
127
-
128
- def get_vocab(self) -> dict[str, int]:
129
- return {}
130
-
131
- def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
132
- return list(text.encode("utf-8"))
133
-
134
- def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str:
135
- return bytes(ids).decode("utf-8", errors="replace")
136
-
137
-
138
-@pytest.fixture(scope="module")
139
-def tokenizer() -> MockCohereTokenizer:
140
- return MockCohereTokenizer()
141
-
142
-
143
-@pytest.fixture
144
-def request_obj():
145
- return ChatCompletionRequest(messages=[], model="test-model")
146
-
147
-
148
-REPLACEMENT_CHAR = "\ufffd"
149
-
150
-
151
-def _token_deltas(tokenizer, text: str) -> list[str]:
152
- """Progressively decode the token sequence and return per-step string
153
- deltas. Incomplete multi-byte sequences (trailing U+FFFD) are buffered
154
- until the next token completes them, matching real streaming behaviour."""
155
- ids = tokenizer.encode(text, add_special_tokens=False)
156
- deltas: list[str] = []
157
- prev = ""
158
- for i in range(1, len(ids) + 1):
159
- current = tokenizer.decode(ids[:i], skip_special_tokens=False)
160
- if current.endswith(REPLACEMENT_CHAR):
161
- continue
162
- delta = current[len(prev) :]
163
- if delta:
164
- deltas.append(delta)
165
- prev = current
166
- return deltas
167
-
168
-
169
-@pytest.mark.parametrize("case", REASONING_CASES)
170
-class TestExtractReasoning:
171
- def test_nonstreaming(self, tokenizer, request_obj, case: ReasoningCase):
172
- parser = case.parser_cls(tokenizer)
173
- reasoning, content = parser.extract_reasoning(case.model_output, request_obj)
174
-
175
- assert reasoning == case.expected_reasoning
176
- assert content == case.expected_content
177
-
178
- def test_streaming(self, tokenizer, case: ReasoningCase):
179
- parser = case.parser_cls(tokenizer)
180
- token_strings = _token_deltas(tokenizer, case.model_output)
181
-
182
- reasoning_parts: list[str] = []
183
- content_parts: list[str] = []
184
- tool_call_deltas: list[dict] = []
185
-
186
- previous_text = ""
187
- previous_token_ids: list[int] = []
188
-
189
- for token_str in token_strings:
190
- current_text = previous_text + token_str
191
- current_token_ids = previous_token_ids + [0]
192
-
193
- delta = parser.extract_reasoning_streaming(
194
- previous_text=previous_text,
195
- current_text=current_text,
196
- delta_text=token_str,
197
- previous_token_ids=previous_token_ids,
198
- current_token_ids=current_token_ids,
199
- delta_token_ids=[0],
200
- )
201
- if delta is not None:
202
- if delta.reasoning is not None:
203
- reasoning_parts.append(delta.reasoning)
204
- if delta.content is not None:
205
- content_parts.append(delta.content)
206
- for tc in delta.tool_calls:
207
- tool_call_deltas.append(
208
- {
209
- "id": tc.id,
210
- "index": tc.index,
211
- "name": tc.function.name if tc.function else None,
212
- "arguments": (
213
- tc.function.arguments if tc.function else None
214
- ),
215
- }
216
- )
217
-
218
- previous_text = current_text
219
- previous_token_ids = current_token_ids
220
-
221
- reasoning = "".join(reasoning_parts) if reasoning_parts else None
222
- assert reasoning == case.expected_reasoning
223
-
224
- content = "".join(content_parts) if content_parts else None
225
- assert content == case.expected_content
226
- assert tool_call_deltas == []
227
-
228
-
229
-class TestIsReasoningEnd:
230
- @pytest.mark.parametrize(
231
- "parser_cls",
232
- [CohereCommand3ReasoningParser, CohereCommand4ReasoningParser],
233
- ids=["cmd3", "cmd4"],
234
- )
235
- def test_is_reasoning_end(self, tokenizer, parser_cls):
236
- parser = parser_cls(tokenizer)
237
- start_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>")
238
- end_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>")
239
- chatbot_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>")
240
- content_ids = [99, 100]
241
-
242
- # Generation-only tokens have no chatbot marker, so the whole sequence
243
- # is considered.
244
- assert parser.is_reasoning_end([end_id])
245
- assert parser.is_reasoning_end([start_id, *content_ids, end_id])
246
- assert not parser.is_reasoning_end([start_id, *content_ids])
247
-
248
- # Full prompt/history tokens are scoped to the latest chatbot marker,
249
- # so stray thinking tokens from the preamble or previous turns are ignored.
250
- assert not parser.is_reasoning_end([start_id, end_id, chatbot_id, *content_ids])
251
- assert parser.is_reasoning_end(
252
- [start_id, end_id, chatbot_id, start_id, *content_ids, end_id]
253
- )
254
-
255
-
256
-@pytest.mark.parametrize(
257
- "parser_cls",
258
- [
259
- CohereCommand3ReasoningParser,
260
- CohereCommand4ReasoningParser,
261
- ],
262
- ids=["cmd3", "cmd4"],
263
-)
264
-def test_count_reasoning_tokens(tokenizer, parser_cls):
265
- parser = parser_cls(tokenizer)
266
- start = parser.start_token_id
267
- end = parser.end_token_id
268
-
269
- assert parser.count_reasoning_tokens([99, start, 11, 12, end, 100]) == 2
270
- assert parser.count_reasoning_tokens([end, start, 1, start, 2, end, 3, end]) == 3
271
- assert parser.count_reasoning_tokens([start, 1, 2]) == 2
272
- assert parser.count_reasoning_tokens([1, 2, end]) == 0
273
-
274
-
275
-SCHEMA_A = {"type": "object", "properties": {"a": {"type": "string"}}}
276
-SCHEMA_B = {"type": "object", "properties": {"b": {"type": "number"}}}
277
-GET_WEATHER_TOOL = {
278
- "type": "function",
279
- "function": {
280
- "name": "get_weather",
281
- "parameters": {
282
- "type": "object",
283
- "properties": {"city": {"type": "string"}},
284
- },
285
- },
286
-}
287
-VALID_STRUCTURAL_TAG = {
288
- "type": "structural_tag",
289
- "format": {
290
- "type": "triggered_tags",
291
- "tags": [
292
- {
293
- "begin": "<tool>",
294
- "content": {"type": "any_text"},
295
- "end": "</tool>",
296
- }
297
- ],
298
- "triggers": ["<tool>"],
299
- },
300
-}
301
-
302
-
303
-def _model_config(arch: str) -> SimpleNamespace:
304
- return SimpleNamespace(
305
- architecture=arch,
306
- architectures=[arch],
307
- hf_text_config=SimpleNamespace(architectures=[arch]),
308
- )
309
-
310
-
311
-def _make_chat_request(**kwargs) -> ChatCompletionRequest:
312
- data = {"messages": [{"role": "user", "content": "hi"}], "model": "m"}
313
- data.update(kwargs)
314
- return ChatCompletionRequest.model_validate(data)
315
-
316
-
317
-def _first_json_schema(tag_json: str) -> dict | None:
318
- outer = json.loads(tag_json)
319
- for t in (outer.get("format") or {}).get("tags") or []:
320
- c = t.get("content") or {}
321
- if c.get("type") == "json_schema":
322
- js = c.get("json_schema")
323
- return js if isinstance(js, dict) else None
324
- return None
325
-
326
-
327
-def _content_types(tag_json: str) -> set[str]:
328
- outer = json.loads(tag_json)
329
- out: set[str] = set()
330
- for t in (outer.get("format") or {}).get("tags") or []:
331
- ty = (t.get("content") or {}).get("type")
332
- if isinstance(ty, str):
333
- out.add(ty)
334
- return out
335
-
336
-
337
-@pytest.fixture(scope="module")
338
-def parser(tokenizer: MockCohereTokenizer) -> CohereCommand4ReasoningParser:
339
- """Parser configured with a supported Cohere architecture."""
340
- return CohereCommand4ReasoningParser(
341
- tokenizer,
342
- model_config=_model_config("Cohere2ForCausalLM"),
343
- )
344
-
345
-
346
-@pytest.fixture(scope="module")
347
-def parser_no_model_config(
348
- tokenizer: MockCohereTokenizer,
349
-) -> CohereCommand4ReasoningParser:
350
- """Parser with no ``model_config`` (cannot resolve architecture)."""
351
- return CohereCommand4ReasoningParser(tokenizer, model_config=None)
352
-
353
-
354
-@pytest.fixture(scope="module")
355
-def parser_unsupported_arch(
356
- tokenizer: MockCohereTokenizer,
357
-) -> CohereCommand4ReasoningParser:
358
- """Parser configured with an architecture that has no structural tag style."""
359
- return CohereCommand4ReasoningParser(
360
- tokenizer,
361
- model_config=_model_config("LlamaForCausalLM"),
362
- )
363
-
364
-
365
-class TestAdjustRequestPassthrough:
366
- def test_structured_outputs_structural_tag_not_modified(self, parser) -> None:
367
- tag = json.dumps(VALID_STRUCTURAL_TAG)
368
- r = _make_chat_request(structured_outputs={"structural_tag": tag})
369
- o = parser.adjust_request(r)
370
- assert o.structured_outputs.structural_tag == tag
371
-
372
- def test_response_format_structural_tag_short_circuit(self, parser) -> None:
373
- # ``ChatCompletionRequest`` validates ``response_format`` as a union;
374
- # bare ``{"type": "structural_tag"}`` is invalid (use pydantic model).
375
- rf = StructuralTagResponseFormat(
376
- type="structural_tag",
377
- format=VALID_STRUCTURAL_TAG["format"],
378
- )
379
- r = _make_chat_request(response_format=rf)
380
- o = parser.adjust_request(r)
381
- assert _response_format_type(o.response_format) == "structural_tag"
382
- assert o.structured_outputs is None
383
-
384
-
385
-class TestAdjustRequestNoOp:
386
- def test_no_schema_no_tools(self, parser) -> None:
387
- o = parser.adjust_request(_make_chat_request())
388
- assert o.structured_outputs is None
389
- assert o.response_format is None
390
-
391
- def test_no_model_config(self, parser_no_model_config) -> None:
392
- inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A)
393
- r = _make_chat_request(
394
- response_format=ResponseFormat(type="json_schema", json_schema=inner),
395
- )
396
- o = parser_no_model_config.adjust_request(r)
397
- assert o.response_format is not None
398
- assert o.structured_outputs is None
399
-
400
-
401
-class TestAdjustRequestUnsupportedArchitecture:
402
- def test_json_schema_raises(self, parser_unsupported_arch) -> None:
403
- inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A)
404
- r = _make_chat_request(
405
- response_format=ResponseFormat(type="json_schema", json_schema=inner),
406
- )
407
- with pytest.raises(ValueError, match="does not support"):
408
- parser_unsupported_arch.adjust_request(r)
409
-
410
-
411
-class TestAdjustRequestFoldFromResponseFormat:
412
- @pytest.mark.parametrize(
413
- "response_format, expected_schema",
414
- [
415
- pytest.param(
416
- ResponseFormat(
417
- type="json_schema",
418
- json_schema=JsonSchemaResponseFormat(
419
- name="n", json_schema=SCHEMA_A
420
- ),
421
- ),
422
- SCHEMA_A,
423
- id="json_schema_pydantic",
424
- ),
425
- pytest.param(
426
- {
427
- "type": "json_schema",
428
- "json_schema": {"name": "n", "schema": SCHEMA_A},
429
- },
430
- SCHEMA_A,
431
- id="json_schema_dict",
432
- ),
433
- pytest.param(
434
- {"type": "json_object"},
435
- {"type": "object"},
436
- id="json_object",
437
- ),
438
- ],
439
- )
440
- def test_response_format_cleared(
441
- self, parser, response_format, expected_schema
442
- ) -> None:
443
- r = _make_chat_request(response_format=response_format)
444
- o = parser.adjust_request(r)
445
- assert o.response_format is None
446
- assert (
447
- _first_json_schema(o.structured_outputs.structural_tag) == expected_schema
448
- )
449
-
450
-
451
-class TestHasEffectiveTools:
452
- @pytest.mark.parametrize(
453
- "tools, expected",
454
- [
455
- pytest.param(None, False, id="none"),
456
- pytest.param([], False, id="empty_list"),
457
- pytest.param(" ", False, id="blank_str"),
458
- pytest.param(
459
- [{"type": "function", "function": {"name": "f"}}],
460
- True,
461
- id="non_empty_list",
462
- ),
463
- pytest.param('{"x": 1}', True, id="non_empty_str"),
464
- ],
465
- )
466
- def test_has_effective_tools(self, tools, expected) -> None:
467
- assert _has_effective_tools(tools) is expected
468
-
469
- def test_convert_schema_json_only_with_empty_tools_list(self) -> None:
470
- tag = convert_schema_to_structural_tags(
471
- schema=SCHEMA_B,
472
- tools=[],
473
- model_architecture="Cohere2ForCausalLM",
474
- )
475
- assert tag is not None
476
- assert _first_json_schema(tag) == SCHEMA_B
477
-
478
-
479
-class TestAdjustRequestFoldFromStructuredOutputs:
480
- @pytest.mark.parametrize(
481
- "structured_outputs, expected_schema",
482
- [
483
- pytest.param({"json": SCHEMA_B}, SCHEMA_B, id="json_dict"),
484
- pytest.param({"json": json.dumps(SCHEMA_B)}, SCHEMA_B, id="json_string"),
485
- pytest.param(
486
- {"json_object": True}, {"type": "object"}, id="json_object_flag"
487
- ),
488
- pytest.param(
489
- StructuredOutputsParams(json=SCHEMA_B),
490
- SCHEMA_B,
491
- id="structured_outputs_dataclass",
492
- ),
493
- pytest.param(
494
- {"json": {"name": "n", "schema": SCHEMA_A}},
495
- SCHEMA_A,
496
- id="openai_wrapper_dict_unwrapped",
497
- ),
498
- ],
499
- )
500
- def test_structured_outputs_folded(
501
- self, parser, structured_outputs, expected_schema
502
- ) -> None:
503
- o = parser.adjust_request(
504
- _make_chat_request(structured_outputs=structured_outputs),
505
- )
506
- assert (
507
- _first_json_schema(o.structured_outputs.structural_tag) == expected_schema
508
- )
509
-
510
- def test_responses_request_default_empty_tools(self, parser) -> None:
511
- """``ResponsesRequest.tools`` defaults to ``[]``, not ``None``."""
512
- r = ResponsesRequest.model_validate(
513
- {
514
- "input": "hi",
515
- "model": "m",
516
- "structured_outputs": {"json": SCHEMA_B},
517
- }
518
- )
519
- assert r.tools == []
520
- o = parser.adjust_request(r)
521
- assert _first_json_schema(o.structured_outputs.structural_tag) == SCHEMA_B
522
-
523
- def test_json_userdict_mapping_unwrapped(self) -> None:
524
- inner = {"type": "object", "properties": {"u": {"type": "number"}}}
525
- so = StructuredOutputsParams(json=UserDict(inner))
526
- assert _schema_dict_from_structured_outputs(so) == inner
527
-
528
- @pytest.mark.parametrize(
529
- "json_value, match",
530
- [
531
- pytest.param("{not json}", "valid JSON", id="invalid_json_string"),
532
- pytest.param(
533
- json.dumps(["a", "b"]), "JSON object", id="non_object_json_string"
534
- ),
535
- pytest.param(" ", "empty", id="empty_json_string"),
536
- ],
537
- )
538
- def test_structured_outputs_json_string_raises(
539
- self, parser, json_value, match
540
- ) -> None:
541
- with pytest.raises(ValueError, match=match):
542
- parser.adjust_request(
543
- _make_chat_request(structured_outputs={"json": json_value}),
544
- )
545
-
546
- @pytest.mark.parametrize(
547
- "construct",
548
- [
549
- pytest.param(
550
- lambda: _make_chat_request(structured_outputs={"json": [1, 2, 3]}),
551
- id="chat_completion_request",
552
- ),
553
- pytest.param(
554
- lambda: StructuredOutputsParams(json=[1, 2, 3]), # type: ignore[arg-type]
555
- id="structured_outputs_params",
556
- ),
557
- ],
558
- )
559
- def test_json_wrong_type_raises(self, construct) -> None:
560
- """Non-str / non-dict ``json`` fails at Pydantic validation."""
561
- with pytest.raises(ValidationError):
562
- construct()
563
-
564
-
565
-class TestAdjustRequestPrecedence:
566
- def test_response_format_over_structured_outputs_json(self, parser) -> None:
567
- s_rf = {"type": "object", "properties": {"rf": {"type": "string"}}}
568
- s_so = {"type": "object", "properties": {"so": {"type": "number"}}}
569
- inner = JsonSchemaResponseFormat(name="n", json_schema=s_rf)
570
- r = _make_chat_request(
571
- response_format=ResponseFormat(type="json_schema", json_schema=inner),
572
- structured_outputs={"json": s_so},
573
- )
574
- o = parser.adjust_request(r)
575
- assert _first_json_schema(o.structured_outputs.structural_tag) == s_rf
576
-
577
-
578
-class TestAdjustRequestTextPlusStructuredOutputs:
579
- def test_text_response_format_preserved(self, parser) -> None:
580
- sch = {"type": "object", "properties": {"k": {"type": "string"}}}
581
- r = _make_chat_request(
582
- response_format=ResponseFormat(type="text"),
583
- structured_outputs={"json": sch},
584
- )
585
- o = parser.adjust_request(r)
586
- assert o.response_format is not None
587
- assert o.response_format.type == "text"
588
- assert _first_json_schema(o.structured_outputs.structural_tag) == sch
589
-
590
-
591
-class TestAdjustRequestTools:
592
- def test_tools_only_command_a_grammar(self, parser) -> None:
593
- o = parser.adjust_request(
594
- _make_chat_request(tools=[GET_WEATHER_TOOL], tool_choice="auto"),
595
- )
596
- assert "grammar" in _content_types(o.structured_outputs.structural_tag)
597
-
598
- def test_tools_plus_json_schema_both_kinds(self, parser) -> None:
599
- inner = JsonSchemaResponseFormat(
600
- name="n",
601
- json_schema={"type": "object", "properties": {"r": {"type": "string"}}},
602
- )
603
- r = _make_chat_request(
604
- response_format=ResponseFormat(type="json_schema", json_schema=inner),
605
- tools=[GET_WEATHER_TOOL],
606
- tool_choice="auto",
607
- )
608
- o = parser.adjust_request(r)
609
- types = _content_types(o.structured_outputs.structural_tag)
610
- assert "grammar" in types
611
- assert "json_schema" in types
612
-
613
-
614
-class TestMelodySourceResolution:
615
- """Pin the parser-side source resolution (see
616
- ``_melody_sources_to_vllm``). The parser receives melody's numeric
617
- ``(tool_call_index, tool_result_indices)`` addressing and resolves
618
- it against a ``position_to_source`` map handed in by
619
- ``CohereServingChatV2._apply_cohere_template_kwargs`` -- the same
620
- map that would otherwise live in the serving layer's resolver.
621
- """
622
-
623
- @staticmethod
624
- def _fake_melody_source(bucket: int, indices: list[int]) -> Any:
625
- return SimpleNamespace(tool_call_index=bucket, tool_result_indices=indices)
626
-
627
- def test_multi_index_source_fans_out(self):
628
- from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
629
-
630
- position_map: dict[tuple[int, int], CitationSource] = {
631
- (0, 0): CitationSource(type="document", id="d0", document={"id": "d0"}),
632
- (0, 1): CitationSource(type="document", id="d1", document={"id": "d1"}),
633
- }
634
- raw = [self._fake_melody_source(0, [0, 1])]
635
- out = _melody_sources_to_vllm(raw, position_map)
636
- assert [s.id for s in out] == ["d0", "d1"]
637
- # Verify type / payload were plumbed through, not just the id.
638
- assert out[0].type == "document"
639
- assert out[0].document == {"id": "d0"}
640
-
641
- def test_unresolvable_position_skipped(self):
642
- from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
643
-
644
- position_map: dict[tuple[int, int], CitationSource] = {
645
- (0, 0): CitationSource(type="document", id="d0"),
646
- }
647
- raw = [self._fake_melody_source(9, [0])]
648
- assert _melody_sources_to_vllm(raw, position_map) == []
649
-
650
- def test_missing_position_map_drops_all_sources(self):
651
- # A parser instance without a position map (parser wired
652
- # outside of ``CohereServingChatV2``) can't attribute anything,
653
- # so every source is dropped. Callers downstream will see the
654
- # citation with empty ``sources`` and drop it entirely.
655
- raw = [self._fake_melody_source(0, [0])]
656
- assert _melody_sources_to_vllm(raw, None) == []
657
-
658
- def test_citations_pass_through_is_thinking_tag(self):
659
- from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
660
-
661
- position_map: dict[tuple[int, int], CitationSource] = {
662
- (0, 0): CitationSource(type="document", id="d0"),
663
- }
664
- raw = [
665
- SimpleNamespace(
666
- start_index=0,
667
- end_index=5,
668
- text="hello",
669
- is_thinking=True,
670
- sources=[self._fake_melody_source(0, [0])],
671
- )
672
- ]
673
- out = _melody_citations_to_vllm(raw, position_map)
674
- assert out is not None
675
- assert out[0].type == "THINKING_CONTENT"
676
- assert out[0].sources[0].id == "d0"
677
-
678
-
679
-class TestParserStreamingEndToEnd:
680
- """End-to-end streaming shape: raw model output text is tokenized
681
- with the mock byte-level tokenizer and fed one token at a time
682
- through :meth:`CohereCommand3ReasoningParser.extract_reasoning_streaming`.
683
- We collect the ordered sequence of ``DeltaMessage``s the parser
684
- hands back -- i.e. what the client would receive on the wire --
685
- and pin the exact shape.
686
-
687
- Combines the three flavors that reach ``delta``: a thinking block
688
- (populates ``reasoning``), a text block (populates ``content``),
689
- and a resolved citation (populates ``citations``) referencing
690
- ``(start, end)`` offsets into the already-streamed content.
691
- """
692
-
693
- _MODEL_OUTPUT = (
694
- "<|START_THINKING|>Let me check.<|END_THINKING|>"
695
- "<|START_RESPONSE|>The capital is <co>Paris</co: 0:[1]>.<|END_RESPONSE|>"
696
- )
697
-
698
- @staticmethod
699
- def _drive_parser(
700
- parser: CohereCommand3ReasoningParser, text: str, tok: MockCohereTokenizer
701
- ) -> list[Any]:
702
- """Feed ``text`` through the parser one token at a time
703
- (byte-level) and return the ordered list of non-``None``
704
- ``DeltaMessage``s.
705
- """
706
- deltas: list[Any] = []
707
- previous_text = ""
708
- previous_token_ids: list[int] = []
709
- for token_str in _token_deltas(tok, text):
710
- current_text = previous_text + token_str
711
- current_token_ids = previous_token_ids + [0]
712
- delta = parser.extract_reasoning_streaming(
713
- previous_text=previous_text,
714
- current_text=current_text,
715
- delta_text=token_str,
716
- previous_token_ids=previous_token_ids,
717
- current_token_ids=current_token_ids,
718
- delta_token_ids=[0],
719
- )
720
- if delta is not None:
721
- deltas.append(delta)
722
- previous_text = current_text
723
- previous_token_ids = current_token_ids
724
- return deltas
725
-
726
- def test_full_wire_stream_shape(self, tokenizer):
727
- from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
728
- from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY
729
-
730
- # Mimic what ``CohereServingChatV2._apply_cohere_template_kwargs``
731
- # installs on a request whose tool_call_index=0 tool result at
732
- # index 1 is the "France" document.
733
- position_to_source = {
734
- (0, 1): CitationSource(
735
- type="document",
736
- id="doc-paris",
737
- document={"id": "doc-paris", "title": "France"},
738
- ),
739
- }
740
-
741
- parser = CohereCommand3ReasoningParser(
742
- tokenizer,
743
- chat_template_kwargs={POSITION_TO_SOURCE_KEY: position_to_source},
744
- )
745
-
746
- deltas = self._drive_parser(parser, self._MODEL_OUTPUT, tokenizer)
747
-
748
- # -- 1. Reasoning and content stream separately and reach
749
- # the wire only as their real payload. Framing tokens
750
- # (``<|START_THINKING|>`` etc.) are consumed by the parser's
751
- # state machine and never produce a delta.
752
- reasoning_stream = [d.reasoning for d in deltas if d.reasoning is not None]
753
- content_stream = [d.content for d in deltas if d.content is not None]
754
- citation_deltas = [d for d in deltas if getattr(d, "citations", None)]
755
-
756
- assert "".join(reasoning_stream) == "Let me check."
757
- assert "".join(content_stream) == "The capital is Paris."
758
-
759
- # -- 2. Exactly one citation reaches the wire, resolved to the
760
- # source from ``position_to_source`` (not the raw
761
- # ``(bucket, idx)`` coordinates).
762
- assert len(citation_deltas) == 1
763
- (cite_delta,) = citation_deltas
764
- assert cite_delta.citations is not None
765
- assert len(cite_delta.citations) == 1
766
- cite = cite_delta.citations[0]
767
- assert cite.text == "Paris"
768
- # Char offsets are into the accumulated ``content``: "The
769
- # capital is " is 15 chars, "Paris" is 5.
770
- assert cite.start == 15
771
- assert cite.end == 20
772
- assert len(cite.sources) == 1
773
- assert cite.sources[0].id == "doc-paris"
774
- assert cite.sources[0].type == "document"
775
-
776
- # -- 3. The citation's ``(start, end)`` offsets point into
777
- # bytes that were already emitted as content-deltas earlier
778
- # in the stream -- i.e. the citation is anchored to real
779
- # already-shipped content.
780
- cite_delta_idx = deltas.index(cite_delta)
781
- earlier_content = "".join(
782
- d.content for d in deltas[:cite_delta_idx] if d.content is not None
783
- )
784
- assert earlier_content[cite.start : cite.end] == "Paris"
785
-
786
- # Everything after the citation delta is the trailing period.
787
- later_content = "".join(
788
- d.content for d in deltas[cite_delta_idx + 1 :] if d.content is not None
789
- )
790
- assert later_content == "."
791
-
792
- # -- 4. Reasoning and content are never mixed on the same
793
- # delta: the parser flips modes on the ``<|END_THINKING|>``
794
- # boundary and never emits a delta with both fields set.
795
- for d in deltas:
796
- assert not (d.reasoning is not None and d.content is not None)
tests/tool_parsers/test_cohere_command_tool_parser.py
deleted
-510
@@ -1,510 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from __future__ import annotations
5
-
6
-import json
7
-from dataclasses import dataclass, field
8
-from typing import Any
9
-
10
-import pytest
11
-
12
-from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
13
-from vllm.tool_parsers.cohere_command_tool_parser import (
14
- CohereCommand3ToolParser,
15
- CohereCommand4ToolParser,
16
-)
17
-
18
-
19
-@dataclass
20
-class ExpectedToolCall:
21
- id: str
22
- name: str
23
- arguments: dict
24
-
25
-
26
-@dataclass
27
-class ToolCallCase:
28
- parser_cls: Any
29
- model_output: str
30
- expected_tool_calls: list[ExpectedToolCall] = field(default_factory=list)
31
- expected_reasoning: str | None = None
32
- expected_content: str | None = None
33
-
34
-
35
-TOOL_CALL_CASES = [
36
- pytest.param(
37
- ToolCallCase(
38
- parser_cls=CohereCommand3ToolParser,
39
- model_output="""\
40
-<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
41
-[
42
- {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
43
-]
44
-<|END_ACTION|>""",
45
- expected_tool_calls=[
46
- ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}),
47
- ],
48
- expected_reasoning="i will call foo with query1",
49
- ),
50
- id="cmd3-single_tool_call",
51
- ),
52
- pytest.param(
53
- ToolCallCase(
54
- parser_cls=CohereCommand4ToolParser,
55
- model_output="""\
56
-<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
57
-[
58
- {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
59
-]
60
-<|END_ACTION|>""",
61
- expected_tool_calls=[
62
- ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}),
63
- ],
64
- expected_reasoning="i will call foo with query1",
65
- ),
66
- id="cmd4-single_tool_call",
67
- ),
68
- pytest.param(
69
- ToolCallCase(
70
- parser_cls=CohereCommand3ToolParser,
71
- model_output="""\
72
-<|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
73
-<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_RESPONSE|>""",
74
- expected_reasoning="This is a rainbow emoji: 🌈",
75
- expected_content="foo bar",
76
- ),
77
- id="cmd3-citations_no_tool_calls",
78
- ),
79
- pytest.param(
80
- ToolCallCase(
81
- parser_cls=CohereCommand4ToolParser,
82
- model_output="""\
83
-<|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
84
-<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_RESPONSE|>""",
85
- expected_reasoning="This is a rainbow emoji: 🌈",
86
- expected_content="foo bar",
87
- ),
88
- id="cmd4-citations_no_tool_calls",
89
- ),
90
- pytest.param(
91
- ToolCallCase(
92
- parser_cls=CohereCommand3ToolParser,
93
- model_output="""\
94
-<|START_THINKING|>first I think about foo<|END_THINKING|><|START_ACTION|>
95
-[
96
- {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}},
97
- {"tool_call_id": "1", "tool_name": "bar", "parameters": {"x": 42}}
98
-]
99
-<|END_ACTION|>""",
100
- expected_tool_calls=[
101
- ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}),
102
- ExpectedToolCall(id="1", name="bar", arguments={"x": 42}),
103
- ],
104
- expected_reasoning="first I think about foo",
105
- ),
106
- id="cmd3-multiple_tool_calls",
107
- ),
108
- pytest.param(
109
- ToolCallCase(
110
- parser_cls=CohereCommand4ToolParser,
111
- model_output="""\
112
-<|START_THINKING|>first I think about foo<|END_THINKING|><|START_ACTION|>
113
-[
114
- {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}},
115
- {"tool_call_id": "1", "tool_name": "bar", "parameters": {"x": 42}}
116
-]
117
-<|END_ACTION|>""",
118
- expected_tool_calls=[
119
- ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}),
120
- ExpectedToolCall(id="1", name="bar", arguments={"x": 42}),
121
- ],
122
- expected_reasoning="first I think about foo",
123
- ),
124
- id="cmd4-multiple_tool_calls",
125
- ),
126
- pytest.param(
127
- ToolCallCase(
128
- parser_cls=CohereCommand3ToolParser,
129
- model_output="""\
130
-<|START_THINKING|>just think, no response<|END_THINKING|>""",
131
- expected_reasoning="just think, no response",
132
- ),
133
- id="cmd3-reasoning_only",
134
- ),
135
- pytest.param(
136
- ToolCallCase(
137
- parser_cls=CohereCommand4ToolParser,
138
- model_output="""\
139
-<|START_THINKING|>just think, no response<|END_THINKING|>""",
140
- expected_reasoning="just think, no response",
141
- ),
142
- id="cmd4-reasoning_only",
143
- ),
144
-]
145
-
146
-
147
-class MockCohereTokenizer:
148
- """Minimal byte-level stand-in for the Cohere tokenizer."""
149
-
150
- def get_vocab(self) -> dict[str, int]:
151
- return {}
152
-
153
- def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
154
- return list(text.encode("utf-8"))
155
-
156
- def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str:
157
- return bytes(ids).decode("utf-8", errors="replace")
158
-
159
-
160
-@pytest.fixture(scope="module")
161
-def tokenizer() -> MockCohereTokenizer:
162
- return MockCohereTokenizer()
163
-
164
-
165
-@pytest.fixture
166
-def request_obj() -> ChatCompletionRequest:
167
- return ChatCompletionRequest(messages=[], model="test-model")
168
-
169
-
170
-REPLACEMENT_CHAR = "\ufffd"
171
-
172
-
173
-def _token_deltas(tokenizer: MockCohereTokenizer, text: str) -> list[str]:
174
- """Decode per-token string deltas, buffering incomplete multi-byte chars."""
175
- ids = tokenizer.encode(text, add_special_tokens=False)
176
- deltas: list[str] = []
177
- prev = ""
178
- for i in range(1, len(ids) + 1):
179
- current = tokenizer.decode(ids[:i], skip_special_tokens=False)
180
- if current.endswith(REPLACEMENT_CHAR):
181
- continue
182
- delta = current[len(prev) :]
183
- if delta:
184
- deltas.append(delta)
185
- prev = current
186
- return deltas
187
-
188
-
189
-@dataclass
190
-class StreamingResult:
191
- tool_calls: dict[int, dict]
192
- reasoning: str | None
193
- content: str | None
194
-
195
-
196
-def _run_streaming_over_deltas(
197
- parser,
198
- deltas: list[str],
199
- request_obj: ChatCompletionRequest,
200
-) -> StreamingResult:
201
- accumulated: dict[int, dict] = {}
202
- reasoning_parts: list[str] = []
203
- content_parts: list[str] = []
204
- previous_text = ""
205
-
206
- for token_str in deltas:
207
- current_text = previous_text + token_str
208
- delta = parser.extract_tool_calls_streaming(
209
- previous_text=previous_text,
210
- current_text=current_text,
211
- delta_text=token_str,
212
- previous_token_ids=[],
213
- current_token_ids=[],
214
- delta_token_ids=[],
215
- request=request_obj,
216
- )
217
- if delta is not None:
218
- if delta.reasoning is not None:
219
- reasoning_parts.append(delta.reasoning)
220
- if delta.content is not None:
221
- content_parts.append(delta.content)
222
- for tc in delta.tool_calls:
223
- idx = tc.index
224
- if idx not in accumulated:
225
- accumulated[idx] = {"id": "", "name": "", "arguments": ""}
226
- if tc.id:
227
- accumulated[idx]["id"] = tc.id
228
- if tc.function and tc.function.name:
229
- accumulated[idx]["name"] = tc.function.name
230
- if tc.function and tc.function.arguments:
231
- accumulated[idx]["arguments"] += tc.function.arguments
232
- previous_text = current_text
233
-
234
- return StreamingResult(
235
- tool_calls=accumulated,
236
- reasoning="".join(reasoning_parts) if reasoning_parts else None,
237
- content="".join(content_parts) if content_parts else None,
238
- )
239
-
240
-
241
-def _run_streaming(
242
- parser,
243
- tokenizer: MockCohereTokenizer,
244
- model_output: str,
245
- request_obj: ChatCompletionRequest,
246
-) -> StreamingResult:
247
- return _run_streaming_over_deltas(
248
- parser,
249
- _token_deltas(tokenizer, model_output),
250
- request_obj,
251
- )
252
-
253
-
254
-@pytest.mark.parametrize("case", TOOL_CALL_CASES)
255
-class TestExtractToolCalls:
256
- def test_streaming(
257
- self,
258
- tokenizer: MockCohereTokenizer,
259
- request_obj: ChatCompletionRequest,
260
- case: ToolCallCase,
261
- ):
262
- parser = case.parser_cls(tokenizer)
263
- streamed = _run_streaming(parser, tokenizer, case.model_output, request_obj)
264
-
265
- assert len(streamed.tool_calls) == len(case.expected_tool_calls)
266
-
267
- for i, expected_tc in enumerate(case.expected_tool_calls):
268
- tc = streamed.tool_calls[i]
269
- assert tc["id"] == expected_tc.id
270
- assert tc["name"] == expected_tc.name
271
- assert json.loads(tc["arguments"]) == expected_tc.arguments
272
-
273
- def test_streaming_reasoning(
274
- self,
275
- tokenizer: MockCohereTokenizer,
276
- request_obj: ChatCompletionRequest,
277
- case: ToolCallCase,
278
- ):
279
- parser = case.parser_cls(tokenizer)
280
- streamed = _run_streaming(parser, tokenizer, case.model_output, request_obj)
281
-
282
- assert streamed.reasoning == case.expected_reasoning
283
-
284
- def test_streaming_content(
285
- self,
286
- tokenizer: MockCohereTokenizer,
287
- request_obj: ChatCompletionRequest,
288
- case: ToolCallCase,
289
- ):
290
- parser = case.parser_cls(tokenizer)
291
- streamed = _run_streaming(parser, tokenizer, case.model_output, request_obj)
292
-
293
- assert streamed.content == case.expected_content
294
-
295
- def test_nonstreaming(
296
- self,
297
- request_obj: ChatCompletionRequest,
298
- tokenizer: MockCohereTokenizer,
299
- case: ToolCallCase,
300
- ):
301
- parser = case.parser_cls(tokenizer)
302
- result = parser.extract_tool_calls(case.model_output, request_obj)
303
-
304
- assert result.tools_called == (len(case.expected_tool_calls) > 0)
305
- assert len(result.tool_calls) == len(case.expected_tool_calls)
306
-
307
- for actual_tc, expected_tc in zip(result.tool_calls, case.expected_tool_calls):
308
- assert actual_tc.type == "function"
309
- assert actual_tc.function.name == expected_tc.name
310
- assert json.loads(actual_tc.function.arguments) == expected_tc.arguments
311
-
312
- def test_streaming_nonstreaming_agree(
313
- self,
314
- tokenizer: MockCohereTokenizer,
315
- request_obj: ChatCompletionRequest,
316
- case: ToolCallCase,
317
- ):
318
- parser_streaming = case.parser_cls(tokenizer)
319
- parser_nonstreaming = case.parser_cls(tokenizer)
320
-
321
- streamed = _run_streaming(
322
- parser_streaming,
323
- tokenizer,
324
- case.model_output,
325
- request_obj,
326
- )
327
- result = parser_nonstreaming.extract_tool_calls(
328
- case.model_output,
329
- request_obj,
330
- )
331
-
332
- assert len(streamed.tool_calls) == len(result.tool_calls)
333
-
334
- for i, actual_tc in enumerate(result.tool_calls):
335
- assert streamed.tool_calls[i]["name"] == actual_tc.function.name
336
- assert json.loads(streamed.tool_calls[i]["arguments"]) == json.loads(
337
- actual_tc.function.arguments
338
- )
339
-
340
-
341
-SPECIAL_TOKEN_MARKERS = (
342
- "<|START_THINKING|>",
343
- "<|END_THINKING|>",
344
- "<|START_RESPONSE|>",
345
- "<|END_RESPONSE|>",
346
- "<|START_ACTION|>",
347
- "<|END_ACTION|>",
348
- "<|START_TEXT|>",
349
- "<|END_TEXT|>",
350
-)
351
-
352
-
353
-def _multi_token_deltas(
354
- tokenizer: MockCohereTokenizer,
355
- text: str,
356
- chunk_size: int,
357
-) -> list[str]:
358
- ids = tokenizer.encode(text, add_special_tokens=False)
359
- deltas: list[str] = []
360
- prev = ""
361
- i = 0
362
- while i < len(ids):
363
- end = min(len(ids), i + chunk_size)
364
- current = tokenizer.decode(ids[:end], skip_special_tokens=False)
365
- i = end
366
- if current.endswith(REPLACEMENT_CHAR):
367
- continue
368
- delta = current[len(prev) :]
369
- if delta:
370
- deltas.append(delta)
371
- prev = current
372
- return deltas
373
-
374
-
375
-class TestSpeculativeDecodingMultiTokenDelta:
376
- MODEL_OUTPUT = (
377
- "<|START_THINKING|> i will call foo with query1<|END_THINKING|>"
378
- "<|START_ACTION|>\n"
379
- '[\n {"tool_call_id": "0", "tool_name": "foo", '
380
- '"parameters": {"query": "query1"}}\n]\n'
381
- "<|END_ACTION|>"
382
- )
383
-
384
- @pytest.mark.parametrize(
385
- "parser_cls",
386
- [CohereCommand3ToolParser, CohereCommand4ToolParser],
387
- ids=["cmd3", "cmd4"],
388
- )
389
- @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6])
390
- def test_no_special_token_leak_in_streaming_deltas(
391
- self,
392
- tokenizer: MockCohereTokenizer,
393
- request_obj: ChatCompletionRequest,
394
- parser_cls,
395
- chunk_size: int,
396
- ):
397
- parser = parser_cls(tokenizer)
398
- chunked = _multi_token_deltas(tokenizer, self.MODEL_OUTPUT, chunk_size)
399
-
400
- previous_text = ""
401
- for token_str in chunked:
402
- current_text = previous_text + token_str
403
- delta = parser.extract_tool_calls_streaming(
404
- previous_text=previous_text,
405
- current_text=current_text,
406
- delta_text=token_str,
407
- previous_token_ids=[],
408
- current_token_ids=[],
409
- delta_token_ids=[],
410
- request=request_obj,
411
- )
412
- previous_text = current_text
413
- if delta is None:
414
- continue
415
-
416
- fields: list[tuple[str, str | None]] = [
417
- ("reasoning", delta.reasoning),
418
- ("content", delta.content),
419
- ]
420
- for tc in delta.tool_calls or []:
421
- if tc.function:
422
- fields.append(("tool_call.name", tc.function.name))
423
- fields.append(("tool_call.arguments", tc.function.arguments))
424
-
425
- for marker in SPECIAL_TOKEN_MARKERS:
426
- for field_name, value in fields:
427
- assert value is None or marker not in value, (
428
- f"special token {marker!r} leaked into {field_name} "
429
- f"with chunk_size={chunk_size} delta={delta!r}"
430
- )
431
-
432
- @pytest.mark.parametrize(
433
- "parser_cls",
434
- [CohereCommand3ToolParser, CohereCommand4ToolParser],
435
- ids=["cmd3", "cmd4"],
436
- )
437
- @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6])
438
- def test_multi_token_chunks_still_produce_correct_tool_call(
439
- self,
440
- tokenizer: MockCohereTokenizer,
441
- request_obj: ChatCompletionRequest,
442
- parser_cls,
443
- chunk_size: int,
444
- ):
445
- parser = parser_cls(tokenizer)
446
- chunked = _multi_token_deltas(tokenizer, self.MODEL_OUTPUT, chunk_size)
447
- streamed = _run_streaming_over_deltas(parser, chunked, request_obj)
448
-
449
- assert len(streamed.tool_calls) == 1
450
- tc = streamed.tool_calls[0]
451
- assert tc["id"] == "0"
452
- assert tc["name"] == "foo"
453
- assert json.loads(tc["arguments"]) == {"query": "query1"}
454
-
455
-
456
-class TestStreamingDeltaShape:
457
- @pytest.mark.parametrize(
458
- "parser_cls",
459
- [CohereCommand3ToolParser, CohereCommand4ToolParser],
460
- ids=["cmd3", "cmd4"],
461
- )
462
- def test_reasoning_and_tool_calls_are_separate_deltas(
463
- self,
464
- tokenizer: MockCohereTokenizer,
465
- request_obj: ChatCompletionRequest,
466
- parser_cls,
467
- ):
468
- parser = parser_cls(tokenizer)
469
- model_output = (
470
- "<|START_THINKING|> i will call foo with query1<|END_THINKING|>"
471
- "<|START_ACTION|>\n"
472
- '[\n {"tool_call_id": "0", "tool_name": "foo", '
473
- '"parameters": {"query": "query1"}}\n]\n'
474
- "<|END_ACTION|>"
475
- )
476
-
477
- token_strings = _token_deltas(tokenizer, model_output)
478
- previous_text = ""
479
- saw_reasoning = False
480
- saw_tool_call = False
481
-
482
- for token_str in token_strings:
483
- current_text = previous_text + token_str
484
- delta = parser.extract_tool_calls_streaming(
485
- previous_text=previous_text,
486
- current_text=current_text,
487
- delta_text=token_str,
488
- previous_token_ids=[],
489
- current_token_ids=[],
490
- delta_token_ids=[],
491
- request=request_obj,
492
- )
493
- if delta is not None:
494
- populated = [
495
- delta.content is not None,
496
- delta.reasoning is not None,
497
- bool(delta.tool_calls),
498
- ]
499
- assert sum(populated) == 1, (
500
- "A single streaming delta must carry exactly one of "
501
- f"content/reasoning/tool_calls, got {delta!r}"
502
- )
503
- if delta.reasoning is not None:
504
- saw_reasoning = True
505
- if delta.tool_calls:
506
- saw_tool_call = True
507
- previous_text = current_text
508
-
509
- assert saw_reasoning, "expected at least one reasoning delta"
510
- assert saw_tool_call, "expected at least one tool-call delta"
vllm/entrypoints/cohere/cohere_chat_message.py
+5
-6
@@ -7,8 +7,8 @@ This module keeps the OpenAI chat completion protocol classes
7
It provides two things:
8
9
* :class:`Citation` / :class:`CitationSource`: the vLLM-internal citation
10
- representation produced by the Cohere reasoning parser
11
- (:mod:`vllm.reasoning.cohere_command_reasoning_parser`) and consumed by
10
+ representation produced by the Cohere parser
11
+ (:mod:`vllm.parser.cohere_command`) and consumed by
12
:mod:`vllm.entrypoints.cohere.serving`. This is *not* the on-wire Cohere
13
SDK shape -- that conversion happens in the serving layer.
14
* :class:`CohereChatMessage` / :class:`CohereDeltaMessage`: :class:`ChatMessage`
@@ -42,9 +42,9 @@ class CitationSource(OpenAIBaseModel):
42
document/tool-output identifier; ``document`` and ``tool_output`` carry
43
the original payload that produced the citation.
44
45
- Sources are fully resolved by the reasoning parser at emit time
45
+ Sources are fully resolved by the parser at emit time
46
(see :func:`_melody_sources_to_vllm` in
47
- :mod:`vllm.reasoning.cohere_command_reasoning_parser`) using the
47
+ :mod:`vllm.parser.cohere_command`) using the
48
``POSITION_TO_SOURCE_KEY`` map forwarded through
49
``chat_template_kwargs`` by
50
:meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._apply_cohere_template_kwargs`.
@@ -121,8 +121,7 @@ class CohereChatMessage(ChatMessage):
121
class CohereDeltaMessage(DeltaMessage):
122
""":class:`DeltaMessage` extension carrying grounding citations for streaming.
123
124
- Emitted by
125
- :class:`vllm.reasoning.cohere_command_reasoning_parser.BaseCohereCommandReasoningParser`
124
+ Emitted by :class:`vllm.parser.cohere_command.CohereCommandParser`
125
on delta events whose payload includes citations. Non-Cohere parsers
126
return plain :class:`DeltaMessage`, so their streamed shape is
127
unchanged.
vllm/entrypoints/cohere/serving.py
+9
-14
@@ -714,8 +714,8 @@ class CohereServingChatV2(OpenAIServingChat):
714
# emits, so citation deltas already carry SDK-shape sources by
715
# the time they reach the streaming loop. See
716
# ``_build_position_to_source`` for the numbering rule and
717
- # ``_melody_sources_to_vllm`` in the reasoning parser for the
718
- # consumer.
717
+ # ``_melody_sources_to_vllm`` in ``vllm.parser.cohere_command`` for
718
+ # the consumer.
719
position_to_source = cls._build_position_to_source(request)
720
if position_to_source:
721
kwargs.setdefault(POSITION_TO_SOURCE_KEY, position_to_source)
@@ -1109,12 +1109,11 @@ class CohereServingChatV2(OpenAIServingChat):
1109
*,
1110
parser: Parser | None,
1111
) -> ChatMessage:
1112
- """Copy grounding citations off the reasoning parser onto the message.
1112
+ """Copy grounding citations off the unified parser onto the message.
1113
1114
- The Cohere reasoning parser
1115
- (:mod:`vllm.reasoning.cohere_command_reasoning_parser`) caches the
1116
- citations produced by its most recent unary ``extract_reasoning``
1117
- call on ``last_unary_citations``. We surface them here on
1114
+ :class:`vllm.parser.cohere_command.CohereCommandParser` caches the
1115
+ citations produced by its most recent unary ``parse`` call on
1116
+ ``last_unary_citations``. We surface them here on
1117
:class:`CohereChatMessage` so downstream response conversion can
1118
pick them up without the base :class:`OpenAIServingChat` having to
1119
know about citations.
@@ -1128,11 +1127,7 @@ class CohereServingChatV2(OpenAIServingChat):
1127
# ``model_dump`` via the extras bucket, so the wire is correct
1128
# either way.
1129
message = cast(CohereChatMessage, message)
1131
- citations = getattr(
1132
- getattr(parser, "reasoning_parser", None),
1133
- "last_unary_citations",
1134
- None,
1135
- )
1130
+ citations = getattr(parser, "last_unary_citations", None)
1131
if citations:
1132
message.citations = citations
1133
return message
@@ -1143,10 +1138,10 @@ class CohereServingChatV2(OpenAIServingChat):
1138
:class:`CohereChatMessage` carries a
1139
``citations: list[vllm...Citation] | None`` field populated by
1140
:meth:`_finalize_response_message` (which reads it off the
1146
- reasoning parser's ``last_unary_citations`` cache). Sources are
1141
+ unified parser's ``last_unary_citations`` cache). Sources are
1142
already fully resolved by the parser (see
1143
:func:`_melody_sources_to_vllm` in
1149
- :mod:`vllm.reasoning.cohere_command_reasoning_parser`) using
1144
+ :mod:`vllm.parser.cohere_command`) using
1145
the position map forwarded via ``chat_template_kwargs``. This
1146
method's remaining job is:
1147
vllm/parser/cohere_command.py
new
+687
@@ -0,0 +1,687 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Cohere Command parser.
4
+
5
+Uses ``cohere_melody`` as a stateful parser filter.
6
+"""
7
+
8
+from __future__ import annotations
9
+
10
+import json
11
+from collections.abc import Mapping, Sequence
12
+from typing import TYPE_CHECKING, Any, NamedTuple, TypedDict, TypeGuard
13
+
14
+import regex as re
15
+import xgrammar as xgr
16
+
17
+try:
18
+ from cohere_melody import (
19
+ AccumulatedToolCall,
20
+ FilterAggregatedResult,
21
+ PyFilter,
22
+ PyFilterOptions,
23
+ )
24
+except ImportError as e:
25
+ raise ImportError(
26
+ "The Cohere Command parser requires the `cohere_melody` package. "
27
+ "Install it with:\n pip install 'cohere-melody>=0.14.0'"
28
+ ) from e
29
+
30
+from vllm.entrypoints.cohere.cohere_chat_message import (
31
+ Citation,
32
+ CitationSource,
33
+ CohereDeltaMessage,
34
+)
35
+from vllm.entrypoints.generate.base.protocol import (
36
+ AnyResponseFormat,
37
+ DeltaFunctionCall,
38
+ DeltaMessage,
39
+ DeltaToolCall,
40
+ FunctionCall,
41
+)
42
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
43
+from vllm.parser.abstract_parser import DelegatingParser
44
+from vllm.reasoning.cohere_command_reasoning_parser import (
45
+ BaseCohereCommandReasoningParser,
46
+)
47
+from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY
48
+from vllm.sampling_params import StructuredOutputsParams
49
+from vllm.tool_parsers.cohere_command_tool_parser import BaseCohereCommandToolParser
50
+
51
+if TYPE_CHECKING:
52
+ # Declared in melody's stub only; not exported at runtime.
53
+ from cohere_melody import FilterCitation, Source
54
+
55
+ from vllm.config import ModelConfig
56
+ from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
57
+ from vllm.tokenizers import TokenizerLike
58
+ from vllm.tool_parsers.utils import Tool
59
+
60
+PositionToSource = Mapping[tuple[int, int], CitationSource]
61
+
62
+
63
+class CohereTagRegistry(NamedTuple):
64
+ """A single ``structural_tag`` trigger / end pair (``begin`` uses ``trigger``)."""
65
+
66
+ trigger: str
67
+ end: str
68
+
69
+
70
+class CohereTagStyle(NamedTuple):
71
+ """The structural tags style for a given model architecture.
72
+
73
+ ``json_tags`` lists every JSON-schema wrapper the model may emit (MOE uses
74
+ both response and text delimiters). ``tools`` is the tool-call wrapper.
75
+ """
76
+
77
+ json_tags: tuple[CohereTagRegistry, ...]
78
+ tools: CohereTagRegistry
79
+
80
+
81
+class CohereNormalizedTool(TypedDict):
82
+ """A tool definition normalized to the shape ``collect_tool_schema`` expects.
83
+
84
+ ``parameters`` is a JSON Schema object (possibly empty) describing the tool's
85
+ call signature.
86
+ """
87
+
88
+ name: str
89
+ parameters: dict[str, Any]
90
+
91
+
92
+COMMAND_A_TOOLS_TAG = CohereTagRegistry(
93
+ trigger="<|START_ACTION|>",
94
+ end="<|END_ACTION|>",
95
+)
96
+COMMAND_A_JSON_TAG = CohereTagRegistry(
97
+ trigger="<|START_RESPONSE|>",
98
+ end="<|END_RESPONSE|>",
99
+)
100
+COMMAND_A_PLUS_JSON_TAG = CohereTagRegistry(
101
+ trigger="<|START_TEXT|>",
102
+ end="<|END_TEXT|>",
103
+)
104
+
105
+MODEL_TO_TAG_STYLE: dict[str, CohereTagStyle] = {
106
+ "Cohere2ForCausalLM": CohereTagStyle(
107
+ json_tags=(COMMAND_A_JSON_TAG,),
108
+ tools=COMMAND_A_TOOLS_TAG,
109
+ ),
110
+ "Cohere2VisionForConditionalGeneration": CohereTagStyle(
111
+ json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG),
112
+ tools=COMMAND_A_TOOLS_TAG,
113
+ ),
114
+ "Cohere2MoeForCausalLM": CohereTagStyle(
115
+ json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG),
116
+ tools=COMMAND_A_TOOLS_TAG,
117
+ ),
118
+}
119
+
120
+
121
+def collect_tool_schema(tool_schema: list[CohereNormalizedTool]) -> str:
122
+ """Build an xgrammar EBNF grammar that matches a JSON array of tool calls.
123
+
124
+ The grammar shape is architecture-independent; callers are responsible for
125
+ wrapping it in the correct structural tag (see ``CohereTagStyle.tools``).
126
+ """
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))
146
+ for match in re.findall(r"\b(\w+)\s*::=", tool_grammar):
147
+ tool_grammar = re.sub(
148
+ rf"\b{re.escape(match)}\b", tool_name + match, tool_grammar
149
+ )
150
+ tool_dictionary[tool_name] = f"{tool_name} ::= {tool_name}root\n{tool_grammar}"
151
+ # Emitted grammar shape:
152
+ # root ::= tools
153
+ # tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws
154
+ # ws ::= (" " | "\t" | "\n")*
155
+ # tool ::= <tool_a> | <tool_b> | ... (one alternative per input)
156
+ # <tool_x> ::= <tool_x>root (per-tool xgrammar rules)
157
+ # <tool_x>root ::= ... (from xgr.Grammar.from_json_schema)
158
+ tool_alternatives = "tool ::= " + " | ".join(tool_dictionary.keys())
159
+ tool_rules = "\n ".join(tool_dictionary.values())
160
+ grammar = f"""root ::= tools
161
+ tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws
162
+ ws ::= (" " | "\\t" | "\\n")*
163
+ {tool_alternatives}
164
+ {tool_rules}
165
+ """
166
+ return grammar
167
+
168
+
169
+def _tool_definitions_to_schema_list(
170
+ tools: str | list[Any],
171
+) -> list[CohereNormalizedTool]:
172
+ """
173
+ Build the list of ``CohereNormalizedTool`` dicts expected by
174
+ ``collect_tool_schema``.
175
+
176
+ Accepts:
177
+ - JSON string
178
+ - list of dicts with top-level ``name`` / ``parameters``
179
+ - list of Chat Completions-style ``{"type": "function", "function": {...}}``
180
+ - list of Pydantic models with ``model_dump()``
181
+ """
182
+ if isinstance(tools, str):
183
+ try:
184
+ parsed = json.loads(tools)
185
+ except json.JSONDecodeError:
186
+ return []
187
+ if not isinstance(parsed, list):
188
+ return []
189
+ else:
190
+ parsed = list(tools)
191
+
192
+ out: list[CohereNormalizedTool] = []
193
+ for raw in parsed:
194
+ t = raw.model_dump() if hasattr(raw, "model_dump") else raw
195
+ if not isinstance(t, dict):
196
+ continue
197
+ # Unwrap Chat Completions' ``{"type": "function", "function": {...}}``
198
+ # shape; otherwise take the dict as-is.
199
+ if t.get("type") == "function" and isinstance(t.get("function"), dict):
200
+ t = t["function"]
201
+ name = t.get("name")
202
+ if not isinstance(name, str):
203
+ continue
204
+ params = t.get("parameters")
205
+ out.append(
206
+ CohereNormalizedTool(
207
+ name=name,
208
+ parameters=params if isinstance(params, dict) else {},
209
+ )
210
+ )
211
+ return out
212
+
213
+
214
+def _has_effective_tools(
215
+ tools: str | list[Any] | None,
216
+) -> TypeGuard[str | list[Any]]:
217
+ """
218
+ True when ``tools`` contains at least one tool definition to convert.
219
+
220
+ ``ResponsesRequest`` defaults ``tools`` to ``[]``; ``ChatCompletionRequest``
221
+ uses ``None``. Both mean "no tools" here. Strings (e.g. a JSON blob) are
222
+ treated as effective only when non-blank.
223
+ """
224
+ if tools is None:
225
+ return False
226
+ if isinstance(tools, str):
227
+ return bool(tools.strip())
228
+ return len(tools) > 0
229
+
230
+
231
+# Builder: produces vLLM response_format in xgrammar's canonical format.
232
+# See xgrammar docs: type "structural_tag" with "format" = triggered_tags
233
+# and tag content type = json_schema | grammar.
234
+def convert_schema_to_structural_tags(
235
+ schema: dict | None = None,
236
+ tools: str | list[Any] | None = None,
237
+ model_architecture: str | None = None,
238
+) -> str | None:
239
+ """
240
+ Returns a response_format string accepted by xgrammar's structural tag format.
241
+ Uses the canonical shape: {"type": "structural_tag", "format": {...}} with
242
+ format.type "triggered_tags" and tag content type "json_schema" or "grammar".
243
+
244
+ Callers that are not on an engine path (e.g. the reasoning parser) must pass
245
+ ``model_architecture`` explicitly.
246
+ """
247
+ if model_architecture is None or model_architecture not in MODEL_TO_TAG_STYLE:
248
+ return None
249
+ style = MODEL_TO_TAG_STYLE[model_architecture]
250
+
251
+ tags: list[dict] = []
252
+ triggers: list[str] = []
253
+
254
+ def _add_tag(tag: CohereTagRegistry, content: dict) -> None:
255
+ tags.append({"begin": tag.trigger, "content": content, "end": tag.end})
256
+ triggers.append(tag.trigger)
257
+
258
+ if schema is not None:
259
+ # One structural tag per JSON wrapper (e.g. MOE: response + text).
260
+ # Same for schema-only and "tools plus JSON mode" (North: schema when
261
+ # the model does not call tools).
262
+ for jt in style.json_tags:
263
+ _add_tag(jt, {"type": "json_schema", "json_schema": schema})
264
+
265
+ if _has_effective_tools(tools):
266
+ # ``tools`` may be a JSON string (poseidon / RESPONSE_FORMAT_TOOL_DEFINITIONS)
267
+ # or a list (Chat Completions ``request.tools`` as Pydantic models or dicts).
268
+ tool_schema_list = _tool_definitions_to_schema_list(tools)
269
+ if not tool_schema_list:
270
+ raise ValueError(
271
+ "No valid tool definitions could be parsed from the request for "
272
+ "structural tag conversion."
273
+ )
274
+ tool_grammar = collect_tool_schema(tool_schema_list)
275
+ _add_tag(style.tools, {"type": "grammar", "grammar": tool_grammar})
276
+
277
+ if not tags:
278
+ return None
279
+ return json.dumps(
280
+ {
281
+ "type": "structural_tag",
282
+ "format": {
283
+ "type": "triggered_tags",
284
+ "triggers": triggers,
285
+ "tags": tags,
286
+ },
287
+ }
288
+ )
289
+
290
+
291
+def _response_format_type(
292
+ response_format: AnyResponseFormat | dict | None,
293
+) -> str | None:
294
+ if response_format is None:
295
+ return None
296
+ if isinstance(response_format, dict):
297
+ t = response_format.get("type")
298
+ return t if isinstance(t, str) else None
299
+ return response_format.type
300
+
301
+
302
+def _maybe_parse_json_dict(value: Any) -> dict | None:
303
+ """If value is a JSON string, parse to dict; otherwise require dict."""
304
+ if isinstance(value, dict):
305
+ return value
306
+ if isinstance(value, str):
307
+ try:
308
+ parsed = json.loads(value)
309
+ except (TypeError, json.JSONDecodeError):
310
+ return None
311
+ return parsed if isinstance(parsed, dict) else None
312
+ return None
313
+
314
+
315
+def _unwrap_nested_schema(candidate: Any) -> dict | None:
316
+ """Return ``candidate`` as a dict, unwrapping a nested ``schema`` if present.
317
+
318
+ Returns ``None`` if ``candidate`` is not (and cannot be parsed into) a dict.
319
+ """
320
+ cand = _maybe_parse_json_dict(candidate)
321
+ if not isinstance(cand, dict):
322
+ return None
323
+ nested = cand.get("schema")
324
+ return nested if isinstance(nested, dict) else cand
325
+
326
+
327
+def _schema_from_json_schema_field(js_wr: Any) -> dict | None:
328
+ """
329
+ Extract the JSON Schema object from Chat Completions ``json_schema`` payload.
330
+
331
+ Accepts:
332
+ - ``JsonSchemaResponseFormat`` (Pydantic) with ``schema`` / ``json_schema`` field
333
+ - dict in OpenAI shape ``{"name": ..., "schema": {...}}``
334
+ - dict with ``json_schema`` key holding either the schema or a nested wrapper
335
+ - dict that is already a JSON Schema document (some clients omit the wrapper)
336
+ - JSON strings for any of the above
337
+ """
338
+ if js_wr is None:
339
+ return None
340
+
341
+ parsed_wr = _maybe_parse_json_dict(js_wr)
342
+ if parsed_wr is not None:
343
+ js_wr = parsed_wr
344
+
345
+ if hasattr(js_wr, "model_dump"):
346
+ for by_alias in (True, False):
347
+ try:
348
+ data = js_wr.model_dump(by_alias=by_alias, exclude_none=False)
349
+ except TypeError:
350
+ data = js_wr.model_dump(by_alias=by_alias)
351
+ out = _unwrap_nested_schema(data.get("schema") or data.get("json_schema"))
352
+ if out is not None:
353
+ return out
354
+ inner_attr = getattr(js_wr, "json_schema", None)
355
+ return inner_attr if isinstance(inner_attr, dict) else None
356
+
357
+ if isinstance(js_wr, dict):
358
+ for key in ("schema", "json_schema"):
359
+ out = _unwrap_nested_schema(js_wr.get(key))
360
+ if out is not None:
361
+ return out
362
+ return js_wr
363
+
364
+ return None
365
+
366
+
367
+def _schema_dict_from_chat_response_format(
368
+ rf: AnyResponseFormat | dict | None,
369
+) -> dict | None:
370
+ """JSON schema dict from Chat Completions ``request.response_format`` only."""
371
+ if rf is None:
372
+ return None
373
+ rf_type = _response_format_type(rf)
374
+ if rf_type == "json_object":
375
+ return {"type": "object"}
376
+ if rf_type != "json_schema":
377
+ return None
378
+ js_wr = (
379
+ rf.get("json_schema")
380
+ if isinstance(rf, dict)
381
+ else getattr(rf, "json_schema", None)
382
+ )
383
+ return _schema_from_json_schema_field(js_wr)
384
+
385
+
386
+def _schema_dict_from_structured_outputs(
387
+ so: StructuredOutputsParams | None,
388
+) -> dict | None:
389
+ """Schema dict from ``structured_outputs`` (``json`` / ``json_object``).
390
+
391
+ Same unwrapping as ``json_schema``. ``json`` is expected to be ``str`` or
392
+ ``dict`` (enforced by ``StructuredOutputsParams`` / request models); other
393
+ types raise ``ValueError`` only if a caller bypasses that validation.
394
+ """
395
+ if so is None:
396
+ return None
397
+ if so.json_object:
398
+ return {"type": "object"}
399
+ raw: Any = so.json
400
+ if raw is None:
401
+ return None
402
+
403
+ if hasattr(raw, "model_dump"):
404
+ out = _schema_from_json_schema_field(raw)
405
+ if out is None:
406
+ raise ValueError(
407
+ "structured_outputs.json model has no extractable JSON Schema."
408
+ )
409
+ return out
410
+
411
+ if isinstance(raw, str):
412
+ if not raw.strip():
413
+ raise ValueError("structured_outputs.json cannot be empty.")
414
+ try:
415
+ raw = json.loads(raw)
416
+ except json.JSONDecodeError as e:
417
+ raise ValueError("structured_outputs.json must be valid JSON.") from e
418
+ if not isinstance(raw, dict):
419
+ raise ValueError("structured_outputs.json must decode to a JSON object.")
420
+
421
+ if isinstance(raw, Mapping):
422
+ body = raw if isinstance(raw, dict) else dict(raw)
423
+ return _schema_from_json_schema_field(body) or body
424
+
425
+ raise ValueError(
426
+ f"structured_outputs.json has unsupported type {type(raw).__name__}."
427
+ )
428
+
429
+
430
+def _melody_sources_to_vllm(
431
+ sources: list[Source], position_to_source: PositionToSource | None
432
+) -> list[CitationSource]:
433
+ """Resolve melody ``Source`` addresses into :class:`CitationSource` objects.
434
+
435
+ Melody reports each source as ``(tool_call_index, tool_result_indices)``.
436
+ ``position_to_source`` is built by
437
+ :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._build_position_to_source`
438
+ and forwarded through ``chat_template_kwargs``; each index fans out to one
439
+ resolved source, and unresolved positions are skipped. Without a map (a
440
+ request that did not come through the Cohere v2 handler) every source is
441
+ dropped and the serving layer discards the citation at wire-coercion time.
442
+ """
443
+ if not position_to_source:
444
+ return []
445
+ out: list[CitationSource] = []
446
+ for source in sources:
447
+ for idx in source.tool_result_indices:
448
+ info = position_to_source.get((source.tool_call_index, idx))
449
+ if info is not None:
450
+ out.append(info)
451
+ return out
452
+
453
+
454
+def _melody_citations_to_vllm(
455
+ citations: list[FilterCitation], position_to_source: PositionToSource | None
456
+) -> list[Citation] | None:
457
+ """Convert melody ``FilterCitation`` objects into :class:`Citation`.
458
+
459
+ Citations whose sources all fail to resolve are still returned with
460
+ ``sources=[]``; the serving layer drops them so the fail-closed policy
461
+ lives in one place.
462
+ """
463
+ if not citations:
464
+ return None
465
+ return [
466
+ Citation(
467
+ start=c.start_index,
468
+ end=c.end_index,
469
+ text=c.text,
470
+ sources=_melody_sources_to_vllm(c.sources, position_to_source),
471
+ type="THINKING_CONTENT" if c.is_thinking else "TEXT_CONTENT",
472
+ )
473
+ for c in citations
474
+ ]
475
+
476
+
477
+class CohereCommandParser(DelegatingParser):
478
+ def __init__(
479
+ self,
480
+ tokenizer: TokenizerLike,
481
+ tools: list[Tool] | None = None,
482
+ *args,
483
+ model_config: ModelConfig | None = None,
484
+ **kwargs,
485
+ ):
486
+ super().__init__(tokenizer, tools, *args, model_config=model_config, **kwargs)
487
+ self._model_config = model_config
488
+ self.start_token_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>")
489
+ self.end_token_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>")
490
+ shim = self._reasoning_parser or self._tool_parser
491
+ assert isinstance(
492
+ shim, (BaseCohereCommandReasoningParser, BaseCohereCommandToolParser)
493
+ )
494
+ options = getattr(PyFilterOptions(), shim.melody_preset)()
495
+ self._streaming = PyFilter(options)
496
+ self._unary = PyFilter(options)
497
+ # Citations from the most recent ``parse`` call, surfaced on
498
+ # ``CohereChatMessage`` by ``CohereServingChatV2``.
499
+ self.last_unary_citations: list[Citation] | None = None
500
+ # Melody may emit a tool-call id before the function name; hold it until
501
+ # the first name delta so clients receive both together.
502
+ self._pending_tool_ids: dict[int, str] = {}
503
+ ctk = kwargs.get("chat_template_kwargs") or {}
504
+ raw_map = ctk.get(POSITION_TO_SOURCE_KEY)
505
+ self._position_to_source: PositionToSource | None = (
506
+ raw_map if isinstance(raw_map, Mapping) else None
507
+ )
508
+
509
+ def adjust_request(
510
+ self, request: ChatCompletionRequest | ResponsesRequest
511
+ ) -> ChatCompletionRequest | ResponsesRequest:
512
+ # keep special tokens to ensure melody can correctly identify boundaries.
513
+ request.skip_special_tokens = False
514
+ request = self._apply_structural_tags(request)
515
+ return super().adjust_request(request)
516
+
517
+ def _apply_structural_tags(
518
+ self, request: ChatCompletionRequest | ResponsesRequest
519
+ ) -> ChatCompletionRequest | ResponsesRequest:
520
+ """Fold ``tools`` and any JSON schema into a Cohere ``structural_tag``."""
521
+ so = request.structured_outputs
522
+ if so is not None and so.structural_tag:
523
+ return request
524
+ # Schema: prefer ``response_format`` (OpenAI Chat Completions), then
525
+ # ``structured_outputs.json`` / ``json_object`` (vLLM direct). Tools stay
526
+ # on ``request.tools``.
527
+ rf = (
528
+ request.response_format
529
+ if isinstance(request, ChatCompletionRequest)
530
+ else None
531
+ )
532
+ if rf is not None and _response_format_type(rf) == "structural_tag":
533
+ return request
534
+ model_architecture = (
535
+ self._model_config.architecture if self._model_config is not None else None
536
+ )
537
+ tools = request.tools
538
+ # ``response_format`` wins if both it and ``structured_outputs`` supply JSON.
539
+ schema = _schema_dict_from_chat_response_format(rf)
540
+ if schema is None:
541
+ schema = _schema_dict_from_structured_outputs(so)
542
+ if schema is None and not _has_effective_tools(tools):
543
+ return request
544
+ if model_architecture is None:
545
+ return request
546
+ result = convert_schema_to_structural_tags(
547
+ schema=schema,
548
+ tools=tools,
549
+ model_architecture=model_architecture,
550
+ )
551
+ if result is None:
552
+ # Unsupported architectures are not in ``MODEL_TO_TAG_STYLE``.
553
+ raise ValueError(
554
+ "Failed to build structural_tag guided decoding constraints from "
555
+ "this request's JSON schema and/or tools. The configured model "
556
+ f"architecture ({model_architecture!r}) does not support Cohere "
557
+ "command structural tags, or the schema cannot be expressed in "
558
+ "that format."
559
+ )
560
+ request.structured_outputs = StructuredOutputsParams(structural_tag=result)
561
+ # Folded JSON constraints into ``structural_tag``; drop ``response_format``
562
+ # when it was the source so ``to_sampling_params`` does not also set ``json`` /
563
+ # ``json_object`` (mutually exclusive in ``StructuredOutputsParams``).
564
+ if isinstance(request, ChatCompletionRequest) and rf is not None:
565
+ rf_type = _response_format_type(rf)
566
+ if rf_type in ("json_schema", "json_object"):
567
+ request.response_format = None
568
+ return request
569
+
570
+ def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
571
+ count = 0
572
+ depth = 0
573
+ for token_id in token_ids:
574
+ if token_id == self.start_token_id:
575
+ depth += 1
576
+ continue
577
+ if token_id == self.end_token_id:
578
+ if depth > 0:
579
+ depth -= 1
580
+ continue
581
+ if depth > 0:
582
+ count += 1
583
+ return count
584
+
585
+ def parse(
586
+ self,
587
+ model_output: str,
588
+ request: ChatCompletionRequest | ResponsesRequest,
589
+ enable_auto_tools: bool = False,
590
+ model_output_token_ids: Sequence[int] = (),
591
+ ) -> tuple[str | None, str | None, list[FunctionCall] | None]:
592
+ result = self._unary.process_full_text(model_output)
593
+ self.last_unary_citations = _melody_citations_to_vllm(
594
+ result.citations, self._position_to_source
595
+ )
596
+ tool_calls: list[FunctionCall] | None = None
597
+ if self._tools_allowed(request):
598
+ tool_calls = [
599
+ FunctionCall(id=tc.id, name=tc.name, arguments=tc.arguments)
600
+ for tc in result.tool_calls
601
+ ] or None
602
+ return result.reasoning, result.content, tool_calls
603
+
604
+ def parse_delta(
605
+ self,
606
+ delta_text: str,
607
+ delta_token_ids: list[int],
608
+ request: ChatCompletionRequest | ResponsesRequest,
609
+ prompt_token_ids: list[int] | None = None,
610
+ *,
611
+ finished: bool,
612
+ ) -> DeltaMessage | None:
613
+ results = [self._streaming.write_decoded(delta_text)]
614
+ if finished:
615
+ results.append(self._streaming.flush_partials())
616
+ return self._to_delta(results, request)
617
+
618
+ def _tools_allowed(self, request: ChatCompletionRequest | ResponsesRequest) -> bool:
619
+ return self._tool_parser is not None and request.tool_choice != "none"
620
+
621
+ def _to_delta(
622
+ self,
623
+ results: list[FilterAggregatedResult],
624
+ request: ChatCompletionRequest | ResponsesRequest,
625
+ ) -> CohereDeltaMessage | None:
626
+ content = self._join(r.content for r in results)
627
+ reasoning = (
628
+ self._join(r.reasoning for r in results)
629
+ if request.include_reasoning
630
+ else None
631
+ )
632
+ tool_calls: list[DeltaToolCall] = []
633
+ citations: list[Citation] = []
634
+ for result in results:
635
+ if result.tool_calls and self._tools_allowed(request):
636
+ tool_calls.extend(self._tool_call_deltas(result.tool_calls))
637
+ citations.extend(
638
+ _melody_citations_to_vllm(result.citations, self._position_to_source)
639
+ or []
640
+ )
641
+ if not (content or reasoning or tool_calls or citations):
642
+ return None
643
+ # ``CohereDeltaMessage`` carries citations to the wire via
644
+ # ``SerializeAsAny[DeltaMessage]``; it drops the field when unset so the
645
+ # shape matches a plain ``DeltaMessage``.
646
+ message = CohereDeltaMessage()
647
+ if content is not None:
648
+ message.content = content
649
+ if reasoning is not None:
650
+ message.reasoning = reasoning
651
+ if tool_calls:
652
+ message.tool_calls = tool_calls
653
+ if citations:
654
+ message.citations = citations
655
+ return message
656
+
657
+ @staticmethod
658
+ def _join(parts: Any) -> str | None:
659
+ present = [p for p in parts if p is not None]
660
+ return "".join(present) if present else None
661
+
662
+ def _tool_call_deltas(
663
+ self, melody_tool_calls: list[AccumulatedToolCall]
664
+ ) -> list[DeltaToolCall]:
665
+ tool_calls: list[DeltaToolCall] = []
666
+ for tc in melody_tool_calls:
667
+ name = tc.name or None
668
+ arguments = tc.arguments or None
669
+ if tc.id:
670
+ self._pending_tool_ids[tc.index] = tc.id
671
+ # Empty strings are melody streaming placeholders; skip fragments
672
+ # that carry neither a name nor argument text.
673
+ if name is None and arguments is None:
674
+ continue
675
+ function = DeltaFunctionCall()
676
+ if name is not None:
677
+ function.name = name
678
+ if arguments is not None:
679
+ function.arguments = arguments
680
+ delta = DeltaToolCall(index=tc.index, function=function)
681
+ if name is not None:
682
+ resolved_id = tc.id or self._pending_tool_ids.pop(tc.index, None)
683
+ if resolved_id is not None:
684
+ delta.id = resolved_id
685
+ delta.type = "function"
686
+ tool_calls.append(delta)
687
+ return tool_calls
vllm/parser/parser_manager.py
+15
@@ -126,6 +126,21 @@ class ParserManager:
126
127
return _KimiK3Parser
128
129
+ if {reasoning_parser_name, tool_parser_name} & {
130
+ "cohere_command3",
131
+ "cohere_command4",
132
+ }:
133
+ from vllm.parser.cohere_command import CohereCommandParser
134
+
135
+ r_cls = reasoning_parser_cls
136
+ t_cls = tool_parser_cls
137
+
138
+ class _CohereCommandParser(CohereCommandParser):
139
+ reasoning_parser_cls = r_cls
140
+ tool_parser_cls = t_cls
141
+
142
+ return _CohereCommandParser
143
+
144
from vllm.parser.abstract_parser import DelegatingParser
145
146
r_cls = reasoning_parser_cls
vllm/reasoning/cohere_command_reasoning_parser.py
+11
-669
@@ -3,531 +3,28 @@
3
4
from __future__ import annotations
5
6
-import json
7
-from collections.abc import Mapping, Sequence
8
-from typing import TYPE_CHECKING, Any, NamedTuple, TypedDict, TypeGuard
6
+from collections.abc import Sequence
7
+from typing import TYPE_CHECKING, ClassVar
8
10
-import regex as re
11
-import xgrammar as xgr
12
-
13
-try:
14
- from cohere_melody import PyFilter, PyFilterOptions
15
-except ImportError as e:
16
- raise ImportError(
17
- "The Cohere reasoning parser requires the `cohere_melody` "
18
- "package, which is not installed. Install it with:\n"
19
- " pip install 'cohere-melody>=0.11.1'"
20
- ) from e
21
-
22
-
23
-from vllm.entrypoints.cohere.cohere_chat_message import (
24
- Citation,
25
- CitationSource,
26
- CohereDeltaMessage,
27
-)
28
-from vllm.entrypoints.generate.base.protocol import (
29
- AnyResponseFormat,
30
- DeltaFunctionCall,
31
- DeltaMessage,
32
- DeltaToolCall,
33
-)
34
-from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
9
from vllm.reasoning import ReasoningParser
36
-from vllm.sampling_params import StructuredOutputsParams
37
-from vllm.tokenizers import TokenizerLike
10
11
if TYPE_CHECKING:
40
- from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
41
-
42
-REPLACEMENT_CHAR = "\ufffd"
43
-
44
-
45
-class CohereTagRegistry(NamedTuple):
46
- """A single ``structural_tag`` trigger / end pair (``begin`` uses ``trigger``)."""
47
-
48
- trigger: str
49
- end: str
50
-
51
-
52
-class CohereTagStyle(NamedTuple):
53
- """The structural tags style for a given model architecture.
54
-
55
- ``json_tags`` lists every JSON-schema wrapper the model may emit (MOE uses
56
- both response and text delimiters). ``tools`` is the tool-call wrapper.
57
- """
58
-
59
- json_tags: tuple[CohereTagRegistry, ...]
60
- tools: CohereTagRegistry
61
-
62
-
63
-class CohereNormalizedTool(TypedDict):
64
- """A tool definition normalized to the shape ``collect_tool_schema`` expects.
65
-
66
- ``parameters`` is a JSON Schema object (possibly empty) describing the tool's
67
- call signature.
68
- """
69
-
70
- name: str
71
- parameters: dict[str, Any]
72
-
73
-
74
-COMMAND_A_TOOLS_TAG = CohereTagRegistry(
75
- trigger="<|START_ACTION|>",
76
- end="<|END_ACTION|>",
77
-)
78
-COMMAND_A_JSON_TAG = CohereTagRegistry(
79
- trigger="<|START_RESPONSE|>",
80
- end="<|END_RESPONSE|>",
81
-)
82
-COMMAND_A_PLUS_JSON_TAG = CohereTagRegistry(
83
- trigger="<|START_TEXT|>",
84
- end="<|END_TEXT|>",
85
-)
86
-
87
-MODEL_TO_TAG_STYLE: dict[str, CohereTagStyle] = {
88
- "Cohere2ForCausalLM": CohereTagStyle(
89
- json_tags=(COMMAND_A_JSON_TAG,),
90
- tools=COMMAND_A_TOOLS_TAG,
91
- ),
92
- "Cohere2VisionForConditionalGeneration": CohereTagStyle(
93
- json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG),
94
- tools=COMMAND_A_TOOLS_TAG,
95
- ),
96
- "Cohere2MoeForCausalLM": CohereTagStyle(
97
- json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG),
98
- tools=COMMAND_A_TOOLS_TAG,
99
- ),
100
-}
101
-
102
-
103
-def collect_tool_schema(tool_schema: list[CohereNormalizedTool]) -> str:
104
- """Build an xgrammar EBNF grammar that matches a JSON array of tool calls.
105
-
106
- The grammar shape is architecture-independent; callers are responsible for
107
- wrapping it in the correct structural tag (see ``CohereTagStyle.tools``).
108
- """
109
- tool_dictionary: dict[str, str] = {}
110
- for tool in tool_schema:
111
- tool_name = tool["name"]
112
- tool_parameters = json.dumps(tool["parameters"])
113
- json_schema = f"""{{
114
- "type": "object",
115
- "properties": {{
116
- "tool_call_id": {{
117
- "type": "string",
118
- "pattern": "^[0-9]+$"
119
- }},
120
- "tool_name": {{
121
- "type": "string",
122
- "const": "{tool_name}"
123
- }},
124
- "parameters": {tool_parameters}
125
- }}
126
- }}"""
127
- tool_grammar = str(xgr.Grammar.from_json_schema(json_schema))
128
- for match in re.findall(r"\b(\w+)\s*::=", tool_grammar):
129
- tool_grammar = re.sub(
130
- rf"\b{re.escape(match)}\b", tool_name + match, tool_grammar
131
- )
132
- tool_dictionary[tool_name] = f"{tool_name} ::= {tool_name}root\n{tool_grammar}"
133
- # Emitted grammar shape:
134
- # root ::= tools
135
- # tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws
136
- # ws ::= (" " | "\t" | "\n")*
137
- # tool ::= <tool_a> | <tool_b> | ... (one alternative per input)
138
- # <tool_x> ::= <tool_x>root (per-tool xgrammar rules)
139
- # <tool_x>root ::= ... (from xgr.Grammar.from_json_schema)
140
- tool_alternatives = "tool ::= " + " | ".join(tool_dictionary.keys())
141
- tool_rules = "\n ".join(tool_dictionary.values())
142
- grammar = f"""root ::= tools
143
- tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws
144
- ws ::= (" " | "\\t" | "\\n")*
145
- {tool_alternatives}
146
- {tool_rules}
147
- """
148
- return grammar
149
-
150
-
151
-def _tool_definitions_to_schema_list(
152
- tools: str | list[Any],
153
-) -> list[CohereNormalizedTool]:
154
- """
155
- Build the list of ``CohereNormalizedTool`` dicts expected by
156
- ``collect_tool_schema``.
157
-
158
- Accepts:
159
- - JSON string
160
- - list of dicts with top-level ``name`` / ``parameters``
161
- - list of Chat Completions-style ``{"type": "function", "function": {...}}``
162
- - list of Pydantic models with ``model_dump()``
163
- """
164
- if isinstance(tools, str):
165
- try:
166
- parsed = json.loads(tools)
167
- except json.JSONDecodeError:
168
- return []
169
- if not isinstance(parsed, list):
170
- return []
171
- else:
172
- parsed = list(tools)
173
-
174
- out: list[CohereNormalizedTool] = []
175
- for raw in parsed:
176
- t = raw.model_dump() if hasattr(raw, "model_dump") else raw
177
- if not isinstance(t, dict):
178
- continue
179
- # Unwrap Chat Completions' ``{"type": "function", "function": {...}}``
180
- # shape; otherwise take the dict as-is.
181
- if t.get("type") == "function" and isinstance(t.get("function"), dict):
182
- t = t["function"]
183
- name = t.get("name")
184
- if not isinstance(name, str):
185
- continue
186
- params = t.get("parameters")
187
- out.append(
188
- CohereNormalizedTool(
189
- name=name,
190
- parameters=params if isinstance(params, dict) else {},
191
- )
192
- )
193
- return out
194
-
195
-
196
-def _has_effective_tools(
197
- tools: str | list[Any] | None,
198
-) -> TypeGuard[str | list[Any]]:
199
- """
200
- True when ``tools`` contains at least one tool definition to convert.
201
-
202
- ``ResponsesRequest`` defaults ``tools`` to ``[]``; ``ChatCompletionRequest``
203
- uses ``None``. Both mean "no tools" here. Strings (e.g. a JSON blob) are
204
- treated as effective only when non-blank.
205
- """
206
- if tools is None:
207
- return False
208
- if isinstance(tools, str):
209
- return bool(tools.strip())
210
- return len(tools) > 0
211
-
212
-
213
-# Builder: produces vLLM response_format in xgrammar's canonical format.
214
-# See xgrammar docs: type "structural_tag" with "format" = triggered_tags
215
-# and tag content type = json_schema | grammar.
216
-def convert_schema_to_structural_tags(
217
- schema: dict | None = None,
218
- tools: str | list[Any] | None = None,
219
- model_architecture: str | None = None,
220
-) -> str | None:
221
- """
222
- Returns a response_format string accepted by xgrammar's structural tag format.
223
- Uses the canonical shape: {"type": "structural_tag", "format": {...}} with
224
- format.type "triggered_tags" and tag content type "json_schema" or "grammar".
225
-
226
- Callers that are not on an engine path (e.g. the reasoning parser) must pass
227
- ``model_architecture`` explicitly.
228
- """
229
- if model_architecture is None or model_architecture not in MODEL_TO_TAG_STYLE:
230
- return None
231
- style = MODEL_TO_TAG_STYLE[model_architecture]
232
-
233
- tags: list[dict] = []
234
- triggers: list[str] = []
235
-
236
- def _add_tag(tag: CohereTagRegistry, content: dict) -> None:
237
- tags.append({"begin": tag.trigger, "content": content, "end": tag.end})
238
- triggers.append(tag.trigger)
239
-
240
- if schema is not None:
241
- # One structural tag per JSON wrapper (e.g. MOE: response + text).
242
- # Same for schema-only and "tools plus JSON mode" (North: schema when
243
- # the model does not call tools).
244
- for jt in style.json_tags:
245
- _add_tag(jt, {"type": "json_schema", "json_schema": schema})
246
-
247
- if _has_effective_tools(tools):
248
- # ``tools`` may be a JSON string (poseidon / RESPONSE_FORMAT_TOOL_DEFINITIONS)
249
- # or a list (Chat Completions ``request.tools`` as Pydantic models or dicts).
250
- tool_schema_list = _tool_definitions_to_schema_list(tools)
251
- if not tool_schema_list:
252
- raise ValueError(
253
- "No valid tool definitions could be parsed from the request for "
254
- "structural tag conversion."
255
- )
256
- tool_grammar = collect_tool_schema(tool_schema_list)
257
- _add_tag(style.tools, {"type": "grammar", "grammar": tool_grammar})
12
+ from vllm.tokenizers import TokenizerLike
13
259
- if not tags:
260
- return None
261
- return json.dumps(
262
- {
263
- "type": "structural_tag",
264
- "format": {
265
- "type": "triggered_tags",
266
- "triggers": triggers,
267
- "tags": tags,
268
- },
269
- }
270
- )
14
272
-
273
-def _response_format_type(
274
- response_format: AnyResponseFormat | dict | None,
275
-) -> str | None:
276
- if response_format is None:
277
- return None
278
- if isinstance(response_format, dict):
279
- t = response_format.get("type")
280
- return t if isinstance(t, str) else None
281
- return response_format.type
282
-
283
-
284
-def _maybe_parse_json_dict(value: Any) -> dict | None:
285
- """If value is a JSON string, parse to dict; otherwise require dict."""
286
- if isinstance(value, dict):
287
- return value
288
- if isinstance(value, str):
289
- try:
290
- parsed = json.loads(value)
291
- except (TypeError, json.JSONDecodeError):
292
- return None
293
- return parsed if isinstance(parsed, dict) else None
294
- return None
295
-
296
-
297
-def _unwrap_nested_schema(candidate: Any) -> dict | None:
298
- """Return ``candidate`` as a dict, unwrapping a nested ``schema`` if present.
299
-
300
- Returns ``None`` if ``candidate`` is not (and cannot be parsed into) a dict.
301
- """
302
- cand = _maybe_parse_json_dict(candidate)
303
- if not isinstance(cand, dict):
304
- return None
305
- nested = cand.get("schema")
306
- return nested if isinstance(nested, dict) else cand
307
-
308
-
309
-def _schema_from_json_schema_field(js_wr: Any) -> dict | None:
310
- """
311
- Extract the JSON Schema object from Chat Completions ``json_schema`` payload.
312
-
313
- Accepts:
314
- - ``JsonSchemaResponseFormat`` (Pydantic) with ``schema`` / ``json_schema`` field
315
- - dict in OpenAI shape ``{"name": ..., "schema": {...}}``
316
- - dict with ``json_schema`` key holding either the schema or a nested wrapper
317
- - dict that is already a JSON Schema document (some clients omit the wrapper)
318
- - JSON strings for any of the above
319
- """
320
- if js_wr is None:
321
- return None
322
-
323
- parsed_wr = _maybe_parse_json_dict(js_wr)
324
- if parsed_wr is not None:
325
- js_wr = parsed_wr
326
-
327
- if hasattr(js_wr, "model_dump"):
328
- for by_alias in (True, False):
329
- try:
330
- data = js_wr.model_dump(by_alias=by_alias, exclude_none=False)
331
- except TypeError:
332
- data = js_wr.model_dump(by_alias=by_alias)
333
- out = _unwrap_nested_schema(data.get("schema") or data.get("json_schema"))
334
- if out is not None:
335
- return out
336
- inner_attr = getattr(js_wr, "json_schema", None)
337
- return inner_attr if isinstance(inner_attr, dict) else None
338
-
339
- if isinstance(js_wr, dict):
340
- for key in ("schema", "json_schema"):
341
- out = _unwrap_nested_schema(js_wr.get(key))
342
- if out is not None:
343
- return out
344
- return js_wr
345
-
346
- return None
347
-
348
-
349
-def _schema_dict_from_chat_response_format(
350
- rf: AnyResponseFormat | dict | None,
351
-) -> dict | None:
352
- """JSON schema dict from Chat Completions ``request.response_format`` only."""
353
- if rf is None:
354
- return None
355
- rf_type = _response_format_type(rf)
356
- if rf_type == "json_object":
357
- return {"type": "object"}
358
- if rf_type != "json_schema":
359
- return None
360
- js_wr = (
361
- rf.get("json_schema")
362
- if isinstance(rf, dict)
363
- else getattr(rf, "json_schema", None)
364
- )
365
- return _schema_from_json_schema_field(js_wr)
366
-
367
-
368
-def _schema_dict_from_structured_outputs(
369
- so: StructuredOutputsParams | None,
370
-) -> dict | None:
371
- """Schema dict from ``structured_outputs`` (``json`` / ``json_object``).
372
-
373
- Same unwrapping as ``json_schema``. ``json`` is expected to be ``str`` or
374
- ``dict`` (enforced by ``StructuredOutputsParams`` / request models); other
375
- types raise ``ValueError`` only if a caller bypasses that validation.
376
- """
377
- if so is None:
378
- return None
379
- if so.json_object:
380
- return {"type": "object"}
381
- raw: Any = so.json
382
- if raw is None:
383
- return None
384
-
385
- if hasattr(raw, "model_dump"):
386
- out = _schema_from_json_schema_field(raw)
387
- if out is None:
388
- raise ValueError(
389
- "structured_outputs.json model has no extractable JSON Schema."
390
- )
391
- return out
392
-
393
- if isinstance(raw, str):
394
- if not raw.strip():
395
- raise ValueError("structured_outputs.json cannot be empty.")
396
- try:
397
- raw = json.loads(raw)
398
- except json.JSONDecodeError as e:
399
- raise ValueError("structured_outputs.json must be valid JSON.") from e
400
- if not isinstance(raw, dict):
401
- raise ValueError("structured_outputs.json must decode to a JSON object.")
402
-
403
- if isinstance(raw, Mapping):
404
- body = raw if isinstance(raw, dict) else dict(raw)
405
- return _schema_from_json_schema_field(body) or body
406
-
407
- raise ValueError(
408
- f"structured_outputs.json has unsupported type {type(raw).__name__}."
409
- )
410
-
411
-
412
-def _melody_sources_to_vllm(
413
- raw_sources: Any,
414
- position_to_source: Mapping[tuple[int, int], CitationSource] | None,
415
-) -> list[CitationSource]:
416
- """Convert melody's ``Source`` objects into resolved :class:`CitationSource`.
417
-
418
- Melody's ``Source`` shape is
419
- ``{tool_call_index, tool_result_indices, document_ids}``. Without
420
- :meth:`PyFilterOptions.with_message_history` the ``document_ids``
421
- list is empty, so we resolve the numeric address ourselves using
422
- ``position_to_source``, which is
423
- :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._build_position_to_source`
424
- applied to the inbound request and forwarded through
425
- ``chat_template_kwargs`` at parser-construction time. Each melody
426
- source with ``tool_result_indices=[i, j, ...]`` fans out to one
427
- :class:`CitationSource` per resolved position, populating the
428
- wire-shape ``type`` / ``id`` / ``document`` / ``tool_output``
429
- fields; unresolved positions are silently skipped (the calling
430
- citation is dropped at wire-coercion time if every source in it
431
- ends up unresolved -- see ``_to_wire_citation`` in
432
- :mod:`vllm.entrypoints.cohere.serving`).
433
-
434
- When ``position_to_source`` is ``None`` (parser wired outside of
435
- :class:`CohereServingChatV2`, e.g. a plain OpenAI chat completion
436
- request that happens to hit a cohere reasoning parser) sources
437
- can't be resolved and get dropped; citations then never reach any
438
- citation-aware handler downstream.
439
- """
440
- if not position_to_source:
441
- return []
442
- out: list[CitationSource] = []
443
- for s in raw_sources or []:
444
- bucket = getattr(s, "tool_call_index", None)
445
- indices = list(getattr(s, "tool_result_indices", None) or [])
446
- if bucket is None or not indices:
447
- continue
448
- for idx in indices:
449
- info = position_to_source.get((bucket, idx))
450
- if info is not None:
451
- out.append(info)
452
- return out
453
-
454
-
455
-def _melody_citations_to_vllm(
456
- raw_citations: Any,
457
- position_to_source: Mapping[tuple[int, int], CitationSource] | None,
458
-) -> list[Citation] | None:
459
- """Convert melody's ``FilterCitation`` objects into :class:`Citation`.
460
-
461
- Resolves each source in-place via ``position_to_source`` (see
462
- :func:`_melody_sources_to_vllm`). Citations whose sources all fail
463
- to resolve are still emitted here with ``sources=[]``; the serving
464
- layer drops them at wire-coercion time so the fail-closed policy
465
- lives in exactly one place.
15
+class BaseCohereCommandReasoningParser(ReasoningParser):
16
+ """Expose is_reasoning_end for the structured-output engine.
17
+ Everything else runs in
18
+ :class:`vllm.parser.cohere_command.CohereCommandParser`.
19
"""
467
- if not raw_citations:
468
- return None
469
- out: list[Citation] = []
470
- for c in raw_citations:
471
- out.append(
472
- Citation(
473
- start=getattr(c, "start_index", None),
474
- end=getattr(c, "end_index", None),
475
- text=getattr(c, "text", None),
476
- sources=_melody_sources_to_vllm(
477
- getattr(c, "sources", None), position_to_source
478
- ),
479
- type=(
480
- "THINKING_CONTENT"
481
- if getattr(c, "is_thinking", False)
482
- else "TEXT_CONTENT"
483
- ),
484
- )
485
- )
486
- return out
20
21
+ melody_preset: ClassVar[str]
22
489
-class BaseCohereCommandReasoningParser(ReasoningParser):
490
- def __init__(
491
- self,
492
- tokenizer: TokenizerLike,
493
- *args,
494
- streaming_opts: PyFilterOptions,
495
- unary_opts: PyFilterOptions,
496
- **kwargs,
497
- ):
23
+ def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
24
super().__init__(tokenizer, *args, **kwargs)
25
self.start_token_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>")
26
self.end_token_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>")
27
self.chatbot_token_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>")
502
- self.unary_opts = unary_opts
503
- self.melody_unary = PyFilter(unary_opts)
504
- self.melody_streaming = PyFilter(streaming_opts)
505
- # Citations extracted by the most recent ``extract_reasoning`` call.
506
- # Citation-aware handlers (e.g. :class:`CohereServingChatV2` via its
507
- # ``_finalize_response_message`` hook) read this back from the
508
- # parser instance (which is constructed per-request) and attach
509
- # the result to :class:`CohereChatMessage.citations`. ``None`` when
510
- # the last parse produced no citations.
511
- self.last_unary_citations: list[Citation] | None = None
512
- # Request-scoped ``(tool_call_index, tool_result_idx) -> resolved
513
- # CitationSource`` map, forwarded from
514
- # :meth:`CohereServingChatV2._apply_cohere_template_kwargs` via
515
- # ``chat_template_kwargs``. See :func:`_melody_sources_to_vllm`
516
- # for how it's consumed. Absent for parser instances that
517
- # weren't created by ``CohereServingChatV2`` (e.g. an OpenAI
518
- # chat completion request routed through a cohere reasoning
519
- # parser); citations from those requests are dropped since we
520
- # have no way to attribute their sources.
521
- ctk = kwargs.get("chat_template_kwargs") or {}
522
- # Lazy import: ``renderers.cohere`` transitively imports
523
- # ``cohere_melody`` and other heavy deps; hoisting the constant
524
- # to module scope would tie parser importability to that graph.
525
- from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY
526
-
527
- raw_map = ctk.get(POSITION_TO_SOURCE_KEY)
528
- self._position_to_source: Mapping[tuple[int, int], CitationSource] | None = (
529
- raw_map if isinstance(raw_map, Mapping) else None
530
- )
28
29
@property
30
def reasoning_start_str(self) -> str | None:
@@ -537,80 +34,6 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
34
def reasoning_end_str(self) -> str | None:
35
return "<|END_THINKING|>"
36
540
- def extract_reasoning_streaming(
541
- self,
542
- previous_text: str,
543
- current_text: str,
544
- delta_text: str,
545
- previous_token_ids: Sequence[int],
546
- current_token_ids: Sequence[int],
547
- delta_token_ids: Sequence[int],
548
- ) -> DeltaMessage | None:
549
- r = self.melody_streaming.write_decoded(delta_text)
550
- citations = _melody_citations_to_vllm(
551
- getattr(r, "citations", None), self._position_to_source
552
- )
553
- if (
554
- r.content is None
555
- and r.reasoning is None
556
- and not r.tool_calls
557
- and not citations
558
- ):
559
- return None
560
- # Always emit CohereDeltaMessage so citations reach the wire via
561
- # ``SerializeAsAny[DeltaMessage]`` on the streaming choice envelope.
562
- # When ``citations`` is unset, ``CohereDeltaMessage._serialize``
563
- # drops the field so the emitted shape matches plain
564
- # :class:`DeltaMessage`
565
- msg = CohereDeltaMessage()
566
- if r.content is not None:
567
- msg.content = r.content
568
- if r.reasoning is not None:
569
- msg.reasoning = r.reasoning
570
- if r.tool_calls:
571
- msg.tool_calls = [
572
- DeltaToolCall(
573
- id=tc.id,
574
- index=tc.index,
575
- type="function",
576
- function=DeltaFunctionCall(name=tc.name, arguments=tc.arguments),
577
- )
578
- for tc in r.tool_calls
579
- ]
580
- if citations:
581
- msg.citations = citations
582
- return msg
583
-
584
- def extract_reasoning(
585
- self, model_output: str, request: ChatCompletionRequest | ResponsesRequest
586
- ) -> tuple[str | None, str | None]:
587
- result = self.melody_unary.process_full_text(model_output)
588
- # Cache citations so citation-aware handlers can surface them
589
- # on :class:`CohereChatMessage.citations`. The base
590
- # :meth:`ReasoningParser.extract_reasoning` contract only
591
- # returns ``(reasoning, content)``, so citations are passed
592
- # back via parser-instance state -- safe because the parser is
593
- # constructed per-request.
594
- self.last_unary_citations = _melody_citations_to_vllm(
595
- getattr(result, "citations", None), self._position_to_source
596
- )
597
- return result.reasoning, result.content
598
-
599
- def extract_content_ids(self, input_ids: list[int]) -> list[int]:
600
- token_buf: list[int] = []
601
- content_ids: list[int] = []
602
- content_filter = PyFilter(self.unary_opts)
603
- for t in input_ids:
604
- token_buf.append(t)
605
- s = self.model_tokenizer.decode(token_buf, skip_special_tokens=False)
606
- if s.endswith(REPLACEMENT_CHAR):
607
- continue
608
- r = content_filter.write_decoded(s)
609
- if r.content is not None:
610
- content_ids.extend(token_buf)
611
- token_buf = []
612
- return content_ids
613
-
37
def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
38
chatbot = self.chatbot_token_id
39
start = self.start_token_id
@@ -628,91 +51,10 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
51
52
return has_end_token
53
631
- def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
632
- count = 0
633
- depth = 0
634
- for token_id in token_ids:
635
- if token_id == self.start_token_id:
636
- depth += 1
637
- continue
638
- if token_id == self.end_token_id:
639
- if depth > 0:
640
- depth -= 1
641
- continue
642
- if depth > 0:
643
- count += 1
644
- return count
645
-
646
- def adjust_request(
647
- self, request: ChatCompletionRequest | ResponsesRequest
648
- ) -> ChatCompletionRequest | ResponsesRequest:
649
- so = request.structured_outputs
650
- if so is not None and so.structural_tag:
651
- return request
652
- # Schema: prefer ``response_format`` (OpenAI Chat Completions), then
653
- # ``structured_outputs.json`` / ``json_object`` (vLLM direct). Tools stay
654
- # on ``request.tools``.
655
- rf = (
656
- request.response_format
657
- if isinstance(request, ChatCompletionRequest)
658
- else None
659
- )
660
- if rf is not None and _response_format_type(rf) == "structural_tag":
661
- return request
662
- model_architecture = (
663
- self._model_config.architecture if self._model_config is not None else None
664
- )
665
- tools = request.tools
666
- # ``response_format`` wins if both it and ``structured_outputs`` supply JSON.
667
- schema = _schema_dict_from_chat_response_format(rf)
668
- if schema is None:
669
- schema = _schema_dict_from_structured_outputs(so)
670
- if schema is None and not _has_effective_tools(tools):
671
- return request
672
- if model_architecture is None:
673
- return request
674
- result = convert_schema_to_structural_tags(
675
- schema=schema,
676
- tools=tools,
677
- model_architecture=model_architecture,
678
- )
679
- if result is None:
680
- # Unsupported architectures are not in ``MODEL_TO_TAG_STYLE``.
681
- raise ValueError(
682
- "Failed to build structural_tag guided decoding constraints from "
683
- "this request's JSON schema and/or tools. The configured model "
684
- f"architecture ({model_architecture!r}) does not support Cohere "
685
- "command structural tags, or the schema cannot be expressed in "
686
- "that format."
687
- )
688
- request.structured_outputs = StructuredOutputsParams(structural_tag=result)
689
- # Folded JSON constraints into ``structural_tag``; drop ``response_format``
690
- # when it was the source so ``to_sampling_params`` does not also set ``json`` /
691
- # ``json_object`` (mutually exclusive in ``StructuredOutputsParams``).
692
- if isinstance(request, ChatCompletionRequest) and rf is not None:
693
- rf_type = _response_format_type(rf)
694
- if rf_type in ("json_schema", "json_object"):
695
- request.response_format = None
696
- return request
697
-
54
55
class CohereCommand3ReasoningParser(BaseCohereCommandReasoningParser):
700
- def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
701
- super().__init__(
702
- tokenizer,
703
- *args,
704
- streaming_opts=PyFilterOptions().cmd3().no_tools(),
705
- unary_opts=PyFilterOptions().cmd3().no_tools(),
706
- **kwargs,
707
- )
56
+ melody_preset = "cmd3"
57
58
59
class CohereCommand4ReasoningParser(BaseCohereCommandReasoningParser):
711
- def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
712
- super().__init__(
713
- tokenizer,
714
- *args,
715
- streaming_opts=PyFilterOptions().cmd4().no_tools(),
716
- unary_opts=PyFilterOptions().cmd4().no_tools(),
717
- **kwargs,
718
- )
60
+ melody_preset = "cmd4"
vllm/renderers/cohere.py
+1
-1
@@ -244,7 +244,7 @@ def _try_import_melody():
244
raise ImportError(
245
"The `cohere` tokenizer/renderer mode requires the "
246
"`cohere_melody` package. Install it via "
247
- "`pip install 'cohere-melody>=0.11.1'` or build from "
247
+ "`pip install 'cohere-melody>=0.14.0'` or build from "
248
"https://github.com/cohere-ai/melody."
249
) from e
250
vllm/tool_parsers/cohere_command_tool_parser.py
+4
-134
@@ -1,148 +1,18 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Sequence
4
+from typing import ClassVar
5
6
-try:
7
- from cohere_melody import PyFilter, PyFilterOptions
8
-except ImportError as e:
9
- raise ImportError(
10
- "The Cohere tool parser requires the `cohere_melody` "
11
- "package, which is not installed. Install it with:\n"
12
- " pip install 'cohere-melody>=0.11.1'"
13
- ) from e
14
-
15
-from vllm.entrypoints.generate.base.protocol import (
16
- DeltaFunctionCall,
17
- DeltaMessage,
18
- DeltaToolCall,
19
- ExtractedToolCallInformation,
20
- FunctionCall,
21
- ToolCall,
22
-)
23
-from vllm.entrypoints.openai.chat_completion.protocol import (
24
- ChatCompletionRequest,
25
-)
26
-from vllm.entrypoints.openai.responses.protocol import (
27
- ResponsesRequest,
28
-)
29
-from vllm.tokenizers import TokenizerLike
6
from vllm.tool_parsers import ToolParser
31
-from vllm.tool_parsers.utils import Tool
7
8
9
class BaseCohereCommandToolParser(ToolParser):
35
- def __init__(
36
- self,
37
- tokenizer: TokenizerLike,
38
- streaming_opts: PyFilterOptions,
39
- unary_opts: PyFilterOptions,
40
- ):
41
- super().__init__(tokenizer)
42
- self.melody_streaming = PyFilter(streaming_opts)
43
- self.melody_unary = PyFilter(unary_opts)
44
- # Melody can emit the tool-call id before the function name. Keep it
45
- # until the first real name delta so clients receive both together.
46
- self._pending_streaming_tool_call_ids: dict[int, str] = {}
47
-
48
- def adjust_request(
49
- self, request: ChatCompletionRequest | ResponsesRequest
50
- ) -> ChatCompletionRequest | ResponsesRequest:
51
- request = super().adjust_request(request)
52
- request.skip_special_tokens = False
53
- return request
54
-
55
- def extract_tool_calls_streaming(
56
- self,
57
- previous_text: str,
58
- current_text: str,
59
- delta_text: str,
60
- previous_token_ids: Sequence[int],
61
- current_token_ids: Sequence[int],
62
- delta_token_ids: Sequence[int],
63
- request: ChatCompletionRequest,
64
- ) -> DeltaMessage | None:
65
- r = self.melody_streaming.write_decoded(delta_text)
66
- if r.content is not None:
67
- return DeltaMessage(content=r.content)
68
- if r.reasoning is not None:
69
- return DeltaMessage(reasoning=r.reasoning)
70
- if r.tool_calls:
71
- tool_calls: list[DeltaToolCall] = []
72
- for tc in r.tool_calls:
73
- if tc.id:
74
- self._pending_streaming_tool_call_ids[tc.index] = tc.id
75
- name = tc.name or None
76
- arguments = tc.arguments or None
77
- # Empty strings are placeholders in Melody's streaming output;
78
- # omit them from OpenAI-compatible deltas instead of sending
79
- # invalid tool names or empty argument fragments.
80
- if name is None and arguments is None:
81
- continue
82
-
83
- function_kwargs = {}
84
- if name is not None:
85
- function_kwargs["name"] = name
86
- if arguments is not None:
87
- function_kwargs["arguments"] = arguments
88
- tool_call_kwargs = {
89
- "index": tc.index,
90
- "function": DeltaFunctionCall(**function_kwargs),
91
- }
92
- if name is not None:
93
- tool_call_id = tc.id or self._pending_streaming_tool_call_ids.pop(
94
- tc.index, None
95
- )
96
- if tool_call_id is not None:
97
- tool_call_kwargs["id"] = tool_call_id
98
- tool_call_kwargs["type"] = "function"
99
- tool_calls.append(DeltaToolCall(**tool_call_kwargs))
100
- if tool_calls:
101
- return DeltaMessage(tool_calls=tool_calls)
102
- return None
103
-
104
- def extract_tool_calls(
105
- self,
106
- model_output: str,
107
- request: ChatCompletionRequest,
108
- ) -> ExtractedToolCallInformation:
109
- result = self.melody_unary.process_full_text(model_output)
110
- tool_calls = [
111
- ToolCall(
112
- id=tc.id,
113
- type="function",
114
- function=FunctionCall(name=tc.name, arguments=tc.arguments),
115
- )
116
- for tc in result.tool_calls
117
- ]
118
- return ExtractedToolCallInformation(
119
- tools_called=len(tool_calls) > 0,
120
- tool_calls=tool_calls,
121
- content=result.content,
122
- )
10
+ melody_preset: ClassVar[str]
11
12
13
class CohereCommand3ToolParser(BaseCohereCommandToolParser):
126
- def __init__(
127
- self,
128
- tokenizer: TokenizerLike,
129
- tools: list[Tool] | None = None,
130
- ):
131
- super().__init__(
132
- tokenizer,
133
- streaming_opts=PyFilterOptions().cmd3().start_in_answer(),
134
- unary_opts=PyFilterOptions().cmd3().start_in_answer(),
135
- )
14
+ melody_preset = "cmd3"
15
16
17
class CohereCommand4ToolParser(BaseCohereCommandToolParser):
139
- def __init__(
140
- self,
141
- tokenizer: TokenizerLike,
142
- tools: list[Tool] | None = None,
143
- ):
144
- super().__init__(
145
- tokenizer,
146
- streaming_opts=PyFilterOptions().cmd4().start_in_answer(),
147
- unary_opts=PyFilterOptions().cmd4().start_in_answer(),
148
- )
18
+ melody_preset = "cmd4"