Skip to content

Commit bd8400f

Browse files
bbartelsyiyeguhuchaunceyjiang
authored and
bringlein
committed
[Bugfix] Fix Kimi-K2 tool parser concatenated tool calls parsing (vllm-project#28831)
Signed-off-by: Thomas Mao <yiyeguhu@gmail.com> Signed-off-by: bbartels <benjamin@bartels.dev> Co-authored-by: Thomas Mao <yiyeguhu@gmail.com> Co-authored-by: Chauncey <chaunceyjiang@gmail.com>
1 parent 630a977 commit bd8400f

File tree

2 files changed

+124
-1
lines changed

2 files changed

+124
-1
lines changed

‎tests/tool_use/test_kimi_k2_tool_parser.py‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,11 @@ def test_extract_tool_calls_no_tools(kimi_k2_tool_parser):
6060
ids=[
6161
"tool_call_with_content_before",
6262
"multi_tool_call_with_content_before",
63+
"concatenated_tool_calls_bug_fix",
64+
"three_concatenated_tool_calls",
65+
"mixed_spacing_tool_calls",
66+
"angle_brackets_in_json",
67+
"newlines_in_json",
6368
],
6469
argnames=["model_output", "expected_tool_calls", "expected_content"],
6570
argvalues=[
@@ -114,6 +119,123 @@ def test_extract_tool_calls_no_tools(kimi_k2_tool_parser):
114119
],
115120
"I'll help you check the weather. ",
116121
),
122+
(
123+
"""I'll get the weather and news for LA today. First, let me get the weather using Los Angeles coordinates, and then get the latest news. <|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{"latitude": 34.0522, "longitude": -118.2437}<|tool_call_end|><|tool_call_begin|>functions.get_news:1<|tool_call_argument_begin|>{"content": "Los Angeles today"}<|tool_call_end|><|tool_calls_section_end|>""",
124+
[
125+
ToolCall(
126+
id="functions.get_weather:0",
127+
function=FunctionCall(
128+
name="get_weather",
129+
arguments=json.dumps(
130+
{"latitude": 34.0522, "longitude": -118.2437}
131+
),
132+
),
133+
type="function",
134+
),
135+
ToolCall(
136+
id="functions.get_news:1",
137+
function=FunctionCall(
138+
name="get_news",
139+
arguments=json.dumps({"content": "Los Angeles today"}),
140+
),
141+
type="function",
142+
),
143+
],
144+
"I'll get the weather and news for LA today. First, let me get the weather using Los Angeles coordinates, and then get the latest news. ",
145+
),
146+
(
147+
"""I'll help you with multiple tasks. <|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{"city": "New York"}<|tool_call_end|><|tool_call_begin|>functions.get_news:1<|tool_call_argument_begin|>{"topic": "technology"}<|tool_call_end|><|tool_call_begin|>functions.send_email:2<|tool_call_argument_begin|>{"to": "user@example.com", "subject": "Daily Update"}<|tool_call_end|><|tool_calls_section_end|>""",
148+
[
149+
ToolCall(
150+
id="functions.get_weather:0",
151+
function=FunctionCall(
152+
name="get_weather",
153+
arguments=json.dumps({"city": "New York"}),
154+
),
155+
type="function",
156+
),
157+
ToolCall(
158+
id="functions.get_news:1",
159+
function=FunctionCall(
160+
name="get_news",
161+
arguments=json.dumps({"topic": "technology"}),
162+
),
163+
type="function",
164+
),
165+
ToolCall(
166+
id="functions.send_email:2",
167+
function=FunctionCall(
168+
name="send_email",
169+
arguments=json.dumps(
170+
{"to": "user@example.com", "subject": "Daily Update"}
171+
),
172+
),
173+
type="function",
174+
),
175+
],
176+
"I'll help you with multiple tasks. ",
177+
),
178+
(
179+
"""Mixed spacing test. <|tool_calls_section_begin|> <|tool_call_begin|> functions.test:0 <|tool_call_argument_begin|> {} <|tool_call_end|><|tool_call_begin|>functions.test2:1<|tool_call_argument_begin|>{}<|tool_call_end|> <|tool_calls_section_end|>""",
180+
[
181+
ToolCall(
182+
id="functions.test:0",
183+
function=FunctionCall(
184+
name="test",
185+
arguments=json.dumps({}),
186+
),
187+
type="function",
188+
),
189+
ToolCall(
190+
id="functions.test2:1",
191+
function=FunctionCall(
192+
name="test2",
193+
arguments=json.dumps({}),
194+
),
195+
type="function",
196+
),
197+
],
198+
"Mixed spacing test. ",
199+
),
200+
(
201+
"""I need to process HTML content. <|tool_calls_section_begin|><|tool_call_begin|>functions.process_html:0<|tool_call_argument_begin|>{"html": "<div>content</div>", "text": "normal text"}<|tool_call_end|><|tool_calls_section_end|>""",
202+
[
203+
ToolCall(
204+
id="functions.process_html:0",
205+
function=FunctionCall(
206+
name="process_html",
207+
arguments=json.dumps(
208+
{"html": "<div>content</div>", "text": "normal text"}
209+
),
210+
),
211+
type="function",
212+
)
213+
],
214+
"I need to process HTML content. ",
215+
),
216+
(
217+
"""I need to process formatted JSON. <|tool_calls_section_begin|><|tool_call_begin|>functions.process_data:0<|tool_call_argument_begin|>{
218+
"name": "test",
219+
"value": 123,
220+
"nested": {
221+
"key": "value"
222+
}
223+
}<|tool_call_end|><|tool_calls_section_end|>""",
224+
[
225+
ToolCall(
226+
id="functions.process_data:0",
227+
function=FunctionCall(
228+
name="process_data",
229+
arguments=json.dumps(
230+
{"name": "test", "value": 123, "nested": {"key": "value"}},
231+
indent=2,
232+
),
233+
),
234+
type="function",
235+
)
236+
],
237+
"I need to process formatted JSON. ",
238+
),
117239
],
118240
)
119241
def test_extract_tool_calls(

‎vllm/entrypoints/openai/tool_parsers/kimi_k2_tool_parser.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,8 @@ def __init__(self, tokenizer: AnyTokenizer):
6060
self.tool_call_end_token: str = "<|tool_call_end|>"
6161

6262
self.tool_call_regex = re.compile(
63-
r"<\|tool_call_begin\|>\s*(?P<tool_call_id>.+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>.*?)\s*<\|tool_call_end\|>"
63+
r"<\|tool_call_begin\|>\s*(?P<tool_call_id>[^<]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>(?:(?!<\|tool_call_begin\|>).)*?)\s*<\|tool_call_end\|>",
64+
re.DOTALL,
6465
)
6566

6667
self.stream_tool_call_portion_regex = re.compile(

0 commit comments

Comments
 (0)