[Bugfix][Tool Parser] Stream tool calls that arrive whole in a single delta in llama3_json, jamba and ernie45 parsers

tests/tool_parsers/test_ernie45_moe_tool_parser.py

@@ -7,6 +7,7 @@
 
 import pytest
 
+from tests.tool_parsers.utils import run_tool_extraction_streaming
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
 from vllm.entrypoints.openai.engine.protocol import (
     DeltaMessage,
@@ -357,3 +358,40 @@ def test_extract_tool_calls_streaming_incremental(
     assert len(actual_tool_calls) > 0
     # check tool call format
     assert_tool_calls(actual_tool_calls, expected_tool_calls)
+
+
+def test_extract_tool_calls_streaming_multiple_tools_single_delta(
+    ernie45_tool_parser,
+):
+    """Multiple complete tool calls in one delta must all be emitted.
+
+    This happens e.g. with async scheduling or stream_interval > 1, where
+    a single streamed delta can contain the entire tool call message.
+    Previously only the first tool call was emitted and the rest were left
+    stranded in the internal buffer.
+    """
+    model_output = (
+        "I will check the weather and the time.\n"
+        "<tool_call>\n"
+        '{"name": "get_current_temperature", "arguments": {"location": "Beijing"}}\n'
+        "</tool_call>\n"
+        "<tool_call>\n"
+        '{"name": "get_current_time", "arguments": {"timezone": "Asia/Shanghai"}}\n'
+        "</tool_call>\n"
+    )
+
+    reconstructor = run_tool_extraction_streaming(
+        ernie45_tool_parser, [model_output], assert_one_tool_per_delta=False
+    )
+
+    assert [call.function.name for call in reconstructor.tool_calls] == [
+        "get_current_temperature",
+        "get_current_time",
+    ]
+    assert json.loads(reconstructor.tool_calls[0].function.arguments) == {
+        "location": "Beijing"
+    }
+    assert json.loads(reconstructor.tool_calls[1].function.arguments) == {
+        "timezone": "Asia/Shanghai"
+    }
+    assert "I will check the weather and the time." in reconstructor.other_content

tests/tool_parsers/test_jamba_tool_parser.py

@@ -8,6 +8,7 @@
 import pytest
 from partial_json_parser.core.options import Allow
 
+from tests.tool_parsers.utils import run_tool_extraction_streaming
 from vllm.entrypoints.openai.engine.protocol import DeltaMessage, FunctionCall, ToolCall
 from vllm.tokenizers import TokenizerLike, get_tokenizer
 from vllm.tokenizers.detokenizer_utils import detokenize_incrementally
@@ -306,3 +307,31 @@ def test_extract_tool_calls_streaming(
         )
     ]
     assert_tool_calls(actual_tool_calls, expected_tool_calls)
+
+
+def test_extract_tool_calls_streaming_single_delta(jamba_tool_parser):
+    """Tool calls arriving whole in the first delta must not be dropped.
+
+    This happens e.g. with async scheduling or stream_interval > 1, where
+    a single streamed delta can contain the entire tool call message.
+    """
+    model_output = """<tool_calls>[\n    {"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}},\n    {"name": "get_current_weather", "arguments": {"city": "Orlando", "state": "FL", "unit": "fahrenheit"}}\n]</tool_calls>"""  # noqa: E501
+
+    reconstructor = run_tool_extraction_streaming(
+        jamba_tool_parser, [model_output], assert_one_tool_per_delta=False
+    )
+
+    assert [call.function.name for call in reconstructor.tool_calls] == [
+        "get_current_weather",
+        "get_current_weather",
+    ]
+    assert json.loads(reconstructor.tool_calls[0].function.arguments) == {
+        "city": "Dallas",
+        "state": "TX",
+        "unit": "fahrenheit",
+    }
+    assert json.loads(reconstructor.tool_calls[1].function.arguments) == {
+        "city": "Orlando",
+        "state": "FL",
+        "unit": "fahrenheit",
+    }

tests/tool_parsers/test_llama3_json_tool_parser.py

@@ -1,11 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
 from unittest.mock import MagicMock, patch
 
 import pytest
 from transformers import AutoTokenizer
 
+from tests.tool_parsers.utils import run_tool_extraction_streaming
 from vllm.entrypoints.openai.engine.protocol import ExtractedToolCallInformation
 from vllm.tool_parsers.llama_tool_parser import Llama3JsonToolParser
 
@@ -267,3 +269,82 @@ def test_regex_timeout_handling(parser):
         assert result.tools_called is False
         assert len(result.tool_calls) == 0
         mock_regex.finditer.assert_called_once()
+
+
+class StubTokenizer:
+    """Minimal tokenizer stub for streaming tests.
+
+    The streaming parser operates on text only, so the tokenizer just needs
+    to expose the bot token in its vocab. Using a stub avoids downloading
+    the gated Llama tokenizer.
+    """
+
+    def get_vocab(self) -> dict[str, int]:
+        return {"<|python_tag|>": 128010}
+
+    def tokenize(self, text: str) -> list[str]:
+        return []
+
+
+def test_streaming_complete_tool_call_single_delta():
+    """A tool call arriving whole in the first delta must not be dropped."""
+    parser = Llama3JsonToolParser(StubTokenizer())
+    model_output = '{"name": "get_weather", "arguments": {"city": "Tokyo"}}'
+
+    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_complete_tool_call_single_delta_with_bot_token():
+    parser = Llama3JsonToolParser(StubTokenizer())
+    model_output = (
+        '<|python_tag|>{"name": "get_weather", "arguments": {"city": "Tokyo"}}'
+    )
+
+    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_complete_tool_call_single_delta_parameters_key():
+    """Llama may emit "parameters" instead of "arguments"; both must work."""
+    parser = Llama3JsonToolParser(StubTokenizer())
+    model_output = '{"name": "get_weather", "parameters": {"city": "Tokyo"}}'
+
+    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_parallel_tool_calls_single_delta():
+    """Multiple complete tool calls in one delta must all be emitted."""
+    parser = Llama3JsonToolParser(StubTokenizer())
+    model_output = (
+        '{"name": "get_weather", "arguments": {"city": "Tokyo"}}; '
+        '{"name": "get_time", "arguments": {"timezone": "Asia/Tokyo"}}'
+    )
+
+    reconstructor = run_tool_extraction_streaming(
+        parser, [model_output], assert_one_tool_per_delta=False
+    )
+
+    assert [call.function.name for call in reconstructor.tool_calls] == [
+        "get_weather",
+        "get_time",
+    ]
+    assert json.loads(reconstructor.tool_calls[0].function.arguments) == {
+        "city": "Tokyo"
+    }
+    assert json.loads(reconstructor.tool_calls[1].function.arguments) == {
+        "timezone": "Asia/Tokyo"
+    }

vllm/tool_parsers/ernie45_tool_parser.py

@@ -163,8 +163,12 @@ def extract_tool_calls_streaming(
 
             return DeltaMessage(content=content if content else None)
         logger.debug("cur_text = %s", cur_text)
-        end_idx = cur_text.find(self.tool_call_end_token)
-        if end_idx != -1:
+        # a single delta may contain several complete tool calls (e.g. the
+        # whole message arrived in one delta), so keep extracting until no
+        # complete tool call is left in the buffer
+        delta_content: str | None = None
+        delta_tool_calls: list[DeltaToolCall] = []
+        while (end_idx := cur_text.find(self.tool_call_end_token)) != -1:
             if self.current_tool_id == -1:
                 self.current_tool_id = 0
                 self.prev_tool_call_arr = []
@@ -180,7 +184,9 @@ def extract_tool_calls_streaming(
 
             if len(extracted_tool_calls.tool_calls) == 0:
                 logger.warning("Failed to extract any tool calls.")
-                return None
+                if not delta_tool_calls:
+                    return None
+                break
             tool_call = extracted_tool_calls.tool_calls[0]
             self.prev_tool_call_arr[self.current_tool_id] = {
                 "name": tool_call.function.name,
@@ -189,23 +195,25 @@ def extract_tool_calls_streaming(
             self.streamed_args_for_tool[self.current_tool_id] = (
                 tool_call.function.arguments
             )
-            delta = DeltaMessage(
-                content=extracted_tool_calls.content,
-                tool_calls=[
-                    DeltaToolCall(
-                        index=self.current_tool_id,
-                        id=tool_call.id,
-                        type=tool_call.type,
-                        function=DeltaFunctionCall(
-                            name=tool_call.function.name,
-                            arguments=tool_call.function.arguments,
-                        ),
-                    )
-                ],
+            if extracted_tool_calls.content:
+                delta_content = (delta_content or "") + extracted_tool_calls.content
+            delta_tool_calls.append(
+                DeltaToolCall(
+                    index=self.current_tool_id,
+                    id=tool_call.id,
+                    type=tool_call.type,
+                    function=DeltaFunctionCall(
+                        name=tool_call.function.name,
+                        arguments=tool_call.function.arguments,
+                    ),
+                )
             )
             self.current_tool_id += 1
-            self._buffer = cur_text[end_idx + len(self.tool_call_end_token) :]
-            return delta
+            cur_text = cur_text[end_idx + len(self.tool_call_end_token) :]
+
+        if delta_tool_calls:
+            self._buffer = cur_text
+            return DeltaMessage(content=delta_content, tool_calls=delta_tool_calls)
 
         self._buffer = cur_text[start_idx:]
         content = cur_text[:start_idx].rstrip("\n")

vllm/tool_parsers/jamba_tool_parser.py

@@ -24,7 +24,7 @@
 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__)
@@ -69,6 +69,39 @@ def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
                 "tokens in the tokenizer!"
             )
 
+    def _emit_initial_complete_tool_calls(
+        self, tool_call_arr: list[dict]
+    ) -> DeltaMessage | None:
+        delta_tool_calls: list[DeltaToolCall] = []
+        streamed_args_for_tool: list[str] = []
+
+        for index, tool_call in enumerate(tool_call_arr):
+            function_name = tool_call.get("name")
+            if not function_name or "arguments" not in tool_call:
+                return None
+            arguments = json.dumps(tool_call["arguments"], ensure_ascii=False)
+            delta_tool_calls.append(
+                DeltaToolCall(
+                    index=index,
+                    type="function",
+                    id=make_tool_call_id(),
+                    function=DeltaFunctionCall(
+                        name=function_name,
+                        arguments=arguments,
+                    ).model_dump(exclude_none=True),
+                )
+            )
+            streamed_args_for_tool.append(arguments)
+
+        if not delta_tool_calls:
+            return None
+
+        self.current_tool_id = len(tool_call_arr) - 1
+        self.current_tool_name_sent = True
+        self.prev_tool_call_arr = tool_call_arr
+        self.streamed_args_for_tool = streamed_args_for_tool
+        return DeltaMessage(tool_calls=delta_tool_calls)
+
     def adjust_request(
         self, request: ChatCompletionRequest | ResponsesRequest
     ) -> ChatCompletionRequest | ResponsesRequest:
@@ -184,11 +217,17 @@ def extract_tool_calls_streaming(
             if len(tool_call_arr) == 0:
                 return None
 
+            # case: the first delta already contains the complete tool call
+            #   array (e.g. the whole message arrived in one delta); emit all
+            #   tool calls now since no further deltas may arrive
+            if self.current_tool_id < 0 and is_complete_json(parsable_arr):
+                initial_delta = self._emit_initial_complete_tool_calls(tool_call_arr)
+                if initial_delta is not None:
+                    return initial_delta
+
             # case: we are starting a new tool in the array
             #   -> array has > 0 length AND length has moved past cursor
-            elif (
-                len(tool_call_arr) > 0 and len(tool_call_arr) > self.current_tool_id + 1
-            ):
+            if len(tool_call_arr) > 0 and len(tool_call_arr) > self.current_tool_id + 1:
                 # if we're moving on to a new call, first make sure we
                 # haven't missed anything in the previous one that was
                 # auto-generated due to JSON completions, but wasn't
@@ -218,7 +257,10 @@ def extract_tool_calls_streaming(
                 # re-set stuff pertaining to progress in the current tool
                 self.current_tool_id = len(tool_call_arr) - 1
                 self.current_tool_name_sent = False
-                self.streamed_args_for_tool.append("")
+                # multiple tools may have appeared in a single delta, so make
+                # sure a slot exists for every tool up to the current one
+                while len(self.streamed_args_for_tool) <= self.current_tool_id:
+                    self.streamed_args_for_tool.append("")
                 logger.debug("starting on new tool %d", self.current_tool_id)
                 return delta