diff --git a/docs/features/reasoning_outputs.md b/docs/features/reasoning_outputs.md index fbd23fcdd5..a58dc6033e 100644 --- a/docs/features/reasoning_outputs.md +++ b/docs/features/reasoning_outputs.md @@ -15,6 +15,7 @@ vLLM currently supports the following reasoning models: | Model Series | Parser Name | Structured Output Support | Tool Calling | | ------------ | ----------- | ---------------- | ----------- | | [Cohere Command A Reasoning](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) | `cohere_command3` | `json`, `regex` | ✅ | +| [Cohere Command A Plus](https://huggingface.co/CohereLabs/command-a-plus-05-2026-bf16) | `cohere_command4` | `json`, `regex` | ✅ | | [DeepSeek R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d) | `deepseek_r1` | `json`, `regex` | ❌ | | [Gemma 4 series](https://huggingface.co/google/gemma-4-26B-A4B-it) | `gemma4` | `json`, `regex` | ✅ | | [DeepSeek-V3.1](https://huggingface.co/collections/deepseek-ai/deepseek-v31-68a491bed32bd77e7fca048f) | `deepseek_v3` | `json`, `regex` | ❌ | diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md index ae65231919..9628bc431a 100644 --- a/docs/features/tool_calling.md +++ b/docs/features/tool_calling.md @@ -366,15 +366,16 @@ Flags: * For non-reasoning: `--tool-call-parser hunyuan_a13b` * For reasoning: `--tool-call-parser hunyuan_a13b --reasoning-parser hunyuan_a13b` -### Cohere Command A Reasoning (`cohere_command3`) +### Cohere Command (`cohere_command3` / `cohere_command4`) Supported models: -* [`CohereLabs/command-a-reasoning-08-2025`](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) +* [`CohereLabs/command-a-reasoning-08-2025`](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) (`cohere_command3`) +* [`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`) -Flags: `--tool-call-parser cohere_command3 --reasoning-parser cohere_command3` +Flags: `--tool-call-parser cohere_command4 --reasoning-parser cohere_command4`. -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. +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. ### LongCat-Flash-Chat Models (`longcat`) diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index 2b1b8d372d..052896c9f2 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -127,7 +127,7 @@ cloudpickle==3.1.2 # via -r requirements/test/../common.txt cohere==7.0.8 # via -r requirements/test/cuda.in -cohere-melody==0.11.1 +cohere-melody==0.14.0 # via -r requirements/test/cuda.in colorama==0.4.6 # via diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 2910f01e9f..db6e359182 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -132,7 +132,7 @@ cloudpickle==3.1.2 # via -r requirements/test/../common.txt cohere==7.0.8 # via -r requirements/test/cuda.in -cohere-melody==0.11.1 +cohere-melody==0.14.0 # via -r requirements/test/cuda.in colorama==0.4.6 # via diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index ed47efded1..a6b9271989 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -134,7 +134,7 @@ cloudpickle==3.1.2 # tilelang cohere==7.0.8 # via -r requirements/test/rocm.in -cohere-melody==0.11.1 +cohere-melody==0.14.0 # via -r requirements/test/rocm.in colorama==0.4.6 # via diff --git a/tests/parser/cohere/__init__.py b/tests/parser/cohere/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/tests/parser/cohere/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/parser/cohere/conftest.py b/tests/parser/cohere/conftest.py new file mode 100644 index 0000000000..4477cacf43 --- /dev/null +++ b/tests/parser/cohere/conftest.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest + +from .utils import MockCohereTokenizer + + +@pytest.fixture(scope="package") +def tokenizer() -> MockCohereTokenizer: + return MockCohereTokenizer() + + +@pytest.fixture +def request_obj() -> ChatCompletionRequest: + return ChatCompletionRequest(messages=[], model="test-model") diff --git a/tests/parser/cohere/test_citations.py b/tests/parser/cohere/test_citations.py new file mode 100644 index 0000000000..14ee1612ea --- /dev/null +++ b/tests/parser/cohere/test_citations.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Citation source resolution and the streamed citation wire shape.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +from vllm.parser.cohere_command import ( + _melody_citations_to_vllm, + _melody_sources_to_vllm, +) + +from .utils import make_parser, stream_parser + + +class TestMelodySourceResolution: + """Pin the parser-side source resolution (see + ``_melody_sources_to_vllm``). The parser receives melody's numeric + ``(tool_call_index, tool_result_indices)`` addressing and resolves + it against a ``position_to_source`` map handed in by + ``CohereServingChatV2._apply_cohere_template_kwargs`` -- the same + map that would otherwise live in the serving layer's resolver. + """ + + @staticmethod + def _fake_melody_source(bucket: int, indices: list[int]) -> Any: + return SimpleNamespace(tool_call_index=bucket, tool_result_indices=indices) + + def test_multi_index_source_fans_out(self): + from vllm.entrypoints.cohere.cohere_chat_message import CitationSource + + position_map: dict[tuple[int, int], CitationSource] = { + (0, 0): CitationSource(type="document", id="d0", document={"id": "d0"}), + (0, 1): CitationSource(type="document", id="d1", document={"id": "d1"}), + } + raw = [self._fake_melody_source(0, [0, 1])] + out = _melody_sources_to_vllm(raw, position_map) + assert [s.id for s in out] == ["d0", "d1"] + # Verify type / payload were plumbed through, not just the id. + assert out[0].type == "document" + assert out[0].document == {"id": "d0"} + + def test_unresolvable_position_skipped(self): + from vllm.entrypoints.cohere.cohere_chat_message import CitationSource + + position_map: dict[tuple[int, int], CitationSource] = { + (0, 0): CitationSource(type="document", id="d0"), + } + raw = [self._fake_melody_source(9, [0])] + assert _melody_sources_to_vllm(raw, position_map) == [] + + def test_missing_position_map_drops_all_sources(self): + # A parser instance without a position map (parser wired + # outside of ``CohereServingChatV2``) can't attribute anything, + # so every source is dropped. Callers downstream will see the + # citation with empty ``sources`` and drop it entirely. + raw = [self._fake_melody_source(0, [0])] + assert _melody_sources_to_vllm(raw, None) == [] + + def test_citations_pass_through_is_thinking_tag(self): + from vllm.entrypoints.cohere.cohere_chat_message import CitationSource + + position_map: dict[tuple[int, int], CitationSource] = { + (0, 0): CitationSource(type="document", id="d0"), + } + raw = [ + SimpleNamespace( + start_index=0, + end_index=5, + text="hello", + is_thinking=True, + sources=[self._fake_melody_source(0, [0])], + ) + ] + out = _melody_citations_to_vllm(raw, position_map) + assert out is not None + assert out[0].type == "THINKING_CONTENT" + assert out[0].sources[0].id == "d0" + + +class TestParserStreamingEndToEnd: + """End-to-end streaming shape: raw model output text is tokenized + with the mock byte-level tokenizer and fed one token at a time + through :meth:`CohereCommand3ReasoningParser.extract_reasoning_streaming`. + We collect the ordered sequence of ``DeltaMessage``s the parser + hands back -- i.e. what the client would receive on the wire -- + and pin the exact shape. + + Combines the three flavors that reach ``delta``: a thinking block + (populates ``reasoning``), a text block (populates ``content``), + and a resolved citation (populates ``citations``) referencing + ``(start, end)`` offsets into the already-streamed content. + """ + + _MODEL_OUTPUT = ( + "<|START_THINKING|>Let me check.<|END_THINKING|>" + "<|START_RESPONSE|>The capital is Paris.<|END_RESPONSE|>" + ) + + def test_full_wire_stream_shape(self, tokenizer, request_obj): + from vllm.entrypoints.cohere.cohere_chat_message import CitationSource + from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY + + # Mimic what ``CohereServingChatV2._apply_cohere_template_kwargs`` + # installs on a request whose tool_call_index=0 tool result at + # index 1 is the "France" document. + position_to_source = { + (0, 1): CitationSource( + type="document", + id="doc-paris", + document={"id": "doc-paris", "title": "France"}, + ), + } + + parser = make_parser( + tokenizer, + "cohere_command3", + chat_template_kwargs={POSITION_TO_SOURCE_KEY: position_to_source}, + ) + deltas = stream_parser(parser, request_obj, tokenizer, self._MODEL_OUTPUT) + + # -- 1. Reasoning and content stream separately and reach + # the wire only as their real payload. Framing tokens + # (``<|START_THINKING|>`` etc.) are consumed by the parser's + # state machine and never produce a delta. + reasoning_stream = [d.reasoning for d in deltas if d.reasoning is not None] + content_stream = [d.content for d in deltas if d.content is not None] + citation_deltas = [d for d in deltas if getattr(d, "citations", None)] + + assert "".join(reasoning_stream) == "Let me check." + assert "".join(content_stream) == "The capital is Paris." + + # -- 2. Exactly one citation reaches the wire, resolved to the + # source from ``position_to_source`` (not the raw + # ``(bucket, idx)`` coordinates). + assert len(citation_deltas) == 1 + (cite_delta,) = citation_deltas + assert cite_delta.citations is not None + assert len(cite_delta.citations) == 1 + cite = cite_delta.citations[0] + assert cite.text == "Paris" + # Char offsets are into the accumulated ``content``: "The + # capital is " is 15 chars, "Paris" is 5. + assert cite.start == 15 + assert cite.end == 20 + assert len(cite.sources) == 1 + assert cite.sources[0].id == "doc-paris" + assert cite.sources[0].type == "document" + + # -- 3. The citation's ``(start, end)`` offsets point into + # bytes that were already emitted as content-deltas earlier + # in the stream -- i.e. the citation is anchored to real + # already-shipped content. + cite_delta_idx = deltas.index(cite_delta) + earlier_content = "".join( + d.content for d in deltas[:cite_delta_idx] if d.content is not None + ) + assert earlier_content[cite.start : cite.end] == "Paris" + + # Everything after the citation delta is the trailing period. + later_content = "".join( + d.content for d in deltas[cite_delta_idx + 1 :] if d.content is not None + ) + assert later_content == "." + + # -- 4. Reasoning and content are never mixed on the same + # delta: the parser flips modes on the ``<|END_THINKING|>`` + # boundary and never emits a delta with both fields set. + for d in deltas: + assert not (d.reasoning is not None and d.content is not None) diff --git a/tests/parser/cohere/test_reasoning.py b/tests/parser/cohere/test_reasoning.py new file mode 100644 index 0000000000..8c3311e838 --- /dev/null +++ b/tests/parser/cohere/test_reasoning.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Reasoning extraction, reasoning-end gating and parser selection for the +unified Cohere Command parser.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.parser import ParserManager +from vllm.reasoning.cohere_command_reasoning_parser import ( + CohereCommand3ReasoningParser, + CohereCommand4ReasoningParser, +) + +from .utils import make_parser, stream_parser + + +@dataclass +class ReasoningCase: + parser_name: str + model_output: str + expected_reasoning: str | None + expected_content: str | None + + +REASONING_CASES = [ + pytest.param( + ReasoningCase( + parser_name="cohere_command3", + model_output="""\ +<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> +[ + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} +] +<|END_ACTION|>""", + expected_reasoning="i will call foo with query1", + expected_content=None, + ), + id="cmd3-single_tool_call", + ), + pytest.param( + ReasoningCase( + parser_name="cohere_command4", + model_output="""\ +<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> +[ + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} +] +<|END_ACTION|>""", + expected_reasoning="i will call foo with query1", + expected_content=None, + ), + id="cmd4-single_tool_call", + ), + pytest.param( + ReasoningCase( + parser_name="cohere_command3", + model_output="""\ +<|START_THINKING|>This is a rainbow emoji: 🌈<|END_THINKING|> +<|START_RESPONSE|>foo bar<|END_RESPONSE|>""", + expected_reasoning="This is a rainbow emoji: 🌈", + expected_content="foo bar", + ), + id="cmd3-citations_with_emoji", + ), + pytest.param( + ReasoningCase( + parser_name="cohere_command4", + model_output="""\ +<|START_THINKING|>This is a rainbow emoji: 🌈<|END_THINKING|> +<|START_RESPONSE|>foo bar<|END_RESPONSE|>""", + expected_reasoning="This is a rainbow emoji: 🌈", + expected_content="foo bar", + ), + id="cmd4-citations_with_emoji", + ), +] + + +@pytest.mark.parametrize("case", REASONING_CASES) +class TestExtractReasoning: + def test_nonstreaming(self, tokenizer, request_obj, case: ReasoningCase): + parser = make_parser(tokenizer, case.parser_name) + reasoning, content, tool_calls = parser.parse( + case.model_output, request_obj, enable_auto_tools=True + ) + assert reasoning == case.expected_reasoning + assert content == case.expected_content + # No tools on the request, so any tool call is dropped, not surfaced. + assert not tool_calls + + def test_streaming(self, tokenizer, request_obj, case: ReasoningCase): + parser = make_parser(tokenizer, case.parser_name) + deltas = stream_parser(parser, request_obj, tokenizer, case.model_output) + + reasoning = "".join(d.reasoning for d in deltas if d.reasoning) or None + content = "".join(d.content for d in deltas if d.content) or None + assert reasoning == case.expected_reasoning + assert content == case.expected_content + assert all(not d.tool_calls for d in deltas) + + +class TestFramingTokensStripped: + @pytest.mark.parametrize( + ("parser_name", "content_tags"), + [ + pytest.param( + "cohere_command4", ("<|START_TEXT|>", "<|END_TEXT|>"), id="cmd4" + ), + pytest.param( + "cohere_command3", ("<|START_RESPONSE|>", "<|END_RESPONSE|>"), id="cmd3" + ), + ], + ) + @pytest.mark.parametrize("chunk_size", [1, 4], ids=["per_token", "batched"]) + @pytest.mark.parametrize("with_tools", [False, True], ids=["no_tools", "tools"]) + def test_content_framing_tokens_stripped( + self, tokenizer, parser_name, content_tags, chunk_size, with_tools + ): + parser = make_parser(tokenizer, parser_name) + tools = [{"type": "function", "function": {"name": "foo"}}] + request = ChatCompletionRequest( + messages=[], model="test-model", tools=tools if with_tools else None + ) + parser.adjust_request(request) + + start_tag, end_tag = content_tags + generation = ( + f"<|START_THINKING|>Think deeply. The user greets us.<|END_THINKING|>" + f"{start_tag}I'm doing well, thank you{end_tag}" + ) + deltas = stream_parser(parser, request, tokenizer, generation, chunk_size) + + assert "".join(d.reasoning or "" for d in deltas) == ( + "Think deeply. The user greets us." + ) + assert not any(d.tool_calls for d in deltas) + assert "".join(d.content or "" for d in deltas) == "I'm doing well, thank you" + + +class TestIsReasoningEnd: + @pytest.mark.parametrize( + "parser_cls", + [CohereCommand3ReasoningParser, CohereCommand4ReasoningParser], + ids=["cmd3", "cmd4"], + ) + def test_is_reasoning_end(self, tokenizer, parser_cls): + parser = parser_cls(tokenizer) + start_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>") + end_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>") + chatbot_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>") + content_ids = [99, 100] + + # Generation-only tokens have no chatbot marker, so the whole sequence + # is considered. + assert parser.is_reasoning_end([end_id]) + assert parser.is_reasoning_end([start_id, *content_ids, end_id]) + assert not parser.is_reasoning_end([start_id, *content_ids]) + + # Full prompt/history tokens are scoped to the latest chatbot marker, + # so stray thinking tokens from the preamble or previous turns are ignored. + assert not parser.is_reasoning_end([start_id, end_id, chatbot_id, *content_ids]) + assert parser.is_reasoning_end( + [start_id, end_id, chatbot_id, start_id, *content_ids, end_id] + ) + + +@pytest.mark.parametrize("parser_name", ["cohere_command3", "cohere_command4"]) +def test_count_reasoning_tokens(tokenizer, parser_name): + parser = make_parser(tokenizer, parser_name) + start = tokenizer.convert_tokens_to_ids("<|START_THINKING|>") + end = tokenizer.convert_tokens_to_ids("<|END_THINKING|>") + + assert parser.count_reasoning_tokens([99, start, 11, 12, end, 100]) == 2 + assert parser.count_reasoning_tokens([end, start, 1, start, 2, end, 3, end]) == 3 + assert parser.count_reasoning_tokens([start, 1, 2]) == 2 + assert parser.count_reasoning_tokens([1, 2, end]) == 0 + + +class TestParserSelection: + @pytest.mark.parametrize("with_tool_parser", [True, False]) + def test_adjust_request_keeps_framing_tokens(self, tokenizer, with_tool_parser): + """Melody needs the special tokens, with or without a tool parser.""" + cls = ParserManager.get_parser( + "cohere_command4" if with_tool_parser else None, + "cohere_command4", + enable_auto_tools=True, + ) + request = ChatCompletionRequest(messages=[], model="m") + assert request.skip_special_tokens + cls(tokenizer).adjust_request(request) + assert request.skip_special_tokens is False + + def test_non_cohere_selection_does_not_import_melody_parser(self, monkeypatch): + """``cohere_melody`` is optional; other parsers must not depend on it.""" + import sys + + monkeypatch.setitem(sys.modules, "vllm.parser.cohere_command", None) + assert ParserManager.get_parser("hermes", "deepseek_r1", True) is not None diff --git a/tests/parser/cohere/test_structural_tags.py b/tests/parser/cohere/test_structural_tags.py new file mode 100644 index 0000000000..27440d52af --- /dev/null +++ b/tests/parser/cohere/test_structural_tags.py @@ -0,0 +1,362 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""``adjust_request``: folding tools and JSON schemas into Cohere structural tags.""" + +from __future__ import annotations + +import json +from collections import UserDict +from types import SimpleNamespace + +import pytest +from pydantic import ValidationError + +from vllm.entrypoints.generate.base.protocol import ( + JsonSchemaResponseFormat, + ResponseFormat, + StructuralTagResponseFormat, +) +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.parser.cohere_command import ( + CohereCommandParser, + _has_effective_tools, + _response_format_type, + _schema_dict_from_structured_outputs, + convert_schema_to_structural_tags, +) +from vllm.sampling_params import StructuredOutputsParams + +from .utils import MockCohereTokenizer, make_parser + +SCHEMA_A = {"type": "object", "properties": {"a": {"type": "string"}}} +SCHEMA_B = {"type": "object", "properties": {"b": {"type": "number"}}} +GET_WEATHER_TOOL = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, +} +VALID_STRUCTURAL_TAG = { + "type": "structural_tag", + "format": { + "type": "triggered_tags", + "tags": [ + { + "begin": "", + "content": {"type": "any_text"}, + "end": "", + } + ], + "triggers": [""], + }, +} + + +def _model_config(arch: str) -> SimpleNamespace: + return SimpleNamespace( + architecture=arch, + architectures=[arch], + hf_text_config=SimpleNamespace(architectures=[arch]), + ) + + +def _make_chat_request(**kwargs) -> ChatCompletionRequest: + data = {"messages": [{"role": "user", "content": "hi"}], "model": "m"} + data.update(kwargs) + return ChatCompletionRequest.model_validate(data) + + +def _first_json_schema(tag_json: str) -> dict | None: + outer = json.loads(tag_json) + for t in (outer.get("format") or {}).get("tags") or []: + c = t.get("content") or {} + if c.get("type") == "json_schema": + js = c.get("json_schema") + return js if isinstance(js, dict) else None + return None + + +def _content_types(tag_json: str) -> set[str]: + outer = json.loads(tag_json) + out: set[str] = set() + for t in (outer.get("format") or {}).get("tags") or []: + ty = (t.get("content") or {}).get("type") + if isinstance(ty, str): + out.add(ty) + return out + + +@pytest.fixture(scope="module") +def parser(tokenizer: MockCohereTokenizer) -> CohereCommandParser: + """Parser configured with a supported Cohere architecture.""" + return make_parser( + tokenizer, "cohere_command4", model_config=_model_config("Cohere2ForCausalLM") + ) + + +@pytest.fixture(scope="module") +def parser_no_model_config(tokenizer: MockCohereTokenizer) -> CohereCommandParser: + """Parser with no ``model_config`` (cannot resolve architecture).""" + return make_parser(tokenizer, "cohere_command4") + + +@pytest.fixture(scope="module") +def parser_unsupported_arch(tokenizer: MockCohereTokenizer) -> CohereCommandParser: + """Parser configured with an architecture that has no structural tag style.""" + return make_parser( + tokenizer, "cohere_command4", model_config=_model_config("LlamaForCausalLM") + ) + + +class TestAdjustRequestPassthrough: + def test_structured_outputs_structural_tag_not_modified(self, parser) -> None: + tag = json.dumps(VALID_STRUCTURAL_TAG) + r = _make_chat_request(structured_outputs={"structural_tag": tag}) + o = parser.adjust_request(r) + assert o.structured_outputs.structural_tag == tag + + def test_response_format_structural_tag_short_circuit(self, parser) -> None: + # ``ChatCompletionRequest`` validates ``response_format`` as a union; + # bare ``{"type": "structural_tag"}`` is invalid (use pydantic model). + rf = StructuralTagResponseFormat( + type="structural_tag", + format=VALID_STRUCTURAL_TAG["format"], + ) + r = _make_chat_request(response_format=rf) + o = parser.adjust_request(r) + assert _response_format_type(o.response_format) == "structural_tag" + assert o.structured_outputs is None + + +class TestAdjustRequestNoOp: + def test_no_schema_no_tools(self, parser) -> None: + o = parser.adjust_request(_make_chat_request()) + assert o.structured_outputs is None + assert o.response_format is None + + def test_no_model_config(self, parser_no_model_config) -> None: + inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A) + r = _make_chat_request( + response_format=ResponseFormat(type="json_schema", json_schema=inner), + ) + o = parser_no_model_config.adjust_request(r) + assert o.response_format is not None + assert o.structured_outputs is None + + +class TestAdjustRequestUnsupportedArchitecture: + def test_json_schema_raises(self, parser_unsupported_arch) -> None: + inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A) + r = _make_chat_request( + response_format=ResponseFormat(type="json_schema", json_schema=inner), + ) + with pytest.raises(ValueError, match="does not support"): + parser_unsupported_arch.adjust_request(r) + + +class TestAdjustRequestFoldFromResponseFormat: + @pytest.mark.parametrize( + "response_format, expected_schema", + [ + pytest.param( + ResponseFormat( + type="json_schema", + json_schema=JsonSchemaResponseFormat( + name="n", json_schema=SCHEMA_A + ), + ), + SCHEMA_A, + id="json_schema_pydantic", + ), + pytest.param( + { + "type": "json_schema", + "json_schema": {"name": "n", "schema": SCHEMA_A}, + }, + SCHEMA_A, + id="json_schema_dict", + ), + pytest.param( + {"type": "json_object"}, + {"type": "object"}, + id="json_object", + ), + ], + ) + def test_response_format_cleared( + self, parser, response_format, expected_schema + ) -> None: + r = _make_chat_request(response_format=response_format) + o = parser.adjust_request(r) + assert o.response_format is None + assert ( + _first_json_schema(o.structured_outputs.structural_tag) == expected_schema + ) + + +class TestHasEffectiveTools: + @pytest.mark.parametrize( + "tools, expected", + [ + pytest.param(None, False, id="none"), + pytest.param([], False, id="empty_list"), + pytest.param(" ", False, id="blank_str"), + pytest.param( + [{"type": "function", "function": {"name": "f"}}], + True, + id="non_empty_list", + ), + pytest.param('{"x": 1}', True, id="non_empty_str"), + ], + ) + def test_has_effective_tools(self, tools, expected) -> None: + assert _has_effective_tools(tools) is expected + + def test_convert_schema_json_only_with_empty_tools_list(self) -> None: + tag = convert_schema_to_structural_tags( + schema=SCHEMA_B, + tools=[], + model_architecture="Cohere2ForCausalLM", + ) + assert tag is not None + assert _first_json_schema(tag) == SCHEMA_B + + +class TestAdjustRequestFoldFromStructuredOutputs: + @pytest.mark.parametrize( + "structured_outputs, expected_schema", + [ + pytest.param({"json": SCHEMA_B}, SCHEMA_B, id="json_dict"), + pytest.param({"json": json.dumps(SCHEMA_B)}, SCHEMA_B, id="json_string"), + pytest.param( + {"json_object": True}, {"type": "object"}, id="json_object_flag" + ), + pytest.param( + StructuredOutputsParams(json=SCHEMA_B), + SCHEMA_B, + id="structured_outputs_dataclass", + ), + pytest.param( + {"json": {"name": "n", "schema": SCHEMA_A}}, + SCHEMA_A, + id="openai_wrapper_dict_unwrapped", + ), + ], + ) + def test_structured_outputs_folded( + self, parser, structured_outputs, expected_schema + ) -> None: + o = parser.adjust_request( + _make_chat_request(structured_outputs=structured_outputs), + ) + assert ( + _first_json_schema(o.structured_outputs.structural_tag) == expected_schema + ) + + def test_responses_request_default_empty_tools(self, parser) -> None: + """``ResponsesRequest.tools`` defaults to ``[]``, not ``None``.""" + r = ResponsesRequest.model_validate( + { + "input": "hi", + "model": "m", + "structured_outputs": {"json": SCHEMA_B}, + } + ) + assert r.tools == [] + o = parser.adjust_request(r) + assert _first_json_schema(o.structured_outputs.structural_tag) == SCHEMA_B + + def test_json_userdict_mapping_unwrapped(self) -> None: + inner = {"type": "object", "properties": {"u": {"type": "number"}}} + so = StructuredOutputsParams(json=UserDict(inner)) + assert _schema_dict_from_structured_outputs(so) == inner + + @pytest.mark.parametrize( + "json_value, match", + [ + pytest.param("{not json}", "valid JSON", id="invalid_json_string"), + pytest.param( + json.dumps(["a", "b"]), "JSON object", id="non_object_json_string" + ), + pytest.param(" ", "empty", id="empty_json_string"), + ], + ) + def test_structured_outputs_json_string_raises( + self, parser, json_value, match + ) -> None: + with pytest.raises(ValueError, match=match): + parser.adjust_request( + _make_chat_request(structured_outputs={"json": json_value}), + ) + + @pytest.mark.parametrize( + "construct", + [ + pytest.param( + lambda: _make_chat_request(structured_outputs={"json": [1, 2, 3]}), + id="chat_completion_request", + ), + pytest.param( + lambda: StructuredOutputsParams(json=[1, 2, 3]), # type: ignore[arg-type] + id="structured_outputs_params", + ), + ], + ) + def test_json_wrong_type_raises(self, construct) -> None: + """Non-str / non-dict ``json`` fails at Pydantic validation.""" + with pytest.raises(ValidationError): + construct() + + +class TestAdjustRequestPrecedence: + def test_response_format_over_structured_outputs_json(self, parser) -> None: + s_rf = {"type": "object", "properties": {"rf": {"type": "string"}}} + s_so = {"type": "object", "properties": {"so": {"type": "number"}}} + inner = JsonSchemaResponseFormat(name="n", json_schema=s_rf) + r = _make_chat_request( + response_format=ResponseFormat(type="json_schema", json_schema=inner), + structured_outputs={"json": s_so}, + ) + o = parser.adjust_request(r) + assert _first_json_schema(o.structured_outputs.structural_tag) == s_rf + + +class TestAdjustRequestTextPlusStructuredOutputs: + def test_text_response_format_preserved(self, parser) -> None: + sch = {"type": "object", "properties": {"k": {"type": "string"}}} + r = _make_chat_request( + response_format=ResponseFormat(type="text"), + structured_outputs={"json": sch}, + ) + o = parser.adjust_request(r) + assert o.response_format is not None + assert o.response_format.type == "text" + assert _first_json_schema(o.structured_outputs.structural_tag) == sch + + +class TestAdjustRequestTools: + def test_tools_only_command_a_grammar(self, parser) -> None: + o = parser.adjust_request( + _make_chat_request(tools=[GET_WEATHER_TOOL], tool_choice="auto"), + ) + assert "grammar" in _content_types(o.structured_outputs.structural_tag) + + def test_tools_plus_json_schema_both_kinds(self, parser) -> None: + inner = JsonSchemaResponseFormat( + name="n", + json_schema={"type": "object", "properties": {"r": {"type": "string"}}}, + ) + r = _make_chat_request( + response_format=ResponseFormat(type="json_schema", json_schema=inner), + tools=[GET_WEATHER_TOOL], + tool_choice="auto", + ) + o = parser.adjust_request(r) + types = _content_types(o.structured_outputs.structural_tag) + assert "grammar" in types + assert "json_schema" in types diff --git a/tests/parser/cohere/test_tool_calls.py b/tests/parser/cohere/test_tool_calls.py new file mode 100644 index 0000000000..a9fe7edb8d --- /dev/null +++ b/tests/parser/cohere/test_tool_calls.py @@ -0,0 +1,288 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tool-call parsing for the unified Cohere Command parser. + +The ``cohere_command3`` / ``cohere_command4`` tool parsers are registry shims; +``CohereCommandParser`` (resolved through ``ParserManager``) does the parsing. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.parser.abstract_parser import Parser + +from .utils import MockCohereTokenizer, drive_parser, make_parser, token_deltas + +SPECIAL_TOKEN_MARKERS = ( + "<|START_THINKING|>", + "<|END_THINKING|>", + "<|START_RESPONSE|>", + "<|END_RESPONSE|>", + "<|START_ACTION|>", + "<|END_ACTION|>", + "<|START_TEXT|>", + "<|END_TEXT|>", +) + +DUMMY_TOOLS = [ + {"type": "function", "function": {"name": name, "parameters": {"type": "object"}}} + for name in ("foo", "bar") +] + +# (registry name, answer framing tag) +VARIANTS = (("cohere_command3", "RESPONSE"), ("cohere_command4", "TEXT")) + + +@dataclass +class ExpectedToolCall: + id: str + name: str + arguments: dict + + +@dataclass +class ToolCallCase: + parser_name: str + model_output: str + expected_tool_calls: list[ExpectedToolCall] = field(default_factory=list) + expected_reasoning: str | None = None + expected_content: str | None = None + + +SINGLE_TOOL_CALL = """\ +<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> +[ + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} +] +<|END_ACTION|>""" +MULTI_TOOL_CALL = """\ +<|START_THINKING|>first I think about foo<|END_THINKING|><|START_ACTION|> +[ + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}, + {"tool_call_id": "1", "tool_name": "bar", "parameters": {"x": 42}} +] +<|END_ACTION|>""" +CITATIONS_NO_TOOL_CALL = ( + "<|START_THINKING|>This is a rainbow emoji: 🌈<|END_THINKING|>\n" + "<|START_{tag}|>foo bar<|END_{tag}|>" +) + +# (case id, model output template, expectations) +_CASE_TABLE = ( + ( + "single_tool_call", + SINGLE_TOOL_CALL, + dict( + expected_tool_calls=[ExpectedToolCall("0", "foo", {"query": "query1"})], + expected_reasoning="i will call foo with query1", + ), + ), + ( + "citations_no_tool_calls", + CITATIONS_NO_TOOL_CALL, + dict( + expected_reasoning="This is a rainbow emoji: 🌈", expected_content="foo bar" + ), + ), + ( + "multiple_tool_calls", + MULTI_TOOL_CALL, + dict( + expected_tool_calls=[ + ExpectedToolCall("0", "foo", {"query": "query1"}), + ExpectedToolCall("1", "bar", {"x": 42}), + ], + expected_reasoning="first I think about foo", + ), + ), + ( + "reasoning_only", + "<|START_THINKING|>just think, no response<|END_THINKING|>", + dict(expected_reasoning="just think, no response"), + ), +) + +TOOL_CALL_CASES = [ + pytest.param( + ToolCallCase(name, output.replace("{tag}", tag), **expected), + id=f"{name}-{case_id}", + ) + for name, tag in VARIANTS + for case_id, output, expected in _CASE_TABLE +] + +PARSER_NAMES = pytest.mark.parametrize("parser_name", [name for name, _ in VARIANTS]) + + +def _tool_parser(tokenizer: MockCohereTokenizer, parser_name: str) -> Parser: + return make_parser(tokenizer, parser_name, DUMMY_TOOLS) + + +def _tool_request() -> ChatCompletionRequest: + return ChatCompletionRequest( + messages=[], model="test-model", tools=DUMMY_TOOLS, tool_choice="auto" + ) + + +@dataclass +class StreamingResult: + tool_calls: dict[int, dict] + reasoning: str | None + content: str | None + + +def run_streaming( + tokenizer: MockCohereTokenizer, + parser_name: str, + model_output: str, + chunk_size: int = 1, +) -> StreamingResult: + deltas = drive_parser( + _tool_parser(tokenizer, parser_name), + _tool_request(), + token_deltas(tokenizer, model_output, chunk_size), + ) + accumulated: dict[int, dict] = {} + for delta in deltas: + for tc in delta.tool_calls or []: + slot = accumulated.setdefault( + tc.index, {"id": "", "name": "", "arguments": ""} + ) + if tc.id: + slot["id"] = tc.id + if tc.function and tc.function.name: + slot["name"] = tc.function.name + if tc.function and tc.function.arguments: + slot["arguments"] += tc.function.arguments + return StreamingResult( + tool_calls=accumulated, + reasoning="".join(d.reasoning for d in deltas if d.reasoning) or None, + content="".join(d.content for d in deltas if d.content) or None, + ) + + +def run_nonstreaming( + tokenizer: MockCohereTokenizer, parser_name: str, model_output: str +) -> tuple[str | None, str | None, list]: + reasoning, content, tool_calls = _tool_parser(tokenizer, parser_name).parse( + model_output, _tool_request(), enable_auto_tools=True + ) + return reasoning, content, tool_calls or [] + + +@pytest.mark.parametrize("case", TOOL_CALL_CASES) +class TestExtractToolCalls: + def test_streaming(self, tokenizer: MockCohereTokenizer, case: ToolCallCase): + streamed = run_streaming(tokenizer, case.parser_name, case.model_output) + assert len(streamed.tool_calls) == len(case.expected_tool_calls) + for i, expected in enumerate(case.expected_tool_calls): + tc = streamed.tool_calls[i] + assert tc["id"] == expected.id + assert tc["name"] == expected.name + assert json.loads(tc["arguments"]) == expected.arguments + + def test_streaming_reasoning( + self, tokenizer: MockCohereTokenizer, case: ToolCallCase + ): + streamed = run_streaming(tokenizer, case.parser_name, case.model_output) + assert streamed.reasoning == case.expected_reasoning + + def test_streaming_content( + self, tokenizer: MockCohereTokenizer, case: ToolCallCase + ): + streamed = run_streaming(tokenizer, case.parser_name, case.model_output) + assert streamed.content == case.expected_content + + def test_nonstreaming(self, tokenizer: MockCohereTokenizer, case: ToolCallCase): + reasoning, content, tool_calls = run_nonstreaming( + tokenizer, case.parser_name, case.model_output + ) + assert reasoning == case.expected_reasoning + assert content == case.expected_content + assert len(tool_calls) == len(case.expected_tool_calls) + for actual, expected in zip(tool_calls, case.expected_tool_calls): + assert actual.name == expected.name + assert json.loads(actual.arguments) == expected.arguments + + def test_streaming_nonstreaming_agree( + self, tokenizer: MockCohereTokenizer, case: ToolCallCase + ): + streamed = run_streaming(tokenizer, case.parser_name, case.model_output) + _, _, tool_calls = run_nonstreaming( + tokenizer, case.parser_name, case.model_output + ) + assert len(streamed.tool_calls) == len(tool_calls) + for i, actual in enumerate(tool_calls): + assert streamed.tool_calls[i]["name"] == actual.name + assert json.loads(streamed.tool_calls[i]["arguments"]) == json.loads( + actual.arguments + ) + + +class TestSpeculativeDecodingMultiTokenDelta: + @PARSER_NAMES + @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6]) + def test_no_special_token_leak_in_streaming_deltas( + self, tokenizer: MockCohereTokenizer, parser_name: str, chunk_size: int + ): + deltas = drive_parser( + _tool_parser(tokenizer, parser_name), + _tool_request(), + token_deltas(tokenizer, SINGLE_TOOL_CALL, chunk_size), + ) + for delta in deltas: + fields: list[tuple[str, str | None]] = [ + ("reasoning", delta.reasoning), + ("content", delta.content), + ] + for tc in delta.tool_calls or []: + if tc.function: + fields.append(("tool_call.name", tc.function.name)) + fields.append(("tool_call.arguments", tc.function.arguments)) + for marker in SPECIAL_TOKEN_MARKERS: + for field_name, value in fields: + assert value is None or marker not in value, ( + f"special token {marker!r} leaked into {field_name} " + f"with chunk_size={chunk_size} delta={delta!r}" + ) + + @PARSER_NAMES + @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6]) + def test_multi_token_chunks_still_produce_correct_tool_call( + self, tokenizer: MockCohereTokenizer, parser_name: str, chunk_size: int + ): + streamed = run_streaming(tokenizer, parser_name, SINGLE_TOOL_CALL, chunk_size) + assert len(streamed.tool_calls) == 1 + tc = streamed.tool_calls[0] + assert tc["id"] == "0" + assert tc["name"] == "foo" + assert json.loads(tc["arguments"]) == {"query": "query1"} + + +class TestStreamingDeltaShape: + @PARSER_NAMES + def test_reasoning_and_tool_calls_are_separate_deltas( + self, tokenizer: MockCohereTokenizer, parser_name: str + ): + deltas = drive_parser( + _tool_parser(tokenizer, parser_name), + _tool_request(), + token_deltas(tokenizer, SINGLE_TOOL_CALL), + ) + for delta in deltas: + populated = [ + delta.content is not None, + delta.reasoning is not None, + bool(delta.tool_calls), + ] + assert sum(populated) == 1, ( + "A single streaming delta must carry exactly one of " + f"content/reasoning/tool_calls, got {delta!r}" + ) + assert any(d.reasoning is not None for d in deltas) + assert any(d.tool_calls for d in deltas) diff --git a/tests/parser/cohere/utils.py b/tests/parser/cohere/utils.py new file mode 100644 index 0000000000..f26df5296c --- /dev/null +++ b/tests/parser/cohere/utils.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared helpers for the unified Cohere Command parser tests.""" + +from __future__ import annotations + +import itertools + +import regex as re + +from vllm.entrypoints.generate.base.protocol import DeltaMessage +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.parser import ParserManager +from vllm.parser.abstract_parser import Parser + +_SPECIAL_TOKEN_RE = re.compile(r"(<\|[A-Z_]+\|>)") +REPLACEMENT_CHAR = "�" + + +class MockCohereTokenizer: + """Byte-level stand-in for the Cohere tokenizer. + + Text round-trips through UTF-8 bytes so a multi-byte character can split + across "tokens" (reproducing trailing U+FFFD buffering), while each + ``<|...|>`` marker is a single id so token-id reasoning-end gating behaves + as with the real tokenizer. + """ + + _SPECIAL_TOKEN_IDS = { + tok: 256 + i + for i, tok in enumerate( + ( + "<|START_THINKING|>", + "<|END_THINKING|>", + "<|CHATBOT_TOKEN|>", + "<|START_RESPONSE|>", + "<|END_RESPONSE|>", + "<|START_TEXT|>", + "<|END_TEXT|>", + "<|START_ACTION|>", + "<|END_ACTION|>", + ) + ) + } + _ID_TO_SPECIAL_TOKEN = {v: k for k, v in _SPECIAL_TOKEN_IDS.items()} + + def convert_tokens_to_ids(self, token: str) -> int: + return self._SPECIAL_TOKEN_IDS.get(token, 0) + + def get_vocab(self) -> dict[str, int]: + return {} + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + ids: list[int] = [] + for part in _SPECIAL_TOKEN_RE.split(text): + if part in self._SPECIAL_TOKEN_IDS: + ids.append(self._SPECIAL_TOKEN_IDS[part]) + else: + ids.extend(part.encode("utf-8")) + return ids + + def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str: + out: list[str] = [] + for special, run in itertools.groupby( + ids, self._ID_TO_SPECIAL_TOKEN.__contains__ + ): + if not special: + out.append(bytes(run).decode("utf-8", errors="replace")) + elif not skip_special_tokens: + out.extend(self._ID_TO_SPECIAL_TOKEN[i] for i in run) + return "".join(out) + + +def token_deltas( + tokenizer: MockCohereTokenizer, text: str, chunk_size: int = 1 +) -> list[tuple[str, list[int]]]: + """Split ``text`` into ``(delta_text, delta_token_ids)`` steps of + ``chunk_size`` tokens, buffering incomplete multi-byte sequences (trailing + U+FFFD) into the next step as real streaming does.""" + ids = tokenizer.encode(text) + deltas: list[tuple[str, list[int]]] = [] + prev = "" + pending: list[int] = [] + for start in range(0, len(ids), chunk_size): + end = start + chunk_size + pending.extend(ids[start:end]) + current = tokenizer.decode(ids[:end]) + if current.endswith(REPLACEMENT_CHAR) and end < len(ids): + continue + deltas.append((current[len(prev) :], pending)) + prev, pending = current, [] + return deltas + + +def make_parser( + tokenizer: MockCohereTokenizer, + name: str, + tools: list[dict] | None = None, + chat_template_kwargs: dict | None = None, + model_config=None, +) -> Parser: + """Resolve the unified Cohere parser through ``ParserManager`` by registry + name (``cohere_command3`` / ``cohere_command4``).""" + cls = ParserManager.get_parser(name, name, enable_auto_tools=True) + assert cls is not None + return cls( + tokenizer, + tools, + model_config=model_config, + chat_template_kwargs=chat_template_kwargs or {}, + ) + + +def drive_parser( + parser: Parser, + request: ChatCompletionRequest, + deltas: list[tuple[str, list[int]]], +) -> list[DeltaMessage]: + """Feed ``deltas`` through ``parse_delta`` (flushing on the last one) and + return the non-``None`` messages in order.""" + out: list[DeltaMessage] = [] + for i, (text, ids) in enumerate(deltas): + delta = parser.parse_delta(text, ids, request, finished=i == len(deltas) - 1) + if delta is not None: + out.append(delta) + return out + + +def stream_parser(parser, request, tokenizer, text: str, chunk_size: int = 1): + """Stream ``text`` through ``parser`` in ``chunk_size``-token deltas.""" + return drive_parser(parser, request, token_deltas(tokenizer, text, chunk_size)) diff --git a/tests/reasoning/test_cohere_command_reasoning_parser.py b/tests/reasoning/test_cohere_command_reasoning_parser.py deleted file mode 100644 index ea9ab63aef..0000000000 --- a/tests/reasoning/test_cohere_command_reasoning_parser.py +++ /dev/null @@ -1,796 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import json -from collections import UserDict -from dataclasses import dataclass -from types import SimpleNamespace -from typing import Any - -import pytest -from pydantic import ValidationError - -from vllm.entrypoints.generate.base.protocol import ( - JsonSchemaResponseFormat, - ResponseFormat, - StructuralTagResponseFormat, -) -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, -) -from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -from vllm.reasoning.cohere_command_reasoning_parser import ( - CohereCommand3ReasoningParser, - CohereCommand4ReasoningParser, - _has_effective_tools, - _melody_citations_to_vllm, - _melody_sources_to_vllm, - _response_format_type, - _schema_dict_from_structured_outputs, - convert_schema_to_structural_tags, -) -from vllm.sampling_params import StructuredOutputsParams - - -@dataclass -class ReasoningCase: - parser_cls: Any - model_output: str - expected_reasoning: str | None - expected_content: str | None - - -REASONING_CASES = [ - pytest.param( - ReasoningCase( - parser_cls=CohereCommand3ReasoningParser, - model_output="""\ -<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> -[ - {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} -] -<|END_ACTION|>""", - expected_reasoning="i will call foo with query1", - expected_content="""\ -<|START_ACTION|> -[ - {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} -] -<|END_ACTION|>""", - ), - id="cmd3-single_tool_call", - ), - pytest.param( - ReasoningCase( - parser_cls=CohereCommand4ReasoningParser, - model_output="""\ -<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> -[ - {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} -] -<|END_ACTION|>""", - expected_reasoning="i will call foo with query1", - expected_content="""\ -<|START_ACTION|> -[ - {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} -] -<|END_ACTION|>""", - ), - id="cmd4-single_tool_call", - ), - pytest.param( - ReasoningCase( - parser_cls=CohereCommand3ReasoningParser, - model_output="""\ -<|START_THINKING|>This is a rainbow emoji: 🌈<|END_THINKING|> -<|START_RESPONSE|>foo bar<|END_RESPONSE|>""", - expected_reasoning="This is a rainbow emoji: 🌈", - expected_content="foo bar", - ), - id="cmd3-citations_with_emoji", - ), - pytest.param( - ReasoningCase( - parser_cls=CohereCommand4ReasoningParser, - model_output="""\ -<|START_THINKING|>This is a rainbow emoji: 🌈<|END_THINKING|> -<|START_RESPONSE|>foo bar<|END_RESPONSE|>""", - expected_reasoning="This is a rainbow emoji: 🌈", - expected_content="foo bar", - ), - id="cmd4-citations_with_emoji", - ), -] - - -class MockCohereTokenizer: - """Minimal byte-level stand-in for the Cohere tokenizer. - - ``encode``/``decode`` round-trip through UTF-8 bytes so splitting a - multi-byte character (e.g. an emoji) across "tokens" reproduces the - trailing U+FFFD buffering that real streaming exhibits. Cohere special - tokens map to distinct synthetic ids; everything else shares a default id. - ``adjust_request`` only needs the token ids, not real tokenization. - """ - - _SPECIAL_TOKEN_IDS = { - "<|START_THINKING|>": -1, - "<|END_THINKING|>": -2, - "<|CHATBOT_TOKEN|>": -3, - } - - def convert_tokens_to_ids(self, token: str) -> int: - return self._SPECIAL_TOKEN_IDS.get(token, 0) - - def get_vocab(self) -> dict[str, int]: - return {} - - def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: - return list(text.encode("utf-8")) - - def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str: - return bytes(ids).decode("utf-8", errors="replace") - - -@pytest.fixture(scope="module") -def tokenizer() -> MockCohereTokenizer: - return MockCohereTokenizer() - - -@pytest.fixture -def request_obj(): - return ChatCompletionRequest(messages=[], model="test-model") - - -REPLACEMENT_CHAR = "\ufffd" - - -def _token_deltas(tokenizer, text: str) -> list[str]: - """Progressively decode the token sequence and return per-step string - deltas. Incomplete multi-byte sequences (trailing U+FFFD) are buffered - until the next token completes them, matching real streaming behaviour.""" - ids = tokenizer.encode(text, add_special_tokens=False) - deltas: list[str] = [] - prev = "" - for i in range(1, len(ids) + 1): - current = tokenizer.decode(ids[:i], skip_special_tokens=False) - if current.endswith(REPLACEMENT_CHAR): - continue - delta = current[len(prev) :] - if delta: - deltas.append(delta) - prev = current - return deltas - - -@pytest.mark.parametrize("case", REASONING_CASES) -class TestExtractReasoning: - def test_nonstreaming(self, tokenizer, request_obj, case: ReasoningCase): - parser = case.parser_cls(tokenizer) - reasoning, content = parser.extract_reasoning(case.model_output, request_obj) - - assert reasoning == case.expected_reasoning - assert content == case.expected_content - - def test_streaming(self, tokenizer, case: ReasoningCase): - parser = case.parser_cls(tokenizer) - token_strings = _token_deltas(tokenizer, case.model_output) - - reasoning_parts: list[str] = [] - content_parts: list[str] = [] - tool_call_deltas: list[dict] = [] - - previous_text = "" - previous_token_ids: list[int] = [] - - for token_str in token_strings: - current_text = previous_text + token_str - current_token_ids = previous_token_ids + [0] - - delta = parser.extract_reasoning_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=token_str, - previous_token_ids=previous_token_ids, - current_token_ids=current_token_ids, - delta_token_ids=[0], - ) - if delta is not None: - if delta.reasoning is not None: - reasoning_parts.append(delta.reasoning) - if delta.content is not None: - content_parts.append(delta.content) - for tc in delta.tool_calls: - tool_call_deltas.append( - { - "id": tc.id, - "index": tc.index, - "name": tc.function.name if tc.function else None, - "arguments": ( - tc.function.arguments if tc.function else None - ), - } - ) - - previous_text = current_text - previous_token_ids = current_token_ids - - reasoning = "".join(reasoning_parts) if reasoning_parts else None - assert reasoning == case.expected_reasoning - - content = "".join(content_parts) if content_parts else None - assert content == case.expected_content - assert tool_call_deltas == [] - - -class TestIsReasoningEnd: - @pytest.mark.parametrize( - "parser_cls", - [CohereCommand3ReasoningParser, CohereCommand4ReasoningParser], - ids=["cmd3", "cmd4"], - ) - def test_is_reasoning_end(self, tokenizer, parser_cls): - parser = parser_cls(tokenizer) - start_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>") - end_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>") - chatbot_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>") - content_ids = [99, 100] - - # Generation-only tokens have no chatbot marker, so the whole sequence - # is considered. - assert parser.is_reasoning_end([end_id]) - assert parser.is_reasoning_end([start_id, *content_ids, end_id]) - assert not parser.is_reasoning_end([start_id, *content_ids]) - - # Full prompt/history tokens are scoped to the latest chatbot marker, - # so stray thinking tokens from the preamble or previous turns are ignored. - assert not parser.is_reasoning_end([start_id, end_id, chatbot_id, *content_ids]) - assert parser.is_reasoning_end( - [start_id, end_id, chatbot_id, start_id, *content_ids, end_id] - ) - - -@pytest.mark.parametrize( - "parser_cls", - [ - CohereCommand3ReasoningParser, - CohereCommand4ReasoningParser, - ], - ids=["cmd3", "cmd4"], -) -def test_count_reasoning_tokens(tokenizer, parser_cls): - parser = parser_cls(tokenizer) - start = parser.start_token_id - end = parser.end_token_id - - assert parser.count_reasoning_tokens([99, start, 11, 12, end, 100]) == 2 - assert parser.count_reasoning_tokens([end, start, 1, start, 2, end, 3, end]) == 3 - assert parser.count_reasoning_tokens([start, 1, 2]) == 2 - assert parser.count_reasoning_tokens([1, 2, end]) == 0 - - -SCHEMA_A = {"type": "object", "properties": {"a": {"type": "string"}}} -SCHEMA_B = {"type": "object", "properties": {"b": {"type": "number"}}} -GET_WEATHER_TOOL = { - "type": "function", - "function": { - "name": "get_weather", - "parameters": { - "type": "object", - "properties": {"city": {"type": "string"}}, - }, - }, -} -VALID_STRUCTURAL_TAG = { - "type": "structural_tag", - "format": { - "type": "triggered_tags", - "tags": [ - { - "begin": "", - "content": {"type": "any_text"}, - "end": "", - } - ], - "triggers": [""], - }, -} - - -def _model_config(arch: str) -> SimpleNamespace: - return SimpleNamespace( - architecture=arch, - architectures=[arch], - hf_text_config=SimpleNamespace(architectures=[arch]), - ) - - -def _make_chat_request(**kwargs) -> ChatCompletionRequest: - data = {"messages": [{"role": "user", "content": "hi"}], "model": "m"} - data.update(kwargs) - return ChatCompletionRequest.model_validate(data) - - -def _first_json_schema(tag_json: str) -> dict | None: - outer = json.loads(tag_json) - for t in (outer.get("format") or {}).get("tags") or []: - c = t.get("content") or {} - if c.get("type") == "json_schema": - js = c.get("json_schema") - return js if isinstance(js, dict) else None - return None - - -def _content_types(tag_json: str) -> set[str]: - outer = json.loads(tag_json) - out: set[str] = set() - for t in (outer.get("format") or {}).get("tags") or []: - ty = (t.get("content") or {}).get("type") - if isinstance(ty, str): - out.add(ty) - return out - - -@pytest.fixture(scope="module") -def parser(tokenizer: MockCohereTokenizer) -> CohereCommand4ReasoningParser: - """Parser configured with a supported Cohere architecture.""" - return CohereCommand4ReasoningParser( - tokenizer, - model_config=_model_config("Cohere2ForCausalLM"), - ) - - -@pytest.fixture(scope="module") -def parser_no_model_config( - tokenizer: MockCohereTokenizer, -) -> CohereCommand4ReasoningParser: - """Parser with no ``model_config`` (cannot resolve architecture).""" - return CohereCommand4ReasoningParser(tokenizer, model_config=None) - - -@pytest.fixture(scope="module") -def parser_unsupported_arch( - tokenizer: MockCohereTokenizer, -) -> CohereCommand4ReasoningParser: - """Parser configured with an architecture that has no structural tag style.""" - return CohereCommand4ReasoningParser( - tokenizer, - model_config=_model_config("LlamaForCausalLM"), - ) - - -class TestAdjustRequestPassthrough: - def test_structured_outputs_structural_tag_not_modified(self, parser) -> None: - tag = json.dumps(VALID_STRUCTURAL_TAG) - r = _make_chat_request(structured_outputs={"structural_tag": tag}) - o = parser.adjust_request(r) - assert o.structured_outputs.structural_tag == tag - - def test_response_format_structural_tag_short_circuit(self, parser) -> None: - # ``ChatCompletionRequest`` validates ``response_format`` as a union; - # bare ``{"type": "structural_tag"}`` is invalid (use pydantic model). - rf = StructuralTagResponseFormat( - type="structural_tag", - format=VALID_STRUCTURAL_TAG["format"], - ) - r = _make_chat_request(response_format=rf) - o = parser.adjust_request(r) - assert _response_format_type(o.response_format) == "structural_tag" - assert o.structured_outputs is None - - -class TestAdjustRequestNoOp: - def test_no_schema_no_tools(self, parser) -> None: - o = parser.adjust_request(_make_chat_request()) - assert o.structured_outputs is None - assert o.response_format is None - - def test_no_model_config(self, parser_no_model_config) -> None: - inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A) - r = _make_chat_request( - response_format=ResponseFormat(type="json_schema", json_schema=inner), - ) - o = parser_no_model_config.adjust_request(r) - assert o.response_format is not None - assert o.structured_outputs is None - - -class TestAdjustRequestUnsupportedArchitecture: - def test_json_schema_raises(self, parser_unsupported_arch) -> None: - inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A) - r = _make_chat_request( - response_format=ResponseFormat(type="json_schema", json_schema=inner), - ) - with pytest.raises(ValueError, match="does not support"): - parser_unsupported_arch.adjust_request(r) - - -class TestAdjustRequestFoldFromResponseFormat: - @pytest.mark.parametrize( - "response_format, expected_schema", - [ - pytest.param( - ResponseFormat( - type="json_schema", - json_schema=JsonSchemaResponseFormat( - name="n", json_schema=SCHEMA_A - ), - ), - SCHEMA_A, - id="json_schema_pydantic", - ), - pytest.param( - { - "type": "json_schema", - "json_schema": {"name": "n", "schema": SCHEMA_A}, - }, - SCHEMA_A, - id="json_schema_dict", - ), - pytest.param( - {"type": "json_object"}, - {"type": "object"}, - id="json_object", - ), - ], - ) - def test_response_format_cleared( - self, parser, response_format, expected_schema - ) -> None: - r = _make_chat_request(response_format=response_format) - o = parser.adjust_request(r) - assert o.response_format is None - assert ( - _first_json_schema(o.structured_outputs.structural_tag) == expected_schema - ) - - -class TestHasEffectiveTools: - @pytest.mark.parametrize( - "tools, expected", - [ - pytest.param(None, False, id="none"), - pytest.param([], False, id="empty_list"), - pytest.param(" ", False, id="blank_str"), - pytest.param( - [{"type": "function", "function": {"name": "f"}}], - True, - id="non_empty_list", - ), - pytest.param('{"x": 1}', True, id="non_empty_str"), - ], - ) - def test_has_effective_tools(self, tools, expected) -> None: - assert _has_effective_tools(tools) is expected - - def test_convert_schema_json_only_with_empty_tools_list(self) -> None: - tag = convert_schema_to_structural_tags( - schema=SCHEMA_B, - tools=[], - model_architecture="Cohere2ForCausalLM", - ) - assert tag is not None - assert _first_json_schema(tag) == SCHEMA_B - - -class TestAdjustRequestFoldFromStructuredOutputs: - @pytest.mark.parametrize( - "structured_outputs, expected_schema", - [ - pytest.param({"json": SCHEMA_B}, SCHEMA_B, id="json_dict"), - pytest.param({"json": json.dumps(SCHEMA_B)}, SCHEMA_B, id="json_string"), - pytest.param( - {"json_object": True}, {"type": "object"}, id="json_object_flag" - ), - pytest.param( - StructuredOutputsParams(json=SCHEMA_B), - SCHEMA_B, - id="structured_outputs_dataclass", - ), - pytest.param( - {"json": {"name": "n", "schema": SCHEMA_A}}, - SCHEMA_A, - id="openai_wrapper_dict_unwrapped", - ), - ], - ) - def test_structured_outputs_folded( - self, parser, structured_outputs, expected_schema - ) -> None: - o = parser.adjust_request( - _make_chat_request(structured_outputs=structured_outputs), - ) - assert ( - _first_json_schema(o.structured_outputs.structural_tag) == expected_schema - ) - - def test_responses_request_default_empty_tools(self, parser) -> None: - """``ResponsesRequest.tools`` defaults to ``[]``, not ``None``.""" - r = ResponsesRequest.model_validate( - { - "input": "hi", - "model": "m", - "structured_outputs": {"json": SCHEMA_B}, - } - ) - assert r.tools == [] - o = parser.adjust_request(r) - assert _first_json_schema(o.structured_outputs.structural_tag) == SCHEMA_B - - def test_json_userdict_mapping_unwrapped(self) -> None: - inner = {"type": "object", "properties": {"u": {"type": "number"}}} - so = StructuredOutputsParams(json=UserDict(inner)) - assert _schema_dict_from_structured_outputs(so) == inner - - @pytest.mark.parametrize( - "json_value, match", - [ - pytest.param("{not json}", "valid JSON", id="invalid_json_string"), - pytest.param( - json.dumps(["a", "b"]), "JSON object", id="non_object_json_string" - ), - pytest.param(" ", "empty", id="empty_json_string"), - ], - ) - def test_structured_outputs_json_string_raises( - self, parser, json_value, match - ) -> None: - with pytest.raises(ValueError, match=match): - parser.adjust_request( - _make_chat_request(structured_outputs={"json": json_value}), - ) - - @pytest.mark.parametrize( - "construct", - [ - pytest.param( - lambda: _make_chat_request(structured_outputs={"json": [1, 2, 3]}), - id="chat_completion_request", - ), - pytest.param( - lambda: StructuredOutputsParams(json=[1, 2, 3]), # type: ignore[arg-type] - id="structured_outputs_params", - ), - ], - ) - def test_json_wrong_type_raises(self, construct) -> None: - """Non-str / non-dict ``json`` fails at Pydantic validation.""" - with pytest.raises(ValidationError): - construct() - - -class TestAdjustRequestPrecedence: - def test_response_format_over_structured_outputs_json(self, parser) -> None: - s_rf = {"type": "object", "properties": {"rf": {"type": "string"}}} - s_so = {"type": "object", "properties": {"so": {"type": "number"}}} - inner = JsonSchemaResponseFormat(name="n", json_schema=s_rf) - r = _make_chat_request( - response_format=ResponseFormat(type="json_schema", json_schema=inner), - structured_outputs={"json": s_so}, - ) - o = parser.adjust_request(r) - assert _first_json_schema(o.structured_outputs.structural_tag) == s_rf - - -class TestAdjustRequestTextPlusStructuredOutputs: - def test_text_response_format_preserved(self, parser) -> None: - sch = {"type": "object", "properties": {"k": {"type": "string"}}} - r = _make_chat_request( - response_format=ResponseFormat(type="text"), - structured_outputs={"json": sch}, - ) - o = parser.adjust_request(r) - assert o.response_format is not None - assert o.response_format.type == "text" - assert _first_json_schema(o.structured_outputs.structural_tag) == sch - - -class TestAdjustRequestTools: - def test_tools_only_command_a_grammar(self, parser) -> None: - o = parser.adjust_request( - _make_chat_request(tools=[GET_WEATHER_TOOL], tool_choice="auto"), - ) - assert "grammar" in _content_types(o.structured_outputs.structural_tag) - - def test_tools_plus_json_schema_both_kinds(self, parser) -> None: - inner = JsonSchemaResponseFormat( - name="n", - json_schema={"type": "object", "properties": {"r": {"type": "string"}}}, - ) - r = _make_chat_request( - response_format=ResponseFormat(type="json_schema", json_schema=inner), - tools=[GET_WEATHER_TOOL], - tool_choice="auto", - ) - o = parser.adjust_request(r) - types = _content_types(o.structured_outputs.structural_tag) - assert "grammar" in types - assert "json_schema" in types - - -class TestMelodySourceResolution: - """Pin the parser-side source resolution (see - ``_melody_sources_to_vllm``). The parser receives melody's numeric - ``(tool_call_index, tool_result_indices)`` addressing and resolves - it against a ``position_to_source`` map handed in by - ``CohereServingChatV2._apply_cohere_template_kwargs`` -- the same - map that would otherwise live in the serving layer's resolver. - """ - - @staticmethod - def _fake_melody_source(bucket: int, indices: list[int]) -> Any: - return SimpleNamespace(tool_call_index=bucket, tool_result_indices=indices) - - def test_multi_index_source_fans_out(self): - from vllm.entrypoints.cohere.cohere_chat_message import CitationSource - - position_map: dict[tuple[int, int], CitationSource] = { - (0, 0): CitationSource(type="document", id="d0", document={"id": "d0"}), - (0, 1): CitationSource(type="document", id="d1", document={"id": "d1"}), - } - raw = [self._fake_melody_source(0, [0, 1])] - out = _melody_sources_to_vllm(raw, position_map) - assert [s.id for s in out] == ["d0", "d1"] - # Verify type / payload were plumbed through, not just the id. - assert out[0].type == "document" - assert out[0].document == {"id": "d0"} - - def test_unresolvable_position_skipped(self): - from vllm.entrypoints.cohere.cohere_chat_message import CitationSource - - position_map: dict[tuple[int, int], CitationSource] = { - (0, 0): CitationSource(type="document", id="d0"), - } - raw = [self._fake_melody_source(9, [0])] - assert _melody_sources_to_vllm(raw, position_map) == [] - - def test_missing_position_map_drops_all_sources(self): - # A parser instance without a position map (parser wired - # outside of ``CohereServingChatV2``) can't attribute anything, - # so every source is dropped. Callers downstream will see the - # citation with empty ``sources`` and drop it entirely. - raw = [self._fake_melody_source(0, [0])] - assert _melody_sources_to_vllm(raw, None) == [] - - def test_citations_pass_through_is_thinking_tag(self): - from vllm.entrypoints.cohere.cohere_chat_message import CitationSource - - position_map: dict[tuple[int, int], CitationSource] = { - (0, 0): CitationSource(type="document", id="d0"), - } - raw = [ - SimpleNamespace( - start_index=0, - end_index=5, - text="hello", - is_thinking=True, - sources=[self._fake_melody_source(0, [0])], - ) - ] - out = _melody_citations_to_vllm(raw, position_map) - assert out is not None - assert out[0].type == "THINKING_CONTENT" - assert out[0].sources[0].id == "d0" - - -class TestParserStreamingEndToEnd: - """End-to-end streaming shape: raw model output text is tokenized - with the mock byte-level tokenizer and fed one token at a time - through :meth:`CohereCommand3ReasoningParser.extract_reasoning_streaming`. - We collect the ordered sequence of ``DeltaMessage``s the parser - hands back -- i.e. what the client would receive on the wire -- - and pin the exact shape. - - Combines the three flavors that reach ``delta``: a thinking block - (populates ``reasoning``), a text block (populates ``content``), - and a resolved citation (populates ``citations``) referencing - ``(start, end)`` offsets into the already-streamed content. - """ - - _MODEL_OUTPUT = ( - "<|START_THINKING|>Let me check.<|END_THINKING|>" - "<|START_RESPONSE|>The capital is Paris.<|END_RESPONSE|>" - ) - - @staticmethod - def _drive_parser( - parser: CohereCommand3ReasoningParser, text: str, tok: MockCohereTokenizer - ) -> list[Any]: - """Feed ``text`` through the parser one token at a time - (byte-level) and return the ordered list of non-``None`` - ``DeltaMessage``s. - """ - deltas: list[Any] = [] - previous_text = "" - previous_token_ids: list[int] = [] - for token_str in _token_deltas(tok, text): - current_text = previous_text + token_str - current_token_ids = previous_token_ids + [0] - delta = parser.extract_reasoning_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=token_str, - previous_token_ids=previous_token_ids, - current_token_ids=current_token_ids, - delta_token_ids=[0], - ) - if delta is not None: - deltas.append(delta) - previous_text = current_text - previous_token_ids = current_token_ids - return deltas - - def test_full_wire_stream_shape(self, tokenizer): - from vllm.entrypoints.cohere.cohere_chat_message import CitationSource - from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY - - # Mimic what ``CohereServingChatV2._apply_cohere_template_kwargs`` - # installs on a request whose tool_call_index=0 tool result at - # index 1 is the "France" document. - position_to_source = { - (0, 1): CitationSource( - type="document", - id="doc-paris", - document={"id": "doc-paris", "title": "France"}, - ), - } - - parser = CohereCommand3ReasoningParser( - tokenizer, - chat_template_kwargs={POSITION_TO_SOURCE_KEY: position_to_source}, - ) - - deltas = self._drive_parser(parser, self._MODEL_OUTPUT, tokenizer) - - # -- 1. Reasoning and content stream separately and reach - # the wire only as their real payload. Framing tokens - # (``<|START_THINKING|>`` etc.) are consumed by the parser's - # state machine and never produce a delta. - reasoning_stream = [d.reasoning for d in deltas if d.reasoning is not None] - content_stream = [d.content for d in deltas if d.content is not None] - citation_deltas = [d for d in deltas if getattr(d, "citations", None)] - - assert "".join(reasoning_stream) == "Let me check." - assert "".join(content_stream) == "The capital is Paris." - - # -- 2. Exactly one citation reaches the wire, resolved to the - # source from ``position_to_source`` (not the raw - # ``(bucket, idx)`` coordinates). - assert len(citation_deltas) == 1 - (cite_delta,) = citation_deltas - assert cite_delta.citations is not None - assert len(cite_delta.citations) == 1 - cite = cite_delta.citations[0] - assert cite.text == "Paris" - # Char offsets are into the accumulated ``content``: "The - # capital is " is 15 chars, "Paris" is 5. - assert cite.start == 15 - assert cite.end == 20 - assert len(cite.sources) == 1 - assert cite.sources[0].id == "doc-paris" - assert cite.sources[0].type == "document" - - # -- 3. The citation's ``(start, end)`` offsets point into - # bytes that were already emitted as content-deltas earlier - # in the stream -- i.e. the citation is anchored to real - # already-shipped content. - cite_delta_idx = deltas.index(cite_delta) - earlier_content = "".join( - d.content for d in deltas[:cite_delta_idx] if d.content is not None - ) - assert earlier_content[cite.start : cite.end] == "Paris" - - # Everything after the citation delta is the trailing period. - later_content = "".join( - d.content for d in deltas[cite_delta_idx + 1 :] if d.content is not None - ) - assert later_content == "." - - # -- 4. Reasoning and content are never mixed on the same - # delta: the parser flips modes on the ``<|END_THINKING|>`` - # boundary and never emits a delta with both fields set. - for d in deltas: - assert not (d.reasoning is not None and d.content is not None) diff --git a/tests/tool_parsers/test_cohere_command_tool_parser.py b/tests/tool_parsers/test_cohere_command_tool_parser.py deleted file mode 100644 index f102474b62..0000000000 --- a/tests/tool_parsers/test_cohere_command_tool_parser.py +++ /dev/null @@ -1,510 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import json -from dataclasses import dataclass, field -from typing import Any - -import pytest - -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest -from vllm.tool_parsers.cohere_command_tool_parser import ( - CohereCommand3ToolParser, - CohereCommand4ToolParser, -) - - -@dataclass -class ExpectedToolCall: - id: str - name: str - arguments: dict - - -@dataclass -class ToolCallCase: - parser_cls: Any - model_output: str - expected_tool_calls: list[ExpectedToolCall] = field(default_factory=list) - expected_reasoning: str | None = None - expected_content: str | None = None - - -TOOL_CALL_CASES = [ - pytest.param( - ToolCallCase( - parser_cls=CohereCommand3ToolParser, - model_output="""\ -<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> -[ - {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} -] -<|END_ACTION|>""", - expected_tool_calls=[ - ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}), - ], - expected_reasoning="i will call foo with query1", - ), - id="cmd3-single_tool_call", - ), - pytest.param( - ToolCallCase( - parser_cls=CohereCommand4ToolParser, - model_output="""\ -<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> -[ - {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} -] -<|END_ACTION|>""", - expected_tool_calls=[ - ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}), - ], - expected_reasoning="i will call foo with query1", - ), - id="cmd4-single_tool_call", - ), - pytest.param( - ToolCallCase( - parser_cls=CohereCommand3ToolParser, - model_output="""\ -<|START_THINKING|>This is a rainbow emoji: 🌈<|END_THINKING|> -<|START_RESPONSE|>foo bar<|END_RESPONSE|>""", - expected_reasoning="This is a rainbow emoji: 🌈", - expected_content="foo bar", - ), - id="cmd3-citations_no_tool_calls", - ), - pytest.param( - ToolCallCase( - parser_cls=CohereCommand4ToolParser, - model_output="""\ -<|START_THINKING|>This is a rainbow emoji: 🌈<|END_THINKING|> -<|START_RESPONSE|>foo bar<|END_RESPONSE|>""", - expected_reasoning="This is a rainbow emoji: 🌈", - expected_content="foo bar", - ), - id="cmd4-citations_no_tool_calls", - ), - pytest.param( - ToolCallCase( - parser_cls=CohereCommand3ToolParser, - model_output="""\ -<|START_THINKING|>first I think about foo<|END_THINKING|><|START_ACTION|> -[ - {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}, - {"tool_call_id": "1", "tool_name": "bar", "parameters": {"x": 42}} -] -<|END_ACTION|>""", - expected_tool_calls=[ - ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}), - ExpectedToolCall(id="1", name="bar", arguments={"x": 42}), - ], - expected_reasoning="first I think about foo", - ), - id="cmd3-multiple_tool_calls", - ), - pytest.param( - ToolCallCase( - parser_cls=CohereCommand4ToolParser, - model_output="""\ -<|START_THINKING|>first I think about foo<|END_THINKING|><|START_ACTION|> -[ - {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}, - {"tool_call_id": "1", "tool_name": "bar", "parameters": {"x": 42}} -] -<|END_ACTION|>""", - expected_tool_calls=[ - ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}), - ExpectedToolCall(id="1", name="bar", arguments={"x": 42}), - ], - expected_reasoning="first I think about foo", - ), - id="cmd4-multiple_tool_calls", - ), - pytest.param( - ToolCallCase( - parser_cls=CohereCommand3ToolParser, - model_output="""\ -<|START_THINKING|>just think, no response<|END_THINKING|>""", - expected_reasoning="just think, no response", - ), - id="cmd3-reasoning_only", - ), - pytest.param( - ToolCallCase( - parser_cls=CohereCommand4ToolParser, - model_output="""\ -<|START_THINKING|>just think, no response<|END_THINKING|>""", - expected_reasoning="just think, no response", - ), - id="cmd4-reasoning_only", - ), -] - - -class MockCohereTokenizer: - """Minimal byte-level stand-in for the Cohere tokenizer.""" - - def get_vocab(self) -> dict[str, int]: - return {} - - def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: - return list(text.encode("utf-8")) - - def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str: - return bytes(ids).decode("utf-8", errors="replace") - - -@pytest.fixture(scope="module") -def tokenizer() -> MockCohereTokenizer: - return MockCohereTokenizer() - - -@pytest.fixture -def request_obj() -> ChatCompletionRequest: - return ChatCompletionRequest(messages=[], model="test-model") - - -REPLACEMENT_CHAR = "\ufffd" - - -def _token_deltas(tokenizer: MockCohereTokenizer, text: str) -> list[str]: - """Decode per-token string deltas, buffering incomplete multi-byte chars.""" - ids = tokenizer.encode(text, add_special_tokens=False) - deltas: list[str] = [] - prev = "" - for i in range(1, len(ids) + 1): - current = tokenizer.decode(ids[:i], skip_special_tokens=False) - if current.endswith(REPLACEMENT_CHAR): - continue - delta = current[len(prev) :] - if delta: - deltas.append(delta) - prev = current - return deltas - - -@dataclass -class StreamingResult: - tool_calls: dict[int, dict] - reasoning: str | None - content: str | None - - -def _run_streaming_over_deltas( - parser, - deltas: list[str], - request_obj: ChatCompletionRequest, -) -> StreamingResult: - accumulated: dict[int, dict] = {} - reasoning_parts: list[str] = [] - content_parts: list[str] = [] - previous_text = "" - - for token_str in deltas: - current_text = previous_text + token_str - delta = parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=token_str, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=request_obj, - ) - if delta is not None: - if delta.reasoning is not None: - reasoning_parts.append(delta.reasoning) - if delta.content is not None: - content_parts.append(delta.content) - for tc in delta.tool_calls: - idx = tc.index - if idx not in accumulated: - accumulated[idx] = {"id": "", "name": "", "arguments": ""} - if tc.id: - accumulated[idx]["id"] = tc.id - if tc.function and tc.function.name: - accumulated[idx]["name"] = tc.function.name - if tc.function and tc.function.arguments: - accumulated[idx]["arguments"] += tc.function.arguments - previous_text = current_text - - return StreamingResult( - tool_calls=accumulated, - reasoning="".join(reasoning_parts) if reasoning_parts else None, - content="".join(content_parts) if content_parts else None, - ) - - -def _run_streaming( - parser, - tokenizer: MockCohereTokenizer, - model_output: str, - request_obj: ChatCompletionRequest, -) -> StreamingResult: - return _run_streaming_over_deltas( - parser, - _token_deltas(tokenizer, model_output), - request_obj, - ) - - -@pytest.mark.parametrize("case", TOOL_CALL_CASES) -class TestExtractToolCalls: - def test_streaming( - self, - tokenizer: MockCohereTokenizer, - request_obj: ChatCompletionRequest, - case: ToolCallCase, - ): - parser = case.parser_cls(tokenizer) - streamed = _run_streaming(parser, tokenizer, case.model_output, request_obj) - - assert len(streamed.tool_calls) == len(case.expected_tool_calls) - - for i, expected_tc in enumerate(case.expected_tool_calls): - tc = streamed.tool_calls[i] - assert tc["id"] == expected_tc.id - assert tc["name"] == expected_tc.name - assert json.loads(tc["arguments"]) == expected_tc.arguments - - def test_streaming_reasoning( - self, - tokenizer: MockCohereTokenizer, - request_obj: ChatCompletionRequest, - case: ToolCallCase, - ): - parser = case.parser_cls(tokenizer) - streamed = _run_streaming(parser, tokenizer, case.model_output, request_obj) - - assert streamed.reasoning == case.expected_reasoning - - def test_streaming_content( - self, - tokenizer: MockCohereTokenizer, - request_obj: ChatCompletionRequest, - case: ToolCallCase, - ): - parser = case.parser_cls(tokenizer) - streamed = _run_streaming(parser, tokenizer, case.model_output, request_obj) - - assert streamed.content == case.expected_content - - def test_nonstreaming( - self, - request_obj: ChatCompletionRequest, - tokenizer: MockCohereTokenizer, - case: ToolCallCase, - ): - parser = case.parser_cls(tokenizer) - result = parser.extract_tool_calls(case.model_output, request_obj) - - assert result.tools_called == (len(case.expected_tool_calls) > 0) - assert len(result.tool_calls) == len(case.expected_tool_calls) - - for actual_tc, expected_tc in zip(result.tool_calls, case.expected_tool_calls): - assert actual_tc.type == "function" - assert actual_tc.function.name == expected_tc.name - assert json.loads(actual_tc.function.arguments) == expected_tc.arguments - - def test_streaming_nonstreaming_agree( - self, - tokenizer: MockCohereTokenizer, - request_obj: ChatCompletionRequest, - case: ToolCallCase, - ): - parser_streaming = case.parser_cls(tokenizer) - parser_nonstreaming = case.parser_cls(tokenizer) - - streamed = _run_streaming( - parser_streaming, - tokenizer, - case.model_output, - request_obj, - ) - result = parser_nonstreaming.extract_tool_calls( - case.model_output, - request_obj, - ) - - assert len(streamed.tool_calls) == len(result.tool_calls) - - for i, actual_tc in enumerate(result.tool_calls): - assert streamed.tool_calls[i]["name"] == actual_tc.function.name - assert json.loads(streamed.tool_calls[i]["arguments"]) == json.loads( - actual_tc.function.arguments - ) - - -SPECIAL_TOKEN_MARKERS = ( - "<|START_THINKING|>", - "<|END_THINKING|>", - "<|START_RESPONSE|>", - "<|END_RESPONSE|>", - "<|START_ACTION|>", - "<|END_ACTION|>", - "<|START_TEXT|>", - "<|END_TEXT|>", -) - - -def _multi_token_deltas( - tokenizer: MockCohereTokenizer, - text: str, - chunk_size: int, -) -> list[str]: - ids = tokenizer.encode(text, add_special_tokens=False) - deltas: list[str] = [] - prev = "" - i = 0 - while i < len(ids): - end = min(len(ids), i + chunk_size) - current = tokenizer.decode(ids[:end], skip_special_tokens=False) - i = end - if current.endswith(REPLACEMENT_CHAR): - continue - delta = current[len(prev) :] - if delta: - deltas.append(delta) - prev = current - return deltas - - -class TestSpeculativeDecodingMultiTokenDelta: - MODEL_OUTPUT = ( - "<|START_THINKING|> i will call foo with query1<|END_THINKING|>" - "<|START_ACTION|>\n" - '[\n {"tool_call_id": "0", "tool_name": "foo", ' - '"parameters": {"query": "query1"}}\n]\n' - "<|END_ACTION|>" - ) - - @pytest.mark.parametrize( - "parser_cls", - [CohereCommand3ToolParser, CohereCommand4ToolParser], - ids=["cmd3", "cmd4"], - ) - @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6]) - def test_no_special_token_leak_in_streaming_deltas( - self, - tokenizer: MockCohereTokenizer, - request_obj: ChatCompletionRequest, - parser_cls, - chunk_size: int, - ): - parser = parser_cls(tokenizer) - chunked = _multi_token_deltas(tokenizer, self.MODEL_OUTPUT, chunk_size) - - previous_text = "" - for token_str in chunked: - current_text = previous_text + token_str - delta = parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=token_str, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=request_obj, - ) - previous_text = current_text - if delta is None: - continue - - fields: list[tuple[str, str | None]] = [ - ("reasoning", delta.reasoning), - ("content", delta.content), - ] - for tc in delta.tool_calls or []: - if tc.function: - fields.append(("tool_call.name", tc.function.name)) - fields.append(("tool_call.arguments", tc.function.arguments)) - - for marker in SPECIAL_TOKEN_MARKERS: - for field_name, value in fields: - assert value is None or marker not in value, ( - f"special token {marker!r} leaked into {field_name} " - f"with chunk_size={chunk_size} delta={delta!r}" - ) - - @pytest.mark.parametrize( - "parser_cls", - [CohereCommand3ToolParser, CohereCommand4ToolParser], - ids=["cmd3", "cmd4"], - ) - @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6]) - def test_multi_token_chunks_still_produce_correct_tool_call( - self, - tokenizer: MockCohereTokenizer, - request_obj: ChatCompletionRequest, - parser_cls, - chunk_size: int, - ): - parser = parser_cls(tokenizer) - chunked = _multi_token_deltas(tokenizer, self.MODEL_OUTPUT, chunk_size) - streamed = _run_streaming_over_deltas(parser, chunked, request_obj) - - assert len(streamed.tool_calls) == 1 - tc = streamed.tool_calls[0] - assert tc["id"] == "0" - assert tc["name"] == "foo" - assert json.loads(tc["arguments"]) == {"query": "query1"} - - -class TestStreamingDeltaShape: - @pytest.mark.parametrize( - "parser_cls", - [CohereCommand3ToolParser, CohereCommand4ToolParser], - ids=["cmd3", "cmd4"], - ) - def test_reasoning_and_tool_calls_are_separate_deltas( - self, - tokenizer: MockCohereTokenizer, - request_obj: ChatCompletionRequest, - parser_cls, - ): - parser = parser_cls(tokenizer) - model_output = ( - "<|START_THINKING|> i will call foo with query1<|END_THINKING|>" - "<|START_ACTION|>\n" - '[\n {"tool_call_id": "0", "tool_name": "foo", ' - '"parameters": {"query": "query1"}}\n]\n' - "<|END_ACTION|>" - ) - - token_strings = _token_deltas(tokenizer, model_output) - previous_text = "" - saw_reasoning = False - saw_tool_call = False - - for token_str in token_strings: - current_text = previous_text + token_str - delta = parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=token_str, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=request_obj, - ) - if delta is not None: - populated = [ - delta.content is not None, - delta.reasoning is not None, - bool(delta.tool_calls), - ] - assert sum(populated) == 1, ( - "A single streaming delta must carry exactly one of " - f"content/reasoning/tool_calls, got {delta!r}" - ) - if delta.reasoning is not None: - saw_reasoning = True - if delta.tool_calls: - saw_tool_call = True - previous_text = current_text - - assert saw_reasoning, "expected at least one reasoning delta" - assert saw_tool_call, "expected at least one tool-call delta" diff --git a/vllm/entrypoints/cohere/cohere_chat_message.py b/vllm/entrypoints/cohere/cohere_chat_message.py index 042856dea6..85925d0f18 100644 --- a/vllm/entrypoints/cohere/cohere_chat_message.py +++ b/vllm/entrypoints/cohere/cohere_chat_message.py @@ -7,8 +7,8 @@ This module keeps the OpenAI chat completion protocol classes It provides two things: * :class:`Citation` / :class:`CitationSource`: the vLLM-internal citation - representation produced by the Cohere reasoning parser - (:mod:`vllm.reasoning.cohere_command_reasoning_parser`) and consumed by + representation produced by the Cohere parser + (:mod:`vllm.parser.cohere_command`) and consumed by :mod:`vllm.entrypoints.cohere.serving`. This is *not* the on-wire Cohere SDK shape -- that conversion happens in the serving layer. * :class:`CohereChatMessage` / :class:`CohereDeltaMessage`: :class:`ChatMessage` @@ -42,9 +42,9 @@ class CitationSource(OpenAIBaseModel): document/tool-output identifier; ``document`` and ``tool_output`` carry the original payload that produced the citation. - Sources are fully resolved by the reasoning parser at emit time + Sources are fully resolved by the parser at emit time (see :func:`_melody_sources_to_vllm` in - :mod:`vllm.reasoning.cohere_command_reasoning_parser`) using the + :mod:`vllm.parser.cohere_command`) using the ``POSITION_TO_SOURCE_KEY`` map forwarded through ``chat_template_kwargs`` by :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._apply_cohere_template_kwargs`. @@ -121,8 +121,7 @@ class CohereChatMessage(ChatMessage): class CohereDeltaMessage(DeltaMessage): """:class:`DeltaMessage` extension carrying grounding citations for streaming. - Emitted by - :class:`vllm.reasoning.cohere_command_reasoning_parser.BaseCohereCommandReasoningParser` + Emitted by :class:`vllm.parser.cohere_command.CohereCommandParser` on delta events whose payload includes citations. Non-Cohere parsers return plain :class:`DeltaMessage`, so their streamed shape is unchanged. diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 8ec4701907..7885586452 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -714,8 +714,8 @@ class CohereServingChatV2(OpenAIServingChat): # emits, so citation deltas already carry SDK-shape sources by # the time they reach the streaming loop. See # ``_build_position_to_source`` for the numbering rule and - # ``_melody_sources_to_vllm`` in the reasoning parser for the - # consumer. + # ``_melody_sources_to_vllm`` in ``vllm.parser.cohere_command`` for + # the consumer. position_to_source = cls._build_position_to_source(request) if position_to_source: kwargs.setdefault(POSITION_TO_SOURCE_KEY, position_to_source) @@ -1109,12 +1109,11 @@ class CohereServingChatV2(OpenAIServingChat): *, parser: Parser | None, ) -> ChatMessage: - """Copy grounding citations off the reasoning parser onto the message. + """Copy grounding citations off the unified parser onto the message. - The Cohere reasoning parser - (:mod:`vllm.reasoning.cohere_command_reasoning_parser`) caches the - citations produced by its most recent unary ``extract_reasoning`` - call on ``last_unary_citations``. We surface them here on + :class:`vllm.parser.cohere_command.CohereCommandParser` caches the + citations produced by its most recent unary ``parse`` call on + ``last_unary_citations``. We surface them here on :class:`CohereChatMessage` so downstream response conversion can pick them up without the base :class:`OpenAIServingChat` having to know about citations. @@ -1128,11 +1127,7 @@ class CohereServingChatV2(OpenAIServingChat): # ``model_dump`` via the extras bucket, so the wire is correct # either way. message = cast(CohereChatMessage, message) - citations = getattr( - getattr(parser, "reasoning_parser", None), - "last_unary_citations", - None, - ) + citations = getattr(parser, "last_unary_citations", None) if citations: message.citations = citations return message @@ -1143,10 +1138,10 @@ class CohereServingChatV2(OpenAIServingChat): :class:`CohereChatMessage` carries a ``citations: list[vllm...Citation] | None`` field populated by :meth:`_finalize_response_message` (which reads it off the - reasoning parser's ``last_unary_citations`` cache). Sources are + unified parser's ``last_unary_citations`` cache). Sources are already fully resolved by the parser (see :func:`_melody_sources_to_vllm` in - :mod:`vllm.reasoning.cohere_command_reasoning_parser`) using + :mod:`vllm.parser.cohere_command`) using the position map forwarded via ``chat_template_kwargs``. This method's remaining job is: diff --git a/vllm/parser/cohere_command.py b/vllm/parser/cohere_command.py new file mode 100644 index 0000000000..cf8d79b10e --- /dev/null +++ b/vllm/parser/cohere_command.py @@ -0,0 +1,687 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cohere Command parser. + +Uses ``cohere_melody`` as a stateful parser filter. +""" + +from __future__ import annotations + +import json +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, NamedTuple, TypedDict, TypeGuard + +import regex as re +import xgrammar as xgr + +try: + from cohere_melody import ( + AccumulatedToolCall, + FilterAggregatedResult, + PyFilter, + PyFilterOptions, + ) +except ImportError as e: + raise ImportError( + "The Cohere Command parser requires the `cohere_melody` package. " + "Install it with:\n pip install 'cohere-melody>=0.14.0'" + ) from e + +from vllm.entrypoints.cohere.cohere_chat_message import ( + Citation, + CitationSource, + CohereDeltaMessage, +) +from vllm.entrypoints.generate.base.protocol import ( + AnyResponseFormat, + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, + FunctionCall, +) +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.parser.abstract_parser import DelegatingParser +from vllm.reasoning.cohere_command_reasoning_parser import ( + BaseCohereCommandReasoningParser, +) +from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY +from vllm.sampling_params import StructuredOutputsParams +from vllm.tool_parsers.cohere_command_tool_parser import BaseCohereCommandToolParser + +if TYPE_CHECKING: + # Declared in melody's stub only; not exported at runtime. + from cohere_melody import FilterCitation, Source + + from vllm.config import ModelConfig + from vllm.entrypoints.openai.responses.protocol import ResponsesRequest + from vllm.tokenizers import TokenizerLike + from vllm.tool_parsers.utils import Tool + +PositionToSource = Mapping[tuple[int, int], CitationSource] + + +class CohereTagRegistry(NamedTuple): + """A single ``structural_tag`` trigger / end pair (``begin`` uses ``trigger``).""" + + trigger: str + end: str + + +class CohereTagStyle(NamedTuple): + """The structural tags style for a given model architecture. + + ``json_tags`` lists every JSON-schema wrapper the model may emit (MOE uses + both response and text delimiters). ``tools`` is the tool-call wrapper. + """ + + json_tags: tuple[CohereTagRegistry, ...] + tools: CohereTagRegistry + + +class CohereNormalizedTool(TypedDict): + """A tool definition normalized to the shape ``collect_tool_schema`` expects. + + ``parameters`` is a JSON Schema object (possibly empty) describing the tool's + call signature. + """ + + name: str + parameters: dict[str, Any] + + +COMMAND_A_TOOLS_TAG = CohereTagRegistry( + trigger="<|START_ACTION|>", + end="<|END_ACTION|>", +) +COMMAND_A_JSON_TAG = CohereTagRegistry( + trigger="<|START_RESPONSE|>", + end="<|END_RESPONSE|>", +) +COMMAND_A_PLUS_JSON_TAG = CohereTagRegistry( + trigger="<|START_TEXT|>", + end="<|END_TEXT|>", +) + +MODEL_TO_TAG_STYLE: dict[str, CohereTagStyle] = { + "Cohere2ForCausalLM": CohereTagStyle( + json_tags=(COMMAND_A_JSON_TAG,), + tools=COMMAND_A_TOOLS_TAG, + ), + "Cohere2VisionForConditionalGeneration": CohereTagStyle( + json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG), + tools=COMMAND_A_TOOLS_TAG, + ), + "Cohere2MoeForCausalLM": CohereTagStyle( + json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG), + tools=COMMAND_A_TOOLS_TAG, + ), +} + + +def collect_tool_schema(tool_schema: list[CohereNormalizedTool]) -> str: + """Build an xgrammar EBNF grammar that matches a JSON array of tool calls. + + The grammar shape is architecture-independent; callers are responsible for + wrapping it in the correct structural tag (see ``CohereTagStyle.tools``). + """ + tool_dictionary: dict[str, str] = {} + for tool in tool_schema: + tool_name = tool["name"] + tool_parameters = json.dumps(tool["parameters"]) + json_schema = f"""{{ + "type": "object", + "properties": {{ + "tool_call_id": {{ + "type": "string", + "pattern": "^[0-9]+$" + }}, + "tool_name": {{ + "type": "string", + "const": "{tool_name}" + }}, + "parameters": {tool_parameters} + }} + }}""" + tool_grammar = str(xgr.Grammar.from_json_schema(json_schema)) + for match in re.findall(r"\b(\w+)\s*::=", tool_grammar): + tool_grammar = re.sub( + rf"\b{re.escape(match)}\b", tool_name + match, tool_grammar + ) + tool_dictionary[tool_name] = f"{tool_name} ::= {tool_name}root\n{tool_grammar}" + # Emitted grammar shape: + # root ::= tools + # tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws + # ws ::= (" " | "\t" | "\n")* + # tool ::= | | ... (one alternative per input) + # ::= root (per-tool xgrammar rules) + # root ::= ... (from xgr.Grammar.from_json_schema) + tool_alternatives = "tool ::= " + " | ".join(tool_dictionary.keys()) + tool_rules = "\n ".join(tool_dictionary.values()) + grammar = f"""root ::= tools + tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws + ws ::= (" " | "\\t" | "\\n")* + {tool_alternatives} + {tool_rules} + """ + return grammar + + +def _tool_definitions_to_schema_list( + tools: str | list[Any], +) -> list[CohereNormalizedTool]: + """ + Build the list of ``CohereNormalizedTool`` dicts expected by + ``collect_tool_schema``. + + Accepts: + - JSON string + - list of dicts with top-level ``name`` / ``parameters`` + - list of Chat Completions-style ``{"type": "function", "function": {...}}`` + - list of Pydantic models with ``model_dump()`` + """ + if isinstance(tools, str): + try: + parsed = json.loads(tools) + except json.JSONDecodeError: + return [] + if not isinstance(parsed, list): + return [] + else: + parsed = list(tools) + + out: list[CohereNormalizedTool] = [] + for raw in parsed: + t = raw.model_dump() if hasattr(raw, "model_dump") else raw + if not isinstance(t, dict): + continue + # Unwrap Chat Completions' ``{"type": "function", "function": {...}}`` + # shape; otherwise take the dict as-is. + if t.get("type") == "function" and isinstance(t.get("function"), dict): + t = t["function"] + name = t.get("name") + if not isinstance(name, str): + continue + params = t.get("parameters") + out.append( + CohereNormalizedTool( + name=name, + parameters=params if isinstance(params, dict) else {}, + ) + ) + return out + + +def _has_effective_tools( + tools: str | list[Any] | None, +) -> TypeGuard[str | list[Any]]: + """ + True when ``tools`` contains at least one tool definition to convert. + + ``ResponsesRequest`` defaults ``tools`` to ``[]``; ``ChatCompletionRequest`` + uses ``None``. Both mean "no tools" here. Strings (e.g. a JSON blob) are + treated as effective only when non-blank. + """ + if tools is None: + return False + if isinstance(tools, str): + return bool(tools.strip()) + return len(tools) > 0 + + +# Builder: produces vLLM response_format in xgrammar's canonical format. +# See xgrammar docs: type "structural_tag" with "format" = triggered_tags +# and tag content type = json_schema | grammar. +def convert_schema_to_structural_tags( + schema: dict | None = None, + tools: str | list[Any] | None = None, + model_architecture: str | None = None, +) -> str | None: + """ + Returns a response_format string accepted by xgrammar's structural tag format. + Uses the canonical shape: {"type": "structural_tag", "format": {...}} with + format.type "triggered_tags" and tag content type "json_schema" or "grammar". + + Callers that are not on an engine path (e.g. the reasoning parser) must pass + ``model_architecture`` explicitly. + """ + if model_architecture is None or model_architecture not in MODEL_TO_TAG_STYLE: + return None + style = MODEL_TO_TAG_STYLE[model_architecture] + + tags: list[dict] = [] + triggers: list[str] = [] + + def _add_tag(tag: CohereTagRegistry, content: dict) -> None: + tags.append({"begin": tag.trigger, "content": content, "end": tag.end}) + triggers.append(tag.trigger) + + if schema is not None: + # One structural tag per JSON wrapper (e.g. MOE: response + text). + # Same for schema-only and "tools plus JSON mode" (North: schema when + # the model does not call tools). + for jt in style.json_tags: + _add_tag(jt, {"type": "json_schema", "json_schema": schema}) + + if _has_effective_tools(tools): + # ``tools`` may be a JSON string (poseidon / RESPONSE_FORMAT_TOOL_DEFINITIONS) + # or a list (Chat Completions ``request.tools`` as Pydantic models or dicts). + tool_schema_list = _tool_definitions_to_schema_list(tools) + if not tool_schema_list: + raise ValueError( + "No valid tool definitions could be parsed from the request for " + "structural tag conversion." + ) + tool_grammar = collect_tool_schema(tool_schema_list) + _add_tag(style.tools, {"type": "grammar", "grammar": tool_grammar}) + + if not tags: + return None + return json.dumps( + { + "type": "structural_tag", + "format": { + "type": "triggered_tags", + "triggers": triggers, + "tags": tags, + }, + } + ) + + +def _response_format_type( + response_format: AnyResponseFormat | dict | None, +) -> str | None: + if response_format is None: + return None + if isinstance(response_format, dict): + t = response_format.get("type") + return t if isinstance(t, str) else None + return response_format.type + + +def _maybe_parse_json_dict(value: Any) -> dict | None: + """If value is a JSON string, parse to dict; otherwise require dict.""" + if isinstance(value, dict): + return value + if isinstance(value, str): + try: + parsed = json.loads(value) + except (TypeError, json.JSONDecodeError): + return None + return parsed if isinstance(parsed, dict) else None + return None + + +def _unwrap_nested_schema(candidate: Any) -> dict | None: + """Return ``candidate`` as a dict, unwrapping a nested ``schema`` if present. + + Returns ``None`` if ``candidate`` is not (and cannot be parsed into) a dict. + """ + cand = _maybe_parse_json_dict(candidate) + if not isinstance(cand, dict): + return None + nested = cand.get("schema") + return nested if isinstance(nested, dict) else cand + + +def _schema_from_json_schema_field(js_wr: Any) -> dict | None: + """ + Extract the JSON Schema object from Chat Completions ``json_schema`` payload. + + Accepts: + - ``JsonSchemaResponseFormat`` (Pydantic) with ``schema`` / ``json_schema`` field + - dict in OpenAI shape ``{"name": ..., "schema": {...}}`` + - dict with ``json_schema`` key holding either the schema or a nested wrapper + - dict that is already a JSON Schema document (some clients omit the wrapper) + - JSON strings for any of the above + """ + if js_wr is None: + return None + + parsed_wr = _maybe_parse_json_dict(js_wr) + if parsed_wr is not None: + js_wr = parsed_wr + + if hasattr(js_wr, "model_dump"): + for by_alias in (True, False): + try: + data = js_wr.model_dump(by_alias=by_alias, exclude_none=False) + except TypeError: + data = js_wr.model_dump(by_alias=by_alias) + out = _unwrap_nested_schema(data.get("schema") or data.get("json_schema")) + if out is not None: + return out + inner_attr = getattr(js_wr, "json_schema", None) + return inner_attr if isinstance(inner_attr, dict) else None + + if isinstance(js_wr, dict): + for key in ("schema", "json_schema"): + out = _unwrap_nested_schema(js_wr.get(key)) + if out is not None: + return out + return js_wr + + return None + + +def _schema_dict_from_chat_response_format( + rf: AnyResponseFormat | dict | None, +) -> dict | None: + """JSON schema dict from Chat Completions ``request.response_format`` only.""" + if rf is None: + return None + rf_type = _response_format_type(rf) + if rf_type == "json_object": + return {"type": "object"} + if rf_type != "json_schema": + return None + js_wr = ( + rf.get("json_schema") + if isinstance(rf, dict) + else getattr(rf, "json_schema", None) + ) + return _schema_from_json_schema_field(js_wr) + + +def _schema_dict_from_structured_outputs( + so: StructuredOutputsParams | None, +) -> dict | None: + """Schema dict from ``structured_outputs`` (``json`` / ``json_object``). + + Same unwrapping as ``json_schema``. ``json`` is expected to be ``str`` or + ``dict`` (enforced by ``StructuredOutputsParams`` / request models); other + types raise ``ValueError`` only if a caller bypasses that validation. + """ + if so is None: + return None + if so.json_object: + return {"type": "object"} + raw: Any = so.json + if raw is None: + return None + + if hasattr(raw, "model_dump"): + out = _schema_from_json_schema_field(raw) + if out is None: + raise ValueError( + "structured_outputs.json model has no extractable JSON Schema." + ) + return out + + if isinstance(raw, str): + if not raw.strip(): + raise ValueError("structured_outputs.json cannot be empty.") + try: + raw = json.loads(raw) + except json.JSONDecodeError as e: + raise ValueError("structured_outputs.json must be valid JSON.") from e + if not isinstance(raw, dict): + raise ValueError("structured_outputs.json must decode to a JSON object.") + + if isinstance(raw, Mapping): + body = raw if isinstance(raw, dict) else dict(raw) + return _schema_from_json_schema_field(body) or body + + raise ValueError( + f"structured_outputs.json has unsupported type {type(raw).__name__}." + ) + + +def _melody_sources_to_vllm( + sources: list[Source], position_to_source: PositionToSource | None +) -> list[CitationSource]: + """Resolve melody ``Source`` addresses into :class:`CitationSource` objects. + + Melody reports each source as ``(tool_call_index, tool_result_indices)``. + ``position_to_source`` is built by + :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._build_position_to_source` + and forwarded through ``chat_template_kwargs``; each index fans out to one + resolved source, and unresolved positions are skipped. Without a map (a + request that did not come through the Cohere v2 handler) every source is + dropped and the serving layer discards the citation at wire-coercion time. + """ + if not position_to_source: + return [] + out: list[CitationSource] = [] + for source in sources: + for idx in source.tool_result_indices: + info = position_to_source.get((source.tool_call_index, idx)) + if info is not None: + out.append(info) + return out + + +def _melody_citations_to_vllm( + citations: list[FilterCitation], position_to_source: PositionToSource | None +) -> list[Citation] | None: + """Convert melody ``FilterCitation`` objects into :class:`Citation`. + + Citations whose sources all fail to resolve are still returned with + ``sources=[]``; the serving layer drops them so the fail-closed policy + lives in one place. + """ + if not citations: + return None + return [ + Citation( + start=c.start_index, + end=c.end_index, + text=c.text, + sources=_melody_sources_to_vllm(c.sources, position_to_source), + type="THINKING_CONTENT" if c.is_thinking else "TEXT_CONTENT", + ) + for c in citations + ] + + +class CohereCommandParser(DelegatingParser): + def __init__( + self, + tokenizer: TokenizerLike, + tools: list[Tool] | None = None, + *args, + model_config: ModelConfig | None = None, + **kwargs, + ): + super().__init__(tokenizer, tools, *args, model_config=model_config, **kwargs) + self._model_config = model_config + self.start_token_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>") + self.end_token_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>") + shim = self._reasoning_parser or self._tool_parser + assert isinstance( + shim, (BaseCohereCommandReasoningParser, BaseCohereCommandToolParser) + ) + options = getattr(PyFilterOptions(), shim.melody_preset)() + self._streaming = PyFilter(options) + self._unary = PyFilter(options) + # Citations from the most recent ``parse`` call, surfaced on + # ``CohereChatMessage`` by ``CohereServingChatV2``. + self.last_unary_citations: list[Citation] | None = None + # Melody may emit a tool-call id before the function name; hold it until + # the first name delta so clients receive both together. + self._pending_tool_ids: dict[int, str] = {} + ctk = kwargs.get("chat_template_kwargs") or {} + raw_map = ctk.get(POSITION_TO_SOURCE_KEY) + self._position_to_source: PositionToSource | None = ( + raw_map if isinstance(raw_map, Mapping) else None + ) + + def adjust_request( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + # keep special tokens to ensure melody can correctly identify boundaries. + request.skip_special_tokens = False + request = self._apply_structural_tags(request) + return super().adjust_request(request) + + def _apply_structural_tags( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + """Fold ``tools`` and any JSON schema into a Cohere ``structural_tag``.""" + so = request.structured_outputs + if so is not None and so.structural_tag: + return request + # Schema: prefer ``response_format`` (OpenAI Chat Completions), then + # ``structured_outputs.json`` / ``json_object`` (vLLM direct). Tools stay + # on ``request.tools``. + rf = ( + request.response_format + if isinstance(request, ChatCompletionRequest) + else None + ) + if rf is not None and _response_format_type(rf) == "structural_tag": + return request + model_architecture = ( + self._model_config.architecture if self._model_config is not None else None + ) + tools = request.tools + # ``response_format`` wins if both it and ``structured_outputs`` supply JSON. + schema = _schema_dict_from_chat_response_format(rf) + if schema is None: + schema = _schema_dict_from_structured_outputs(so) + if schema is None and not _has_effective_tools(tools): + return request + if model_architecture is None: + return request + result = convert_schema_to_structural_tags( + schema=schema, + tools=tools, + model_architecture=model_architecture, + ) + if result is None: + # Unsupported architectures are not in ``MODEL_TO_TAG_STYLE``. + raise ValueError( + "Failed to build structural_tag guided decoding constraints from " + "this request's JSON schema and/or tools. The configured model " + f"architecture ({model_architecture!r}) does not support Cohere " + "command structural tags, or the schema cannot be expressed in " + "that format." + ) + request.structured_outputs = StructuredOutputsParams(structural_tag=result) + # Folded JSON constraints into ``structural_tag``; drop ``response_format`` + # when it was the source so ``to_sampling_params`` does not also set ``json`` / + # ``json_object`` (mutually exclusive in ``StructuredOutputsParams``). + if isinstance(request, ChatCompletionRequest) and rf is not None: + rf_type = _response_format_type(rf) + if rf_type in ("json_schema", "json_object"): + request.response_format = None + return request + + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: + count = 0 + depth = 0 + for token_id in token_ids: + if token_id == self.start_token_id: + depth += 1 + continue + if token_id == self.end_token_id: + if depth > 0: + depth -= 1 + continue + if depth > 0: + count += 1 + return count + + def parse( + self, + model_output: str, + request: ChatCompletionRequest | ResponsesRequest, + enable_auto_tools: bool = False, + model_output_token_ids: Sequence[int] = (), + ) -> tuple[str | None, str | None, list[FunctionCall] | None]: + result = self._unary.process_full_text(model_output) + self.last_unary_citations = _melody_citations_to_vllm( + result.citations, self._position_to_source + ) + tool_calls: list[FunctionCall] | None = None + if self._tools_allowed(request): + tool_calls = [ + FunctionCall(id=tc.id, name=tc.name, arguments=tc.arguments) + for tc in result.tool_calls + ] or None + return result.reasoning, result.content, tool_calls + + def parse_delta( + self, + delta_text: str, + delta_token_ids: list[int], + request: ChatCompletionRequest | ResponsesRequest, + prompt_token_ids: list[int] | None = None, + *, + finished: bool, + ) -> DeltaMessage | None: + results = [self._streaming.write_decoded(delta_text)] + if finished: + results.append(self._streaming.flush_partials()) + return self._to_delta(results, request) + + def _tools_allowed(self, request: ChatCompletionRequest | ResponsesRequest) -> bool: + return self._tool_parser is not None and request.tool_choice != "none" + + def _to_delta( + self, + results: list[FilterAggregatedResult], + request: ChatCompletionRequest | ResponsesRequest, + ) -> CohereDeltaMessage | None: + content = self._join(r.content for r in results) + reasoning = ( + self._join(r.reasoning for r in results) + if request.include_reasoning + else None + ) + tool_calls: list[DeltaToolCall] = [] + citations: list[Citation] = [] + for result in results: + if result.tool_calls and self._tools_allowed(request): + tool_calls.extend(self._tool_call_deltas(result.tool_calls)) + citations.extend( + _melody_citations_to_vllm(result.citations, self._position_to_source) + or [] + ) + if not (content or reasoning or tool_calls or citations): + return None + # ``CohereDeltaMessage`` carries citations to the wire via + # ``SerializeAsAny[DeltaMessage]``; it drops the field when unset so the + # shape matches a plain ``DeltaMessage``. + message = CohereDeltaMessage() + if content is not None: + message.content = content + if reasoning is not None: + message.reasoning = reasoning + if tool_calls: + message.tool_calls = tool_calls + if citations: + message.citations = citations + return message + + @staticmethod + def _join(parts: Any) -> str | None: + present = [p for p in parts if p is not None] + return "".join(present) if present else None + + def _tool_call_deltas( + self, melody_tool_calls: list[AccumulatedToolCall] + ) -> list[DeltaToolCall]: + tool_calls: list[DeltaToolCall] = [] + for tc in melody_tool_calls: + name = tc.name or None + arguments = tc.arguments or None + if tc.id: + self._pending_tool_ids[tc.index] = tc.id + # Empty strings are melody streaming placeholders; skip fragments + # that carry neither a name nor argument text. + if name is None and arguments is None: + continue + function = DeltaFunctionCall() + if name is not None: + function.name = name + if arguments is not None: + function.arguments = arguments + delta = DeltaToolCall(index=tc.index, function=function) + if name is not None: + resolved_id = tc.id or self._pending_tool_ids.pop(tc.index, None) + if resolved_id is not None: + delta.id = resolved_id + delta.type = "function" + tool_calls.append(delta) + return tool_calls diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py index 420873a58d..cc9dd3a6c6 100644 --- a/vllm/parser/parser_manager.py +++ b/vllm/parser/parser_manager.py @@ -126,6 +126,21 @@ class ParserManager: return _KimiK3Parser + if {reasoning_parser_name, tool_parser_name} & { + "cohere_command3", + "cohere_command4", + }: + from vllm.parser.cohere_command import CohereCommandParser + + r_cls = reasoning_parser_cls + t_cls = tool_parser_cls + + class _CohereCommandParser(CohereCommandParser): + reasoning_parser_cls = r_cls + tool_parser_cls = t_cls + + return _CohereCommandParser + from vllm.parser.abstract_parser import DelegatingParser r_cls = reasoning_parser_cls diff --git a/vllm/reasoning/cohere_command_reasoning_parser.py b/vllm/reasoning/cohere_command_reasoning_parser.py index 4dada76ee7..d02091e600 100644 --- a/vllm/reasoning/cohere_command_reasoning_parser.py +++ b/vllm/reasoning/cohere_command_reasoning_parser.py @@ -3,531 +3,28 @@ from __future__ import annotations -import json -from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, NamedTuple, TypedDict, TypeGuard +from collections.abc import Sequence +from typing import TYPE_CHECKING, ClassVar -import regex as re -import xgrammar as xgr - -try: - from cohere_melody import PyFilter, PyFilterOptions -except ImportError as e: - raise ImportError( - "The Cohere reasoning parser requires the `cohere_melody` " - "package, which is not installed. Install it with:\n" - " pip install 'cohere-melody>=0.11.1'" - ) from e - - -from vllm.entrypoints.cohere.cohere_chat_message import ( - Citation, - CitationSource, - CohereDeltaMessage, -) -from vllm.entrypoints.generate.base.protocol import ( - AnyResponseFormat, - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, -) -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.reasoning import ReasoningParser -from vllm.sampling_params import StructuredOutputsParams -from vllm.tokenizers import TokenizerLike if TYPE_CHECKING: - from vllm.entrypoints.openai.responses.protocol import ResponsesRequest - -REPLACEMENT_CHAR = "\ufffd" - - -class CohereTagRegistry(NamedTuple): - """A single ``structural_tag`` trigger / end pair (``begin`` uses ``trigger``).""" - - trigger: str - end: str - - -class CohereTagStyle(NamedTuple): - """The structural tags style for a given model architecture. - - ``json_tags`` lists every JSON-schema wrapper the model may emit (MOE uses - both response and text delimiters). ``tools`` is the tool-call wrapper. - """ - - json_tags: tuple[CohereTagRegistry, ...] - tools: CohereTagRegistry - - -class CohereNormalizedTool(TypedDict): - """A tool definition normalized to the shape ``collect_tool_schema`` expects. - - ``parameters`` is a JSON Schema object (possibly empty) describing the tool's - call signature. - """ - - name: str - parameters: dict[str, Any] - - -COMMAND_A_TOOLS_TAG = CohereTagRegistry( - trigger="<|START_ACTION|>", - end="<|END_ACTION|>", -) -COMMAND_A_JSON_TAG = CohereTagRegistry( - trigger="<|START_RESPONSE|>", - end="<|END_RESPONSE|>", -) -COMMAND_A_PLUS_JSON_TAG = CohereTagRegistry( - trigger="<|START_TEXT|>", - end="<|END_TEXT|>", -) - -MODEL_TO_TAG_STYLE: dict[str, CohereTagStyle] = { - "Cohere2ForCausalLM": CohereTagStyle( - json_tags=(COMMAND_A_JSON_TAG,), - tools=COMMAND_A_TOOLS_TAG, - ), - "Cohere2VisionForConditionalGeneration": CohereTagStyle( - json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG), - tools=COMMAND_A_TOOLS_TAG, - ), - "Cohere2MoeForCausalLM": CohereTagStyle( - json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG), - tools=COMMAND_A_TOOLS_TAG, - ), -} - - -def collect_tool_schema(tool_schema: list[CohereNormalizedTool]) -> str: - """Build an xgrammar EBNF grammar that matches a JSON array of tool calls. - - The grammar shape is architecture-independent; callers are responsible for - wrapping it in the correct structural tag (see ``CohereTagStyle.tools``). - """ - tool_dictionary: dict[str, str] = {} - for tool in tool_schema: - tool_name = tool["name"] - tool_parameters = json.dumps(tool["parameters"]) - json_schema = f"""{{ - "type": "object", - "properties": {{ - "tool_call_id": {{ - "type": "string", - "pattern": "^[0-9]+$" - }}, - "tool_name": {{ - "type": "string", - "const": "{tool_name}" - }}, - "parameters": {tool_parameters} - }} - }}""" - tool_grammar = str(xgr.Grammar.from_json_schema(json_schema)) - for match in re.findall(r"\b(\w+)\s*::=", tool_grammar): - tool_grammar = re.sub( - rf"\b{re.escape(match)}\b", tool_name + match, tool_grammar - ) - tool_dictionary[tool_name] = f"{tool_name} ::= {tool_name}root\n{tool_grammar}" - # Emitted grammar shape: - # root ::= tools - # tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws - # ws ::= (" " | "\t" | "\n")* - # tool ::= | | ... (one alternative per input) - # ::= root (per-tool xgrammar rules) - # root ::= ... (from xgr.Grammar.from_json_schema) - tool_alternatives = "tool ::= " + " | ".join(tool_dictionary.keys()) - tool_rules = "\n ".join(tool_dictionary.values()) - grammar = f"""root ::= tools - tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws - ws ::= (" " | "\\t" | "\\n")* - {tool_alternatives} - {tool_rules} - """ - return grammar - - -def _tool_definitions_to_schema_list( - tools: str | list[Any], -) -> list[CohereNormalizedTool]: - """ - Build the list of ``CohereNormalizedTool`` dicts expected by - ``collect_tool_schema``. - - Accepts: - - JSON string - - list of dicts with top-level ``name`` / ``parameters`` - - list of Chat Completions-style ``{"type": "function", "function": {...}}`` - - list of Pydantic models with ``model_dump()`` - """ - if isinstance(tools, str): - try: - parsed = json.loads(tools) - except json.JSONDecodeError: - return [] - if not isinstance(parsed, list): - return [] - else: - parsed = list(tools) - - out: list[CohereNormalizedTool] = [] - for raw in parsed: - t = raw.model_dump() if hasattr(raw, "model_dump") else raw - if not isinstance(t, dict): - continue - # Unwrap Chat Completions' ``{"type": "function", "function": {...}}`` - # shape; otherwise take the dict as-is. - if t.get("type") == "function" and isinstance(t.get("function"), dict): - t = t["function"] - name = t.get("name") - if not isinstance(name, str): - continue - params = t.get("parameters") - out.append( - CohereNormalizedTool( - name=name, - parameters=params if isinstance(params, dict) else {}, - ) - ) - return out - - -def _has_effective_tools( - tools: str | list[Any] | None, -) -> TypeGuard[str | list[Any]]: - """ - True when ``tools`` contains at least one tool definition to convert. - - ``ResponsesRequest`` defaults ``tools`` to ``[]``; ``ChatCompletionRequest`` - uses ``None``. Both mean "no tools" here. Strings (e.g. a JSON blob) are - treated as effective only when non-blank. - """ - if tools is None: - return False - if isinstance(tools, str): - return bool(tools.strip()) - return len(tools) > 0 - - -# Builder: produces vLLM response_format in xgrammar's canonical format. -# See xgrammar docs: type "structural_tag" with "format" = triggered_tags -# and tag content type = json_schema | grammar. -def convert_schema_to_structural_tags( - schema: dict | None = None, - tools: str | list[Any] | None = None, - model_architecture: str | None = None, -) -> str | None: - """ - Returns a response_format string accepted by xgrammar's structural tag format. - Uses the canonical shape: {"type": "structural_tag", "format": {...}} with - format.type "triggered_tags" and tag content type "json_schema" or "grammar". - - Callers that are not on an engine path (e.g. the reasoning parser) must pass - ``model_architecture`` explicitly. - """ - if model_architecture is None or model_architecture not in MODEL_TO_TAG_STYLE: - return None - style = MODEL_TO_TAG_STYLE[model_architecture] - - tags: list[dict] = [] - triggers: list[str] = [] - - def _add_tag(tag: CohereTagRegistry, content: dict) -> None: - tags.append({"begin": tag.trigger, "content": content, "end": tag.end}) - triggers.append(tag.trigger) - - if schema is not None: - # One structural tag per JSON wrapper (e.g. MOE: response + text). - # Same for schema-only and "tools plus JSON mode" (North: schema when - # the model does not call tools). - for jt in style.json_tags: - _add_tag(jt, {"type": "json_schema", "json_schema": schema}) - - if _has_effective_tools(tools): - # ``tools`` may be a JSON string (poseidon / RESPONSE_FORMAT_TOOL_DEFINITIONS) - # or a list (Chat Completions ``request.tools`` as Pydantic models or dicts). - tool_schema_list = _tool_definitions_to_schema_list(tools) - if not tool_schema_list: - raise ValueError( - "No valid tool definitions could be parsed from the request for " - "structural tag conversion." - ) - tool_grammar = collect_tool_schema(tool_schema_list) - _add_tag(style.tools, {"type": "grammar", "grammar": tool_grammar}) - - if not tags: - return None - return json.dumps( - { - "type": "structural_tag", - "format": { - "type": "triggered_tags", - "triggers": triggers, - "tags": tags, - }, - } - ) - - -def _response_format_type( - response_format: AnyResponseFormat | dict | None, -) -> str | None: - if response_format is None: - return None - if isinstance(response_format, dict): - t = response_format.get("type") - return t if isinstance(t, str) else None - return response_format.type - - -def _maybe_parse_json_dict(value: Any) -> dict | None: - """If value is a JSON string, parse to dict; otherwise require dict.""" - if isinstance(value, dict): - return value - if isinstance(value, str): - try: - parsed = json.loads(value) - except (TypeError, json.JSONDecodeError): - return None - return parsed if isinstance(parsed, dict) else None - return None - - -def _unwrap_nested_schema(candidate: Any) -> dict | None: - """Return ``candidate`` as a dict, unwrapping a nested ``schema`` if present. - - Returns ``None`` if ``candidate`` is not (and cannot be parsed into) a dict. - """ - cand = _maybe_parse_json_dict(candidate) - if not isinstance(cand, dict): - return None - nested = cand.get("schema") - return nested if isinstance(nested, dict) else cand - - -def _schema_from_json_schema_field(js_wr: Any) -> dict | None: - """ - Extract the JSON Schema object from Chat Completions ``json_schema`` payload. - - Accepts: - - ``JsonSchemaResponseFormat`` (Pydantic) with ``schema`` / ``json_schema`` field - - dict in OpenAI shape ``{"name": ..., "schema": {...}}`` - - dict with ``json_schema`` key holding either the schema or a nested wrapper - - dict that is already a JSON Schema document (some clients omit the wrapper) - - JSON strings for any of the above - """ - if js_wr is None: - return None - - parsed_wr = _maybe_parse_json_dict(js_wr) - if parsed_wr is not None: - js_wr = parsed_wr - - if hasattr(js_wr, "model_dump"): - for by_alias in (True, False): - try: - data = js_wr.model_dump(by_alias=by_alias, exclude_none=False) - except TypeError: - data = js_wr.model_dump(by_alias=by_alias) - out = _unwrap_nested_schema(data.get("schema") or data.get("json_schema")) - if out is not None: - return out - inner_attr = getattr(js_wr, "json_schema", None) - return inner_attr if isinstance(inner_attr, dict) else None - - if isinstance(js_wr, dict): - for key in ("schema", "json_schema"): - out = _unwrap_nested_schema(js_wr.get(key)) - if out is not None: - return out - return js_wr - - return None - - -def _schema_dict_from_chat_response_format( - rf: AnyResponseFormat | dict | None, -) -> dict | None: - """JSON schema dict from Chat Completions ``request.response_format`` only.""" - if rf is None: - return None - rf_type = _response_format_type(rf) - if rf_type == "json_object": - return {"type": "object"} - if rf_type != "json_schema": - return None - js_wr = ( - rf.get("json_schema") - if isinstance(rf, dict) - else getattr(rf, "json_schema", None) - ) - return _schema_from_json_schema_field(js_wr) - - -def _schema_dict_from_structured_outputs( - so: StructuredOutputsParams | None, -) -> dict | None: - """Schema dict from ``structured_outputs`` (``json`` / ``json_object``). - - Same unwrapping as ``json_schema``. ``json`` is expected to be ``str`` or - ``dict`` (enforced by ``StructuredOutputsParams`` / request models); other - types raise ``ValueError`` only if a caller bypasses that validation. - """ - if so is None: - return None - if so.json_object: - return {"type": "object"} - raw: Any = so.json - if raw is None: - return None - - if hasattr(raw, "model_dump"): - out = _schema_from_json_schema_field(raw) - if out is None: - raise ValueError( - "structured_outputs.json model has no extractable JSON Schema." - ) - return out - - if isinstance(raw, str): - if not raw.strip(): - raise ValueError("structured_outputs.json cannot be empty.") - try: - raw = json.loads(raw) - except json.JSONDecodeError as e: - raise ValueError("structured_outputs.json must be valid JSON.") from e - if not isinstance(raw, dict): - raise ValueError("structured_outputs.json must decode to a JSON object.") - - if isinstance(raw, Mapping): - body = raw if isinstance(raw, dict) else dict(raw) - return _schema_from_json_schema_field(body) or body - - raise ValueError( - f"structured_outputs.json has unsupported type {type(raw).__name__}." - ) - - -def _melody_sources_to_vllm( - raw_sources: Any, - position_to_source: Mapping[tuple[int, int], CitationSource] | None, -) -> list[CitationSource]: - """Convert melody's ``Source`` objects into resolved :class:`CitationSource`. - - Melody's ``Source`` shape is - ``{tool_call_index, tool_result_indices, document_ids}``. Without - :meth:`PyFilterOptions.with_message_history` the ``document_ids`` - list is empty, so we resolve the numeric address ourselves using - ``position_to_source``, which is - :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._build_position_to_source` - applied to the inbound request and forwarded through - ``chat_template_kwargs`` at parser-construction time. Each melody - source with ``tool_result_indices=[i, j, ...]`` fans out to one - :class:`CitationSource` per resolved position, populating the - wire-shape ``type`` / ``id`` / ``document`` / ``tool_output`` - fields; unresolved positions are silently skipped (the calling - citation is dropped at wire-coercion time if every source in it - ends up unresolved -- see ``_to_wire_citation`` in - :mod:`vllm.entrypoints.cohere.serving`). - - When ``position_to_source`` is ``None`` (parser wired outside of - :class:`CohereServingChatV2`, e.g. a plain OpenAI chat completion - request that happens to hit a cohere reasoning parser) sources - can't be resolved and get dropped; citations then never reach any - citation-aware handler downstream. - """ - if not position_to_source: - return [] - out: list[CitationSource] = [] - for s in raw_sources or []: - bucket = getattr(s, "tool_call_index", None) - indices = list(getattr(s, "tool_result_indices", None) or []) - if bucket is None or not indices: - continue - for idx in indices: - info = position_to_source.get((bucket, idx)) - if info is not None: - out.append(info) - return out - - -def _melody_citations_to_vllm( - raw_citations: Any, - position_to_source: Mapping[tuple[int, int], CitationSource] | None, -) -> list[Citation] | None: - """Convert melody's ``FilterCitation`` objects into :class:`Citation`. - - Resolves each source in-place via ``position_to_source`` (see - :func:`_melody_sources_to_vllm`). Citations whose sources all fail - to resolve are still emitted here with ``sources=[]``; the serving - layer drops them at wire-coercion time so the fail-closed policy - lives in exactly one place. - """ - if not raw_citations: - return None - out: list[Citation] = [] - for c in raw_citations: - out.append( - Citation( - start=getattr(c, "start_index", None), - end=getattr(c, "end_index", None), - text=getattr(c, "text", None), - sources=_melody_sources_to_vllm( - getattr(c, "sources", None), position_to_source - ), - type=( - "THINKING_CONTENT" - if getattr(c, "is_thinking", False) - else "TEXT_CONTENT" - ), - ) - ) - return out + from vllm.tokenizers import TokenizerLike class BaseCohereCommandReasoningParser(ReasoningParser): - def __init__( - self, - tokenizer: TokenizerLike, - *args, - streaming_opts: PyFilterOptions, - unary_opts: PyFilterOptions, - **kwargs, - ): + """Expose is_reasoning_end for the structured-output engine. + Everything else runs in + :class:`vllm.parser.cohere_command.CohereCommandParser`. + """ + + melody_preset: ClassVar[str] + + def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): super().__init__(tokenizer, *args, **kwargs) self.start_token_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>") self.end_token_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>") self.chatbot_token_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>") - self.unary_opts = unary_opts - self.melody_unary = PyFilter(unary_opts) - self.melody_streaming = PyFilter(streaming_opts) - # Citations extracted by the most recent ``extract_reasoning`` call. - # Citation-aware handlers (e.g. :class:`CohereServingChatV2` via its - # ``_finalize_response_message`` hook) read this back from the - # parser instance (which is constructed per-request) and attach - # the result to :class:`CohereChatMessage.citations`. ``None`` when - # the last parse produced no citations. - self.last_unary_citations: list[Citation] | None = None - # Request-scoped ``(tool_call_index, tool_result_idx) -> resolved - # CitationSource`` map, forwarded from - # :meth:`CohereServingChatV2._apply_cohere_template_kwargs` via - # ``chat_template_kwargs``. See :func:`_melody_sources_to_vllm` - # for how it's consumed. Absent for parser instances that - # weren't created by ``CohereServingChatV2`` (e.g. an OpenAI - # chat completion request routed through a cohere reasoning - # parser); citations from those requests are dropped since we - # have no way to attribute their sources. - ctk = kwargs.get("chat_template_kwargs") or {} - # Lazy import: ``renderers.cohere`` transitively imports - # ``cohere_melody`` and other heavy deps; hoisting the constant - # to module scope would tie parser importability to that graph. - from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY - - raw_map = ctk.get(POSITION_TO_SOURCE_KEY) - self._position_to_source: Mapping[tuple[int, int], CitationSource] | None = ( - raw_map if isinstance(raw_map, Mapping) else None - ) @property def reasoning_start_str(self) -> str | None: @@ -537,80 +34,6 @@ class BaseCohereCommandReasoningParser(ReasoningParser): def reasoning_end_str(self) -> str | None: return "<|END_THINKING|>" - def extract_reasoning_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - ) -> DeltaMessage | None: - r = self.melody_streaming.write_decoded(delta_text) - citations = _melody_citations_to_vllm( - getattr(r, "citations", None), self._position_to_source - ) - if ( - r.content is None - and r.reasoning is None - and not r.tool_calls - and not citations - ): - return None - # Always emit CohereDeltaMessage so citations reach the wire via - # ``SerializeAsAny[DeltaMessage]`` on the streaming choice envelope. - # When ``citations`` is unset, ``CohereDeltaMessage._serialize`` - # drops the field so the emitted shape matches plain - # :class:`DeltaMessage` - msg = CohereDeltaMessage() - if r.content is not None: - msg.content = r.content - if r.reasoning is not None: - msg.reasoning = r.reasoning - if r.tool_calls: - msg.tool_calls = [ - DeltaToolCall( - id=tc.id, - index=tc.index, - type="function", - function=DeltaFunctionCall(name=tc.name, arguments=tc.arguments), - ) - for tc in r.tool_calls - ] - if citations: - msg.citations = citations - return msg - - def extract_reasoning( - self, model_output: str, request: ChatCompletionRequest | ResponsesRequest - ) -> tuple[str | None, str | None]: - result = self.melody_unary.process_full_text(model_output) - # Cache citations so citation-aware handlers can surface them - # on :class:`CohereChatMessage.citations`. The base - # :meth:`ReasoningParser.extract_reasoning` contract only - # returns ``(reasoning, content)``, so citations are passed - # back via parser-instance state -- safe because the parser is - # constructed per-request. - self.last_unary_citations = _melody_citations_to_vllm( - getattr(result, "citations", None), self._position_to_source - ) - return result.reasoning, result.content - - def extract_content_ids(self, input_ids: list[int]) -> list[int]: - token_buf: list[int] = [] - content_ids: list[int] = [] - content_filter = PyFilter(self.unary_opts) - for t in input_ids: - token_buf.append(t) - s = self.model_tokenizer.decode(token_buf, skip_special_tokens=False) - if s.endswith(REPLACEMENT_CHAR): - continue - r = content_filter.write_decoded(s) - if r.content is not None: - content_ids.extend(token_buf) - token_buf = [] - return content_ids - def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: chatbot = self.chatbot_token_id start = self.start_token_id @@ -628,91 +51,10 @@ class BaseCohereCommandReasoningParser(ReasoningParser): return has_end_token - def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: - count = 0 - depth = 0 - for token_id in token_ids: - if token_id == self.start_token_id: - depth += 1 - continue - if token_id == self.end_token_id: - if depth > 0: - depth -= 1 - continue - if depth > 0: - count += 1 - return count - - def adjust_request( - self, request: ChatCompletionRequest | ResponsesRequest - ) -> ChatCompletionRequest | ResponsesRequest: - so = request.structured_outputs - if so is not None and so.structural_tag: - return request - # Schema: prefer ``response_format`` (OpenAI Chat Completions), then - # ``structured_outputs.json`` / ``json_object`` (vLLM direct). Tools stay - # on ``request.tools``. - rf = ( - request.response_format - if isinstance(request, ChatCompletionRequest) - else None - ) - if rf is not None and _response_format_type(rf) == "structural_tag": - return request - model_architecture = ( - self._model_config.architecture if self._model_config is not None else None - ) - tools = request.tools - # ``response_format`` wins if both it and ``structured_outputs`` supply JSON. - schema = _schema_dict_from_chat_response_format(rf) - if schema is None: - schema = _schema_dict_from_structured_outputs(so) - if schema is None and not _has_effective_tools(tools): - return request - if model_architecture is None: - return request - result = convert_schema_to_structural_tags( - schema=schema, - tools=tools, - model_architecture=model_architecture, - ) - if result is None: - # Unsupported architectures are not in ``MODEL_TO_TAG_STYLE``. - raise ValueError( - "Failed to build structural_tag guided decoding constraints from " - "this request's JSON schema and/or tools. The configured model " - f"architecture ({model_architecture!r}) does not support Cohere " - "command structural tags, or the schema cannot be expressed in " - "that format." - ) - request.structured_outputs = StructuredOutputsParams(structural_tag=result) - # Folded JSON constraints into ``structural_tag``; drop ``response_format`` - # when it was the source so ``to_sampling_params`` does not also set ``json`` / - # ``json_object`` (mutually exclusive in ``StructuredOutputsParams``). - if isinstance(request, ChatCompletionRequest) and rf is not None: - rf_type = _response_format_type(rf) - if rf_type in ("json_schema", "json_object"): - request.response_format = None - return request - class CohereCommand3ReasoningParser(BaseCohereCommandReasoningParser): - def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): - super().__init__( - tokenizer, - *args, - streaming_opts=PyFilterOptions().cmd3().no_tools(), - unary_opts=PyFilterOptions().cmd3().no_tools(), - **kwargs, - ) + melody_preset = "cmd3" class CohereCommand4ReasoningParser(BaseCohereCommandReasoningParser): - def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): - super().__init__( - tokenizer, - *args, - streaming_opts=PyFilterOptions().cmd4().no_tools(), - unary_opts=PyFilterOptions().cmd4().no_tools(), - **kwargs, - ) + melody_preset = "cmd4" diff --git a/vllm/renderers/cohere.py b/vllm/renderers/cohere.py index 8e6c4a5ac8..e22ebf1ebf 100644 --- a/vllm/renderers/cohere.py +++ b/vllm/renderers/cohere.py @@ -244,7 +244,7 @@ def _try_import_melody(): raise ImportError( "The `cohere` tokenizer/renderer mode requires the " "`cohere_melody` package. Install it via " - "`pip install 'cohere-melody>=0.11.1'` or build from " + "`pip install 'cohere-melody>=0.14.0'` or build from " "https://github.com/cohere-ai/melody." ) from e diff --git a/vllm/tool_parsers/cohere_command_tool_parser.py b/vllm/tool_parsers/cohere_command_tool_parser.py index 3739b222e5..eca5c54967 100644 --- a/vllm/tool_parsers/cohere_command_tool_parser.py +++ b/vllm/tool_parsers/cohere_command_tool_parser.py @@ -1,148 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Sequence +from typing import ClassVar -try: - from cohere_melody import PyFilter, PyFilterOptions -except ImportError as e: - raise ImportError( - "The Cohere tool parser requires the `cohere_melody` " - "package, which is not installed. Install it with:\n" - " pip install 'cohere-melody>=0.11.1'" - ) from e - -from vllm.entrypoints.generate.base.protocol import ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, - ExtractedToolCallInformation, - FunctionCall, - ToolCall, -) -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, -) -from vllm.entrypoints.openai.responses.protocol import ( - ResponsesRequest, -) -from vllm.tokenizers import TokenizerLike from vllm.tool_parsers import ToolParser -from vllm.tool_parsers.utils import Tool class BaseCohereCommandToolParser(ToolParser): - def __init__( - self, - tokenizer: TokenizerLike, - streaming_opts: PyFilterOptions, - unary_opts: PyFilterOptions, - ): - super().__init__(tokenizer) - self.melody_streaming = PyFilter(streaming_opts) - self.melody_unary = PyFilter(unary_opts) - # Melody can emit the tool-call id before the function name. Keep it - # until the first real name delta so clients receive both together. - self._pending_streaming_tool_call_ids: dict[int, str] = {} - - def adjust_request( - self, request: ChatCompletionRequest | ResponsesRequest - ) -> ChatCompletionRequest | ResponsesRequest: - request = super().adjust_request(request) - request.skip_special_tokens = False - return request - - def extract_tool_calls_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - request: ChatCompletionRequest, - ) -> DeltaMessage | None: - r = self.melody_streaming.write_decoded(delta_text) - if r.content is not None: - return DeltaMessage(content=r.content) - if r.reasoning is not None: - return DeltaMessage(reasoning=r.reasoning) - if r.tool_calls: - tool_calls: list[DeltaToolCall] = [] - for tc in r.tool_calls: - if tc.id: - self._pending_streaming_tool_call_ids[tc.index] = tc.id - name = tc.name or None - arguments = tc.arguments or None - # Empty strings are placeholders in Melody's streaming output; - # omit them from OpenAI-compatible deltas instead of sending - # invalid tool names or empty argument fragments. - if name is None and arguments is None: - continue - - function_kwargs = {} - if name is not None: - function_kwargs["name"] = name - if arguments is not None: - function_kwargs["arguments"] = arguments - tool_call_kwargs = { - "index": tc.index, - "function": DeltaFunctionCall(**function_kwargs), - } - if name is not None: - tool_call_id = tc.id or self._pending_streaming_tool_call_ids.pop( - tc.index, None - ) - if tool_call_id is not None: - tool_call_kwargs["id"] = tool_call_id - tool_call_kwargs["type"] = "function" - tool_calls.append(DeltaToolCall(**tool_call_kwargs)) - if tool_calls: - return DeltaMessage(tool_calls=tool_calls) - return None - - def extract_tool_calls( - self, - model_output: str, - request: ChatCompletionRequest, - ) -> ExtractedToolCallInformation: - result = self.melody_unary.process_full_text(model_output) - tool_calls = [ - ToolCall( - id=tc.id, - type="function", - function=FunctionCall(name=tc.name, arguments=tc.arguments), - ) - for tc in result.tool_calls - ] - return ExtractedToolCallInformation( - tools_called=len(tool_calls) > 0, - tool_calls=tool_calls, - content=result.content, - ) + melody_preset: ClassVar[str] class CohereCommand3ToolParser(BaseCohereCommandToolParser): - def __init__( - self, - tokenizer: TokenizerLike, - tools: list[Tool] | None = None, - ): - super().__init__( - tokenizer, - streaming_opts=PyFilterOptions().cmd3().start_in_answer(), - unary_opts=PyFilterOptions().cmd3().start_in_answer(), - ) + melody_preset = "cmd3" class CohereCommand4ToolParser(BaseCohereCommandToolParser): - def __init__( - self, - tokenizer: TokenizerLike, - tools: list[Tool] | None = None, - ): - super().__init__( - tokenizer, - streaming_opts=PyFilterOptions().cmd4().start_in_answer(), - unary_opts=PyFilterOptions().cmd4().start_in_answer(), - ) + melody_preset = "cmd4"