@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 419 lines 12.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import json
5 from types import SimpleNamespace
6
7 import pytest
8
9 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
10 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
11 from vllm.parser.abstract_parser import DelegatingParser
12 from vllm.parser.utils import count_history_tool_calls
13 from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
14 from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
15
16
17 @pytest.fixture(autouse=True)
18 def enable_hermes_required_named_parsing(monkeypatch):
19 # With VLLM_ENFORCE_STRICT_TOOL_CALLING (default on), Hermes sets
20 # supports_required_and_named=False and uses structural-tag guided tool
21 # calling. Force it True to test the non-guided JSON required/named parse path.
22 monkeypatch.setattr(Hermes2ProToolParser, "supports_required_and_named", True)
23
24
25 class ThinkReasoningParser(BaseThinkingReasoningParser):
26 @property
27 def start_token(self) -> str:
28 return "<think>"
29
30 @property
31 def end_token(self) -> str:
32 return "</think>"
33
34
35 MODEL_OUTPUT = (
36 "<think>let me think about this</think>"
37 '<tool_call>\n{"name": "get_weather", '
38 '"arguments": {"city": "Dallas"}}\n</tool_call>'
39 )
40
41 PLAIN_TEXT = "The weather in Dallas is sunny and 75°F."
42
43 TOOL_CALL_ONLY = (
44 '<tool_call>\n{"name": "get_weather", '
45 '"arguments": {"city": "Dallas"}}\n</tool_call>'
46 )
47
48 TOOL_ARGUMENTS = '{"city": "Dallas"}'
49
50
51 @pytest.fixture(scope="module")
52 def tokenizer():
53 from vllm.tokenizers import get_tokenizer
54
55 return get_tokenizer("Qwen/Qwen3-32B")
56
57
58 def make_request(**overrides):
59 base = {
60 "model": "test-model",
61 "messages": [{"role": "user", "content": "hi"}],
62 }
63 base.update(overrides)
64 return ChatCompletionRequest.model_validate(base)
65
66
67 TOOLS = [
68 {
69 "type": "function",
70 "function": {
71 "name": "get_weather",
72 "parameters": {"type": "object", "properties": {}},
73 },
74 }
75 ]
76
77
78 KIMI_K2_MODEL_CONFIG = SimpleNamespace(
79 hf_text_config=SimpleNamespace(model_type="kimi_k2"),
80 hf_overrides=None,
81 )
82
83 HISTORY_MESSAGES = [
84 {"role": "user", "content": "first"},
85 {
86 "role": "assistant",
87 "content": None,
88 "tool_calls": [
89 {
90 "id": "functions.get_current_weather:0",
91 "type": "function",
92 "function": {
93 "name": "get_current_weather",
94 "arguments": "{}",
95 },
96 },
97 {
98 "id": "functions.get_forecast:1",
99 "type": "function",
100 "function": {
101 "name": "get_forecast",
102 "arguments": "{}",
103 },
104 },
105 ],
106 },
107 {
108 "role": "tool",
109 "tool_call_id": "functions.get_current_weather:0",
110 "content": "{}",
111 },
112 {
113 "role": "tool",
114 "tool_call_id": "functions.get_forecast:1",
115 "content": "{}",
116 },
117 {"role": "user", "content": "again"},
118 ]
119
120
121 def make_parser(tokenizer, reasoning=False, tool=False, **kwargs):
122 class TestParser(DelegatingParser):
123 reasoning_parser_cls = ThinkReasoningParser if reasoning else None
124 tool_parser_cls = Hermes2ProToolParser if tool else None
125
126 return TestParser(tokenizer, **kwargs)
127
128
129 @pytest.mark.parametrize(
130 "reasoning,tool",
131 [(False, False), (False, True)],
132 ids=["neither", "tool-only"],
133 )
134 def test_parse_plain_text_no_reasoning_parser(tokenizer, reasoning, tool):
135 parser = make_parser(tokenizer, reasoning=reasoning, tool=tool)
136 request = make_request()
137 r, content, tool_calls = parser.parse(PLAIN_TEXT, request)
138
139 assert r is None
140 assert content == PLAIN_TEXT
141 assert tool_calls is not None
142 assert len(tool_calls) == 0
143
144
145 @pytest.mark.parametrize(
146 "reasoning,tool",
147 [(True, False), (True, True)],
148 ids=["reasoning-only", "both"],
149 )
150 def test_parse_plain_text_with_reasoning_parser(tokenizer, reasoning, tool):
151 parser = make_parser(tokenizer, reasoning=reasoning, tool=tool)
152 request = make_request()
153 r, content, tool_calls = parser.parse(PLAIN_TEXT, request)
154
155 assert r == PLAIN_TEXT
156 assert content is None
157 assert tool_calls is not None
158 assert len(tool_calls) == 0
159
160
161 def test_parse_both_parsers(tokenizer):
162 parser = make_parser(tokenizer, reasoning=True, tool=True)
163 request = make_request(tools=TOOLS)
164 reasoning, content, tool_calls = parser.parse(
165 MODEL_OUTPUT, request, enable_auto_tools=True
166 )
167
168 assert reasoning is not None
169 assert "let me think about this" in reasoning
170 assert tool_calls is not None
171 assert len(tool_calls) == 1
172 assert tool_calls[0].name == "get_weather"
173 assert json.loads(tool_calls[0].arguments) == {"city": "Dallas"}
174 assert not content or content.strip() == ""
175
176
177 def test_parse_reasoning_only(tokenizer):
178 parser = make_parser(tokenizer, reasoning=True, tool=False)
179 request = make_request()
180 reasoning, content, tool_calls = parser.parse(MODEL_OUTPUT, request)
181
182 assert reasoning is not None
183 assert "let me think about this" in reasoning
184 assert content is not None
185 assert "<tool_call>" in content
186 assert "get_weather" in content
187 assert tool_calls is not None
188 assert len(tool_calls) == 0
189
190
191 def test_parse_tool_only(tokenizer):
192 parser = make_parser(tokenizer, reasoning=False, tool=True)
193 request = make_request(tools=TOOLS)
194 reasoning, content, tool_calls = parser.parse(
195 MODEL_OUTPUT, request, enable_auto_tools=True
196 )
197
198 assert reasoning is None
199 assert tool_calls is not None
200 assert len(tool_calls) == 1
201 assert tool_calls[0].name == "get_weather"
202 assert json.loads(tool_calls[0].arguments) == {"city": "Dallas"}
203
204
205 def test_parse_named_tool_choice(tokenizer):
206 parser = make_parser(tokenizer, reasoning=False, tool=True)
207 request = make_request(
208 tools=TOOLS,
209 tool_choice={
210 "type": "function",
211 "function": {"name": "get_weather"},
212 },
213 )
214 reasoning, content, tool_calls = parser.parse(
215 TOOL_ARGUMENTS, request, enable_auto_tools=True
216 )
217
218 assert reasoning is None
219 assert content is None
220 assert tool_calls is not None
221 assert len(tool_calls) == 1
222 assert tool_calls[0].name == "get_weather"
223 assert tool_calls[0].arguments == TOOL_ARGUMENTS
224
225
226 def test_parse_named_tool_choice_with_reasoning(tokenizer):
227 parser = make_parser(tokenizer, reasoning=True, tool=True)
228 model_output = f"<think>thinking</think>{TOOL_ARGUMENTS}"
229 request = make_request(
230 tools=TOOLS,
231 tool_choice={
232 "type": "function",
233 "function": {"name": "get_weather"},
234 },
235 )
236 reasoning, content, tool_calls = parser.parse(
237 model_output, request, enable_auto_tools=True
238 )
239
240 assert reasoning is not None
241 assert "thinking" in reasoning
242 assert content is None
243 assert len(tool_calls) == 1
244 assert tool_calls[0].name == "get_weather"
245 assert tool_calls[0].arguments == TOOL_ARGUMENTS
246
247
248 def test_parse_required_tool_choice(tokenizer):
249 parser = make_parser(tokenizer, reasoning=False, tool=True)
250 functions_json = json.dumps(
251 [
252 {"name": "get_weather", "parameters": {"city": "Dallas"}},
253 {"name": "get_time", "parameters": {"timezone": "UTC"}},
254 ]
255 )
256 request = make_request(tools=TOOLS, tool_choice="required")
257 reasoning, content, tool_calls = parser.parse(
258 functions_json, request, enable_auto_tools=True
259 )
260
261 assert reasoning is None
262 assert content is None
263 assert tool_calls is not None
264 assert len(tool_calls) == 2
265 assert tool_calls[0].name == "get_weather"
266 assert json.loads(tool_calls[0].arguments) == {"city": "Dallas"}
267 assert tool_calls[1].name == "get_time"
268 assert json.loads(tool_calls[1].arguments) == {"timezone": "UTC"}
269
270
271 def test_parse_required_tool_choice_kimi_k2_ids(tokenizer):
272 parser = make_parser(
273 tokenizer, reasoning=False, tool=True, model_config=KIMI_K2_MODEL_CONFIG
274 )
275 functions_json = json.dumps(
276 [
277 {"name": "get_current_weather", "parameters": {"city": "Dallas"}},
278 {"name": "get_forecast", "parameters": {"city": "Dallas", "days": 2}},
279 ]
280 )
281 request = make_request(tools=TOOLS, tool_choice="required")
282 _, content, tool_calls = parser.parse(
283 functions_json, request, enable_auto_tools=True
284 )
285
286 assert content is None
287 assert tool_calls is not None
288 assert [tc.id for tc in tool_calls] == [
289 "functions.get_current_weather:0",
290 "functions.get_forecast:1",
291 ]
292
293
294 def test_parse_required_tool_choice_kimi_k2_ids_after_history(tokenizer):
295 parser = make_parser(
296 tokenizer, reasoning=False, tool=True, model_config=KIMI_K2_MODEL_CONFIG
297 )
298 functions_json = json.dumps(
299 [{"name": "get_current_weather", "parameters": {"city": "Dallas"}}]
300 )
301 request = make_request(
302 messages=HISTORY_MESSAGES,
303 tools=TOOLS,
304 tool_choice="required",
305 )
306 _, _, tool_calls = parser.parse(functions_json, request, enable_auto_tools=True)
307
308 assert tool_calls is not None
309 assert tool_calls[0].id == "functions.get_current_weather:2"
310
311
312 def test_count_history_tool_calls_responses_request():
313 request = ResponsesRequest.model_validate(
314 {
315 "model": "test-model",
316 "input": [
317 {
318 "type": "function_call",
319 "call_id": "call_0",
320 "name": "get_current_weather",
321 "arguments": "{}",
322 },
323 {
324 "type": "function_call",
325 "call_id": "call_1",
326 "name": "get_forecast",
327 "arguments": "{}",
328 },
329 ],
330 }
331 )
332
333 assert count_history_tool_calls(request) == 2
334
335
336 def test_parse_required_tool_choice_random_ids_deferred(tokenizer):
337 parser = make_parser(tokenizer, reasoning=False, tool=True)
338 functions_json = json.dumps(
339 [{"name": "get_current_weather", "parameters": {"city": "Dallas"}}]
340 )
341 request = make_request(
342 messages=HISTORY_MESSAGES,
343 tools=TOOLS,
344 tool_choice="required",
345 )
346 _, _, tool_calls = parser.parse(functions_json, request, enable_auto_tools=True)
347
348 assert tool_calls is not None
349 assert tool_calls[0].id is None
350
351
352 def test_parse_named_tool_choice_kimi_k2_id(tokenizer):
353 parser = make_parser(
354 tokenizer, reasoning=False, tool=True, model_config=KIMI_K2_MODEL_CONFIG
355 )
356 request = make_request(
357 tools=TOOLS,
358 tool_choice={
359 "type": "function",
360 "function": {"name": "get_weather"},
361 },
362 )
363 _, content, tool_calls = parser.parse(
364 TOOL_ARGUMENTS, request, enable_auto_tools=True
365 )
366
367 assert content is None
368 assert tool_calls is not None
369 assert tool_calls[0].id == "functions.get_weather:0"
370
371
372 def test_parse_named_tool_choice_content_none(tokenizer):
373 parser = make_parser(tokenizer, reasoning=False, tool=True)
374 request = make_request(
375 tools=TOOLS,
376 tool_choice={
377 "type": "function",
378 "function": {"name": "get_weather"},
379 },
380 )
381 reasoning, content, tool_calls = parser.parse("", request, enable_auto_tools=True)
382 assert reasoning is None
383 assert content is None
384 assert tool_calls is not None
385
386
387 def test_parse_required_tool_choice_content_none(tokenizer):
388 parser = make_parser(tokenizer, reasoning=False, tool=True)
389 request = make_request(tools=TOOLS, tool_choice="required")
390 reasoning, content, tool_calls = parser.parse("", request, enable_auto_tools=True)
391 assert reasoning is None
392 assert content is None
393 assert tool_calls is not None
394 assert len(tool_calls) == 0
395
396
397 def test_parse_auto_tools_no_parser(tokenizer):
398 parser = make_parser(tokenizer, reasoning=False, tool=False)
399 request = make_request()
400 reasoning, content, tool_calls = parser.parse(
401 TOOL_CALL_ONLY, request, enable_auto_tools=True
402 )
403
404 assert reasoning is None
405 assert content == TOOL_CALL_ONLY
406 assert tool_calls is not None
407 assert len(tool_calls) == 0
408
409
410 def test_parse_auto_tools_no_calls_returns_none(tokenizer):
411 parser = make_parser(tokenizer, reasoning=False, tool=True)
412 request = make_request(tools=TOOLS)
413 reasoning, content, tool_calls = parser.parse(
414 PLAIN_TEXT, request, enable_auto_tools=True
415 )
416
417 assert reasoning is None
418 assert content == PLAIN_TEXT
419 assert tool_calls is None