[Bugfix][Tool Parser] Fix dropped streaming arguments in Jamba and InternLM2 parsers (#48852)

Signed-off-by: mosya415 <[email protected]>
Co-authored-by: mosya415 <[email protected]>

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

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 = [
+        '<tool_calls>[{"name": "get_current_weather"',
+        ",",
+        ' "arguments": {"city": "Dallas", "state": "TX"}}]',
+        "</tool_calls>",
+    ]
+
+    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"}

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(

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
                     )