[Bugfix][Tool Parser] Fix dropped streaming arguments in Jamba and InternLM2 parsers (#48852)
Signed-off-by: mosya415 <263250241+mosya415@users.noreply.github.com> Co-authored-by: mosya415 <263250241+mosya415@users.noreply.github.com>
mosya415 committed
Jul 25, 2026 at 16:34 UTC
1423569ff51418fdf2089ed295eeae86f2ea2cd9
4 files changed
+99
-8
tests/tool_parsers/test_internlm2_tool_parser.py
+43
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import json
5
from unittest.mock import MagicMock
6
7
import pytest
@@ -10,6 +11,7 @@ from tests.tool_parsers.common_tests import (
11
ToolParserTests,
12
)
13
from vllm.tokenizers import TokenizerLike
14
+from vllm.tool_parsers.internlm2_tool_parser import Internlm2ToolParser
15
16
17
class TestInternLM2ToolParser(ToolParserTests):
@@ -120,3 +122,44 @@ class TestInternLM2ToolParser(ToolParserTests):
122
),
123
},
124
)
125
+
126
+
127
+def test_streaming_arguments_in_single_delta(default_tokenizer: TokenizerLike) -> None:
128
+ """Arguments arriving whole in one delta must not be dropped."""
129
+ tokenizer_vocab = default_tokenizer.get_vocab()
130
+ default_tokenizer.get_vocab = MagicMock()
131
+ tokenizer_vocab.update(
132
+ {
133
+ "<|action_start|>": 92540,
134
+ "<|plugin|>": 92541,
135
+ "<|action_end|>": 92542,
136
+ }
137
+ )
138
+ default_tokenizer.get_vocab.return_value = tokenizer_vocab
139
+ parser = Internlm2ToolParser(default_tokenizer)
140
+
141
+ deltas = [
142
+ '<|action_start|><|plugin|>{"name": "get_weather"',
143
+ ', "parameters": {"city": "Dallas", "state": "TX"}}<|action_end|>',
144
+ ]
145
+
146
+ streamed = ""
147
+ current_text = ""
148
+ for delta_text in deltas:
149
+ previous_text = current_text
150
+ current_text += delta_text
151
+ delta_message = parser.extract_tool_calls_streaming(
152
+ previous_text=previous_text,
153
+ current_text=current_text,
154
+ delta_text=delta_text,
155
+ previous_token_ids=[],
156
+ current_token_ids=[],
157
+ delta_token_ids=[],
158
+ request=None,
159
+ )
160
+ if delta_message and delta_message.tool_calls:
161
+ arguments = delta_message.tool_calls[0].function.arguments
162
+ if arguments:
163
+ streamed += arguments
164
+
165
+ assert json.loads(streamed) == {"city": "Dallas", "state": "TX"}
tests/tool_parsers/test_jamba_tool_parser.py
+31
@@ -306,3 +306,34 @@ def test_extract_tool_calls_streaming(
306
)
307
]
308
assert_tool_calls(actual_tool_calls, expected_tool_calls)
309
+
310
+
311
+def test_extract_tool_calls_streaming_arguments_in_single_delta(jamba_tool_parser):
312
+ """Arguments delivered whole in one coarse delta must not be dropped."""
313
+ deltas = [
314
+ '<tool_calls>[{"name": "get_current_weather"',
315
+ ",",
316
+ ' "arguments": {"city": "Dallas", "state": "TX"}}]',
317
+ "</tool_calls>",
318
+ ]
319
+
320
+ streamed_arguments = ""
321
+ current_text = ""
322
+ for delta_text in deltas:
323
+ previous_text = current_text
324
+ current_text += delta_text
325
+ delta_message = jamba_tool_parser.extract_tool_calls_streaming(
326
+ previous_text=previous_text,
327
+ current_text=current_text,
328
+ delta_text=delta_text,
329
+ previous_token_ids=[],
330
+ current_token_ids=[],
331
+ delta_token_ids=[],
332
+ request=None,
333
+ )
334
+ if delta_message and delta_message.tool_calls:
335
+ arguments = delta_message.tool_calls[0].function.arguments
336
+ if arguments:
337
+ streamed_arguments += arguments
338
+
339
+ assert json.loads(streamed_arguments) == {"city": "Dallas", "state": "TX"}
vllm/tool_parsers/internlm2_tool_parser.py
+12
-4
@@ -26,7 +26,7 @@ from vllm.tool_parsers.abstract_tool_parser import (
26
Tool,
27
ToolParser,
28
)
29
-from vllm.tool_parsers.utils import extract_intermediate_diff
29
+from vllm.tool_parsers.utils import extract_intermediate_diff, is_complete_json
30
31
logger = init_logger(__name__)
32
@@ -146,9 +146,17 @@ class Internlm2ToolParser(ToolParser):
146
elif cur_arguments and not prev_arguments:
147
cur_arguments_json = json.dumps(cur_arguments, ensure_ascii=False)
148
149
- arguments_delta = cur_arguments_json[
150
- : cur_arguments_json.index(delta_text) + len(delta_text)
151
- ]
149
+ match_start = cur_arguments_json.find(delta_text)
150
+ if match_start != -1:
151
+ arguments_delta = cur_arguments_json[
152
+ : match_start + len(delta_text)
153
+ ]
154
+ elif is_complete_json(parsable_arr):
155
+ # Complete in this delta: send whole, don't drop.
156
+ arguments_delta = cur_arguments_json
157
+ else:
158
+ # Still partial: wait for more text.
159
+ return None
160
delta = DeltaMessage(
161
tool_calls=[
162
DeltaToolCall(
vllm/tool_parsers/jamba_tool_parser.py
+13
-4
@@ -24,7 +24,7 @@ from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
24
from vllm.logger import init_logger
25
from vllm.tokenizers import TokenizerLike
26
from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser
27
-from vllm.tool_parsers.utils import extract_intermediate_diff
27
+from vllm.tool_parsers.utils import extract_intermediate_diff, is_complete_json
28
from vllm.utils.mistral import is_mistral_tokenizer
29
30
logger = init_logger(__name__)
@@ -266,9 +266,18 @@ class JambaToolParser(ToolParser):
266
cur_arguments_json = json.dumps(cur_arguments, ensure_ascii=False)
267
logger.debug("finding %s in %s", new_text, cur_arguments_json)
268
269
- arguments_delta = cur_arguments_json[
270
- : cur_arguments_json.index(new_text) + len(new_text)
271
- ]
269
+ # `new_text` may not appear verbatim in the re-serialized JSON.
270
+ match_start = cur_arguments_json.find(new_text)
271
+ if match_start != -1:
272
+ arguments_delta = cur_arguments_json[
273
+ : match_start + len(new_text)
274
+ ]
275
+ elif is_complete_json(parsable_arr):
276
+ # Complete in this delta: send whole, don't drop.
277
+ arguments_delta = cur_arguments_json
278
+ else:
279
+ # Still partial: wait for more text.
280
+ return None
281
logger.debug(
282
"First tokens in arguments received: %s", arguments_delta
283
)