From 1423569ff51418fdf2089ed295eeae86f2ea2cd9 Mon Sep 17 00:00:00 2001 From: mosya415 Date: Sat, 25 Jul 2026 16:34:43 +0300 Subject: [PATCH] [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> --- .../test_internlm2_tool_parser.py | 43 +++++++++++++++++++ tests/tool_parsers/test_jamba_tool_parser.py | 31 +++++++++++++ vllm/tool_parsers/internlm2_tool_parser.py | 16 +++++-- vllm/tool_parsers/jamba_tool_parser.py | 17 ++++++-- 4 files changed, 99 insertions(+), 8 deletions(-) diff --git a/tests/tool_parsers/test_internlm2_tool_parser.py b/tests/tool_parsers/test_internlm2_tool_parser.py index 2e5069dbed9..7fd3860ef71 100644 --- a/tests/tool_parsers/test_internlm2_tool_parser.py +++ b/tests/tool_parsers/test_internlm2_tool_parser.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json from unittest.mock import MagicMock import pytest @@ -10,6 +11,7 @@ from tests.tool_parsers.common_tests import ( ToolParserTests, ) from vllm.tokenizers import TokenizerLike +from vllm.tool_parsers.internlm2_tool_parser import Internlm2ToolParser class TestInternLM2ToolParser(ToolParserTests): @@ -120,3 +122,44 @@ class TestInternLM2ToolParser(ToolParserTests): ), }, ) + + +def test_streaming_arguments_in_single_delta(default_tokenizer: TokenizerLike) -> None: + """Arguments arriving whole in one delta must not be dropped.""" + tokenizer_vocab = default_tokenizer.get_vocab() + default_tokenizer.get_vocab = MagicMock() + tokenizer_vocab.update( + { + "<|action_start|>": 92540, + "<|plugin|>": 92541, + "<|action_end|>": 92542, + } + ) + default_tokenizer.get_vocab.return_value = tokenizer_vocab + parser = Internlm2ToolParser(default_tokenizer) + + deltas = [ + '<|action_start|><|plugin|>{"name": "get_weather"', + ', "parameters": {"city": "Dallas", "state": "TX"}}<|action_end|>', + ] + + streamed = "" + current_text = "" + for delta_text in deltas: + previous_text = current_text + current_text += delta_text + delta_message = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=None, + ) + if delta_message and delta_message.tool_calls: + arguments = delta_message.tool_calls[0].function.arguments + if arguments: + streamed += arguments + + assert json.loads(streamed) == {"city": "Dallas", "state": "TX"} diff --git a/tests/tool_parsers/test_jamba_tool_parser.py b/tests/tool_parsers/test_jamba_tool_parser.py index f0e7899c8aa..9eb7404209d 100644 --- a/tests/tool_parsers/test_jamba_tool_parser.py +++ b/tests/tool_parsers/test_jamba_tool_parser.py @@ -306,3 +306,34 @@ def test_extract_tool_calls_streaming( ) ] assert_tool_calls(actual_tool_calls, expected_tool_calls) + + +def test_extract_tool_calls_streaming_arguments_in_single_delta(jamba_tool_parser): + """Arguments delivered whole in one coarse delta must not be dropped.""" + deltas = [ + '[{"name": "get_current_weather"', + ",", + ' "arguments": {"city": "Dallas", "state": "TX"}}]', + "", + ] + + streamed_arguments = "" + current_text = "" + for delta_text in deltas: + previous_text = current_text + current_text += delta_text + delta_message = jamba_tool_parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=None, + ) + if delta_message and delta_message.tool_calls: + arguments = delta_message.tool_calls[0].function.arguments + if arguments: + streamed_arguments += arguments + + assert json.loads(streamed_arguments) == {"city": "Dallas", "state": "TX"} diff --git a/vllm/tool_parsers/internlm2_tool_parser.py b/vllm/tool_parsers/internlm2_tool_parser.py index f4aaeef71a0..7cbffcb1c0b 100644 --- a/vllm/tool_parsers/internlm2_tool_parser.py +++ b/vllm/tool_parsers/internlm2_tool_parser.py @@ -26,7 +26,7 @@ from vllm.tool_parsers.abstract_tool_parser import ( Tool, ToolParser, ) -from vllm.tool_parsers.utils import extract_intermediate_diff +from vllm.tool_parsers.utils import extract_intermediate_diff, is_complete_json logger = init_logger(__name__) @@ -146,9 +146,17 @@ class Internlm2ToolParser(ToolParser): elif cur_arguments and not prev_arguments: cur_arguments_json = json.dumps(cur_arguments, ensure_ascii=False) - arguments_delta = cur_arguments_json[ - : cur_arguments_json.index(delta_text) + len(delta_text) - ] + match_start = cur_arguments_json.find(delta_text) + if match_start != -1: + arguments_delta = cur_arguments_json[ + : match_start + len(delta_text) + ] + elif is_complete_json(parsable_arr): + # Complete in this delta: send whole, don't drop. + arguments_delta = cur_arguments_json + else: + # Still partial: wait for more text. + return None delta = DeltaMessage( tool_calls=[ DeltaToolCall( diff --git a/vllm/tool_parsers/jamba_tool_parser.py b/vllm/tool_parsers/jamba_tool_parser.py index dec3c88d934..193a51faa3d 100644 --- a/vllm/tool_parsers/jamba_tool_parser.py +++ b/vllm/tool_parsers/jamba_tool_parser.py @@ -24,7 +24,7 @@ from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser -from vllm.tool_parsers.utils import extract_intermediate_diff +from vllm.tool_parsers.utils import extract_intermediate_diff, is_complete_json from vllm.utils.mistral import is_mistral_tokenizer logger = init_logger(__name__) @@ -266,9 +266,18 @@ class JambaToolParser(ToolParser): cur_arguments_json = json.dumps(cur_arguments, ensure_ascii=False) logger.debug("finding %s in %s", new_text, cur_arguments_json) - arguments_delta = cur_arguments_json[ - : cur_arguments_json.index(new_text) + len(new_text) - ] + # `new_text` may not appear verbatim in the re-serialized JSON. + match_start = cur_arguments_json.find(new_text) + if match_start != -1: + arguments_delta = cur_arguments_json[ + : match_start + len(new_text) + ] + elif is_complete_json(parsable_arr): + # Complete in this delta: send whole, don't drop. + arguments_delta = cur_arguments_json + else: + # Still partial: wait for more text. + return None logger.debug( "First tokens in arguments received: %s", arguments_delta )