@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 363528de27
Fork 0

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 15, 03:30 UTC

[Feature] Support MiniMax-M1 function calls features (#20297)

Signed-off-by: QscQ <qscqesze@gmail.com> Signed-off-by: qingjun <qingjun@minimaxi.com>

qscqesze committed Jul 3, 2025 at 14:48 UTC 363528de27db30c2cde6c9abdf3da1544e8c1fe4
5 files changed +842 -1
docs/features/tool_calling.md
+9
@@ -264,6 +264,15 @@ For Qwen2.5, the chat template in tokenizer_config.json has already included sup
264
265 Flags: `--tool-call-parser hermes`
266
267 +### MiniMax Models (`minimax_m1`)
268 +
269 +Supported models:
270 +
271 +* `MiniMaxAi/MiniMax-M1-40k` (use with <gh-file:examples/tool_chat_template_minimax.jinja>)
272 +* `MiniMaxAi/MiniMax-M1-80k` (use with <gh-file:examples/tool_chat_template_minimax.jinja>)
273 +
274 +Flags: `--tool-call-parser minimax --chat-template examples/tool_chat_template_minimax.jinja`
275 +
276 ### DeepSeek-V3 Models (`deepseek_v3`)
277
278 Supported models:
examples/tool_chat_template_minimax_m1.jinja new
+91
@@ -0,0 +1,91 @@
1 +{{ '<begin_of_document>' -}}
2 +{%- if custom_tools is defined %}
3 + {%- set tools = custom_tools %}
4 +{%- endif %}
5 +{%- if not tools is defined %}
6 + {%- set tools = none %}
7 +{%- endif %}
8 +
9 +{#- Extract system message #}
10 +{% set ns = namespace(system_prompt='') -%}
11 +{%- if messages[0]['role'] == 'system' %}
12 + {%- if messages[0]['content'] is string %}
13 + {%- set ns.system_prompt = messages[0]['content']|trim %}
14 + {%- else %}
15 + {%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %}
16 + {%- endif %}
17 + {%- set messages = messages[1:] %}
18 +{%- else %}
19 + {%- if tools is not none %}
20 + {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
21 + {%- else %}
22 + {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
23 + {%- endif %}
24 +{%- endif %}
25 +
26 +{#- System message #}
27 +{%- if ns.system_prompt != '' %}
28 +{{ '<beginning_of_sentence>system ai_setting=assistant\n' + ns.system_prompt + '<end_of_sentence>\n' -}}
29 +{%- endif %}
30 +
31 +{#- Tools configuration #}
32 +{%- if tools is not none %}
33 +{{ '<beginning_of_sentence>system tool_setting=tools\nYou are provided with these tools:\n<tools>\n' -}}
34 +{%- for tool in tools %}
35 +{{ tool | tojson ~ '\n' -}}
36 +{%- endfor %}
37 +{{ '</tools>\n\nIf you need to call tools, please respond with <tool_calls></tool_calls> XML tags, and provide tool-name and json-object of arguments, following the format below:\n<tool_calls>\n{"name": <tool-name>, "arguments": <args-json-object>}\n...\n</tool_calls><end_of_sentence>\n' -}}
38 +{%- endif %}
39 +
40 +{#- Process messages #}
41 +{%- for message in messages %}
42 + {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
43 + {%- if message['role'] == 'user' %}
44 +{{ '<beginning_of_sentence>user name=user\n' -}}
45 +{%- if message['content'] is string %}
46 +{{ message['content']|trim -}}
47 +{%- else %}
48 +{%- for content in message['content'] %}
49 +{%- if content['type'] == 'text' %}
50 +{{ content['text']|trim -}}
51 +{%- endif %}
52 +{%- endfor %}
53 +{%- endif %}
54 +{{ '<end_of_sentence>\n' -}}
55 + {%- elif message['role'] == 'assistant' %}
56 +{{ '<beginning_of_sentence>ai name=assistant\n' -}}
57 +{%- if message['content'] is string %}
58 +{{ message['content']|trim -}}
59 +{%- else %}
60 +{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %}
61 +{{ content['text']|trim -}}
62 +{%- endfor %}
63 +{%- endif %}
64 +{{ '<end_of_sentence>\n' -}}
65 + {%- endif %}
66 + {%- elif 'tool_calls' in message %}
67 +{{ '<beginning_of_sentence>ai name=assistant\n<tool_calls>\n' -}}
68 +{%- for tool_call in message.tool_calls %}
69 +{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}}
70 +{%- endfor %}
71 +{{ '</tool_calls><end_of_sentence>\n' -}}
72 + {%- elif message.role == "tool" or message.role == "ipython" %}
73 +{{ '<beginning_of_sentence>tool name=tools\n' -}}
74 +{%- if message.content is string %}
75 +{{ 'tool result: ' + message.content + '\n\n' -}}
76 +{%- else %}
77 +{%- for content in message['content'] %}
78 +{%- if content['type'] == 'text' %}
79 +{{ 'tool result: ' + content['text'] + '\n\n' -}}
80 +{%- elif content.get('name') %}
81 +{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}}
82 +{%- endif %}
83 +{%- endfor %}
84 +{%- endif %}
85 +{{ '<end_of_sentence>\n' -}}
86 + {%- endif %}
87 +{%- endfor %}
88 +
89 +{%- if add_generation_prompt %}
90 +{{ '<beginning_of_sentence>ai name=assistant\n' -}}
91 +{%- endif %}
\ No newline at end of file
tests/tool_use/test_minimax_tool_parser.py new
+371
@@ -0,0 +1,371 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# ruff: noqa: E501
3 +
4 +import json
5 +
6 +import pytest
7 +
8 +from vllm.entrypoints.openai.protocol import FunctionCall, ToolCall
9 +from vllm.entrypoints.openai.tool_parsers import MinimaxToolParser
10 +from vllm.transformers_utils.tokenizer import get_tokenizer
11 +
12 +# Use a common model that is likely to be available
13 +MODEL = "MiniMaxAi/MiniMax-M1-40k"
14 +
15 +
16 +@pytest.fixture(scope="module")
17 +def minimax_tokenizer():
18 + return get_tokenizer(tokenizer_name=MODEL)
19 +
20 +
21 +@pytest.fixture
22 +def minimax_tool_parser(minimax_tokenizer):
23 + return MinimaxToolParser(minimax_tokenizer)
24 +
25 +
26 +def assert_tool_calls(actual_tool_calls: list[ToolCall],
27 + expected_tool_calls: list[ToolCall]):
28 + assert len(actual_tool_calls) == len(expected_tool_calls)
29 +
30 + for actual_tool_call, expected_tool_call in zip(actual_tool_calls,
31 + expected_tool_calls):
32 + assert isinstance(actual_tool_call.id, str)
33 + assert len(actual_tool_call.id) > 16
34 +
35 + assert actual_tool_call.type == "function"
36 + assert actual_tool_call.function == expected_tool_call.function
37 +
38 +
39 +def test_extract_tool_calls_no_tools(minimax_tool_parser):
40 + model_output = "This is a test"
41 + extracted_tool_calls = minimax_tool_parser.extract_tool_calls(
42 + model_output, request=None) # type: ignore[arg-type]
43 + assert not extracted_tool_calls.tools_called
44 + assert extracted_tool_calls.tool_calls == []
45 + assert extracted_tool_calls.content == model_output
46 +
47 +
48 +@pytest.mark.parametrize(
49 + ids=[
50 + "single_tool_call",
51 + "multiple_tool_calls",
52 + "tool_call_with_content_before",
53 + "tool_call_with_single_line_json",
54 + "tool_call_incomplete_tag",
55 + ],
56 + argnames=["model_output", "expected_tool_calls", "expected_content"],
57 + argvalues=[
58 + (
59 + """<tool_calls>
60 +{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}
61 +</tool_calls>""",
62 + [
63 + ToolCall(function=FunctionCall(
64 + name="get_current_weather",
65 + arguments=json.dumps({
66 + "city": "Dallas",
67 + "state": "TX",
68 + "unit": "fahrenheit",
69 + }),
70 + ))
71 + ],
72 + None,
73 + ),
74 + (
75 + """<tool_calls>
76 +{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}
77 +{"name": "get_current_weather", "arguments": {"city": "Orlando", "state": "FL", "unit": "fahrenheit"}}
78 +</tool_calls>""",
79 + [
80 + ToolCall(function=FunctionCall(
81 + name="get_current_weather",
82 + arguments=json.dumps({
83 + "city": "Dallas",
84 + "state": "TX",
85 + "unit": "fahrenheit",
86 + }),
87 + )),
88 + ToolCall(function=FunctionCall(
89 + name="get_current_weather",
90 + arguments=json.dumps({
91 + "city": "Orlando",
92 + "state": "FL",
93 + "unit": "fahrenheit",
94 + }),
95 + )),
96 + ],
97 + None,
98 + ),
99 + (
100 + """I'll help you check the weather. <tool_calls>
101 +{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}
102 +</tool_calls>""",
103 + [
104 + ToolCall(function=FunctionCall(
105 + name="get_current_weather",
106 + arguments=json.dumps({
107 + "city": "Seattle",
108 + "state": "WA",
109 + "unit": "celsius",
110 + }),
111 + ))
112 + ],
113 + "I'll help you check the weather.",
114 + ),
115 + (
116 + """<tool_calls>
117 +{"name": "get_current_weather", "arguments": {"city": "New York", "state": "NY", "unit": "celsius"}}
118 +</tool_calls>""",
119 + [
120 + ToolCall(function=FunctionCall(
121 + name="get_current_weather",
122 + arguments=json.dumps({
123 + "city": "New York",
124 + "state": "NY",
125 + "unit": "celsius",
126 + }),
127 + ))
128 + ],
129 + None,
130 + ),
131 + (
132 + """<tool_calls>
133 +{"name": "get_current_weather", "arguments": {"city": "Boston", "state": "MA"}}""",
134 + [
135 + ToolCall(function=FunctionCall(
136 + name="get_current_weather",
137 + arguments=json.dumps({
138 + "city": "Boston",
139 + "state": "MA",
140 + }),
141 + ))
142 + ],
143 + None,
144 + ),
145 + ],
146 +)
147 +def test_extract_tool_calls(minimax_tool_parser, model_output,
148 + expected_tool_calls, expected_content):
149 + extracted_tool_calls = minimax_tool_parser.extract_tool_calls(
150 + model_output, request=None) # type: ignore[arg-type]
151 + assert extracted_tool_calls.tools_called
152 +
153 + assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls)
154 +
155 + assert extracted_tool_calls.content == expected_content
156 +
157 +
158 +def test_preprocess_model_output_with_thinking_tags(minimax_tool_parser):
159 + """Test that tool calls within thinking tags are removed during preprocessing."""
160 + model_output = """<think>Let me think about this. <tool_calls>
161 +{"name": "fake_tool", "arguments": {"param": "value"}}
162 +</tool_calls> This should be removed.</think>
163 +
164 +I'll help you with that. <tool_calls>
165 +{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA"}}
166 +</tool_calls>"""
167 +
168 + processed_output = minimax_tool_parser.preprocess_model_output(
169 + model_output)
170 +
171 + # The tool call within thinking tags should be removed
172 + assert "fake_tool" not in processed_output
173 + # But the thinking tag itself should remain
174 + assert "<think>" in processed_output
175 + assert "</think>" in processed_output
176 + # The actual tool call outside thinking tags should remain
177 + assert "get_current_weather" in processed_output
178 +
179 +
180 +def test_extract_tool_calls_with_thinking_tags(minimax_tool_parser):
181 + """Test tool extraction when thinking tags contain tool calls that should be ignored."""
182 + model_output = """<think>I should use a tool. <tool_calls>
183 +{"name": "ignored_tool", "arguments": {"should": "ignore"}}
184 +</tool_calls></think>
185 +
186 +Let me help you with the weather. <tool_calls>
187 +{"name": "get_current_weather", "arguments": {"city": "Miami", "state": "FL", "unit": "fahrenheit"}}
188 +</tool_calls>"""
189 +
190 + extracted_tool_calls = minimax_tool_parser.extract_tool_calls(
191 + model_output, request=None) # type: ignore[arg-type]
192 +
193 + assert extracted_tool_calls.tools_called
194 + assert len(extracted_tool_calls.tool_calls) == 1
195 + assert extracted_tool_calls.tool_calls[
196 + 0].function.name == "get_current_weather"
197 +
198 + # Content extraction is based on the position of the first <tool_calls> in the original model_output
199 + # Since preprocessing removes tool calls within thinking tags, the actual first <tool_calls> is the external one
200 + expected_content = """<think>I should use a tool. <tool_calls>
201 +{"name": "ignored_tool", "arguments": {"should": "ignore"}}
202 +</tool_calls></think>
203 +
204 +Let me help you with the weather."""
205 + assert extracted_tool_calls.content == expected_content
206 +
207 +
208 +def test_extract_tool_calls_invalid_json(minimax_tool_parser):
209 + """Test that invalid JSON in tool calls is handled gracefully."""
210 + model_output = """<tool_calls>
211 +{"name": "valid_tool", "arguments": {"city": "Seattle"}}
212 +{invalid json here}
213 +{"name": "another_valid_tool", "arguments": {"param": "value"}}
214 +</tool_calls>"""
215 +
216 + extracted_tool_calls = minimax_tool_parser.extract_tool_calls(
217 + model_output, request=None) # type: ignore[arg-type]
218 +
219 + assert extracted_tool_calls.tools_called
220 + # Should extract only the valid JSON tool calls
221 + assert len(extracted_tool_calls.tool_calls) == 2
222 + assert extracted_tool_calls.tool_calls[0].function.name == "valid_tool"
223 + assert extracted_tool_calls.tool_calls[
224 + 1].function.name == "another_valid_tool"
225 +
226 +
227 +def test_extract_tool_calls_missing_name_or_arguments(minimax_tool_parser):
228 + """Test that tool calls missing name or arguments are filtered out."""
229 + model_output = """<tool_calls>
230 +{"name": "valid_tool", "arguments": {"city": "Seattle"}}
231 +{"name": "missing_args"}
232 +{"arguments": {"city": "Portland"}}
233 +{"name": "another_valid_tool", "arguments": {"param": "value"}}
234 +</tool_calls>"""
235 +
236 + extracted_tool_calls = minimax_tool_parser.extract_tool_calls(
237 + model_output, request=None) # type: ignore[arg-type]
238 +
239 + assert extracted_tool_calls.tools_called
240 + # Should extract only the valid tool calls with both name and arguments
241 + assert len(extracted_tool_calls.tool_calls) == 2
242 + assert extracted_tool_calls.tool_calls[0].function.name == "valid_tool"
243 + assert extracted_tool_calls.tool_calls[
244 + 1].function.name == "another_valid_tool"
245 +
246 +
247 +def test_streaming_basic_functionality(minimax_tool_parser):
248 + """Test basic streaming functionality."""
249 + # Reset streaming state
250 + minimax_tool_parser.current_tool_name_sent = False
251 + minimax_tool_parser.prev_tool_call_arr = []
252 + minimax_tool_parser.current_tool_id = -1
253 + minimax_tool_parser.streamed_args_for_tool = []
254 +
255 + # Test with a simple tool call
256 + current_text = """<tool_calls>
257 +{"name": "get_current_weather", "arguments": {"city": "Seattle"}}
258 +</tool_calls>"""
259 +
260 + # First call should handle the initial setup
261 + result = minimax_tool_parser.extract_tool_calls_streaming(
262 + previous_text="",
263 + current_text=current_text,
264 + delta_text="</tool_calls>",
265 + previous_token_ids=[],
266 + current_token_ids=[],
267 + delta_token_ids=[],
268 + request=None,
269 + )
270 +
271 + # The result might be None or contain tool call information
272 + # This depends on the internal state management
273 + if result is not None and hasattr(result,
274 + 'tool_calls') and result.tool_calls:
275 + assert len(result.tool_calls) >= 0
276 +
277 +
278 +def test_streaming_with_content_before_tool_calls(minimax_tool_parser):
279 + """Test streaming when there's content before tool calls."""
280 + # Reset streaming state
281 + minimax_tool_parser.current_tool_name_sent = False
282 + minimax_tool_parser.prev_tool_call_arr = []
283 + minimax_tool_parser.current_tool_id = -1
284 + minimax_tool_parser.streamed_args_for_tool = []
285 +
286 + current_text = "I'll help you with that. <tool_calls>"
287 +
288 + # When there's content before tool calls, it should be returned as content
289 + result = minimax_tool_parser.extract_tool_calls_streaming(
290 + previous_text="I'll help you",
291 + current_text=current_text,
292 + delta_text=" with that. <tool_calls>",
293 + previous_token_ids=[],
294 + current_token_ids=[],
295 + delta_token_ids=[],
296 + request=None,
297 + )
298 +
299 + if result is not None and hasattr(result, 'content'):
300 + # Should contain some content
301 + assert result.content is not None
302 +
303 +
304 +def test_streaming_no_tool_calls(minimax_tool_parser):
305 + """Test streaming when there are no tool calls."""
306 + current_text = "This is just regular text without any tool calls."
307 +
308 + result = minimax_tool_parser.extract_tool_calls_streaming(
309 + previous_text="This is just regular text",
310 + current_text=current_text,
311 + delta_text=" without any tool calls.",
312 + previous_token_ids=[],
313 + current_token_ids=[],
314 + delta_token_ids=[],
315 + request=None,
316 + )
317 +
318 + # Should return the delta text as content
319 + assert result is not None
320 + assert hasattr(result, 'content')
321 + assert result.content == " without any tool calls."
322 +
323 +
324 +def test_streaming_with_thinking_tags(minimax_tool_parser):
325 + """Test streaming with thinking tags that contain tool calls."""
326 + # Reset streaming state
327 + minimax_tool_parser.current_tool_name_sent = False
328 + minimax_tool_parser.prev_tool_call_arr = []
329 + minimax_tool_parser.current_tool_id = -1
330 + minimax_tool_parser.streamed_args_for_tool = []
331 +
332 + current_text = """<think><tool_calls>{"name": "ignored", "arguments": {}}</tool_calls></think><tool_calls>{"name": "real_tool", "arguments": {"param": "value"}}</tool_calls>"""
333 +
334 + result = minimax_tool_parser.extract_tool_calls_streaming(
335 + previous_text="",
336 + current_text=current_text,
337 + delta_text=current_text,
338 + previous_token_ids=[],
339 + current_token_ids=[],
340 + delta_token_ids=[],
341 + request=None,
342 + )
343 +
344 + # The preprocessing should remove tool calls from thinking tags
345 + # and only process the real tool call
346 + if result is not None and hasattr(result,
347 + 'tool_calls') and result.tool_calls:
348 + for tool_call in result.tool_calls:
349 + assert tool_call.function.name != "ignored"
350 +
351 +
352 +def test_extract_tool_calls_multiline_json_not_supported(minimax_tool_parser):
353 + """Test that multiline JSON in tool calls is not currently supported."""
354 + model_output = """<tool_calls>
355 +{
356 + "name": "get_current_weather",
357 + "arguments": {
358 + "city": "New York",
359 + "state": "NY",
360 + "unit": "celsius"
361 + }
362 +}
363 +</tool_calls>"""
364 +
365 + extracted_tool_calls = minimax_tool_parser.extract_tool_calls(
366 + model_output, request=None) # type: ignore[arg-type]
367 +
368 + # Multiline JSON is currently not supported, should return no tools called
369 + assert not extracted_tool_calls.tools_called
370 + assert extracted_tool_calls.tool_calls == []
371 + assert extracted_tool_calls.content is None
vllm/entrypoints/openai/tool_parsers/__init__.py
+2 -1
@@ -10,6 +10,7 @@ from .internlm2_tool_parser import Internlm2ToolParser
10 from .jamba_tool_parser import JambaToolParser
11 from .llama4_pythonic_tool_parser import Llama4PythonicToolParser
12 from .llama_tool_parser import Llama3JsonToolParser
13 +from .minimax_tool_parser import MinimaxToolParser
14 from .mistral_tool_parser import MistralToolParser
15 from .phi4mini_tool_parser import Phi4MiniJsonToolParser
16 from .pythonic_tool_parser import PythonicToolParser
@@ -20,5 +21,5 @@ __all__ = [
21 "GraniteToolParser", "Hermes2ProToolParser", "MistralToolParser",
22 "Internlm2ToolParser", "Llama3JsonToolParser", "JambaToolParser",
23 "Llama4PythonicToolParser", "PythonicToolParser", "Phi4MiniJsonToolParser",
23 - "DeepSeekV3ToolParser", "xLAMToolParser"
24 + "DeepSeekV3ToolParser", "xLAMToolParser", "MinimaxToolParser"
25 ]
vllm/entrypoints/openai/tool_parsers/minimax_tool_parser.py new
+369
@@ -0,0 +1,369 @@
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 Sequence
6 +from typing import Union
7 +
8 +import partial_json_parser
9 +import regex as re
10 +from partial_json_parser.core.options import Allow
11 +
12 +from vllm.entrypoints.chat_utils import random_tool_call_id
13 +from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
14 + DeltaFunctionCall, DeltaMessage,
15 + DeltaToolCall,
16 + ExtractedToolCallInformation,
17 + FunctionCall, ToolCall)
18 +from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import (
19 + ToolParser, ToolParserManager)
20 +from vllm.logger import init_logger
21 +from vllm.transformers_utils.tokenizer import AnyTokenizer
22 +
23 +logger = init_logger(__name__)
24 +
25 +
26 +@ToolParserManager.register_module("minimax")
27 +class MinimaxToolParser(ToolParser):
28 +
29 + def __init__(self, tokenizer: AnyTokenizer):
30 + super().__init__(tokenizer)
31 +
32 + self.current_tool_name_sent: bool = False
33 + self.prev_tool_call_arr: list[dict] = []
34 + self.current_tool_id: int = -1
35 + self.streamed_args_for_tool: list[str] = []
36 +
37 + self.tool_call_start_token: str = "<tool_calls>"
38 + self.tool_call_end_token: str = "</tool_calls>"
39 +
40 + self.tool_call_regex = re.compile(
41 + r"<tool_calls>(.*?)</tool_calls>|<tool_calls>(.*)", re.DOTALL)
42 +
43 + # Add regex pattern for thinking tag
44 + self.thinking_tag_pattern = r"<think>(.*?)</think>"
45 +
46 + if not self.model_tokenizer:
47 + raise ValueError(
48 + "The model tokenizer must be passed to the ToolParser "
49 + "constructor during construction.")
50 +
51 + self.tool_call_start_token_id = self.vocab.get(
52 + self.tool_call_start_token)
53 + self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)
54 +
55 + if (self.tool_call_start_token_id is None
56 + or self.tool_call_end_token_id is None):
57 + logger.warning(
58 + "Minimax Tool parser could not locate tool call start/end "
59 + "tokens in the tokenizer. Falling back to string matching.")
60 +
61 + def preprocess_model_output(self, model_output: str) -> str:
62 + """
63 + Remove tool calls from within thinking tags to avoid processing them.
64 + """
65 +
66 + def remove_tool_calls_from_think(match):
67 + think_content = match.group(1)
68 + # Remove tool_calls from within the think tag
69 + cleaned_content = re.sub(r"<tool_calls>.*?</tool_calls>",
70 + "",
71 + think_content,
72 + flags=re.DOTALL)
73 + return f"<think>{cleaned_content}</think>"
74 +
75 + # Process thinking tags and remove tool_calls from within them
76 + processed_output = re.sub(self.thinking_tag_pattern,
77 + remove_tool_calls_from_think,
78 + model_output,
79 + flags=re.DOTALL)
80 +
81 + return processed_output
82 +
83 + def extract_tool_calls(
84 + self,
85 + model_output: str,
86 + request: ChatCompletionRequest,
87 + ) -> ExtractedToolCallInformation:
88 +
89 + # Preprocess to remove tool calls from thinking tags
90 + processed_output = self.preprocess_model_output(model_output)
91 +
92 + if self.tool_call_start_token not in processed_output:
93 + return ExtractedToolCallInformation(tools_called=False,
94 + tool_calls=[],
95 + content=model_output)
96 +
97 + try:
98 + function_call_tuples = (
99 + self.tool_call_regex.findall(processed_output))
100 +
101 + raw_function_calls = []
102 + for match in function_call_tuples:
103 + tool_call_content = match[0] if match[0] else match[1]
104 + if tool_call_content.strip():
105 + lines = tool_call_content.strip().split('\n')
106 + for line in lines:
107 + line = line.strip()
108 + if line and line.startswith('{') and line.endswith(
109 + '}'):
110 + try:
111 + parsed_call = json.loads(line)
112 + raw_function_calls.append(parsed_call)
113 + except json.JSONDecodeError:
114 + continue
115 +
116 + tool_calls = []
117 + for function_call in raw_function_calls:
118 + if "name" in function_call and "arguments" in function_call:
119 + tool_calls.append(
120 + ToolCall(type="function",
121 + function=FunctionCall(
122 + name=function_call["name"],
123 + arguments=json.dumps(
124 + function_call["arguments"],
125 + ensure_ascii=False))))
126 +
127 + # Extract content before the first valid tool call
128 + # Find the position in processed output, then map back to original
129 + processed_pos = processed_output.find(self.tool_call_start_token)
130 + if processed_pos != -1:
131 + # Get the content before tool calls in processed output
132 + processed_content = processed_output[:processed_pos].strip()
133 +
134 + if processed_content:
135 + # Find the end of this content in the original output
136 + # Look for the last non-empty line of processed content
137 + lines = processed_content.split('\n')
138 + for line in reversed(lines):
139 + line = line.strip()
140 + if line:
141 + # Find this line in original output
142 + pos = model_output.find(line)
143 + if pos != -1:
144 + content = model_output[:pos + len(line)]
145 + break
146 + else:
147 + content = ""
148 + else:
149 + content = ""
150 + else:
151 + content = model_output
152 +
153 + return ExtractedToolCallInformation(
154 + tools_called=len(tool_calls) > 0,
155 + tool_calls=tool_calls,
156 + content=content.strip() if content.strip() else None)
157 +
158 + except Exception:
159 + logger.exception(
160 + "An unexpected error occurred during tool call extraction.")
161 + return ExtractedToolCallInformation(tools_called=False,
162 + tool_calls=[],
163 + content=model_output)
164 +
165 + def extract_tool_calls_streaming(
166 + self,
167 + previous_text: str,
168 + current_text: str,
169 + delta_text: str,
170 + previous_token_ids: Sequence[int],
171 + current_token_ids: Sequence[int],
172 + delta_token_ids: Sequence[int],
173 + request: ChatCompletionRequest,
174 + ) -> Union[DeltaMessage, None]:
175 + logger.debug("delta_text: %s", delta_text)
176 + logger.debug("delta_token_ids: %s", delta_token_ids)
177 +
178 + # Preprocess to remove tool calls from thinking tags
179 + processed_current_text = self.preprocess_model_output(current_text)
180 +
181 + if self.tool_call_start_token not in processed_current_text:
182 + return DeltaMessage(content=delta_text)
183 +
184 + if (self.tool_call_start_token_id is not None
185 + and self.tool_call_start_token_id in delta_token_ids
186 + and len(delta_token_ids) == 1):
187 + return None
188 +
189 + original_tool_call_start_pos = current_text.find(
190 + self.tool_call_start_token)
191 + if original_tool_call_start_pos > 0:
192 + delta_start_pos = len(current_text) - len(delta_text)
193 + if delta_start_pos < original_tool_call_start_pos:
194 + content_part = delta_text
195 + if delta_start_pos + len(
196 + delta_text) > original_tool_call_start_pos:
197 + content_part = delta_text[:original_tool_call_start_pos -
198 + delta_start_pos]
199 + if content_part:
200 + return DeltaMessage(content=content_part)
201 +
202 + flags = Allow.ALL if self.current_tool_name_sent \
203 + else Allow.ALL & ~Allow.STR
204 +
205 + try:
206 + parsable_content = processed_current_text.split(
207 + self.tool_call_start_token)[-1].split(
208 + self.tool_call_end_token)[0]
209 +
210 + tool_call_arr = []
211 + if parsable_content.strip():
212 + lines = parsable_content.strip().split('\n')
213 + for line in lines:
214 + line = line.strip()
215 + if line and (line.startswith('{') or '"name"' in line):
216 + try:
217 + if line.endswith('}'):
218 + parsed_call = json.loads(line)
219 + tool_call_arr.append(parsed_call)
220 + else:
221 + parsed_call = partial_json_parser.loads(
222 + line, flags)
223 + if parsed_call and isinstance(
224 + parsed_call, dict):
225 + tool_call_arr.append(parsed_call)
226 + except (json.JSONDecodeError, partial_json_parser.core.
227 + exceptions.MalformedJSON):
228 + continue
229 +
230 + current_tool_call: dict = tool_call_arr[self.current_tool_id] \
231 + if len(tool_call_arr) > self.current_tool_id >= 0 else {}
232 +
233 + if len(tool_call_arr) == 0:
234 + return None
235 +
236 + # Starting a new tool in the array
237 + elif (len(tool_call_arr) > 0
238 + and len(tool_call_arr) > self.current_tool_id + 1):
239 +
240 + # Handle any missed arguments from previous tool
241 + if self.current_tool_id >= 0 and self.current_tool_id < len(
242 + self.prev_tool_call_arr):
243 + prev_tool_call = self.prev_tool_call_arr[
244 + self.current_tool_id]
245 + diff_arguments = prev_tool_call.get("arguments")
246 +
247 + if diff_arguments:
248 + diff_arguments_json = json.dumps(diff_arguments,
249 + ensure_ascii=False)
250 + already_streamed = self.streamed_args_for_tool[
251 + self.
252 + current_tool_id] if self.current_tool_id < len(
253 + self.streamed_args_for_tool) else ""
254 +
255 + if diff_arguments_json != already_streamed:
256 + diff = diff_arguments_json[len(already_streamed):]
257 + delta = DeltaMessage(tool_calls=[
258 + DeltaToolCall(index=self.current_tool_id,
259 + function=DeltaFunctionCall(
260 + arguments=diff).model_dump(
261 + exclude_none=True))
262 + ])
263 + if self.current_tool_id < len(
264 + self.streamed_args_for_tool):
265 + self.streamed_args_for_tool[
266 + self.current_tool_id] = diff_arguments_json
267 + else:
268 + delta = None
269 + else:
270 + delta = None
271 + else:
272 + delta = None
273 +
274 + self.current_tool_id = len(tool_call_arr) - 1
275 + self.current_tool_name_sent = False
276 + self.streamed_args_for_tool.append("")
277 + logger.debug("starting on new tool %d", self.current_tool_id)
278 + return delta
279 +
280 + # Send tool name if not sent yet
281 + if not self.current_tool_name_sent:
282 + function_name = current_tool_call.get("name")
283 + if function_name:
284 + delta = DeltaMessage(tool_calls=[
285 + DeltaToolCall(index=self.current_tool_id,
286 + type="function",
287 + id=random_tool_call_id(),
288 + function=DeltaFunctionCall(
289 + name=function_name).model_dump(
290 + exclude_none=True))
291 + ])
292 + self.current_tool_name_sent = True
293 + else:
294 + delta = None
295 +
296 + # Stream arguments
297 + else:
298 + prev_arguments = None
299 + if (self.current_tool_id < len(self.prev_tool_call_arr)
300 + and self.prev_tool_call_arr[self.current_tool_id]):
301 + prev_arguments = self.prev_tool_call_arr[
302 + self.current_tool_id].get("arguments")
303 +
304 + cur_arguments = current_tool_call.get("arguments")
305 +
306 + if not cur_arguments and not prev_arguments:
307 + delta = None
308 + elif not cur_arguments and prev_arguments:
309 + logger.error(
310 + "Arguments reset mid-call, skipping streaming")
311 + delta = None
312 + elif cur_arguments and not prev_arguments:
313 + cur_arguments_json = json.dumps(cur_arguments,
314 + ensure_ascii=False)
315 + logger.debug("First tokens in arguments received: %s",
316 + cur_arguments_json)
317 +
318 + delta = DeltaMessage(tool_calls=[
319 + DeltaToolCall(index=self.current_tool_id,
320 + function=DeltaFunctionCall(
321 + arguments=cur_arguments_json).
322 + model_dump(exclude_none=True))
323 + ])
324 + self.streamed_args_for_tool[
325 + self.current_tool_id] = cur_arguments_json
326 +
327 + elif cur_arguments and prev_arguments:
328 + cur_args_json = json.dumps(cur_arguments,
329 + ensure_ascii=False)
330 + prev_args_json = json.dumps(prev_arguments,
331 + ensure_ascii=False)
332 +
333 + logger.debug("Searching for diff between \n%s\n%s",
334 + cur_args_json, prev_args_json)
335 +
336 + already_streamed = self.streamed_args_for_tool[
337 + self.current_tool_id] if self.current_tool_id < len(
338 + self.streamed_args_for_tool) else ""
339 +
340 + if cur_args_json.startswith(already_streamed):
341 + argument_diff = cur_args_json[len(already_streamed):]
342 + elif cur_args_json != already_streamed:
343 + argument_diff = cur_args_json
344 + self.streamed_args_for_tool[self.current_tool_id] = ""
345 + else:
346 + argument_diff = ""
347 +
348 + if argument_diff:
349 + logger.debug("got arguments diff: %s", argument_diff)
350 + delta = DeltaMessage(tool_calls=[
351 + DeltaToolCall(index=self.current_tool_id,
352 + function=DeltaFunctionCall(
353 + arguments=argument_diff).
354 + model_dump(exclude_none=True))
355 + ])
356 + self.streamed_args_for_tool[
357 + self.current_tool_id] += argument_diff
358 + else:
359 + delta = None
360 + else:
361 + delta = None
362 +
363 + self.prev_tool_call_arr = tool_call_arr
364 + return delta
365 +
366 + except Exception:
367 + logger.exception("An unexpected error occurred",
368 + "during streaming tool call handling.")
369 + return None