@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 264 lines 8.42 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import json
5
6 import pytest
7
8 from tests.parser.engine.conftest import make_mock_tokenizer
9 from tests.parser.engine.streaming_helpers import (
10 collect_function_name,
11 collect_tool_arguments,
12 simulate_tool_streaming,
13 )
14 from vllm.entrypoints.openai.chat_completion.protocol import (
15 ChatCompletionToolsParam,
16 FunctionDefinition,
17 )
18 from vllm.parser.minimax_m2 import (
19 THINK_END,
20 TOOL_CALL_END,
21 TOOL_CALL_START,
22 MinimaxM2Parser,
23 )
24
25
26 @pytest.fixture
27 def mock_tokenizer():
28 return make_mock_tokenizer(
29 {
30 THINK_END: 99,
31 TOOL_CALL_START: 100,
32 TOOL_CALL_END: 101,
33 }
34 )
35
36
37 @pytest.fixture
38 def parser(mock_tokenizer):
39 return MinimaxM2Parser(mock_tokenizer)
40
41
42 def make_tools(*names: str):
43 return [
44 ChatCompletionToolsParam(
45 function=FunctionDefinition(
46 name=name,
47 parameters={
48 "type": "object",
49 "properties": {},
50 },
51 ),
52 )
53 for name in names
54 ]
55
56
57 class TestNonStreaming:
58 def test_no_tool_calls(self, parser, mock_request):
59 result = parser.extract_tool_calls(
60 "</think>This is a regular response without tool calls.",
61 mock_request,
62 )
63 assert result.tools_called is False
64 assert result.tool_calls == []
65 assert result.content == "This is a regular response without tool calls."
66
67 def test_single_tool_call(self, parser, mock_request):
68 result = parser.extract_tool_calls(
69 '<minimax:tool_call><invoke name="get_weather">'
70 '<parameter name="city">Seattle</parameter>'
71 "</invoke></minimax:tool_call>",
72 mock_request,
73 )
74
75 assert result.tools_called is True
76 assert len(result.tool_calls) == 1
77 assert result.tool_calls[0].function.name == "get_weather"
78 assert json.loads(result.tool_calls[0].function.arguments) == {
79 "city": "Seattle",
80 }
81
82 def test_multiple_invokes(self, parser, mock_request):
83 result = parser.extract_tool_calls(
84 "<minimax:tool_call>"
85 '<invoke name="search"><parameter name="q">OpenAI</parameter></invoke>'
86 '<invoke name="search"><parameter name="q">vLLM</parameter></invoke>'
87 "</minimax:tool_call>",
88 mock_request,
89 )
90
91 assert result.tools_called is True
92 assert [tc.function.name for tc in result.tool_calls] == ["search", "search"]
93 assert json.loads(result.tool_calls[0].function.arguments) == {"q": "OpenAI"}
94 assert json.loads(result.tool_calls[1].function.arguments) == {"q": "vLLM"}
95
96 def test_schema_type_coercion(self, mock_tokenizer, mock_request):
97 tools = [
98 ChatCompletionToolsParam(
99 function=FunctionDefinition(
100 name="forecast",
101 parameters={
102 "type": "object",
103 "properties": {
104 "days": {"type": "integer"},
105 "include_hourly": {"type": "boolean"},
106 },
107 },
108 ),
109 )
110 ]
111 parser = MinimaxM2Parser(mock_tokenizer, tools=tools)
112 mock_request.tools = tools
113
114 result = parser.extract_tool_calls(
115 '<minimax:tool_call><invoke name="forecast">'
116 '<parameter name="days">5</parameter>'
117 '<parameter name="include_hourly">true</parameter>'
118 "</invoke></minimax:tool_call>",
119 mock_request,
120 )
121
122 assert json.loads(result.tool_calls[0].function.arguments) == {
123 "days": 5,
124 "include_hourly": True,
125 }
126
127 def test_invalid_tool_name_is_rejected(self, mock_tokenizer, mock_request):
128 tools = make_tools("search")
129 parser = MinimaxM2Parser(mock_tokenizer)
130 mock_request.tools = tools
131
132 result = parser.extract_tool_calls(
133 '<minimax:tool_call><invoke name="img_gen">'
134 '<parameter name="prompt">a cat</parameter>'
135 "</invoke></minimax:tool_call>",
136 mock_request,
137 )
138
139 assert result.tools_called is False
140 assert result.tool_calls == []
141
142 def test_mixed_tool_names_only_return_valid(self, mock_tokenizer, mock_request):
143 tools = make_tools("search")
144 parser = MinimaxM2Parser(mock_tokenizer)
145 mock_request.tools = tools
146
147 result = parser.extract_tool_calls(
148 "<minimax:tool_call>"
149 '<invoke name="img_gen"><parameter name="prompt">cat</parameter></invoke>'
150 '<invoke name="search"><parameter name="query">news</parameter></invoke>'
151 "</minimax:tool_call>",
152 mock_request,
153 )
154
155 assert result.tools_called is True
156 assert [tc.function.name for tc in result.tool_calls] == ["search"]
157 assert json.loads(result.tool_calls[0].function.arguments) == {
158 "query": "news",
159 }
160
161
162 class TestStreaming:
163 def test_streaming_single_tool_call(self, parser, mock_request):
164 results = simulate_tool_streaming(
165 parser,
166 mock_request,
167 [
168 "<minimax:tool_call>",
169 '<invoke name="get_weather">',
170 '<parameter name="city">Seattle</parameter>',
171 "</invoke></minimax:tool_call>",
172 ],
173 )
174
175 assert collect_function_name(results) == "get_weather"
176 assert json.loads(collect_tool_arguments(results)) == {
177 "city": "Seattle",
178 }
179
180 def test_streaming_multiple_invokes(self, parser, mock_request):
181 results = simulate_tool_streaming(
182 parser,
183 mock_request,
184 [
185 "<minimax:tool_call>",
186 '<invoke name="a"><parameter name="x">1</parameter></invoke>',
187 '<invoke name="b"><parameter name="y">2</parameter></invoke>',
188 "</minimax:tool_call>",
189 ],
190 )
191
192 tool_names = [
193 tc.function.name
194 for delta, _ in results
195 if delta and delta.tool_calls
196 for tc in delta.tool_calls
197 if tc.function and tc.function.name
198 ]
199 assert tool_names == ["a", "b"]
200
201 def test_streaming_invoke_prefix_split_before_quote(self, parser, mock_request):
202 results = simulate_tool_streaming(
203 parser,
204 mock_request,
205 [
206 "<minimax:tool_call>",
207 "<invoke name=",
208 '"get_weather">',
209 '<parameter name="city">Seattle</parameter>',
210 "</invoke></minimax:tool_call>",
211 ],
212 )
213
214 assert collect_function_name(results) == "get_weather"
215 assert json.loads(collect_tool_arguments(results)) == {
216 "city": "Seattle",
217 }
218
219 def test_streaming_invalid_tool_name_is_rejected(
220 self, mock_tokenizer, mock_request
221 ):
222 tools = make_tools("search")
223 parser = MinimaxM2Parser(mock_tokenizer)
224 mock_request.tools = tools
225
226 results = simulate_tool_streaming(
227 parser,
228 mock_request,
229 [
230 "<minimax:tool_call>",
231 '<invoke name="img_gen">',
232 '<parameter name="prompt">cat</parameter>',
233 "</invoke></minimax:tool_call>",
234 ],
235 )
236
237 assert collect_function_name(results) is None
238 assert collect_tool_arguments(results) == ""
239
240
241 class TestReasoning:
242 def test_extract_reasoning_without_start_token(self, parser, mock_request):
243 reasoning, content = parser.extract_reasoning(
244 "This is reasoning</think>This is content",
245 mock_request,
246 )
247
248 assert reasoning == "This is reasoning"
249 assert content == "This is content"
250
251 def test_extract_reasoning_without_end_token(self, parser, mock_request):
252 reasoning, content = parser.extract_reasoning(
253 "This is still reasoning",
254 mock_request,
255 )
256
257 assert reasoning == "This is still reasoning"
258 assert content is None
259
260 def test_extract_content_ids_without_end_token(self, parser):
261 assert parser.extract_content_ids([1, 2, 3]) == []
262
263 def test_extract_content_ids_after_end_token(self, parser):
264 assert parser.extract_content_ids([1, 99, 2, 3]) == [2, 3]