[Bugfix][Tool Parser] InternLM2: stream tool calls that arrive whole in a single delta

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"}

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