[Bugfix] Fix InternLM2 tool parser dropping characters from streamed arguments (#60709)

Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Flora Feng <4florafeng@gmail.com>

tests/tool_parsers/test_internlm2_tool_parser.py

@@ -10,6 +10,7 @@ from tests.tool_parsers.common_tests import (
     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
 
@@ -162,3 +163,30 @@ def test_streaming_arguments_in_single_delta(default_tokenizer: TokenizerLike) -
                 streamed += arguments
 
     assert json.loads(streamed) == {"city": "Dallas", "state": "TX"}
+
+
+@pytest.mark.parametrize(
+    "arguments_deltas",
+    [
+        ['{"location": "Beijing', '", "days": 3}}<|action_end|>'],
+        [*'{"location": "Beijing", "days": 3}}', "<|action_end|>"],
+    ],
+    ids=["split_inside_string", "char_by_char"],
+)
+def test_streaming_arguments_match_final_arguments(
+    default_tokenizer: TokenizerLike, arguments_deltas: list[str]
+) -> None:
+    """Streamed argument deltas must add up to the final arguments, wherever
+    the chunk boundaries fall."""
+    parser = Internlm2ToolParser(default_tokenizer)
+    deltas = [
+        '<|action_start|><|plugin|>{"name": "get_weather", "parameters": ',
+        *arguments_deltas,
+    ]
+
+    reconstructor = run_tool_extraction_streaming(parser, deltas)
+
+    assert len(reconstructor.tool_calls) == 1
+    streamed = reconstructor.tool_calls[0].function.arguments
+    streamed += parser.get_remaining_unstreamed_args()
+    assert streamed == '{"location": "Beijing", "days": 3}'

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, is_complete_json
+from vllm.tool_parsers.utils import find_common_prefix, is_complete_json
 
 logger = init_logger(__name__)
 
@@ -172,10 +172,17 @@ class Internlm2ToolParser(ToolParser):
                 elif cur_arguments and prev_arguments:
                     cur_args_json = json.dumps(cur_arguments, ensure_ascii=False)
                     prev_args_json = json.dumps(prev_arguments, ensure_ascii=False)
+                    sent = len(self.streamed_args_for_tool[self.current_tool_id])
 
-                    argument_diff = extract_intermediate_diff(
-                        cur_args_json, prev_args_json
-                    )
+                    # Partial JSON is serialized with synthetic closers, so
+                    # only the part shared with the previous parse is final.
+                    if is_complete_json(parsable_arr):
+                        argument_diff = cur_args_json[sent:]
+                    elif cur_args_json != prev_args_json:
+                        prefix = find_common_prefix(prev_args_json, cur_args_json)
+                        argument_diff = prefix[sent:]
+                    else:
+                        argument_diff = ""
 
                     delta = DeltaMessage(
                         tool_calls=[