@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 203 lines 7.75 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Reasoning extraction, reasoning-end gating and parser selection for the
4 unified Cohere Command parser."""
5
6 from __future__ import annotations
7
8 from dataclasses import dataclass
9
10 import pytest
11
12 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
13 from vllm.parser import ParserManager
14 from vllm.reasoning.cohere_command_reasoning_parser import (
15 CohereCommand3ReasoningParser,
16 CohereCommand4ReasoningParser,
17 )
18
19 from .utils import make_parser, stream_parser
20
21
22 @dataclass
23 class ReasoningCase:
24 parser_name: str
25 model_output: str
26 expected_reasoning: str | None
27 expected_content: str | None
28
29
30 REASONING_CASES = [
31 pytest.param(
32 ReasoningCase(
33 parser_name="cohere_command3",
34 model_output="""\
35 <|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
36 [
37 {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
38 ]
39 <|END_ACTION|>""",
40 expected_reasoning="i will call foo with query1",
41 expected_content=None,
42 ),
43 id="cmd3-single_tool_call",
44 ),
45 pytest.param(
46 ReasoningCase(
47 parser_name="cohere_command4",
48 model_output="""\
49 <|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
50 [
51 {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
52 ]
53 <|END_ACTION|>""",
54 expected_reasoning="i will call foo with query1",
55 expected_content=None,
56 ),
57 id="cmd4-single_tool_call",
58 ),
59 pytest.param(
60 ReasoningCase(
61 parser_name="cohere_command3",
62 model_output="""\
63 <|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
64 <|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_RESPONSE|>""",
65 expected_reasoning="This is a rainbow emoji: 🌈",
66 expected_content="foo bar",
67 ),
68 id="cmd3-citations_with_emoji",
69 ),
70 pytest.param(
71 ReasoningCase(
72 parser_name="cohere_command4",
73 model_output="""\
74 <|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
75 <|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_RESPONSE|>""",
76 expected_reasoning="This is a rainbow emoji: 🌈",
77 expected_content="foo bar",
78 ),
79 id="cmd4-citations_with_emoji",
80 ),
81 ]
82
83
84 @pytest.mark.parametrize("case", REASONING_CASES)
85 class TestExtractReasoning:
86 def test_nonstreaming(self, tokenizer, request_obj, case: ReasoningCase):
87 parser = make_parser(tokenizer, case.parser_name)
88 reasoning, content, tool_calls = parser.parse(
89 case.model_output, request_obj, enable_auto_tools=True
90 )
91 assert reasoning == case.expected_reasoning
92 assert content == case.expected_content
93 # No tools on the request, so any tool call is dropped, not surfaced.
94 assert not tool_calls
95
96 def test_streaming(self, tokenizer, request_obj, case: ReasoningCase):
97 parser = make_parser(tokenizer, case.parser_name)
98 deltas = stream_parser(parser, request_obj, tokenizer, case.model_output)
99
100 reasoning = "".join(d.reasoning for d in deltas if d.reasoning) or None
101 content = "".join(d.content for d in deltas if d.content) or None
102 assert reasoning == case.expected_reasoning
103 assert content == case.expected_content
104 assert all(not d.tool_calls for d in deltas)
105
106
107 class TestFramingTokensStripped:
108 @pytest.mark.parametrize(
109 ("parser_name", "content_tags"),
110 [
111 pytest.param(
112 "cohere_command4", ("<|START_TEXT|>", "<|END_TEXT|>"), id="cmd4"
113 ),
114 pytest.param(
115 "cohere_command3", ("<|START_RESPONSE|>", "<|END_RESPONSE|>"), id="cmd3"
116 ),
117 ],
118 )
119 @pytest.mark.parametrize("chunk_size", [1, 4], ids=["per_token", "batched"])
120 @pytest.mark.parametrize("with_tools", [False, True], ids=["no_tools", "tools"])
121 def test_content_framing_tokens_stripped(
122 self, tokenizer, parser_name, content_tags, chunk_size, with_tools
123 ):
124 parser = make_parser(tokenizer, parser_name)
125 tools = [{"type": "function", "function": {"name": "foo"}}]
126 request = ChatCompletionRequest(
127 messages=[], model="test-model", tools=tools if with_tools else None
128 )
129 parser.adjust_request(request)
130
131 start_tag, end_tag = content_tags
132 generation = (
133 f"<|START_THINKING|>Think deeply. The user greets us.<|END_THINKING|>"
134 f"{start_tag}I'm doing well, thank you{end_tag}"
135 )
136 deltas = stream_parser(parser, request, tokenizer, generation, chunk_size)
137
138 assert "".join(d.reasoning or "" for d in deltas) == (
139 "Think deeply. The user greets us."
140 )
141 assert not any(d.tool_calls for d in deltas)
142 assert "".join(d.content or "" for d in deltas) == "I'm doing well, thank you"
143
144
145 class TestIsReasoningEnd:
146 @pytest.mark.parametrize(
147 "parser_cls",
148 [CohereCommand3ReasoningParser, CohereCommand4ReasoningParser],
149 ids=["cmd3", "cmd4"],
150 )
151 def test_is_reasoning_end(self, tokenizer, parser_cls):
152 parser = parser_cls(tokenizer)
153 start_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>")
154 end_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>")
155 chatbot_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>")
156 content_ids = [99, 100]
157
158 # Generation-only tokens have no chatbot marker, so the whole sequence
159 # is considered.
160 assert parser.is_reasoning_end([end_id])
161 assert parser.is_reasoning_end([start_id, *content_ids, end_id])
162 assert not parser.is_reasoning_end([start_id, *content_ids])
163
164 # Full prompt/history tokens are scoped to the latest chatbot marker,
165 # so stray thinking tokens from the preamble or previous turns are ignored.
166 assert not parser.is_reasoning_end([start_id, end_id, chatbot_id, *content_ids])
167 assert parser.is_reasoning_end(
168 [start_id, end_id, chatbot_id, start_id, *content_ids, end_id]
169 )
170
171
172 @pytest.mark.parametrize("parser_name", ["cohere_command3", "cohere_command4"])
173 def test_count_reasoning_tokens(tokenizer, parser_name):
174 parser = make_parser(tokenizer, parser_name)
175 start = tokenizer.convert_tokens_to_ids("<|START_THINKING|>")
176 end = tokenizer.convert_tokens_to_ids("<|END_THINKING|>")
177
178 assert parser.count_reasoning_tokens([99, start, 11, 12, end, 100]) == 2
179 assert parser.count_reasoning_tokens([end, start, 1, start, 2, end, 3, end]) == 3
180 assert parser.count_reasoning_tokens([start, 1, 2]) == 2
181 assert parser.count_reasoning_tokens([1, 2, end]) == 0
182
183
184 class TestParserSelection:
185 @pytest.mark.parametrize("with_tool_parser", [True, False])
186 def test_adjust_request_keeps_framing_tokens(self, tokenizer, with_tool_parser):
187 """Melody needs the special tokens, with or without a tool parser."""
188 cls = ParserManager.get_parser(
189 "cohere_command4" if with_tool_parser else None,
190 "cohere_command4",
191 enable_auto_tools=True,
192 )
193 request = ChatCompletionRequest(messages=[], model="m")
194 assert request.skip_special_tokens
195 cls(tokenizer).adjust_request(request)
196 assert request.skip_special_tokens is False
197
198 def test_non_cohere_selection_does_not_import_melody_parser(self, monkeypatch):
199 """``cohere_melody`` is optional; other parsers must not depend on it."""
200 import sys
201
202 monkeypatch.setitem(sys.modules, "vllm.parser.cohere_command", None)
203 assert ParserManager.get_parser("hermes", "deepseek_r1", True) is not None