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