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, 23:45 UTC
[Bugfix] Improve JSON extraction in LlamaToolParser (#19024)
Signed-off-by: keru <keyang.ru@oracle.com> Co-authored-by: keru <keyang.ru@oracle.com>
Keyang Ru committed
Jul 28, 2025 at 05:36 UTC
9ace2eaf352b4339685db59d5367784e59fd5df1
2 files changed
+174
-40
tests/entrypoints/openai/tool_parsers/test_llama3_json_tool_parser.py
new
+132
@@ -0,0 +1,132 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+from transformers import AutoTokenizer
6
+
7
+from vllm.entrypoints.openai.protocol import ExtractedToolCallInformation
8
+from vllm.entrypoints.openai.tool_parsers.llama_tool_parser import (
9
+ Llama3JsonToolParser)
10
+
11
+
12
+@pytest.fixture
13
+def parser():
14
+ # Use a small tokenizer for testing
15
+ tokenizer = AutoTokenizer.from_pretrained("gpt2")
16
+ return Llama3JsonToolParser(tokenizer)
17
+
18
+
19
+def test_extract_tool_calls_simple(parser):
20
+ # Test with a simple tool call
21
+ model_output = ('Here is the result: {"name": "getOpenIncidentsTool", '
22
+ '"parameters": {}} Would you like to know more?')
23
+ result = parser.extract_tool_calls(model_output, None)
24
+
25
+ assert isinstance(result, ExtractedToolCallInformation)
26
+ assert result.tools_called is True
27
+ assert len(result.tool_calls) == 1
28
+ assert result.tool_calls[0].type == "function"
29
+ assert result.tool_calls[0].function.name == "getOpenIncidentsTool"
30
+ assert result.tool_calls[0].function.arguments == "{}"
31
+ assert result.content is None
32
+
33
+
34
+def test_extract_tool_calls_with_arguments(parser):
35
+ # Test with a tool call that has arguments
36
+ model_output = (
37
+ '{"name": "searchTool", "parameters": {"query": "test query", '
38
+ '"limit": 10}}')
39
+ result = parser.extract_tool_calls(model_output, None)
40
+
41
+ assert result.tools_called is True
42
+ assert len(result.tool_calls) == 1
43
+ assert result.tool_calls[0].function.name == "searchTool"
44
+ assert '"query": "test query"' in result.tool_calls[0].function.arguments
45
+ assert '"limit": 10' in result.tool_calls[0].function.arguments
46
+
47
+
48
+def test_extract_tool_calls_no_json(parser):
49
+ # Test with text that doesn't contain a JSON object
50
+ model_output = "This is just some text without any tool calls"
51
+ result = parser.extract_tool_calls(model_output, None)
52
+
53
+ assert result.tools_called is False
54
+ assert len(result.tool_calls) == 0
55
+ assert result.content == model_output
56
+
57
+
58
+def test_extract_tool_calls_invalid_json(parser):
59
+ # Test with invalid JSON
60
+ model_output = '{"name": "invalidTool", "parameters": {invalid json}'
61
+ result = parser.extract_tool_calls(model_output, None)
62
+
63
+ assert result.tools_called is False
64
+ assert len(result.tool_calls) == 0
65
+ assert result.content == model_output
66
+
67
+
68
+def test_extract_tool_calls_with_arguments_key(parser):
69
+ # Test with a tool call that uses "arguments" instead of "parameters"
70
+ model_output = '{"name": "searchTool", "arguments": {"query": "test"}}'
71
+ result = parser.extract_tool_calls(model_output, None)
72
+
73
+ assert result.tools_called is True
74
+ assert len(result.tool_calls) == 1
75
+ assert result.tool_calls[0].function.name == "searchTool"
76
+ assert '"query": "test"' in result.tool_calls[0].function.arguments
77
+
78
+
79
+def test_extract_tool_calls_multiple_json(parser):
80
+ # Test with multiple JSONs separated by semicolons
81
+ model_output = (
82
+ '{"name": "searchTool", "parameters": {"query": "test1"}}; '
83
+ '{"name": "getOpenIncidentsTool", "parameters": {}}; '
84
+ '{"name": "searchTool", "parameters": {"query": "test2"}}')
85
+ result = parser.extract_tool_calls(model_output, None)
86
+
87
+ assert result.tools_called is True
88
+ assert len(result.tool_calls) == 3
89
+
90
+ # Check first tool call
91
+ assert result.tool_calls[0].function.name == "searchTool"
92
+ assert '"query": "test1"' in result.tool_calls[0].function.arguments
93
+
94
+ # Check second tool call
95
+ assert result.tool_calls[1].function.name == "getOpenIncidentsTool"
96
+ assert result.tool_calls[1].function.arguments == "{}"
97
+
98
+ # Check third tool call
99
+ assert result.tool_calls[2].function.name == "searchTool"
100
+ assert '"query": "test2"' in result.tool_calls[2].function.arguments
101
+
102
+
103
+def test_extract_tool_calls_multiple_json_with_whitespace(parser):
104
+ # Test with multiple JSONs separated by semicolons and extra whitespace
105
+ model_output = (
106
+ '{"name": "searchTool", "parameters": {"query": "test1"}} ; '
107
+ '{"name": "getOpenIncidentsTool", "parameters": {}} ; '
108
+ '{"name": "searchTool", "parameters": {"query": "test2"}}')
109
+ result = parser.extract_tool_calls(model_output, None)
110
+
111
+ assert result.tools_called is True
112
+ assert len(result.tool_calls) == 3
113
+ assert result.tool_calls[0].function.name == "searchTool"
114
+ assert result.tool_calls[1].function.name == "getOpenIncidentsTool"
115
+ assert result.tool_calls[2].function.name == "searchTool"
116
+
117
+
118
+def test_extract_tool_calls_multiple_json_with_surrounding_text(parser):
119
+ # Test with multiple JSONs and surrounding text
120
+ model_output = (
121
+ 'Here are the results: '
122
+ '{"name": "searchTool", "parameters": {"query": "test1"}}; '
123
+ '{"name": "getOpenIncidentsTool", "parameters": {}}; '
124
+ '{"name": "searchTool", "parameters": {"query": "test2"}} '
125
+ 'Would you like to know more?')
126
+ result = parser.extract_tool_calls(model_output, None)
127
+
128
+ assert result.tools_called is True
129
+ assert len(result.tool_calls) == 3
130
+ assert result.tool_calls[0].function.name == "searchTool"
131
+ assert result.tool_calls[1].function.name == "getOpenIncidentsTool"
132
+ assert result.tool_calls[2].function.name == "searchTool"
vllm/entrypoints/openai/tool_parsers/llama_tool_parser.py
+42
-40
@@ -3,7 +3,6 @@
3
4
import json
5
from collections.abc import Sequence
6
-from json import JSONDecoder
6
from typing import Union
7
8
import partial_json_parser
@@ -31,11 +30,11 @@ logger = init_logger(__name__)
30
@ToolParserManager.register_module("llama4_json")
31
class Llama3JsonToolParser(ToolParser):
32
"""
34
- Tool call parser for Llama 3.1 models intended for use with the
33
+ Tool call parser for Llama 3.x and 4 models intended for use with the
34
examples/tool_chat_template_llama.jinja template.
35
37
- Used when --enable-auto-tool-choice --tool-call-parser llama3_json
38
- are all set
36
+ Used when --enable-auto-tool-choice --tool-call-parser llama3_json or
37
+ llama4_json are set.
38
"""
39
40
def __init__(self, tokenizer: PreTrainedTokenizerBase):
@@ -51,54 +50,57 @@ class Llama3JsonToolParser(ToolParser):
50
self.bot_token = "<|python_tag|>"
51
self.bot_token_id = tokenizer.encode(self.bot_token,
52
add_special_tokens=False)[0]
54
- self.tool_call_regex = re.compile(r"\[{.*?}\]", re.DOTALL)
53
+ # Updated regex to match multiple JSONs separated by semicolons
54
+ # This pattern is more robust and can handle nested JSON objects
55
+ self.tool_call_regex = re.compile(
56
+ r'{[^{}]*(?:{[^{}]*}[^{}]*)*}(?:\s*;\s*{[^{}]*(?:{[^{}]*}[^{}]*)*})*',
57
+ re.DOTALL)
58
59
def extract_tool_calls(
60
self, model_output: str,
61
request: ChatCompletionRequest) -> ExtractedToolCallInformation:
62
"""
63
Extract the tool calls from a complete model response.
64
+ Only extracts JSON content and ignores any surrounding plain text.
65
+ Supports both single JSON and multiple JSONs separated by semicolons.
66
"""
62
- # case -- if a tool call token is not present, return a text response
63
- if not (model_output.startswith(self.bot_token)
64
- or model_output.startswith('{')):
67
+ # Quick check before running regex
68
+ if not (self.bot_token in model_output or '{' in model_output):
69
+ return ExtractedToolCallInformation(tools_called=False,
70
+ tool_calls=[],
71
+ content=model_output)
72
+
73
+ # Find JSON object(s) in the text using regex
74
+ match = self.tool_call_regex.search(model_output)
75
+ if not match:
76
return ExtractedToolCallInformation(tools_called=False,
77
tool_calls=[],
78
content=model_output)
79
80
try:
70
- # load the JSON, and then use it to build the Function and
71
- # Tool Call
72
- dec = JSONDecoder()
73
- function_call_arr = []
74
-
75
- # depending on the prompt format the Llama model may or may not
76
- # prefix the output with the <|python_tag|> token
77
- start_idx = len(self.bot_token) if model_output.startswith(
78
- self.bot_token) else 0
79
- while start_idx < len(model_output):
80
- (obj, end_idx) = dec.raw_decode(model_output[start_idx:])
81
- start_idx += end_idx + len('; ')
82
- function_call_arr.append(obj)
83
-
84
- tool_calls: list[ToolCall] = [
85
- ToolCall(
86
- type="function",
87
- function=FunctionCall(
88
- name=raw_function_call["name"],
89
- # function call args are JSON but as a string
90
- arguments=json.dumps(raw_function_call["arguments"] \
91
- if "arguments" in raw_function_call \
92
- else raw_function_call["parameters"],
93
- ensure_ascii=False)))
94
- for raw_function_call in function_call_arr
95
- ]
96
-
97
- # get any content before the tool call
98
- ret = ExtractedToolCallInformation(tools_called=True,
99
- tool_calls=tool_calls,
100
- content=None)
101
- return ret
81
+ json_str = match.group(0)
82
+ # Split by semicolon and strip whitespace
83
+ json_objects = [obj.strip() for obj in json_str.split(';')]
84
+
85
+ tool_calls: list[ToolCall] = []
86
+ for json_obj in json_objects:
87
+ if not json_obj: # Skip empty strings
88
+ continue
89
+ obj = json.loads(json_obj)
90
+ tool_calls.append(
91
+ ToolCall(
92
+ type="function",
93
+ function=FunctionCall(
94
+ name=obj["name"],
95
+ # function call args are JSON but as a string
96
+ arguments=json.dumps(
97
+ obj["arguments"]
98
+ if "arguments" in obj else obj["parameters"],
99
+ ensure_ascii=False))))
100
+
101
+ return ExtractedToolCallInformation(tools_called=True,
102
+ tool_calls=tool_calls,
103
+ content=None)
104
105
except Exception:
106
logger.exception("Error in extracting tool call from response.")