main
py 507 lines 15.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import json
5 from dataclasses import dataclass, field
6
7 import openai
8 import pytest
9
10 from tests.tool_use.utils import (
11 MESSAGES_ASKING_FOR_PARALLEL_TOOLS,
12 MESSAGES_ASKING_FOR_TOOLS,
13 MESSAGES_WITH_TOOL_RESPONSE,
14 MESSAGES_WITHOUT_TOOLS,
15 SEARCH_TOOL,
16 SEED,
17 WEATHER_TOOL,
18 ensure_system_prompt,
19 )
20
21 from .utils import ServerConfig
22
23
24 def _requires_tool_parser(server_config: ServerConfig) -> None:
25 r"""Skip test if server was not started with --tool-call-parser."""
26 if "--tool-call-parser" not in server_config.get("arguments", []):
27 pytest.skip(
28 f"Skipping: {server_config['model']} not configured with --tool-call-parser"
29 )
30
31
32 def _is_pre_v11(server_config: ServerConfig) -> bool:
33 r"""Pre-v11 Mistral models lack grammar-based tool call enforcement."""
34 return "7B" in server_config.get("model", "")
35
36
37 @dataclass
38 class StreamedToolCallResult:
39 r"""Accumulated result from streaming a single tool call."""
40
41 function_name: str | None = None
42 function_args_str: str = ""
43 tool_call_id: str | None = None
44 role_name: str | None = None
45 finish_reason_count: int = 0
46 finish_reason: str | None = None
47
48
49 async def _collect_streamed_tool_call(
50 stream: openai.AsyncStream,
51 *,
52 expected_finish_reason: str = "tool_calls",
53 ) -> StreamedToolCallResult:
54 result = StreamedToolCallResult()
55
56 async for chunk in stream:
57 if chunk.choices[0].finish_reason:
58 result.finish_reason_count += 1
59 result.finish_reason = chunk.choices[0].finish_reason
60 assert chunk.choices[0].finish_reason == expected_finish_reason
61
62 if chunk.choices[0].delta.role:
63 assert not result.role_name or result.role_name == "assistant"
64 result.role_name = "assistant"
65
66 streamed_tool_calls = chunk.choices[0].delta.tool_calls
67 if streamed_tool_calls and len(streamed_tool_calls) > 0:
68 assert len(streamed_tool_calls) == 1
69 tool_call = streamed_tool_calls[0]
70
71 if tool_call.id:
72 assert not result.tool_call_id
73 result.tool_call_id = tool_call.id
74
75 if tool_call.function:
76 if tool_call.function.name:
77 assert result.function_name is None
78 result.function_name = tool_call.function.name
79 if tool_call.function.arguments:
80 result.function_args_str += tool_call.function.arguments
81
82 return result
83
84
85 @dataclass
86 class StreamedContentResult:
87 r"""Accumulated result from streaming a content-only response."""
88
89 chunks: list[str] = field(default_factory=list)
90 finish_reason_count: int = 0
91 finish_reason: str | None = None
92 role_sent: bool = False
93
94
95 async def _collect_streamed_content(
96 stream: openai.AsyncStream,
97 *,
98 expected_finish_reason: str | None = None,
99 no_tool_calls: bool = True,
100 ) -> StreamedContentResult:
101 r"""Consume a streaming response and collect text content."""
102 result = StreamedContentResult()
103
104 async for chunk in stream:
105 delta = chunk.choices[0].delta
106
107 if delta.role:
108 assert not result.role_sent
109 assert delta.role == "assistant"
110 result.role_sent = True
111
112 if delta.content:
113 result.chunks.append(delta.content)
114
115 if chunk.choices[0].finish_reason is not None:
116 result.finish_reason_count += 1
117 result.finish_reason = chunk.choices[0].finish_reason
118 if expected_finish_reason is not None:
119 assert result.finish_reason == expected_finish_reason
120
121 if no_tool_calls:
122 assert not delta.tool_calls or len(delta.tool_calls) == 0
123
124 return result
125
126
127 @dataclass
128 class StreamedParallelToolCallResult:
129 r"""Accumulated result from streaming parallel tool calls."""
130
131 function_names: list[str] = field(default_factory=list)
132 function_args_strs: list[str] = field(default_factory=list)
133 tool_call_ids: list[str] = field(default_factory=list)
134 role_name: str | None = None
135 finish_reason_count: int = 0
136
137
138 async def _collect_streamed_parallel_tool_calls(
139 stream: openai.AsyncStream,
140 ) -> StreamedParallelToolCallResult:
141 r"""Consume a streaming response and collect parallel tool calls."""
142 result = StreamedParallelToolCallResult()
143 tool_call_idx: int = -1
144
145 async for chunk in stream:
146 if chunk.choices[0].finish_reason:
147 result.finish_reason_count += 1
148 assert chunk.choices[0].finish_reason == "tool_calls"
149
150 if chunk.choices[0].delta.role:
151 assert not result.role_name or result.role_name == "assistant"
152 result.role_name = "assistant"
153
154 streamed_tool_calls = chunk.choices[0].delta.tool_calls
155 if streamed_tool_calls and len(streamed_tool_calls) > 0:
156 assert len(streamed_tool_calls) == 1
157 tool_call = streamed_tool_calls[0]
158
159 if tool_call.index != tool_call_idx:
160 tool_call_idx = tool_call.index
161 result.function_args_strs.append("")
162 result.tool_call_ids.append("")
163
164 if tool_call.id:
165 result.tool_call_ids[tool_call.index] = tool_call.id
166
167 if tool_call.function:
168 if tool_call.function.name:
169 result.function_names.append(tool_call.function.name)
170 if tool_call.function.arguments:
171 result.function_args_strs[tool_call.index] += (
172 tool_call.function.arguments
173 )
174
175 return result
176
177
178 # test: a tool_choice with mistral-tokenizer results in an ID of length 9
179 @pytest.mark.asyncio
180 async def test_tool_call_with_tool_choice(
181 client: openai.AsyncOpenAI, server_config: ServerConfig
182 ) -> None:
183 _requires_tool_parser(server_config)
184
185 models = await client.models.list()
186 model_name: str = models.data[0].id
187 chat_completion = await client.chat.completions.create(
188 messages=ensure_system_prompt(MESSAGES_ASKING_FOR_TOOLS, server_config),
189 temperature=0,
190 max_completion_tokens=100,
191 model=model_name,
192 tools=[WEATHER_TOOL],
193 tool_choice=WEATHER_TOOL,
194 logprobs=False,
195 seed=SEED,
196 )
197
198 choice = chat_completion.choices[0]
199
200 assert choice.finish_reason != "tool_calls" # "stop" or "length"
201 assert choice.message.role == "assistant"
202 assert choice.message.tool_calls is None or len(choice.message.tool_calls) == 1
203 assert len(choice.message.tool_calls[0].id) == 9 # length of 9 for mistral
204
205
206 _NOT_SET = object()
207
208
209 @pytest.mark.asyncio
210 @pytest.mark.parametrize(
211 "tools, tool_choice, streaming_id_len_pre_v11",
212 [
213 pytest.param(
214 [WEATHER_TOOL, SEARCH_TOOL],
215 _NOT_SET,
216 9,
217 id="auto",
218 ),
219 pytest.param(
220 [WEATHER_TOOL],
221 "required",
222 9,
223 id="required",
224 ),
225 ],
226 )
227 async def test_tool_call_auto_or_required(
228 client: openai.AsyncOpenAI,
229 server_config: ServerConfig,
230 tools: list,
231 tool_choice: object,
232 streaming_id_len_pre_v11: int,
233 ) -> None:
234 _requires_tool_parser(server_config)
235
236 models = await client.models.list()
237 model_name: str = models.data[0].id
238
239 create_kwargs: dict = {
240 "messages": ensure_system_prompt(MESSAGES_ASKING_FOR_TOOLS, server_config),
241 "temperature": 0,
242 "max_completion_tokens": 100,
243 "model": model_name,
244 "tools": tools,
245 "logprobs": False,
246 "seed": SEED,
247 }
248 if tool_choice is not _NOT_SET:
249 create_kwargs["tool_choice"] = tool_choice
250
251 # --- non-streaming ---
252 chat_completion = await client.chat.completions.create(**create_kwargs)
253
254 choice = chat_completion.choices[0]
255 tool_calls = choice.message.tool_calls
256
257 assert choice.finish_reason == "tool_calls"
258 assert tool_calls is not None and len(tool_calls) >= 1
259 assert tool_calls[0].function.name == "get_current_weather"
260 parsed_arguments = json.loads(tool_calls[0].function.arguments)
261 assert "city" in parsed_arguments
262 assert len(tool_calls[0].id) == 9
263
264 # --- streaming ---
265 stream = await client.chat.completions.create(**create_kwargs, stream=True)
266
267 result = await _collect_streamed_tool_call(stream)
268
269 assert result.finish_reason_count == 1
270 assert result.role_name == "assistant"
271 assert result.function_name == "get_current_weather"
272 streamed_args = json.loads(result.function_args_str)
273 assert isinstance(result.tool_call_id, str)
274 if _is_pre_v11(server_config):
275 assert len(result.tool_call_id) == streaming_id_len_pre_v11
276 else:
277 assert len(result.tool_call_id) == 9
278 assert parsed_arguments == streamed_args
279
280
281 @pytest.mark.asyncio
282 async def test_tool_call_none_with_tools(
283 client: openai.AsyncOpenAI, server_config: ServerConfig
284 ) -> None:
285 _requires_tool_parser(server_config)
286
287 models = await client.models.list()
288 model_name: str = models.data[0].id
289
290 # --- non-streaming ---
291 chat_completion = await client.chat.completions.create(
292 messages=ensure_system_prompt(MESSAGES_ASKING_FOR_TOOLS, server_config),
293 temperature=0,
294 max_completion_tokens=100,
295 model=model_name,
296 tools=[WEATHER_TOOL],
297 tool_choice="none",
298 logprobs=False,
299 seed=SEED,
300 )
301
302 choice = chat_completion.choices[0]
303
304 assert choice.finish_reason != "tool_calls"
305 assert choice.message.tool_calls is None or len(choice.message.tool_calls) == 0
306 assert choice.message.content is not None
307 # Without grammar enforcement, pre-v11 models may still emit [TOOL_CALLS]
308 if not _is_pre_v11(server_config):
309 assert "[TOOL_CALLS]" not in choice.message.content
310
311 non_streaming_content = choice.message.content
312
313 # --- streaming ---
314 stream = await client.chat.completions.create(
315 messages=ensure_system_prompt(MESSAGES_ASKING_FOR_TOOLS, server_config),
316 temperature=0,
317 max_completion_tokens=100,
318 model=model_name,
319 tools=[WEATHER_TOOL],
320 tool_choice="none",
321 logprobs=False,
322 seed=SEED,
323 stream=True,
324 )
325
326 # Pre-v11 models lack grammar enforcement, so the model may still
327 # emit tool calls even with tool_choice="none".
328 pre_v11 = _is_pre_v11(server_config)
329 result = await _collect_streamed_content(stream, no_tool_calls=not pre_v11)
330
331 assert result.finish_reason_count == 1
332 if not pre_v11:
333 assert result.finish_reason != "tool_calls"
334 streamed_content = "".join(result.chunks)
335 if not pre_v11:
336 assert "[TOOL_CALLS]" not in streamed_content
337 assert streamed_content == non_streaming_content
338
339
340 @pytest.mark.asyncio
341 async def test_chat_without_tools(
342 client: openai.AsyncOpenAI, server_config: ServerConfig
343 ) -> None:
344 models = await client.models.list()
345 model_name: str = models.data[0].id
346
347 # --- non-streaming ---
348 chat_completion = await client.chat.completions.create(
349 messages=ensure_system_prompt(MESSAGES_WITHOUT_TOOLS, server_config),
350 temperature=0,
351 max_completion_tokens=150,
352 model=model_name,
353 logprobs=False,
354 seed=SEED,
355 )
356
357 choice = chat_completion.choices[0]
358 output_text = choice.message.content
359
360 assert output_text is not None and len(output_text) > 0
361 assert choice.finish_reason != "tool_calls"
362 assert choice.message.tool_calls is None or len(choice.message.tool_calls) == 0
363
364 # --- streaming ---
365 stream = await client.chat.completions.create(
366 messages=ensure_system_prompt(MESSAGES_WITHOUT_TOOLS, server_config),
367 temperature=0,
368 max_completion_tokens=150,
369 model=model_name,
370 logprobs=False,
371 seed=SEED,
372 stream=True,
373 )
374
375 result = await _collect_streamed_content(
376 stream, expected_finish_reason=choice.finish_reason
377 )
378
379 assert result.role_sent
380 assert result.finish_reason_count == 1
381 assert len(result.chunks)
382 assert "".join(result.chunks) == output_text
383
384
385 @pytest.mark.asyncio
386 async def test_tool_call_with_results(
387 client: openai.AsyncOpenAI, server_config: ServerConfig
388 ) -> None:
389 _requires_tool_parser(server_config)
390
391 models = await client.models.list()
392 model_name: str = models.data[0].id
393
394 # --- non-streaming ---
395 chat_completion = await client.chat.completions.create(
396 messages=ensure_system_prompt(MESSAGES_WITH_TOOL_RESPONSE, server_config),
397 temperature=0,
398 max_completion_tokens=100,
399 model=model_name,
400 tools=[WEATHER_TOOL, SEARCH_TOOL],
401 logprobs=False,
402 seed=SEED,
403 )
404
405 choice = chat_completion.choices[0]
406
407 assert choice.finish_reason != "tool_calls"
408 assert choice.message.tool_calls is None or len(choice.message.tool_calls) == 0
409 assert choice.message.content is not None
410 assert "98" in choice.message.content
411
412 # --- streaming ---
413 stream = await client.chat.completions.create(
414 messages=ensure_system_prompt(MESSAGES_WITH_TOOL_RESPONSE, server_config),
415 temperature=0,
416 max_completion_tokens=100,
417 model=model_name,
418 tools=[WEATHER_TOOL, SEARCH_TOOL],
419 logprobs=False,
420 seed=SEED,
421 stream=True,
422 )
423
424 result = await _collect_streamed_content(
425 stream, expected_finish_reason=choice.finish_reason
426 )
427
428 assert result.role_sent
429 assert result.finish_reason_count == 1
430 assert len(result.chunks)
431 assert "".join(result.chunks) == choice.message.content
432
433
434 def _requires_parallel(server_config: ServerConfig) -> None:
435 r"""Skip test if the model does not support parallel tool calls."""
436 if not server_config.get("supports_parallel"):
437 pytest.skip(
438 f"Skipping: {server_config['model']} does not support parallel tool calls"
439 )
440
441
442 @pytest.mark.asyncio
443 async def test_tool_call_parallel(
444 client: openai.AsyncOpenAI, server_config: ServerConfig
445 ) -> None:
446 _requires_tool_parser(server_config)
447 _requires_parallel(server_config)
448
449 models = await client.models.list()
450 model_name: str = models.data[0].id
451
452 # --- non-streaming ---
453 chat_completion = await client.chat.completions.create(
454 messages=ensure_system_prompt(
455 MESSAGES_ASKING_FOR_PARALLEL_TOOLS, server_config
456 ),
457 temperature=0,
458 max_completion_tokens=200,
459 model=model_name,
460 tools=[WEATHER_TOOL],
461 logprobs=False,
462 seed=SEED,
463 )
464
465 choice = chat_completion.choices[0]
466 tool_calls = choice.message.tool_calls
467
468 assert choice.finish_reason == "tool_calls"
469 assert tool_calls is not None and len(tool_calls) >= 2
470 for tc in tool_calls:
471 assert tc.type == "function"
472 assert tc.function.name == "get_current_weather"
473 assert isinstance(tc.function.arguments, str)
474 parsed = json.loads(tc.function.arguments)
475 assert "city" in parsed
476 assert len(tc.id) == 9
477
478 non_streaming_tool_calls = tool_calls
479
480 # --- streaming ---
481 stream = await client.chat.completions.create(
482 messages=ensure_system_prompt(
483 MESSAGES_ASKING_FOR_PARALLEL_TOOLS, server_config
484 ),
485 temperature=0,
486 max_completion_tokens=200,
487 model=model_name,
488 tools=[WEATHER_TOOL],
489 logprobs=False,
490 seed=SEED,
491 stream=True,
492 )
493
494 result = await _collect_streamed_parallel_tool_calls(stream)
495
496 assert result.finish_reason_count == 1
497 assert result.role_name == "assistant"
498 assert len(result.function_names) >= 2
499 assert all(name == "get_current_weather" for name in result.function_names)
500 assert len(result.tool_call_ids) >= 2
501 assert all(isinstance(tid, str) and len(tid) == 9 for tid in result.tool_call_ids)
502
503 for args_str in result.function_args_strs:
504 streamed_args = json.loads(args_str)
505 assert "city" in streamed_args
506
507 assert len(result.function_names) == len(non_streaming_tool_calls)