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, 02:15 UTC
Add kimi-k2 tool parser (#20789)
Signed-off-by: wangzhengtao <wangzhengtao@moonshot.cn> Co-authored-by: wangzhengtao <wangzhengtao@moonshot.cn> Co-authored-by: wangzhengtao <wangzhengtao@msh.team>
bigmoyan committed
Jul 11, 2025 at 10:36 UTC
0cf893cae174d634d03c5b399bc6787bf2d5a6cc
4 files changed
+576
-2
tests/tool_use/test_kimi_k2_tool_parser.py
new
+195
@@ -0,0 +1,195 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# ruff: noqa: E501
4
+
5
+import json
6
+
7
+import pytest
8
+
9
+from vllm.entrypoints.openai.protocol import FunctionCall, ToolCall
10
+from vllm.entrypoints.openai.tool_parsers import KimiK2ToolParser
11
+from vllm.transformers_utils.tokenizer import get_tokenizer
12
+
13
+pytest.skip("skip kimi_k2 parser test", allow_module_level=True)
14
+
15
+# Use a common model that is likely to be available
16
+MODEL = "moonshotai/Kimi-K2-Instruct"
17
+
18
+
19
+@pytest.fixture(scope="module")
20
+def kimi_k2_tokenizer():
21
+ return get_tokenizer(tokenizer_name=MODEL, trust_remote_code=True)
22
+
23
+
24
+@pytest.fixture
25
+def kimi_k2_tool_parser(kimi_k2_tokenizer):
26
+ return KimiK2ToolParser(kimi_k2_tokenizer)
27
+
28
+
29
+def assert_tool_calls(actual_tool_calls: list[ToolCall],
30
+ expected_tool_calls: list[ToolCall]):
31
+ assert len(actual_tool_calls) == len(expected_tool_calls)
32
+
33
+ for actual_tool_call, expected_tool_call in zip(actual_tool_calls,
34
+ expected_tool_calls):
35
+
36
+ assert actual_tool_call.type == "function"
37
+ assert actual_tool_call.function == expected_tool_call.function
38
+
39
+ # assert tool call id format
40
+ assert actual_tool_call.id.startswith("functions.")
41
+ assert actual_tool_call.id.split(':')[-1].isdigit()
42
+ assert actual_tool_call.id.split('.')[1].split(
43
+ ':')[0] == expected_tool_call.function.name
44
+
45
+
46
+def test_extract_tool_calls_no_tools(kimi_k2_tool_parser):
47
+ model_output = "This is a test"
48
+ extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls(
49
+ model_output, request=None) # type: ignore[arg-type]
50
+ assert not extracted_tool_calls.tools_called
51
+ assert extracted_tool_calls.tool_calls == []
52
+ assert extracted_tool_calls.content == model_output
53
+
54
+
55
+@pytest.mark.parametrize(
56
+ ids=[
57
+ "tool_call_with_content_before",
58
+ "multi_tool_call_with_content_before",
59
+ ],
60
+ argnames=["model_output", "expected_tool_calls", "expected_content"],
61
+ argvalues=[
62
+ (
63
+ """I'll help you check the weather. <|tool_calls_section_begin|> <|tool_call_begin|>
64
+functions.get_weather:0 <|tool_call_argument_begin|> {"city": "Beijing"} <|tool_call_end|> <|tool_calls_section_end|>""",
65
+ [
66
+ ToolCall(id='functions.get_weather:0',
67
+ function=FunctionCall(
68
+ name="get_weather",
69
+ arguments=json.dumps({
70
+ "city": "Beijing",
71
+ }, ),
72
+ ),
73
+ type='function')
74
+ ],
75
+ "I'll help you check the weather. ",
76
+ ),
77
+ (
78
+ """I'll help you check the weather. <|tool_calls_section_begin|> <|tool_call_begin|>
79
+functions.get_weather:0 <|tool_call_argument_begin|> {"city": "Beijing"} <|tool_call_end|> <|tool_call_begin|>
80
+functions.get_weather:1 <|tool_call_argument_begin|> {"city": "Shanghai"} <|tool_call_end|> <|tool_calls_section_end|>""",
81
+ [
82
+ ToolCall(id='functions.get_weather:0',
83
+ function=FunctionCall(
84
+ name="get_weather",
85
+ arguments=json.dumps({
86
+ "city": "Beijing",
87
+ }, ),
88
+ ),
89
+ type='function'),
90
+ ToolCall(id='functions.get_weather:1',
91
+ function=FunctionCall(
92
+ name="get_weather",
93
+ arguments=json.dumps({
94
+ "city": "Shanghai",
95
+ }, ),
96
+ ),
97
+ type='function')
98
+ ],
99
+ "I'll help you check the weather. ",
100
+ ),
101
+ ],
102
+)
103
+def test_extract_tool_calls(kimi_k2_tool_parser, model_output,
104
+ expected_tool_calls, expected_content):
105
+ extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls(
106
+ model_output, request=None) # type: ignore[arg-type]
107
+ assert extracted_tool_calls.tools_called
108
+
109
+ assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls)
110
+
111
+ assert extracted_tool_calls.content == expected_content
112
+
113
+
114
+def test_extract_tool_calls_invalid_json(kimi_k2_tool_parser):
115
+ """we'll return every funcall result"""
116
+ model_output = """I'll help you check the weather. <|tool_calls_section_begin|> <|tool_call_begin|>
117
+functions.invalid_get_weather:0 <|tool_call_argument_begin|> {"city": "Beijing" <|tool_call_end|> <|tool_call_begin|>
118
+functions.valid_get_weather:1 <|tool_call_argument_begin|> {"city": "Shanghai"} <|tool_call_end|> <|tool_calls_section_end|>"""
119
+
120
+ extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls(
121
+ model_output, request=None) # type: ignore[arg-type]
122
+
123
+ assert extracted_tool_calls.tools_called
124
+ # Should extract only the valid JSON tool calls
125
+ assert len(extracted_tool_calls.tool_calls) == 2
126
+ assert extracted_tool_calls.tool_calls[
127
+ 0].function.name == "invalid_get_weather"
128
+ assert extracted_tool_calls.tool_calls[
129
+ 1].function.name == "valid_get_weather"
130
+
131
+
132
+def test_extract_tool_calls_invalid_funcall(kimi_k2_tool_parser):
133
+ """we'll return every funcall result"""
134
+ model_output = """I'll help you check the weather. <|tool_calls_section_begin|> <|tool_call_begin|>
135
+functions.invalid_get_weather.0 <|tool_call_argument_begin|> {"city": "Beijing"} <|tool_call_end|> <|tool_call_begin|>
136
+functions.valid_get_weather:1 <|tool_call_argument_begin|> {"city": "Shanghai"} <|tool_call_end|> <|tool_calls_section_end|>"""
137
+
138
+ extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls(
139
+ model_output, request=None) # type: ignore[arg-type]
140
+
141
+ assert extracted_tool_calls.tools_called
142
+ # Should extract only the valid JSON tool calls
143
+ assert len(extracted_tool_calls.tool_calls) == 1
144
+ assert extracted_tool_calls.tool_calls[
145
+ 0].function.name == "valid_get_weather"
146
+
147
+
148
+def test_streaming_basic_functionality(kimi_k2_tool_parser):
149
+ """Test basic streaming functionality."""
150
+ # Reset streaming state
151
+ kimi_k2_tool_parser.current_tool_name_sent = False
152
+ kimi_k2_tool_parser.prev_tool_call_arr = []
153
+ kimi_k2_tool_parser.current_tool_id = -1
154
+ kimi_k2_tool_parser.streamed_args_for_tool = []
155
+
156
+ # Test with a simple tool call
157
+ current_text = """ check the weather. <|tool_calls_section_begin|> <|tool_call_begin|>
158
+functions.get_weather:0 <|tool_call_argument_begin|> {"city": "Beijing"} <|tool_call_end|> <|tool_calls_section_end|>"""
159
+
160
+ # First call should handle the initial setup
161
+ result = kimi_k2_tool_parser.extract_tool_calls_streaming(
162
+ previous_text="I'll help you",
163
+ current_text=current_text,
164
+ delta_text="<|tool_calls_section_end|>",
165
+ previous_token_ids=[],
166
+ current_token_ids=[],
167
+ delta_token_ids=[],
168
+ request=None,
169
+ )
170
+
171
+ # The result might be None or contain tool call information
172
+ # This depends on the internal state management
173
+ if result is not None and hasattr(result,
174
+ 'tool_calls') and result.tool_calls:
175
+ assert len(result.tool_calls) >= 0
176
+
177
+
178
+def test_streaming_no_tool_calls(kimi_k2_tool_parser):
179
+ """Test streaming when there are no tool calls."""
180
+ current_text = "This is just regular text without any tool calls."
181
+
182
+ result = kimi_k2_tool_parser.extract_tool_calls_streaming(
183
+ previous_text="This is just regular text",
184
+ current_text=current_text,
185
+ delta_text=" without any tool calls.",
186
+ previous_token_ids=[],
187
+ current_token_ids=[],
188
+ delta_token_ids=[],
189
+ request=None,
190
+ )
191
+
192
+ # Should return the delta text as content
193
+ assert result is not None
194
+ assert hasattr(result, 'content')
195
+ assert result.content == " without any tool calls."
vllm/config.py
+1
-1
@@ -1143,7 +1143,7 @@ class ModelConfig:
1143
if not hasattr(self.hf_text_config, "model_type"):
1144
return False
1145
elif self.hf_text_config.model_type in \
1146
- ('deepseek_v2', 'deepseek_v3', 'deepseek_mtp'):
1146
+ ('deepseek_v2', 'deepseek_v3', 'deepseek_mtp', 'kimi_k2'):
1147
return self.hf_text_config.kv_lora_rank is not None
1148
elif self.hf_text_config.model_type == 'eagle':
1149
# if the model is an EAGLE module, check for the
vllm/entrypoints/openai/tool_parsers/__init__.py
+3
-1
@@ -8,6 +8,7 @@ from .granite_tool_parser import GraniteToolParser
8
from .hermes_tool_parser import Hermes2ProToolParser
9
from .internlm2_tool_parser import Internlm2ToolParser
10
from .jamba_tool_parser import JambaToolParser
11
+from .kimi_k2_tool_parser import KimiK2ToolParser
12
from .llama4_pythonic_tool_parser import Llama4PythonicToolParser
13
from .llama_tool_parser import Llama3JsonToolParser
14
from .minimax_tool_parser import MinimaxToolParser
@@ -21,5 +22,6 @@ __all__ = [
22
"GraniteToolParser", "Hermes2ProToolParser", "MistralToolParser",
23
"Internlm2ToolParser", "Llama3JsonToolParser", "JambaToolParser",
24
"Llama4PythonicToolParser", "PythonicToolParser", "Phi4MiniJsonToolParser",
24
- "DeepSeekV3ToolParser", "xLAMToolParser", "MinimaxToolParser"
25
+ "DeepSeekV3ToolParser", "xLAMToolParser", "MinimaxToolParser",
26
+ "KimiK2ToolParser"
27
]
vllm/entrypoints/openai/tool_parsers/kimi_k2_tool_parser.py
new
+377
@@ -0,0 +1,377 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# code modified from deepseekv3_tool_parser.py
4
+
5
+from collections.abc import Sequence
6
+from typing import Union
7
+
8
+import regex as re
9
+
10
+from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
11
+ DeltaFunctionCall, DeltaMessage,
12
+ DeltaToolCall,
13
+ ExtractedToolCallInformation,
14
+ FunctionCall, ToolCall)
15
+from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import (
16
+ ToolParser, ToolParserManager)
17
+from vllm.logger import init_logger
18
+from vllm.transformers_utils.tokenizer import AnyTokenizer
19
+
20
+logger = init_logger(__name__)
21
+
22
+
23
+@ToolParserManager.register_module(["kimi_k2"])
24
+class KimiK2ToolParser(ToolParser):
25
+
26
+ def __init__(self, tokenizer: AnyTokenizer):
27
+ super().__init__(tokenizer)
28
+ self.current_tool_name_sent: bool = False
29
+ self.prev_tool_call_arr: list[dict] = []
30
+ self.current_tool_id: int = -1
31
+ self.streamed_args_for_tool: list[str] = (
32
+ []) # map what has been streamed for each tool so far to a list
33
+
34
+ self.tool_calls_start_token: str = "<|tool_calls_section_begin|>"
35
+ self.tool_calls_end_token: str = "<|tool_calls_section_end|>"
36
+
37
+ self.tool_call_start_token: str = "<|tool_call_begin|>"
38
+ self.tool_call_end_token: str = "<|tool_call_end|>"
39
+
40
+ self.tool_call_regex = re.compile(
41
+ r"<\|tool_call_begin\|>\s*(?P<tool_call_id>[\w\.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>.*?)\s*<\|tool_call_end\|>"
42
+ )
43
+
44
+ self.stream_tool_call_portion_regex = re.compile(
45
+ r"(?P<tool_call_id>[\w\.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>.*)"
46
+ )
47
+
48
+ self.stream_tool_call_name_regex = re.compile(
49
+ r"(?P<tool_call_id>[\w\.]+:\d+)\s*")
50
+
51
+ if not self.model_tokenizer:
52
+ raise ValueError(
53
+ "The model tokenizer must be passed to the ToolParser "
54
+ "constructor during construction.")
55
+ self.tool_calls_start_token_id = self.vocab.get(
56
+ self.tool_calls_start_token)
57
+ self.tool_calls_end_token_id = self.vocab.get(
58
+ self.tool_calls_end_token)
59
+
60
+ self.tool_call_start_token_id = self.vocab.get(
61
+ self.tool_call_start_token)
62
+ self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)
63
+
64
+ if (self.tool_calls_start_token_id is None
65
+ or self.tool_calls_end_token_id is None):
66
+ raise RuntimeError(
67
+ "Kimi-K2 Tool parser could not locate tool call start/end "
68
+ "tokens in the tokenizer!")
69
+
70
+ def extract_tool_calls(
71
+ self,
72
+ model_output: str,
73
+ request: ChatCompletionRequest,
74
+ ) -> ExtractedToolCallInformation:
75
+
76
+ # sanity check; avoid unnecessary processing
77
+ if self.tool_calls_start_token not in model_output:
78
+ return ExtractedToolCallInformation(tools_called=False,
79
+ tool_calls=[],
80
+ content=model_output)
81
+
82
+ else:
83
+ try:
84
+ # there are two possible captures - between tags, or between a
85
+ # tag and end-of-string so the result of
86
+ # findall is an array of tuples where one is a function call and
87
+ # the other is None
88
+ function_call_tuples = self.tool_call_regex.findall(
89
+ model_output)
90
+
91
+ logger.debug("function_call_tuples: %s", function_call_tuples)
92
+
93
+ tool_calls = []
94
+ for match in function_call_tuples:
95
+ function_id, function_args = match
96
+ # function_id: functions.get_weather:0
97
+ function_name = function_id.split('.')[1].split(':')[0]
98
+ tool_calls.append(
99
+ ToolCall(
100
+ id=function_id,
101
+ type='function',
102
+ function=FunctionCall(name=function_name,
103
+ arguments=function_args),
104
+ ))
105
+
106
+ content = model_output[:model_output.
107
+ find(self.tool_calls_start_token)]
108
+ return ExtractedToolCallInformation(
109
+ tools_called=True,
110
+ tool_calls=tool_calls,
111
+ content=content if content else None,
112
+ )
113
+
114
+ except Exception:
115
+ logger.exception(
116
+ "Error in extracting tool call from response.")
117
+ return ExtractedToolCallInformation(tools_called=False,
118
+ tool_calls=[],
119
+ content=model_output)
120
+
121
+ def extract_tool_calls_streaming(
122
+ self,
123
+ previous_text: str,
124
+ current_text: str,
125
+ delta_text: str,
126
+ previous_token_ids: Sequence[int],
127
+ current_token_ids: Sequence[int],
128
+ delta_token_ids: Sequence[int],
129
+ request: ChatCompletionRequest,
130
+ ) -> Union[DeltaMessage, None]:
131
+
132
+ logger.debug("delta_text: %s", delta_text)
133
+ logger.debug("delta_token_ids: %s", delta_token_ids)
134
+ # check to see if we should be streaming a tool call - is there a
135
+ if self.tool_calls_start_token_id not in current_token_ids:
136
+ logger.debug("No tool call tokens found!")
137
+ return DeltaMessage(content=delta_text)
138
+ delta_text = delta_text.replace(self.tool_calls_start_token,
139
+ "").replace(self.tool_calls_end_token,
140
+ "")
141
+ try:
142
+
143
+ # figure out where we are in the parsing by counting tool call
144
+ # start & end tags
145
+ prev_tool_start_count = previous_token_ids.count(
146
+ self.tool_call_start_token_id)
147
+ prev_tool_end_count = previous_token_ids.count(
148
+ self.tool_call_end_token_id)
149
+ cur_tool_start_count = current_token_ids.count(
150
+ self.tool_call_start_token_id)
151
+ cur_tool_end_count = current_token_ids.count(
152
+ self.tool_call_end_token_id)
153
+ tool_call_portion = None
154
+ text_portion = None
155
+
156
+ # case: if we're generating text, OR rounding out a tool call
157
+ if (cur_tool_start_count == cur_tool_end_count
158
+ and prev_tool_end_count == cur_tool_end_count
159
+ and self.tool_call_end_token not in delta_text):
160
+ logger.debug("Generating text content! skipping tool parsing.")
161
+ return DeltaMessage(content=delta_text)
162
+
163
+ if self.tool_call_end_token in delta_text:
164
+ logger.debug("tool_call_end_token in delta_text")
165
+ full_text = current_text + delta_text
166
+ tool_call_portion = full_text.split(
167
+ self.tool_call_start_token)[-1].split(
168
+ self.tool_call_end_token)[0].rstrip()
169
+ delta_text = delta_text.split(
170
+ self.tool_call_end_token)[0].rstrip()
171
+ text_portion = delta_text.split(
172
+ self.tool_call_end_token)[-1].lstrip()
173
+
174
+ # case -- we're starting a new tool call
175
+ if (cur_tool_start_count > cur_tool_end_count
176
+ and cur_tool_start_count > prev_tool_start_count):
177
+ if len(delta_token_ids) > 1:
178
+ tool_call_portion = current_text.split(
179
+ self.tool_call_start_token)[-1]
180
+ else:
181
+ tool_call_portion = None
182
+ delta = None
183
+
184
+ text_portion = None
185
+
186
+ # set cursors and state appropriately
187
+ self.current_tool_id += 1
188
+ self.current_tool_name_sent = False
189
+ self.streamed_args_for_tool.append("")
190
+ logger.debug("Starting on a new tool %s", self.current_tool_id)
191
+
192
+ # case -- we're updating an existing tool call
193
+ elif (cur_tool_start_count > cur_tool_end_count
194
+ and cur_tool_start_count == prev_tool_start_count):
195
+
196
+ # get the portion of the text that's the tool call
197
+ tool_call_portion = current_text.split(
198
+ self.tool_call_start_token)[-1]
199
+ text_portion = None
200
+
201
+ # case -- the current tool call is being closed.
202
+ elif (cur_tool_start_count == cur_tool_end_count
203
+ and cur_tool_end_count >= prev_tool_end_count):
204
+ if self.prev_tool_call_arr is None or len(
205
+ self.prev_tool_call_arr) == 0:
206
+ logger.debug(
207
+ "attempting to close tool call, but no tool call")
208
+ return None
209
+ diff = self.prev_tool_call_arr[self.current_tool_id].get(
210
+ "arguments")
211
+ if diff:
212
+ diff = (diff.encode("utf-8").decode("unicode_escape")
213
+ if diff is str else diff)
214
+ if '"}' not in delta_text:
215
+ return None
216
+ end_loc = delta_text.rindex('"}')
217
+ diff = delta_text[:end_loc] + '"}'
218
+ logger.debug(
219
+ "Finishing tool and found diff that had not "
220
+ "been streamed yet: %s",
221
+ diff,
222
+ )
223
+ self.streamed_args_for_tool[self.current_tool_id] += diff
224
+ return DeltaMessage(tool_calls=[
225
+ DeltaToolCall(
226
+ index=self.current_tool_id,
227
+ function=DeltaFunctionCall(
228
+ arguments=diff).model_dump(exclude_none=True),
229
+ )
230
+ ])
231
+
232
+ # case -- otherwise we're just generating text
233
+ else:
234
+ text = delta_text.replace(self.tool_call_start_token, "")
235
+ text = text.replace(self.tool_call_end_token, "")
236
+ delta = DeltaMessage(tool_calls=[], content=text)
237
+ return delta
238
+
239
+ current_tool_call = dict()
240
+ if tool_call_portion:
241
+ current_tool_call_matches = (
242
+ self.stream_tool_call_portion_regex.match(
243
+ tool_call_portion))
244
+ if current_tool_call_matches:
245
+ tool_id, tool_args = (current_tool_call_matches.groups())
246
+ tool_name = tool_id.split('.')[1].split(':')[0]
247
+ current_tool_call['id'] = tool_id
248
+ current_tool_call["name"] = tool_name
249
+ current_tool_call["arguments"] = tool_args
250
+ else:
251
+ current_tool_call_name_matches = (
252
+ self.stream_tool_call_name_regex.match(
253
+ tool_call_portion))
254
+ if current_tool_call_name_matches:
255
+ tool_id_str, = current_tool_call_name_matches.groups()
256
+ tool_name = tool_id_str.split('.')[1].split(':')[0]
257
+ current_tool_call['id'] = tool_id_str
258
+ current_tool_call["name"] = tool_name
259
+ current_tool_call["arguments"] = ""
260
+ else:
261
+ logger.debug("Not enough token")
262
+ return None
263
+
264
+ # case - we haven't sent the tool name yet. If it's available, send
265
+ # it. otherwise, wait until it's available.
266
+ if not self.current_tool_name_sent:
267
+ if current_tool_call is None:
268
+ return None
269
+ function_name: Union[str, None] = current_tool_call.get("name")
270
+ tool_id = current_tool_call.get("id")
271
+ if function_name:
272
+ self.current_tool_name_sent = True
273
+ return DeltaMessage(tool_calls=[
274
+ DeltaToolCall(
275
+ index=self.current_tool_id,
276
+ type="function",
277
+ id=tool_id,
278
+ function=DeltaFunctionCall(
279
+ name=function_name).model_dump(
280
+ exclude_none=True),
281
+ )
282
+ ])
283
+ else:
284
+ return None
285
+
286
+ # case -- otherwise, send the tool call delta
287
+
288
+ # if the tool call portion is None, send the delta as text
289
+ if tool_call_portion is None:
290
+ # if there's text but not tool calls, send that -
291
+ # otherwise None to skip chunk
292
+ delta = (DeltaMessage(
293
+ content=delta_text) if text_portion is not None else None)
294
+ return delta
295
+
296
+ # now, the nitty-gritty of tool calls
297
+ # now we have the portion to parse as tool call.
298
+
299
+ logger.debug("Trying to parse current tool call with ID %s",
300
+ self.current_tool_id)
301
+
302
+ # if we're starting a new tool call, push an empty object in as
303
+ # a placeholder for the arguments
304
+ if len(self.prev_tool_call_arr) <= self.current_tool_id:
305
+ self.prev_tool_call_arr.append({})
306
+
307
+ # main logic for tool parsing here - compare prev. partially-parsed
308
+ # JSON to the current partially-parsed JSON
309
+ prev_arguments = self.prev_tool_call_arr[self.current_tool_id].get(
310
+ "arguments")
311
+ cur_arguments = current_tool_call.get("arguments")
312
+
313
+ logger.debug("diffing old arguments: %s", prev_arguments)
314
+ logger.debug("against new ones: %s", cur_arguments)
315
+
316
+ # case -- no arguments have been created yet. skip sending a delta.
317
+ if not cur_arguments and not prev_arguments:
318
+ logger.debug("Skipping text %s - no arguments", delta_text)
319
+ delta = None
320
+
321
+ # case -- prev arguments are defined, but non are now.
322
+ # probably impossible, but not a fatal error - just keep going
323
+ elif not cur_arguments and prev_arguments:
324
+ logger.error("should be impossible to have arguments reset "
325
+ "mid-call. skipping streaming anything.")
326
+ delta = None
327
+
328
+ # case -- we now have the first info about arguments available from
329
+ # autocompleting the JSON
330
+ elif cur_arguments and not prev_arguments:
331
+
332
+ delta = DeltaMessage(tool_calls=[
333
+ DeltaToolCall(
334
+ index=self.current_tool_id,
335
+ function=DeltaFunctionCall(
336
+ arguments=cur_arguments).model_dump(
337
+ exclude_none=True),
338
+ )
339
+ ])
340
+ self.streamed_args_for_tool[
341
+ self.current_tool_id] = cur_arguments
342
+
343
+ # last case -- we have an update to existing arguments.
344
+ elif cur_arguments and prev_arguments:
345
+ if (isinstance(delta_text, str)
346
+ and cur_arguments != prev_arguments
347
+ and len(cur_arguments) > len(prev_arguments)
348
+ and cur_arguments.startswith(prev_arguments)):
349
+ delta_arguments = cur_arguments[len(prev_arguments):]
350
+ logger.debug("got diff %s", delta_text)
351
+
352
+ delta = DeltaMessage(tool_calls=[
353
+ DeltaToolCall(
354
+ index=self.current_tool_id,
355
+ function=DeltaFunctionCall(
356
+ arguments=delta_arguments).model_dump(
357
+ exclude_none=True),
358
+ )
359
+ ])
360
+ self.streamed_args_for_tool[
361
+ self.current_tool_id] = cur_arguments
362
+ else:
363
+ delta = None
364
+
365
+ # handle saving the state for the current tool into
366
+ # the "prev" list for use in diffing for the next iteration
367
+ if self.current_tool_id == len(self.prev_tool_call_arr) - 1:
368
+ self.prev_tool_call_arr[
369
+ self.current_tool_id] = current_tool_call
370
+ else:
371
+ self.prev_tool_call_arr.append(current_tool_call)
372
+
373
+ return delta
374
+
375
+ except Exception:
376
+ logger.exception("Error trying to handle streaming tool call.")
377
+ return None # do not stream a delta. skip this token ID.
\ No newline at end of file