From a1ed5f230a4113e2b31655f1b70897521ad00189 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Sat, 11 Jul 2026 17:42:41 +0530 Subject: [PATCH] [Bugfix][Tool Parser] InternLM2: stream tool calls that arrive whole in a single delta When an entire tool call arrives in one streamed delta (async scheduling, stream_interval > 1, or the whole message emitted at once), the InternLM2 streaming parser lost information in two ways: - the name-sending branch emitted only the function name and returned; with no further deltas arriving, the arguments were never streamed and the client received a tool call with empty arguments, while the non-streaming parser extracted the full call. - when assistant content preceded the tool call in the same delta, the parser emitted the content and returned immediately, losing the tool call entirely. Fix, following the approach of #47955 (gigachat3/granite) and #48295 (llama3_json/jamba/ernie45): when the tool-call JSON in the delta is already complete, emit the name and the full arguments in one DeltaMessage, carrying any preceding content in the same message instead of returning early. Add streaming regression tests for both scenarios; they fail on the previous code. Existing xfail-marked streaming tests are unaffected (still xfail, no behavior change for per-token streaming). Co-authored-by: Claude Signed-off-by: Abhay Singh --- .../test_internlm2_tool_parser.py | 38 +++++++++++++++++++ vllm/tool_parsers/internlm2_tool_parser.py | 38 ++++++++++++++----- 2 files changed, 66 insertions(+), 10 deletions(-) diff --git a/tests/tool_parsers/test_internlm2_tool_parser.py b/tests/tool_parsers/test_internlm2_tool_parser.py index 2e5069dbed94..b3e1e6ebc8b6 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 @@ -9,7 +10,9 @@ ToolParserTestConfig, ToolParserTests, ) +from tests.tool_parsers.utils import run_tool_extraction_streaming from vllm.tokenizers import TokenizerLike +from vllm.tool_parsers.internlm2_tool_parser import Internlm2ToolParser class TestInternLM2ToolParser(ToolParserTests): @@ -120,3 +123,38 @@ def test_config(self) -> ToolParserTestConfig: ), }, ) + + def test_streaming_complete_tool_call_single_delta(self, tokenizer): + """A tool call arriving whole in one delta (e.g. async scheduling or + stream_interval > 1) must emit both the name and the arguments; + previously only the name was emitted and the arguments were lost.""" + parser = Internlm2ToolParser(tokenizer) + model_output = ( + '<|action_start|><|plugin|>{"name": "get_weather", ' + '"parameters": {"city": "Tokyo"}}<|action_end|>' + ) + + reconstructor = run_tool_extraction_streaming(parser, [model_output]) + + assert len(reconstructor.tool_calls) == 1 + call = reconstructor.tool_calls[0] + assert call.function.name == "get_weather" + assert json.loads(call.function.arguments) == {"city": "Tokyo"} + + def test_streaming_content_then_complete_tool_call_single_delta(self, tokenizer): + """Content preceding a complete tool call in the same delta must not + cause the tool call to be dropped; previously the parser emitted the + content and returned, losing the tool call entirely.""" + parser = Internlm2ToolParser(tokenizer) + model_output = ( + 'Sure, let me check.<|action_start|><|plugin|>{"name": "get_weather", ' + '"parameters": {"city": "Tokyo"}}<|action_end|>' + ) + + reconstructor = run_tool_extraction_streaming(parser, [model_output]) + + assert reconstructor.other_content == "Sure, let me check." + assert len(reconstructor.tool_calls) == 1 + call = reconstructor.tool_calls[0] + assert call.function.name == "get_weather" + assert json.loads(call.function.arguments) == {"city": "Tokyo"} diff --git a/vllm/tool_parsers/internlm2_tool_parser.py b/vllm/tool_parsers/internlm2_tool_parser.py index f4aaeef71a0f..4fe36a88d07b 100644 --- a/vllm/tool_parsers/internlm2_tool_parser.py +++ b/vllm/tool_parsers/internlm2_tool_parser.py @@ -26,7 +26,7 @@ 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__) @@ -79,13 +79,20 @@ def extract_tool_calls_streaming( new_delta = current_text[last_pos:] text, action = new_delta.split("<|action_start|><|plugin|>") - if len(text) > 0: - self.position = self.position + len(text) - return DeltaMessage(content=text) - action = action.strip() action = action.split("<|action_end|>".strip())[0] + content_prefix: str | None = None + if len(text) > 0: + self.position = self.position + len(text) + # if the tool call is not yet complete, emit the content now and + # keep parsing the tool call from later deltas; otherwise fall + # through, so that a tool call arriving whole in this same delta + # is not lost when no further deltas arrive + if not is_complete_json(action): + return DeltaMessage(content=text) + content_prefix = text + # bit mask flags for partial JSON parsing. If the name hasn't been # sent yet, don't allow sending # an incomplete string since OpenAI only ever (as far as I have @@ -109,20 +116,31 @@ def extract_tool_calls_streaming( function_name = tool_call_arr.get("name") if function_name: self.current_tool_id = self.current_tool_id + 1 + delta_function_call = DeltaFunctionCall(name=function_name) + streamed_args = "" + cur_arguments = self.get_arguments(tool_call_arr) + # if the whole tool call arrived in a single delta (e.g. + # with async scheduling or stream_interval > 1), emit the + # complete arguments together with the name, since no + # further deltas may arrive + if cur_arguments is not None and is_complete_json(action): + streamed_args = json.dumps(cur_arguments, ensure_ascii=False) + delta_function_call.arguments = streamed_args delta = DeltaMessage( + content=content_prefix, tool_calls=[ DeltaToolCall( index=self.current_tool_id, type="function", id=make_tool_call_id(), - function=DeltaFunctionCall( - name=function_name - ).model_dump(exclude_none=True), + function=delta_function_call.model_dump( + exclude_none=True + ), ) - ] + ], ) self.current_tool_name_sent = True - self.streamed_args_for_tool.append("") + self.streamed_args_for_tool.append(streamed_args) else: delta = None # now we know we're on the same tool call and we're streaming