[Frontend] create unified Cohere parser (#56392)
Signed-off-by: walterbm <walter.beller.morales@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com>
This commit is contained in:
Walter Beller-Morales
co-authored by
Flora Feng
Flora Feng
parent
d406f09eb8
commit
756794a9a7
@@ -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` | ❌ |
|
||||
|
||||
@@ -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`)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -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")
|
||||
@@ -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 <co>Paris</co: 0:[1]>.<|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)
|
||||
@@ -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 <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
|
||||
<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|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 <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
|
||||
<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|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
|
||||
@@ -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": "<tool>",
|
||||
"content": {"type": "any_text"},
|
||||
"end": "</tool>",
|
||||
}
|
||||
],
|
||||
"triggers": ["<tool>"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
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
|
||||
@@ -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 <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>\n"
|
||||
"<|START_{tag}|>foo <co>bar</co: 0:[1,2],1:[3,4]><|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)
|
||||
@@ -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))
|
||||
@@ -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 <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
|
||||
<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|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 <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
|
||||
<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|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": "<tool>",
|
||||
"content": {"type": "any_text"},
|
||||
"end": "</tool>",
|
||||
}
|
||||
],
|
||||
"triggers": ["<tool>"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
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 <co>Paris</co: 0:[1]>.<|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)
|
||||
@@ -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 <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
|
||||
<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|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 <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
|
||||
<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|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"
|
||||
@@ -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.
|
||||
|
||||
@@ -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:
|
||||
|
||||
|
||||
@@ -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 ::= <tool_a> | <tool_b> | ... (one alternative per input)
|
||||
# <tool_x> ::= <tool_x>root (per-tool xgrammar rules)
|
||||
# <tool_x>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
|
||||
@@ -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
|
||||
|
||||
@@ -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 ::= <tool_a> | <tool_b> | ... (one alternative per input)
|
||||
# <tool_x> ::= <tool_x>root (per-tool xgrammar rules)
|
||||
# <tool_x>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"
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user