| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import json |
| 5 | from collections.abc import Generator |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | from vllm.entrypoints.generate.base.protocol import ( |
| 10 | DeltaMessage, |
| 11 | FunctionCall, |
| 12 | ToolCall, |
| 13 | ) |
| 14 | from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest |
| 15 | from vllm.tokenizers import TokenizerLike, get_tokenizer |
| 16 | from vllm.tokenizers.detokenizer_utils import detokenize_incrementally |
| 17 | from vllm.tool_parsers.xlam_tool_parser import xLAMToolParser |
| 18 | |
| 19 | # Use a common model that is likely to be available |
| 20 | MODEL = "Salesforce/Llama-xLAM-2-8B-fc-r" |
| 21 | |
| 22 | |
| 23 | @pytest.fixture(scope="module") |
| 24 | def xlam_tokenizer(): |
| 25 | return get_tokenizer(tokenizer_name=MODEL) |
| 26 | |
| 27 | |
| 28 | @pytest.fixture |
| 29 | def xlam_tool_parser(xlam_tokenizer): |
| 30 | return xLAMToolParser(xlam_tokenizer) |
| 31 | |
| 32 | |
| 33 | def assert_tool_calls( |
| 34 | actual_tool_calls: list[ToolCall], expected_tool_calls: list[ToolCall] |
| 35 | ): |
| 36 | assert len(actual_tool_calls) == len(expected_tool_calls) |
| 37 | |
| 38 | for actual_tool_call, expected_tool_call in zip( |
| 39 | actual_tool_calls, expected_tool_calls |
| 40 | ): |
| 41 | assert isinstance(actual_tool_call.id, str) |
| 42 | assert len(actual_tool_call.id) > 16 |
| 43 | |
| 44 | assert actual_tool_call.type == "function" |
| 45 | assert actual_tool_call.function == expected_tool_call.function |
| 46 | |
| 47 | |
| 48 | def stream_delta_message_generator( |
| 49 | xlam_tool_parser: xLAMToolParser, |
| 50 | xlam_tokenizer: TokenizerLike, |
| 51 | model_output: str, |
| 52 | request: ChatCompletionRequest | None = None, |
| 53 | ) -> Generator[DeltaMessage, None, None]: |
| 54 | all_token_ids = xlam_tokenizer.encode(model_output, add_special_tokens=False) |
| 55 | |
| 56 | previous_text = "" |
| 57 | previous_tokens = None |
| 58 | prefix_offset = 0 |
| 59 | read_offset = 0 |
| 60 | for i, delta_token in enumerate(all_token_ids): |
| 61 | delta_token_ids = [delta_token] |
| 62 | previous_token_ids = all_token_ids[:i] |
| 63 | current_token_ids = all_token_ids[: i + 1] |
| 64 | |
| 65 | (new_tokens, delta_text, new_prefix_offset, new_read_offset) = ( |
| 66 | detokenize_incrementally( |
| 67 | tokenizer=xlam_tokenizer, |
| 68 | all_input_ids=current_token_ids, |
| 69 | prev_tokens=previous_tokens, |
| 70 | prefix_offset=prefix_offset, |
| 71 | read_offset=read_offset, |
| 72 | skip_special_tokens=False, |
| 73 | spaces_between_special_tokens=True, |
| 74 | ) |
| 75 | ) |
| 76 | |
| 77 | current_text = previous_text + delta_text |
| 78 | |
| 79 | delta_message = xlam_tool_parser.extract_tool_calls_streaming( |
| 80 | previous_text, |
| 81 | current_text, |
| 82 | delta_text, |
| 83 | previous_token_ids, |
| 84 | current_token_ids, |
| 85 | delta_token_ids, |
| 86 | request=request, |
| 87 | ) |
| 88 | if delta_message: |
| 89 | yield delta_message |
| 90 | |
| 91 | previous_text = current_text |
| 92 | previous_tokens = ( |
| 93 | previous_tokens + new_tokens if previous_tokens else new_tokens |
| 94 | ) |
| 95 | prefix_offset = new_prefix_offset |
| 96 | read_offset = new_read_offset |
| 97 | |
| 98 | |
| 99 | def test_extract_tool_calls_no_tools(xlam_tool_parser): |
| 100 | model_output = "This is a test" |
| 101 | extracted_tool_calls = xlam_tool_parser.extract_tool_calls( |
| 102 | model_output, request=None |
| 103 | ) |
| 104 | assert not extracted_tool_calls.tools_called |
| 105 | assert extracted_tool_calls.tool_calls == [] |
| 106 | assert extracted_tool_calls.content == model_output |
| 107 | |
| 108 | |
| 109 | @pytest.mark.parametrize( |
| 110 | ids=[ |
| 111 | "parallel_tool_calls", |
| 112 | "single_tool_with_think_tag", |
| 113 | "single_tool_with_json_code_block", |
| 114 | "single_tool_with_tool_calls_tag", |
| 115 | "single_tool_with_tool_call_xml_tags", |
| 116 | ], |
| 117 | argnames=["model_output", "expected_tool_calls", "expected_content"], |
| 118 | argvalues=[ |
| 119 | ( |
| 120 | """[{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}, {"name": "get_current_weather", "arguments": {"city": "Orlando", "state": "FL", "unit": "fahrenheit"}}]""", # noqa: E501 |
| 121 | [ |
| 122 | ToolCall( |
| 123 | function=FunctionCall( |
| 124 | name="get_current_weather", |
| 125 | arguments=json.dumps( |
| 126 | { |
| 127 | "city": "Dallas", |
| 128 | "state": "TX", |
| 129 | "unit": "fahrenheit", |
| 130 | } |
| 131 | ), |
| 132 | ) |
| 133 | ), |
| 134 | ToolCall( |
| 135 | function=FunctionCall( |
| 136 | name="get_current_weather", |
| 137 | arguments=json.dumps( |
| 138 | { |
| 139 | "city": "Orlando", |
| 140 | "state": "FL", |
| 141 | "unit": "fahrenheit", |
| 142 | } |
| 143 | ), |
| 144 | ) |
| 145 | ), |
| 146 | ], |
| 147 | None, |
| 148 | ), |
| 149 | ( |
| 150 | """<think>I'll help you with that.</think>[{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}]""", # noqa: E501 |
| 151 | [ |
| 152 | ToolCall( |
| 153 | function=FunctionCall( |
| 154 | name="get_current_weather", |
| 155 | arguments=json.dumps( |
| 156 | { |
| 157 | "city": "Dallas", |
| 158 | "state": "TX", |
| 159 | "unit": "fahrenheit", |
| 160 | } |
| 161 | ), |
| 162 | ) |
| 163 | ) |
| 164 | ], |
| 165 | "<think>I'll help you with that.</think>", |
| 166 | ), |
| 167 | ( |
| 168 | """I'll help you with that.\n```json\n[{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}]\n```""", # noqa: E501 |
| 169 | [ |
| 170 | ToolCall( |
| 171 | function=FunctionCall( |
| 172 | name="get_current_weather", |
| 173 | arguments=json.dumps( |
| 174 | { |
| 175 | "city": "Dallas", |
| 176 | "state": "TX", |
| 177 | "unit": "fahrenheit", |
| 178 | } |
| 179 | ), |
| 180 | ) |
| 181 | ) |
| 182 | ], |
| 183 | "I'll help you with that.", |
| 184 | ), |
| 185 | ( |
| 186 | """I'll check the weather for you.[TOOL_CALLS][{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}]""", # noqa: E501 |
| 187 | [ |
| 188 | ToolCall( |
| 189 | function=FunctionCall( |
| 190 | name="get_current_weather", |
| 191 | arguments=json.dumps( |
| 192 | { |
| 193 | "city": "Dallas", |
| 194 | "state": "TX", |
| 195 | "unit": "fahrenheit", |
| 196 | } |
| 197 | ), |
| 198 | ) |
| 199 | ) |
| 200 | ], |
| 201 | "I'll check the weather for you.", |
| 202 | ), |
| 203 | ( |
| 204 | """I'll help you check the weather.<tool_call>[{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}]</tool_call>""", # noqa: E501 |
| 205 | [ |
| 206 | ToolCall( |
| 207 | function=FunctionCall( |
| 208 | name="get_current_weather", |
| 209 | arguments=json.dumps( |
| 210 | { |
| 211 | "city": "Dallas", |
| 212 | "state": "TX", |
| 213 | "unit": "fahrenheit", |
| 214 | } |
| 215 | ), |
| 216 | ) |
| 217 | ) |
| 218 | ], |
| 219 | "I'll help you check the weather.", |
| 220 | ), |
| 221 | ], |
| 222 | ) |
| 223 | def test_extract_tool_calls( |
| 224 | xlam_tool_parser, model_output, expected_tool_calls, expected_content |
| 225 | ): |
| 226 | extracted_tool_calls = xlam_tool_parser.extract_tool_calls( |
| 227 | model_output, request=None |
| 228 | ) |
| 229 | assert extracted_tool_calls.tools_called |
| 230 | |
| 231 | assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls) |
| 232 | |
| 233 | assert extracted_tool_calls.content == expected_content |
| 234 | |
| 235 | |
| 236 | @pytest.mark.parametrize( |
| 237 | ids=["list_structured_tool_call"], |
| 238 | argnames=["model_output", "expected_tool_calls", "expected_content"], |
| 239 | argvalues=[ |
| 240 | ( |
| 241 | """[{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}]""", # noqa: E501 |
| 242 | [ |
| 243 | ToolCall( |
| 244 | function=FunctionCall( |
| 245 | name="get_current_weather", |
| 246 | arguments=json.dumps( |
| 247 | { |
| 248 | "city": "Seattle", |
| 249 | "state": "WA", |
| 250 | "unit": "celsius", |
| 251 | } |
| 252 | ), |
| 253 | ) |
| 254 | ) |
| 255 | ], |
| 256 | None, |
| 257 | ), |
| 258 | ], |
| 259 | ) |
| 260 | def test_extract_tool_calls_list_structure( |
| 261 | xlam_tool_parser, model_output, expected_tool_calls, expected_content |
| 262 | ): |
| 263 | """Test extraction of tool calls when the model outputs a list-structured tool call.""" # noqa: E501 |
| 264 | extracted_tool_calls = xlam_tool_parser.extract_tool_calls( |
| 265 | model_output, request=None |
| 266 | ) |
| 267 | assert extracted_tool_calls.tools_called |
| 268 | |
| 269 | assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls) |
| 270 | |
| 271 | assert extracted_tool_calls.content == expected_content |
| 272 | |
| 273 | |
| 274 | # Test for preprocess_model_output method |
| 275 | def test_preprocess_model_output(xlam_tool_parser): |
| 276 | # Test with list structure |
| 277 | model_output = ( |
| 278 | """[{"name": "get_current_weather", "arguments": {"city": "Seattle"}}]""" # noqa: E501 |
| 279 | ) |
| 280 | content, potential_tool_calls = xlam_tool_parser.preprocess_model_output( |
| 281 | model_output |
| 282 | ) |
| 283 | assert content is None |
| 284 | assert potential_tool_calls == model_output |
| 285 | |
| 286 | # Test with thinking tag |
| 287 | model_output = """<think>I'll help you with that.</think>[{"name": "get_current_weather", "arguments": {"city": "Seattle"}}]""" # noqa: E501 |
| 288 | content, potential_tool_calls = xlam_tool_parser.preprocess_model_output( |
| 289 | model_output |
| 290 | ) |
| 291 | assert content == "<think>I'll help you with that.</think>" |
| 292 | assert ( |
| 293 | potential_tool_calls |
| 294 | == '[{"name": "get_current_weather", "arguments": {"city": "Seattle"}}]' |
| 295 | ) |
| 296 | |
| 297 | # Test with JSON code block |
| 298 | model_output = """I'll help you with that. |
| 299 | ```json |
| 300 | [{"name": "get_current_weather", "arguments": {"city": "Seattle"}}] |
| 301 | ```""" |
| 302 | content, potential_tool_calls = xlam_tool_parser.preprocess_model_output( |
| 303 | model_output |
| 304 | ) |
| 305 | assert content == "I'll help you with that." |
| 306 | assert "get_current_weather" in potential_tool_calls |
| 307 | |
| 308 | # Test with no tool calls |
| 309 | model_output = """I'll help you with that.""" |
| 310 | content, potential_tool_calls = xlam_tool_parser.preprocess_model_output( |
| 311 | model_output |
| 312 | ) |
| 313 | assert content == model_output |
| 314 | assert potential_tool_calls is None |
| 315 | |
| 316 | |
| 317 | # Simulate streaming to test extract_tool_calls_streaming |
| 318 | def test_streaming_with_list_structure(xlam_tool_parser): |
| 319 | # Reset streaming state |
| 320 | xlam_tool_parser.prev_tool_calls = [] |
| 321 | xlam_tool_parser.current_tools_sent = [] |
| 322 | xlam_tool_parser.streamed_args = [] |
| 323 | xlam_tool_parser.current_tool_id = -1 |
| 324 | |
| 325 | # Simulate receiving a message with list structure |
| 326 | current_text = ( |
| 327 | """[{"name": "get_current_weather", "arguments": {"city": "Seattle"}}]""" # noqa: E501 |
| 328 | ) |
| 329 | |
| 330 | # First call to set up the tool |
| 331 | xlam_tool_parser.extract_tool_calls_streaming( |
| 332 | previous_text="", |
| 333 | current_text=current_text, |
| 334 | delta_text="]", |
| 335 | previous_token_ids=[], |
| 336 | current_token_ids=[], |
| 337 | delta_token_ids=[], |
| 338 | request=None, |
| 339 | ) |
| 340 | |
| 341 | # Make sure the tool is set up correctly |
| 342 | assert xlam_tool_parser.current_tool_id >= 0, "Tool index should be initialized" |
| 343 | |
| 344 | # Manually set up the state for sending the tool name |
| 345 | xlam_tool_parser.current_tools_sent = [False] |
| 346 | |
| 347 | # Call to send the function name |
| 348 | result = xlam_tool_parser.extract_tool_calls_streaming( |
| 349 | previous_text=current_text, |
| 350 | current_text=current_text, |
| 351 | delta_text="", |
| 352 | previous_token_ids=[], |
| 353 | current_token_ids=[], |
| 354 | delta_token_ids=[], |
| 355 | request=None, |
| 356 | ) |
| 357 | |
| 358 | # Check that we get a result with the proper tool call |
| 359 | if result is not None: |
| 360 | assert hasattr(result, "tool_calls") |
| 361 | assert len(result.tool_calls) == 1 |
| 362 | assert result.tool_calls[0].function.name == "get_current_weather" |
| 363 | |
| 364 | |
| 365 | @pytest.mark.parametrize( |
| 366 | ids=[ |
| 367 | "parallel_tool_calls", |
| 368 | "single_tool_with_think_tag", |
| 369 | "single_tool_with_json_code_block", |
| 370 | "single_tool_with_tool_calls_tag", |
| 371 | "single_tool_with_tool_call_xml_tags", |
| 372 | ], |
| 373 | argnames=["model_output", "expected_tool_calls", "expected_content"], |
| 374 | argvalues=[ |
| 375 | ( |
| 376 | """[{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}, {"name": "get_current_weather", "arguments": {"city": "Orlando", "state": "FL", "unit": "fahrenheit"}}]""", # noqa: E501 |
| 377 | [ |
| 378 | ToolCall( |
| 379 | function=FunctionCall( |
| 380 | name="get_current_weather", |
| 381 | arguments=json.dumps( |
| 382 | { |
| 383 | "city": "Dallas", |
| 384 | "state": "TX", |
| 385 | "unit": "fahrenheit", |
| 386 | } |
| 387 | ), |
| 388 | ) |
| 389 | ), |
| 390 | ToolCall( |
| 391 | function=FunctionCall( |
| 392 | name="get_current_weather", |
| 393 | arguments=json.dumps( |
| 394 | { |
| 395 | "city": "Orlando", |
| 396 | "state": "FL", |
| 397 | "unit": "fahrenheit", |
| 398 | } |
| 399 | ), |
| 400 | ) |
| 401 | ), |
| 402 | ], |
| 403 | "", |
| 404 | ), |
| 405 | ( |
| 406 | """<think>I'll help you with that.</think>[{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}]""", # noqa: E501 |
| 407 | [ |
| 408 | ToolCall( |
| 409 | function=FunctionCall( |
| 410 | name="get_current_weather", |
| 411 | arguments=json.dumps( |
| 412 | { |
| 413 | "city": "Dallas", |
| 414 | "state": "TX", |
| 415 | "unit": "fahrenheit", |
| 416 | } |
| 417 | ), |
| 418 | ) |
| 419 | ) |
| 420 | ], |
| 421 | "<think>I'll help you with that.</think>", |
| 422 | ), |
| 423 | ( |
| 424 | """```json\n[{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}]\n```""", # noqa: E501 |
| 425 | [ |
| 426 | ToolCall( |
| 427 | function=FunctionCall( |
| 428 | name="get_current_weather", |
| 429 | arguments=json.dumps( |
| 430 | { |
| 431 | "city": "Dallas", |
| 432 | "state": "TX", |
| 433 | "unit": "fahrenheit", |
| 434 | } |
| 435 | ), |
| 436 | ) |
| 437 | ) |
| 438 | ], |
| 439 | "", |
| 440 | ), |
| 441 | ( |
| 442 | """[TOOL_CALLS][{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}]""", # noqa: E501 |
| 443 | [ |
| 444 | ToolCall( |
| 445 | function=FunctionCall( |
| 446 | name="get_current_weather", |
| 447 | arguments=json.dumps( |
| 448 | { |
| 449 | "city": "Dallas", |
| 450 | "state": "TX", |
| 451 | "unit": "fahrenheit", |
| 452 | } |
| 453 | ), |
| 454 | ) |
| 455 | ) |
| 456 | ], |
| 457 | "", |
| 458 | ), |
| 459 | ( |
| 460 | """I can help with that.<tool_call>[{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}]</tool_call>""", # noqa: E501 |
| 461 | [ |
| 462 | ToolCall( |
| 463 | function=FunctionCall( |
| 464 | name="get_current_weather", |
| 465 | arguments=json.dumps( |
| 466 | { |
| 467 | "city": "Dallas", |
| 468 | "state": "TX", |
| 469 | "unit": "fahrenheit", |
| 470 | } |
| 471 | ), |
| 472 | ) |
| 473 | ) |
| 474 | ], |
| 475 | "I can help with that.", |
| 476 | ), |
| 477 | ], |
| 478 | ) |
| 479 | def test_extract_tool_calls_streaming_incremental( |
| 480 | xlam_tool_parser, |
| 481 | xlam_tokenizer, |
| 482 | model_output, |
| 483 | expected_tool_calls, |
| 484 | expected_content, |
| 485 | ): |
| 486 | """Verify the XLAM Parser streaming behavior by verifying each chunk is as expected.""" # noqa: E501 |
| 487 | request = ChatCompletionRequest(model=MODEL, messages=[]) |
| 488 | |
| 489 | chunks = [] |
| 490 | for delta_message in stream_delta_message_generator( |
| 491 | xlam_tool_parser, xlam_tokenizer, model_output, request |
| 492 | ): |
| 493 | chunks.append(delta_message) |
| 494 | |
| 495 | # Should have multiple chunks |
| 496 | assert len(chunks) >= 3 |
| 497 | |
| 498 | # Should have a chunk with tool header (id, name, type) for the first tool call # noqa: E501 |
| 499 | header_found = False |
| 500 | expected_first_tool = expected_tool_calls[0] |
| 501 | for chunk in chunks: |
| 502 | if chunk.tool_calls and chunk.tool_calls[0].id: |
| 503 | header_found = True |
| 504 | assert ( |
| 505 | chunk.tool_calls[0].function.name == expected_first_tool.function.name |
| 506 | ) |
| 507 | assert chunk.tool_calls[0].type == "function" |
| 508 | # Arguments may be empty initially or None |
| 509 | if chunk.tool_calls[0].function.arguments is not None: |
| 510 | # If present, should be empty string initially |
| 511 | assert chunk.tool_calls[0].function.arguments == "" |
| 512 | break |
| 513 | assert header_found |
| 514 | |
| 515 | # Should have chunks with incremental arguments |
| 516 | arg_chunks = [] |
| 517 | for chunk in chunks: |
| 518 | if ( |
| 519 | chunk.tool_calls |
| 520 | and chunk.tool_calls[0].function.arguments |
| 521 | and chunk.tool_calls[0].function.arguments != "" |
| 522 | and chunk.tool_calls[0].index |
| 523 | == 0 # Only collect arguments from the first tool call |
| 524 | ): |
| 525 | arg_chunks.append(chunk.tool_calls[0].function.arguments) |
| 526 | |
| 527 | # Arguments should be streamed incrementally |
| 528 | assert len(arg_chunks) > 1 |
| 529 | |
| 530 | # Concatenated arguments should form valid JSON for the first tool call |
| 531 | full_args = "".join(arg_chunks) |
| 532 | parsed_args = json.loads(full_args) |
| 533 | expected_args = json.loads(expected_first_tool.function.arguments) |
| 534 | assert parsed_args == expected_args |
| 535 | |
| 536 | |
| 537 | @pytest.mark.parametrize("streaming", [False, True]) |
| 538 | def test_extract_tool_calls_non_ascii(xlam_tool_parser, xlam_tokenizer, streaming): |
| 539 | # Use parallel tool calls so the streaming path re-serializes arguments |
| 540 | # (the ensure_ascii fix only runs when tool_count > 1). |
| 541 | model_output = """[{"name": "get_current_weather", "arguments": {"city": "北京"}}, {"name": "get_current_weather", "arguments": {"city": "上海"}}]""" # noqa: E501 |
| 542 | |
| 543 | if streaming: |
| 544 | request = ChatCompletionRequest(model=MODEL, messages=[]) |
| 545 | args = "".join( |
| 546 | delta.tool_calls[0].function.arguments |
| 547 | for delta in stream_delta_message_generator( |
| 548 | xlam_tool_parser, xlam_tokenizer, model_output, request |
| 549 | ) |
| 550 | if delta.tool_calls and delta.tool_calls[0].function.arguments |
| 551 | ) |
| 552 | else: |
| 553 | extracted = xlam_tool_parser.extract_tool_calls(model_output, request=None) |
| 554 | args = "".join(tc.function.arguments for tc in extracted.tool_calls) |
| 555 | |
| 556 | assert "北京" in args |
| 557 | assert "\\u" not in args |