@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d706dec904

fix: Correct reasoning-end detection for prompt history (#44551)

Signed-off-by: jwzheng96 <jianweizheng@pku.edu.cn> Signed-off-by: JianweiZheng <32029023+jwzheng96@users.noreply.github.com> Signed-off-by: Jason Ozuzu <jasonozuzu@cohere.com> Signed-off-by: walterbm <walter.beller.morales@gmail.com> Co-authored-by: JianweiZheng <32029023+jwzheng96@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: walterbm <walter.beller.morales@gmail.com> Co-authored-by: Walter Beller-Morales <walterbm@users.noreply.github.com> Co-authored-by: Flora Feng <4florafeng@gmail.com>

JasonCohere committed Jun 27, 2026 at 06:15 UTC d706dec904e89e4067efd9535b96705aa61f3935
4 files changed +645 -1
requirements/test/cuda.in
+1
@@ -73,6 +73,7 @@ gpt-oss >= 0.0.7; python_version > '3.11'
73
74 perceptron # required for isaac test
75 kaldi-native-fbank >= 1.18.7 # required for fireredasr2 test
76 +cohere_melody>=0.9.0 # required for cohere command reasoning parser test
77
78 # Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library.
79 # Older versions are in conflict with teerratorch requirements.
requirements/test/cuda.txt
+2
@@ -101,6 +101,8 @@ click==8.1.7
101 # schemathesis
102 # typer
103 # uvicorn
104 +cohere-melody==0.9.0
105 + # via -r requirements/test/cuda.in
106 colorama==0.4.6
107 # via
108 # perceptron
tests/reasoning/test_cohere_command_reasoning_parser.py new
+625
@@ -0,0 +1,625 @@
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, field
9 +from types import SimpleNamespace
10 +from typing import Any
11 +
12 +import pytest
13 +from pydantic import ValidationError
14 +
15 +from vllm.entrypoints.openai.chat_completion.protocol import (
16 + ChatCompletionRequest,
17 +)
18 +from vllm.entrypoints.openai.engine.protocol import (
19 + JsonSchemaResponseFormat,
20 + ResponseFormat,
21 + StructuralTagResponseFormat,
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 + _response_format_type,
29 + _schema_dict_from_structured_outputs,
30 + convert_schema_to_structural_tags,
31 +)
32 +from vllm.sampling_params import StructuredOutputsParams
33 +
34 +
35 +@dataclass
36 +class ExpectedToolCall:
37 + id: str
38 + name: str
39 + arguments: dict
40 +
41 +
42 +@dataclass
43 +class ReasoningCase:
44 + parser_cls: Any
45 + model_output: str
46 + expected_reasoning: str | None
47 + expected_content: str | None
48 + expected_tool_calls: list[ExpectedToolCall] = field(default_factory=list)
49 +
50 +
51 +REASONING_CASES = [
52 + pytest.param(
53 + ReasoningCase(
54 + parser_cls=CohereCommand3ReasoningParser,
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_reasoning="i will call foo with query1",
62 + expected_content="""\
63 +<|START_ACTION|>
64 +[
65 + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
66 +]
67 +<|END_ACTION|>""",
68 + expected_tool_calls=[
69 + ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}),
70 + ],
71 + ),
72 + id="cmd3-single_tool_call",
73 + ),
74 + pytest.param(
75 + ReasoningCase(
76 + parser_cls=CohereCommand4ReasoningParser,
77 + model_output="""\
78 +<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
79 +[
80 + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
81 +]
82 +<|END_ACTION|>""",
83 + expected_reasoning="i will call foo with query1",
84 + expected_content="""\
85 +<|START_ACTION|>
86 +[
87 + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
88 +]
89 +<|END_ACTION|>""",
90 + expected_tool_calls=[
91 + ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}),
92 + ],
93 + ),
94 + id="cmd4-single_tool_call",
95 + ),
96 + pytest.param(
97 + ReasoningCase(
98 + parser_cls=CohereCommand3ReasoningParser,
99 + model_output="""\
100 +<|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
101 +<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_RESPONSE|>""",
102 + expected_reasoning="This is a rainbow emoji: 🌈",
103 + expected_content="foo bar",
104 + ),
105 + id="cmd3-citations_with_emoji",
106 + ),
107 + pytest.param(
108 + ReasoningCase(
109 + parser_cls=CohereCommand4ReasoningParser,
110 + model_output="""\
111 +<|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
112 +<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_RESPONSE|>""",
113 + expected_reasoning="This is a rainbow emoji: 🌈",
114 + expected_content="foo bar",
115 + ),
116 + id="cmd4-citations_with_emoji",
117 + ),
118 +]
119 +
120 +
121 +class MockCohereTokenizer:
122 + """Minimal byte-level stand-in for the Cohere tokenizer.
123 +
124 + ``encode``/``decode`` round-trip through UTF-8 bytes so splitting a
125 + multi-byte character (e.g. an emoji) across "tokens" reproduces the
126 + trailing U+FFFD buffering that real streaming exhibits. Cohere special
127 + tokens map to distinct synthetic ids; everything else shares a default id.
128 + ``adjust_request`` only needs the token ids, not real tokenization.
129 + """
130 +
131 + _SPECIAL_TOKEN_IDS = {
132 + "<|START_THINKING|>": -1,
133 + "<|END_THINKING|>": -2,
134 + "<|CHATBOT_TOKEN|>": -3,
135 + }
136 +
137 + def convert_tokens_to_ids(self, token: str) -> int:
138 + return self._SPECIAL_TOKEN_IDS.get(token, 0)
139 +
140 + def get_vocab(self) -> dict[str, int]:
141 + return {}
142 +
143 + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
144 + return list(text.encode("utf-8"))
145 +
146 + def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str:
147 + return bytes(ids).decode("utf-8", errors="replace")
148 +
149 +
150 +@pytest.fixture(scope="module")
151 +def tokenizer() -> MockCohereTokenizer:
152 + return MockCohereTokenizer()
153 +
154 +
155 +@pytest.fixture
156 +def request_obj():
157 + return ChatCompletionRequest(messages=[], model="test-model")
158 +
159 +
160 +REPLACEMENT_CHAR = "\ufffd"
161 +
162 +
163 +def _token_deltas(tokenizer, text: str) -> list[str]:
164 + """Progressively decode the token sequence and return per-step string
165 + deltas. Incomplete multi-byte sequences (trailing U+FFFD) are buffered
166 + until the next token completes them, matching real streaming behaviour."""
167 + ids = tokenizer.encode(text, add_special_tokens=False)
168 + deltas: list[str] = []
169 + prev = ""
170 + for i in range(1, len(ids) + 1):
171 + current = tokenizer.decode(ids[:i], skip_special_tokens=False)
172 + if current.endswith(REPLACEMENT_CHAR):
173 + continue
174 + delta = current[len(prev) :]
175 + if delta:
176 + deltas.append(delta)
177 + prev = current
178 + return deltas
179 +
180 +
181 +@pytest.mark.parametrize("case", REASONING_CASES)
182 +class TestExtractReasoning:
183 + def test_nonstreaming(self, tokenizer, request_obj, case: ReasoningCase):
184 + parser = case.parser_cls(tokenizer)
185 + reasoning, content = parser.extract_reasoning(case.model_output, request_obj)
186 +
187 + assert reasoning == case.expected_reasoning
188 + assert content == case.expected_content
189 +
190 + def test_streaming(self, tokenizer, case: ReasoningCase):
191 + parser = case.parser_cls(tokenizer)
192 + token_strings = _token_deltas(tokenizer, case.model_output)
193 +
194 + reasoning_parts: list[str] = []
195 + content_parts: list[str] = []
196 + tool_call_deltas: list[dict] = []
197 +
198 + previous_text = ""
199 + previous_token_ids: list[int] = []
200 +
201 + for token_str in token_strings:
202 + current_text = previous_text + token_str
203 + current_token_ids = previous_token_ids + [0]
204 +
205 + delta = parser.extract_reasoning_streaming(
206 + previous_text=previous_text,
207 + current_text=current_text,
208 + delta_text=token_str,
209 + previous_token_ids=previous_token_ids,
210 + current_token_ids=current_token_ids,
211 + delta_token_ids=[0],
212 + )
213 + if delta is not None:
214 + if delta.reasoning is not None:
215 + reasoning_parts.append(delta.reasoning)
216 + if delta.content is not None:
217 + content_parts.append(delta.content)
218 + for tc in delta.tool_calls:
219 + tool_call_deltas.append(
220 + {
221 + "id": tc.id,
222 + "index": tc.index,
223 + "name": tc.function.name if tc.function else None,
224 + "arguments": (
225 + tc.function.arguments if tc.function else None
226 + ),
227 + }
228 + )
229 +
230 + previous_text = current_text
231 + previous_token_ids = current_token_ids
232 +
233 + reasoning = "".join(reasoning_parts) if reasoning_parts else None
234 + assert reasoning == case.expected_reasoning
235 +
236 + content = "".join(content_parts) if content_parts else None
237 + if case.expected_tool_calls:
238 + assert content is None or content == ""
239 + else:
240 + assert content == case.expected_content
241 +
242 + accumulated: dict[int, dict] = {}
243 + for d in tool_call_deltas:
244 + idx = d["index"]
245 + if idx not in accumulated:
246 + accumulated[idx] = {"id": "", "name": "", "arguments": ""}
247 + if d["id"]:
248 + accumulated[idx]["id"] = d["id"]
249 + if d["name"]:
250 + accumulated[idx]["name"] = d["name"]
251 + if d["arguments"]:
252 + accumulated[idx]["arguments"] += d["arguments"]
253 +
254 + assert len(accumulated) == len(case.expected_tool_calls)
255 + for i, expected_tc in enumerate(case.expected_tool_calls):
256 + tc = accumulated[i]
257 + assert tc["id"] == expected_tc.id
258 + assert tc["name"] == expected_tc.name
259 + assert json.loads(tc["arguments"]) == expected_tc.arguments
260 +
261 +
262 +class TestIsReasoningEnd:
263 + @pytest.mark.parametrize(
264 + "parser_cls",
265 + [CohereCommand3ReasoningParser, CohereCommand4ReasoningParser],
266 + ids=["cmd3", "cmd4"],
267 + )
268 + def test_is_reasoning_end(self, tokenizer, parser_cls):
269 + parser = parser_cls(tokenizer)
270 + start_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>")
271 + end_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>")
272 + chatbot_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>")
273 + content_ids = [99, 100]
274 +
275 + # Generation-only tokens have no chatbot marker, so the whole sequence
276 + # is considered.
277 + assert parser.is_reasoning_end([end_id])
278 + assert parser.is_reasoning_end([start_id, *content_ids, end_id])
279 + assert not parser.is_reasoning_end([start_id, *content_ids])
280 +
281 + # Full prompt/history tokens are scoped to the latest chatbot marker,
282 + # so stray thinking tokens from the preamble or previous turns are ignored.
283 + assert not parser.is_reasoning_end([start_id, end_id, chatbot_id, *content_ids])
284 + assert parser.is_reasoning_end(
285 + [start_id, end_id, chatbot_id, start_id, *content_ids, end_id]
286 + )
287 +
288 +
289 +SCHEMA_A = {"type": "object", "properties": {"a": {"type": "string"}}}
290 +SCHEMA_B = {"type": "object", "properties": {"b": {"type": "number"}}}
291 +GET_WEATHER_TOOL = {
292 + "type": "function",
293 + "function": {
294 + "name": "get_weather",
295 + "parameters": {
296 + "type": "object",
297 + "properties": {"city": {"type": "string"}},
298 + },
299 + },
300 +}
301 +VALID_STRUCTURAL_TAG = {
302 + "type": "structural_tag",
303 + "format": {
304 + "type": "triggered_tags",
305 + "tags": [
306 + {
307 + "begin": "<tool>",
308 + "content": {"type": "any_text"},
309 + "end": "</tool>",
310 + }
311 + ],
312 + "triggers": ["<tool>"],
313 + },
314 +}
315 +
316 +
317 +def _model_config(arch: str) -> SimpleNamespace:
318 + return SimpleNamespace(
319 + architecture=arch,
320 + architectures=[arch],
321 + hf_text_config=SimpleNamespace(architectures=[arch]),
322 + )
323 +
324 +
325 +def _make_chat_request(**kwargs) -> ChatCompletionRequest:
326 + data = {"messages": [{"role": "user", "content": "hi"}], "model": "m"}
327 + data.update(kwargs)
328 + return ChatCompletionRequest.model_validate(data)
329 +
330 +
331 +def _first_json_schema(tag_json: str) -> dict | None:
332 + outer = json.loads(tag_json)
333 + for t in (outer.get("format") or {}).get("tags") or []:
334 + c = t.get("content") or {}
335 + if c.get("type") == "json_schema":
336 + js = c.get("json_schema")
337 + return js if isinstance(js, dict) else None
338 + return None
339 +
340 +
341 +def _content_types(tag_json: str) -> set[str]:
342 + outer = json.loads(tag_json)
343 + out: set[str] = set()
344 + for t in (outer.get("format") or {}).get("tags") or []:
345 + ty = (t.get("content") or {}).get("type")
346 + if isinstance(ty, str):
347 + out.add(ty)
348 + return out
349 +
350 +
351 +@pytest.fixture(scope="module")
352 +def parser(tokenizer: MockCohereTokenizer) -> CohereCommand4ReasoningParser:
353 + """Parser configured with a supported Cohere architecture."""
354 + return CohereCommand4ReasoningParser(
355 + tokenizer,
356 + model_config=_model_config("Cohere2ForCausalLM"),
357 + )
358 +
359 +
360 +@pytest.fixture(scope="module")
361 +def parser_no_model_config(
362 + tokenizer: MockCohereTokenizer,
363 +) -> CohereCommand4ReasoningParser:
364 + """Parser with no ``model_config`` (cannot resolve architecture)."""
365 + return CohereCommand4ReasoningParser(tokenizer, model_config=None)
366 +
367 +
368 +@pytest.fixture(scope="module")
369 +def parser_unsupported_arch(
370 + tokenizer: MockCohereTokenizer,
371 +) -> CohereCommand4ReasoningParser:
372 + """Parser configured with an architecture that has no structural tag style."""
373 + return CohereCommand4ReasoningParser(
374 + tokenizer,
375 + model_config=_model_config("LlamaForCausalLM"),
376 + )
377 +
378 +
379 +class TestAdjustRequestPassthrough:
380 + def test_structured_outputs_structural_tag_not_modified(self, parser) -> None:
381 + tag = json.dumps(VALID_STRUCTURAL_TAG)
382 + r = _make_chat_request(structured_outputs={"structural_tag": tag})
383 + o = parser.adjust_request(r)
384 + assert o.structured_outputs.structural_tag == tag
385 +
386 + def test_response_format_structural_tag_short_circuit(self, parser) -> None:
387 + # ``ChatCompletionRequest`` validates ``response_format`` as a union;
388 + # bare ``{"type": "structural_tag"}`` is invalid (use pydantic model).
389 + rf = StructuralTagResponseFormat(
390 + type="structural_tag",
391 + format=VALID_STRUCTURAL_TAG["format"],
392 + )
393 + r = _make_chat_request(response_format=rf)
394 + o = parser.adjust_request(r)
395 + assert _response_format_type(o.response_format) == "structural_tag"
396 + assert o.structured_outputs is None
397 +
398 +
399 +class TestAdjustRequestNoOp:
400 + def test_no_schema_no_tools(self, parser) -> None:
401 + o = parser.adjust_request(_make_chat_request())
402 + assert o.structured_outputs is None
403 + assert o.response_format is None
404 +
405 + def test_no_model_config(self, parser_no_model_config) -> None:
406 + inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A)
407 + r = _make_chat_request(
408 + response_format=ResponseFormat(type="json_schema", json_schema=inner),
409 + )
410 + o = parser_no_model_config.adjust_request(r)
411 + assert o.response_format is not None
412 + assert o.structured_outputs is None
413 +
414 +
415 +class TestAdjustRequestUnsupportedArchitecture:
416 + def test_json_schema_raises(self, parser_unsupported_arch) -> None:
417 + inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A)
418 + r = _make_chat_request(
419 + response_format=ResponseFormat(type="json_schema", json_schema=inner),
420 + )
421 + with pytest.raises(ValueError, match="does not support"):
422 + parser_unsupported_arch.adjust_request(r)
423 +
424 +
425 +class TestAdjustRequestFoldFromResponseFormat:
426 + @pytest.mark.parametrize(
427 + "response_format, expected_schema",
428 + [
429 + pytest.param(
430 + ResponseFormat(
431 + type="json_schema",
432 + json_schema=JsonSchemaResponseFormat(
433 + name="n", json_schema=SCHEMA_A
434 + ),
435 + ),
436 + SCHEMA_A,
437 + id="json_schema_pydantic",
438 + ),
439 + pytest.param(
440 + {
441 + "type": "json_schema",
442 + "json_schema": {"name": "n", "schema": SCHEMA_A},
443 + },
444 + SCHEMA_A,
445 + id="json_schema_dict",
446 + ),
447 + pytest.param(
448 + {"type": "json_object"},
449 + {"type": "object"},
450 + id="json_object",
451 + ),
452 + ],
453 + )
454 + def test_response_format_cleared(
455 + self, parser, response_format, expected_schema
456 + ) -> None:
457 + r = _make_chat_request(response_format=response_format)
458 + o = parser.adjust_request(r)
459 + assert o.response_format is None
460 + assert (
461 + _first_json_schema(o.structured_outputs.structural_tag) == expected_schema
462 + )
463 +
464 +
465 +class TestHasEffectiveTools:
466 + @pytest.mark.parametrize(
467 + "tools, expected",
468 + [
469 + pytest.param(None, False, id="none"),
470 + pytest.param([], False, id="empty_list"),
471 + pytest.param(" ", False, id="blank_str"),
472 + pytest.param(
473 + [{"type": "function", "function": {"name": "f"}}],
474 + True,
475 + id="non_empty_list",
476 + ),
477 + pytest.param('{"x": 1}', True, id="non_empty_str"),
478 + ],
479 + )
480 + def test_has_effective_tools(self, tools, expected) -> None:
481 + assert _has_effective_tools(tools) is expected
482 +
483 + def test_convert_schema_json_only_with_empty_tools_list(self) -> None:
484 + tag = convert_schema_to_structural_tags(
485 + schema=SCHEMA_B,
486 + tools=[],
487 + model_architecture="Cohere2ForCausalLM",
488 + )
489 + assert tag is not None
490 + assert _first_json_schema(tag) == SCHEMA_B
491 +
492 +
493 +class TestAdjustRequestFoldFromStructuredOutputs:
494 + @pytest.mark.parametrize(
495 + "structured_outputs, expected_schema",
496 + [
497 + pytest.param({"json": SCHEMA_B}, SCHEMA_B, id="json_dict"),
498 + pytest.param({"json": json.dumps(SCHEMA_B)}, SCHEMA_B, id="json_string"),
499 + pytest.param(
500 + {"json_object": True}, {"type": "object"}, id="json_object_flag"
501 + ),
502 + pytest.param(
503 + StructuredOutputsParams(json=SCHEMA_B),
504 + SCHEMA_B,
505 + id="structured_outputs_dataclass",
506 + ),
507 + pytest.param(
508 + {"json": {"name": "n", "schema": SCHEMA_A}},
509 + SCHEMA_A,
510 + id="openai_wrapper_dict_unwrapped",
511 + ),
512 + ],
513 + )
514 + def test_structured_outputs_folded(
515 + self, parser, structured_outputs, expected_schema
516 + ) -> None:
517 + o = parser.adjust_request(
518 + _make_chat_request(structured_outputs=structured_outputs),
519 + )
520 + assert (
521 + _first_json_schema(o.structured_outputs.structural_tag) == expected_schema
522 + )
523 +
524 + def test_responses_request_default_empty_tools(self, parser) -> None:
525 + """``ResponsesRequest.tools`` defaults to ``[]``, not ``None``."""
526 + r = ResponsesRequest.model_validate(
527 + {
528 + "input": "hi",
529 + "model": "m",
530 + "structured_outputs": {"json": SCHEMA_B},
531 + }
532 + )
533 + assert r.tools == []
534 + o = parser.adjust_request(r)
535 + assert _first_json_schema(o.structured_outputs.structural_tag) == SCHEMA_B
536 +
537 + def test_json_userdict_mapping_unwrapped(self) -> None:
538 + inner = {"type": "object", "properties": {"u": {"type": "number"}}}
539 + so = StructuredOutputsParams(json=UserDict(inner))
540 + assert _schema_dict_from_structured_outputs(so) == inner
541 +
542 + @pytest.mark.parametrize(
543 + "json_value, match",
544 + [
545 + pytest.param("{not json}", "valid JSON", id="invalid_json_string"),
546 + pytest.param(
547 + json.dumps(["a", "b"]), "JSON object", id="non_object_json_string"
548 + ),
549 + pytest.param(" ", "empty", id="empty_json_string"),
550 + ],
551 + )
552 + def test_structured_outputs_json_string_raises(
553 + self, parser, json_value, match
554 + ) -> None:
555 + with pytest.raises(ValueError, match=match):
556 + parser.adjust_request(
557 + _make_chat_request(structured_outputs={"json": json_value}),
558 + )
559 +
560 + @pytest.mark.parametrize(
561 + "construct",
562 + [
563 + pytest.param(
564 + lambda: _make_chat_request(structured_outputs={"json": [1, 2, 3]}),
565 + id="chat_completion_request",
566 + ),
567 + pytest.param(
568 + lambda: StructuredOutputsParams(json=[1, 2, 3]), # type: ignore[arg-type]
569 + id="structured_outputs_params",
570 + ),
571 + ],
572 + )
573 + def test_json_wrong_type_raises(self, construct) -> None:
574 + """Non-str / non-dict ``json`` fails at Pydantic validation."""
575 + with pytest.raises(ValidationError):
576 + construct()
577 +
578 +
579 +class TestAdjustRequestPrecedence:
580 + def test_response_format_over_structured_outputs_json(self, parser) -> None:
581 + s_rf = {"type": "object", "properties": {"rf": {"type": "string"}}}
582 + s_so = {"type": "object", "properties": {"so": {"type": "number"}}}
583 + inner = JsonSchemaResponseFormat(name="n", json_schema=s_rf)
584 + r = _make_chat_request(
585 + response_format=ResponseFormat(type="json_schema", json_schema=inner),
586 + structured_outputs={"json": s_so},
587 + )
588 + o = parser.adjust_request(r)
589 + assert _first_json_schema(o.structured_outputs.structural_tag) == s_rf
590 +
591 +
592 +class TestAdjustRequestTextPlusStructuredOutputs:
593 + def test_text_response_format_preserved(self, parser) -> None:
594 + sch = {"type": "object", "properties": {"k": {"type": "string"}}}
595 + r = _make_chat_request(
596 + response_format=ResponseFormat(type="text"),
597 + structured_outputs={"json": sch},
598 + )
599 + o = parser.adjust_request(r)
600 + assert o.response_format is not None
601 + assert o.response_format.type == "text"
602 + assert _first_json_schema(o.structured_outputs.structural_tag) == sch
603 +
604 +
605 +class TestAdjustRequestTools:
606 + def test_tools_only_command_a_grammar(self, parser) -> None:
607 + o = parser.adjust_request(
608 + _make_chat_request(tools=[GET_WEATHER_TOOL], tool_choice="auto"),
609 + )
610 + assert "grammar" in _content_types(o.structured_outputs.structural_tag)
611 +
612 + def test_tools_plus_json_schema_both_kinds(self, parser) -> None:
613 + inner = JsonSchemaResponseFormat(
614 + name="n",
615 + json_schema={"type": "object", "properties": {"r": {"type": "string"}}},
616 + )
617 + r = _make_chat_request(
618 + response_format=ResponseFormat(type="json_schema", json_schema=inner),
619 + tools=[GET_WEATHER_TOOL],
620 + tool_choice="auto",
621 + )
622 + o = parser.adjust_request(r)
623 + types = _content_types(o.structured_outputs.structural_tag)
624 + assert "grammar" in types
625 + assert "json_schema" in types
vllm/reasoning/cohere_command_reasoning_parser.py
+17 -1
@@ -414,7 +414,9 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
414 **kwargs,
415 ):
416 super().__init__(tokenizer, *args, **kwargs)
417 + self.start_token_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>")
418 self.end_token_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>")
419 + self.chatbot_token_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>")
420 self.unary_opts = unary_opts
421 self.melody_unary = PyFilter(unary_opts)
422 self.melody_streaming = PyFilter(streaming_opts)
@@ -478,7 +480,21 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
480 return content_ids
481
482 def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
481 - return any(tid == self.end_token_id for tid in reversed(input_ids))
483 + chatbot = self.chatbot_token_id
484 + start = self.start_token_id
485 + end = self.end_token_id
486 + has_end_token = False
487 +
488 + for i in reversed(range(len(input_ids))):
489 + tid = input_ids[i]
490 + if tid == start:
491 + return has_end_token
492 + if tid == chatbot:
493 + return False
494 + if tid == end:
495 + has_end_token = True
496 +
497 + return has_end_token
498
499 def adjust_request(
500 self, request: ChatCompletionRequest | ResponsesRequest