A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Bugfix] Add support for `<tool_call>` format in streaming mode for XLAM Tool Parser (#22769)
Signed-off-by: Devon Peroutky <devon@kindo.ai>
Code Jesus committed
Aug 31, 2025 at 23:07 UTC
422e793fa6c4381809ca236946ab5c2206ad59c1
2 files changed
+297
-25
tests/tool_use/test_xlam_tool_parser.py
+216
-2
@@ -2,12 +2,17 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import json
5
+from collections.abc import Generator
6
+from typing import Optional
7
8
import pytest
9
8
-from vllm.entrypoints.openai.protocol import FunctionCall, ToolCall
10
+from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
11
+ DeltaMessage, FunctionCall,
12
+ ToolCall)
13
from vllm.entrypoints.openai.tool_parsers import xLAMToolParser
10
-from vllm.transformers_utils.tokenizer import get_tokenizer
14
+from vllm.transformers_utils.detokenizer import detokenize_incrementally
15
+from vllm.transformers_utils.tokenizer import AnyTokenizer, get_tokenizer
16
17
# Use a common model that is likely to be available
18
MODEL = "Salesforce/Llama-xLAM-2-8B-fc-r"
@@ -36,6 +41,56 @@ def assert_tool_calls(actual_tool_calls: list[ToolCall],
41
assert actual_tool_call.function == expected_tool_call.function
42
43
44
+def stream_delta_message_generator(
45
+ xlam_tool_parser: xLAMToolParser,
46
+ xlam_tokenizer: AnyTokenizer,
47
+ model_output: str,
48
+ request: Optional[ChatCompletionRequest] = None,
49
+) -> Generator[DeltaMessage, None, None]:
50
+ all_token_ids = xlam_tokenizer.encode(model_output,
51
+ add_special_tokens=False)
52
+
53
+ previous_text = ""
54
+ previous_tokens = None
55
+ prefix_offset = 0
56
+ read_offset = 0
57
+ for i, delta_token in enumerate(all_token_ids):
58
+ delta_token_ids = [delta_token]
59
+ previous_token_ids = all_token_ids[:i]
60
+ current_token_ids = all_token_ids[:i + 1]
61
+
62
+ (new_tokens, delta_text, new_prefix_offset,
63
+ new_read_offset) = (detokenize_incrementally(
64
+ tokenizer=xlam_tokenizer,
65
+ all_input_ids=current_token_ids,
66
+ prev_tokens=previous_tokens,
67
+ prefix_offset=prefix_offset,
68
+ read_offset=read_offset,
69
+ skip_special_tokens=False,
70
+ spaces_between_special_tokens=True,
71
+ ))
72
+
73
+ current_text = previous_text + delta_text
74
+
75
+ delta_message = xlam_tool_parser.extract_tool_calls_streaming(
76
+ previous_text,
77
+ current_text,
78
+ delta_text,
79
+ previous_token_ids,
80
+ current_token_ids,
81
+ delta_token_ids,
82
+ request=request,
83
+ )
84
+ if delta_message:
85
+ yield delta_message
86
+
87
+ previous_text = current_text
88
+ previous_tokens = (previous_tokens +
89
+ new_tokens if previous_tokens else new_tokens)
90
+ prefix_offset = new_prefix_offset
91
+ read_offset = new_read_offset
92
+
93
+
94
def test_extract_tool_calls_no_tools(xlam_tool_parser):
95
model_output = "This is a test"
96
extracted_tool_calls = xlam_tool_parser.extract_tool_calls(
@@ -51,6 +106,7 @@ def test_extract_tool_calls_no_tools(xlam_tool_parser):
106
"single_tool_with_think_tag",
107
"single_tool_with_json_code_block",
108
"single_tool_with_tool_calls_tag",
109
+ "single_tool_with_tool_call_xml_tags",
110
],
111
argnames=["model_output", "expected_tool_calls", "expected_content"],
112
argvalues=[
@@ -118,6 +174,20 @@ def test_extract_tool_calls_no_tools(xlam_tool_parser):
174
],
175
"I'll check the weather for you.",
176
),
177
+ (
178
+ """I'll help you check the weather.<tool_call>[{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}]</tool_call>""", # noqa: E501
179
+ [
180
+ ToolCall(function=FunctionCall(
181
+ name="get_current_weather",
182
+ arguments=json.dumps({
183
+ "city": "Dallas",
184
+ "state": "TX",
185
+ "unit": "fahrenheit",
186
+ }),
187
+ ))
188
+ ],
189
+ "I'll help you check the weather.",
190
+ ),
191
],
192
)
193
def test_extract_tool_calls(xlam_tool_parser, model_output,
@@ -245,3 +315,147 @@ def test_streaming_with_list_structure(xlam_tool_parser):
315
assert hasattr(result, "tool_calls")
316
assert len(result.tool_calls) == 1
317
assert result.tool_calls[0].function.name == "get_current_weather"
318
+
319
+
320
+@pytest.mark.parametrize(
321
+ ids=[
322
+ "parallel_tool_calls",
323
+ "single_tool_with_think_tag",
324
+ "single_tool_with_json_code_block",
325
+ "single_tool_with_tool_calls_tag",
326
+ "single_tool_with_tool_call_xml_tags",
327
+ ],
328
+ argnames=["model_output", "expected_tool_calls", "expected_content"],
329
+ argvalues=[
330
+ (
331
+ """[{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}, {"name": "get_current_weather", "arguments": {"city": "Orlando", "state": "FL", "unit": "fahrenheit"}}]""", # noqa: E501
332
+ [
333
+ ToolCall(function=FunctionCall(
334
+ name="get_current_weather",
335
+ arguments=json.dumps({
336
+ "city": "Dallas",
337
+ "state": "TX",
338
+ "unit": "fahrenheit",
339
+ }),
340
+ )),
341
+ ToolCall(function=FunctionCall(
342
+ name="get_current_weather",
343
+ arguments=json.dumps({
344
+ "city": "Orlando",
345
+ "state": "FL",
346
+ "unit": "fahrenheit",
347
+ }),
348
+ )),
349
+ ],
350
+ "",
351
+ ),
352
+ (
353
+ """<think>I'll help you with that.</think>[{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}]""", # noqa: E501
354
+ [
355
+ ToolCall(function=FunctionCall(
356
+ name="get_current_weather",
357
+ arguments=json.dumps({
358
+ "city": "Dallas",
359
+ "state": "TX",
360
+ "unit": "fahrenheit",
361
+ }),
362
+ ))
363
+ ],
364
+ "<think>I'll help you with that.</think>",
365
+ ),
366
+ (
367
+ """```json\n[{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}]\n```""", # noqa: E501
368
+ [
369
+ ToolCall(function=FunctionCall(
370
+ name="get_current_weather",
371
+ arguments=json.dumps({
372
+ "city": "Dallas",
373
+ "state": "TX",
374
+ "unit": "fahrenheit",
375
+ }),
376
+ ))
377
+ ],
378
+ "",
379
+ ),
380
+ (
381
+ """[TOOL_CALLS][{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}]""", # noqa: E501
382
+ [
383
+ ToolCall(function=FunctionCall(
384
+ name="get_current_weather",
385
+ arguments=json.dumps({
386
+ "city": "Dallas",
387
+ "state": "TX",
388
+ "unit": "fahrenheit",
389
+ }),
390
+ ))
391
+ ],
392
+ "",
393
+ ),
394
+ (
395
+ """I can help with that.<tool_call>[{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}]</tool_call>""", # noqa: E501
396
+ [
397
+ ToolCall(function=FunctionCall(
398
+ name="get_current_weather",
399
+ arguments=json.dumps({
400
+ "city": "Dallas",
401
+ "state": "TX",
402
+ "unit": "fahrenheit",
403
+ }),
404
+ ))
405
+ ],
406
+ "I can help with that.",
407
+ ),
408
+ ],
409
+)
410
+def test_extract_tool_calls_streaming_incremental(
411
+ xlam_tool_parser,
412
+ xlam_tokenizer,
413
+ model_output,
414
+ expected_tool_calls,
415
+ expected_content,
416
+):
417
+ """Verify the XLAM Parser streaming behavior by verifying each chunk is as expected.""" # noqa: E501
418
+ request = ChatCompletionRequest(model=MODEL, messages=[], tools=[])
419
+
420
+ chunks = []
421
+ for delta_message in stream_delta_message_generator(
422
+ xlam_tool_parser, xlam_tokenizer, model_output, request):
423
+ chunks.append(delta_message)
424
+
425
+ # Should have multiple chunks
426
+ assert len(chunks) >= 3
427
+
428
+ # Should have a chunk with tool header (id, name, type) for the first tool call # noqa: E501
429
+ header_found = False
430
+ expected_first_tool = expected_tool_calls[0]
431
+ for chunk in chunks:
432
+ if chunk.tool_calls and chunk.tool_calls[0].id:
433
+ header_found = True
434
+ assert (chunk.tool_calls[0].function.name ==
435
+ expected_first_tool.function.name)
436
+ assert chunk.tool_calls[0].type == "function"
437
+ # Arguments may be empty initially or None
438
+ if chunk.tool_calls[0].function.arguments is not None:
439
+ # If present, should be empty string initially
440
+ assert chunk.tool_calls[0].function.arguments == ""
441
+ break
442
+ assert header_found
443
+
444
+ # Should have chunks with incremental arguments
445
+ arg_chunks = []
446
+ for chunk in chunks:
447
+ if (chunk.tool_calls and chunk.tool_calls[0].function.arguments
448
+ and chunk.tool_calls[0].function.arguments != ""
449
+ and chunk.tool_calls[0].index ==
450
+ 0 # Only collect arguments from the first tool call
451
+ ):
452
+ arg_chunks.append(chunk.tool_calls[0].function.arguments)
453
+
454
+ # Arguments should be streamed incrementally
455
+ assert len(arg_chunks) > 1
456
+
457
+ # Concatenated arguments should form valid JSON for the first tool call
458
+ full_args = "".join(arg_chunks)
459
+ parsed_args = json.loads(full_args)
460
+ expected_args = json.loads(expected_first_tool.function.arguments)
461
+ assert parsed_args == expected_args
vllm/entrypoints/openai/tool_parsers/xlam_tool_parser.py
+81
-23
@@ -186,11 +186,31 @@ class xLAMToolParser(ToolParser):
186
"""
187
Extract tool calls for streaming mode.
188
"""
189
- # Simplify detection: if it begins with "[" treat it as a function call
190
- is_function_call = (current_text.strip().startswith("["))
191
-
192
- # If not a function call, return normal content
193
- if not is_function_call:
189
+ # First, check for a definitive start of a tool call block.
190
+ # This prevents premature parsing of incomplete output.
191
+ stripped_text = current_text.strip()
192
+ preprocessed_content, preprocessed_tool_calls = (
193
+ self.preprocess_model_output(current_text))
194
+
195
+ # For JSON code blocks, we need to detect them earlier, even if incomplete
196
+ has_potential_json_block = ("```json" in current_text
197
+ or "```\n[" in current_text
198
+ or "[TOOL_CALLS]" in current_text
199
+ or "<tool_call>" in current_text)
200
+
201
+ is_tool_call_block = (
202
+ stripped_text.startswith("[")
203
+ or stripped_text.startswith("<tool_call>")
204
+ or stripped_text.startswith("[TOOL_CALLS]") or
205
+ # Check if we have thinking tags with JSON-like content following
206
+ ("</think>[" in current_text) or
207
+ # Check if the text contains a JSON array after preprocessing
208
+ preprocessed_tool_calls is not None or
209
+ # For JSON code blocks, detect early if we see enough structure
210
+ (has_potential_json_block and '"name"' in current_text
211
+ and '"arguments"' in current_text))
212
+
213
+ if not is_tool_call_block:
214
return DeltaMessage(content=delta_text)
215
216
try:
@@ -204,7 +224,10 @@ class xLAMToolParser(ToolParser):
224
225
# Try parsing as JSON to check for complete tool calls
226
try:
207
- parsed_tools = json.loads(current_text)
227
+ # Use preprocessed tool calls if available
228
+ tool_calls_text = (preprocessed_tool_calls if
229
+ preprocessed_tool_calls else current_text)
230
+ parsed_tools = json.loads(tool_calls_text)
231
if isinstance(parsed_tools, list):
232
# Update our tool array for next time
233
self.prev_tool_call_arr = parsed_tools
@@ -257,13 +280,40 @@ class xLAMToolParser(ToolParser):
280
return delta
281
282
# Use regex to identify tool calls in the output
283
+ # Use preprocessed tool calls text for better parsing, but also try to extract from incomplete JSON blocks
284
+ search_text = (preprocessed_tool_calls
285
+ if preprocessed_tool_calls else current_text)
286
+
287
+ # For JSON code blocks that aren't complete yet, try to extract the JSON content
288
+ if not preprocessed_tool_calls and has_potential_json_block:
289
+ # Try to extract the JSON array from within the code block
290
+ json_match = re.search(r"```(?:json)?\s*([\s\S]*?)(?:```|$)",
291
+ current_text)
292
+ if json_match:
293
+ potential_json = json_match.group(1).strip()
294
+ # Use this as search text even if it's incomplete
295
+ if potential_json.startswith("[") and (
296
+ '"name"' in potential_json
297
+ and '"arguments"' in potential_json):
298
+ search_text = potential_json
299
+
300
+ # Try to find complete tool names first
301
name_pattern = r'"name"\s*:\s*"([^"]+)"'
261
- name_matches = list(re.finditer(name_pattern, current_text))
302
+ name_matches = list(re.finditer(name_pattern, search_text))
303
tool_count = len(name_matches)
304
264
- # If no tools found yet, return
305
+ # If no complete tool names found, check for partial tool names
306
if tool_count == 0:
266
- return None
307
+ # Check if we're in the middle of parsing a tool name
308
+ partial_name_pattern = r'"name"\s*:\s*"([^"]*)'
309
+ partial_matches = list(
310
+ re.finditer(partial_name_pattern, search_text))
311
+ if partial_matches:
312
+ # We have a partial tool name - not ready to emit yet
313
+ return None
314
+ else:
315
+ # No tools found at all
316
+ return None
317
318
# Ensure our state arrays are large enough
319
while len(self.streaming_state["sent_tools"]) < tool_count:
@@ -332,7 +382,7 @@ class xLAMToolParser(ToolParser):
382
# First, check for the empty arguments case: "arguments": {}
383
empty_args_pattern = (
384
r'"name"\s*:\s*"[^"]+"\s*,\s*"arguments"\s*:\s*\{\s*\}')
335
- empty_args_match = re.search(empty_args_pattern, current_text)
385
+ empty_args_match = re.search(empty_args_pattern, search_text)
386
387
# Check if this tool has empty arguments
388
if empty_args_match and empty_args_match.start() > 0:
@@ -376,7 +426,7 @@ class xLAMToolParser(ToolParser):
426
427
# Extract arguments for current tool using regex for non-empty arguments
428
args_pattern = r'"name"\s*:\s*"[^"]+"\s*,\s*"arguments"\s*:\s*(\{(?:[^{}]|(?:\{[^{}]*\}))*\})'
379
- args_matches = list(re.finditer(args_pattern, current_text))
429
+ args_matches = list(re.finditer(args_pattern, search_text))
430
431
if current_idx < len(args_matches):
432
args_text = args_matches[current_idx].group(1)
@@ -384,17 +434,25 @@ class xLAMToolParser(ToolParser):
434
# Handle transition between tools
435
is_last_tool = current_idx == tool_count - 1
436
387
- # Find where the arguments for our current tool end
388
- if not is_last_tool:
389
- # If we have more tools after this one, try to find the complete argument block
390
- next_tool_pos = current_text.find(
391
- "},{", args_matches[current_idx].start())
392
- if next_tool_pos != -1:
393
- args_end_pos = (next_tool_pos + 1
394
- ) # +1 to include the '}'
395
- args_text = (current_text[args_matches[current_idx]
396
- .start():args_end_pos].
397
- split('"arguments":')[1].strip())
437
+ # For multiple tools, extract only the arguments for the current tool
438
+ if tool_count > 1:
439
+ # Parse the entire JSON structure to properly extract arguments for each tool
440
+ try:
441
+ parsed_tools = json.loads(search_text)
442
+ if isinstance(
443
+ parsed_tools,
444
+ list) and current_idx < len(parsed_tools):
445
+ current_tool = parsed_tools[current_idx]
446
+ if isinstance(current_tool.get("arguments"),
447
+ dict):
448
+ args_text = json.dumps(
449
+ current_tool["arguments"])
450
+ else:
451
+ args_text = str(
452
+ current_tool.get("arguments", "{}"))
453
+ except (json.JSONDecodeError, KeyError, IndexError):
454
+ # Fallback to regex-based extraction
455
+ pass
456
457
# If arguments haven't been sent yet
458
sent_args = self.streaming_state["sent_tools"][
@@ -419,7 +477,7 @@ class xLAMToolParser(ToolParser):
477
index=current_idx,
478
function=DeltaFunctionCall(
479
arguments="{").model_dump(
422
- exclude_none=True), # type: ignore
480
+ exclude_none=True), # type: ignore
481
)
482
])
483
return delta