diff --git a/tests/tool_parsers/test_phi4mini_tool_parser.py b/tests/tool_parsers/test_phi4mini_tool_parser.py index eff9fa9bb8ff..978ea2cd0a8b 100644 --- a/tests/tool_parsers/test_phi4mini_tool_parser.py +++ b/tests/tool_parsers/test_phi4mini_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 import ToolParser, ToolParserManager class TestPhi4MiniToolParser(ToolParserTests): @@ -85,18 +88,7 @@ def test_config(self) -> ToolParserTestConfig: parallel_tool_calls_names=["get_weather", "get_time"], parallel_tool_calls_expected_content=None, # xfail markers - xfail_streaming={ - "test_no_tool_calls": "Phi4 Mini streaming not implemented", - "test_single_tool_call_simple_args": ( - "Phi4 Mini streaming not implemented" - ), - "test_parallel_tool_calls": "Phi4 Mini streaming not implemented", - "test_various_data_types": "Phi4 Mini streaming not implemented", - "test_empty_arguments": "Phi4 Mini streaming not implemented", - "test_surrounding_text": "Phi4 Mini streaming not implemented", - "test_escaped_strings": "Phi4 Mini streaming not implemented", - "test_streaming_reconstruction": "Phi4 Mini streaming not implemented", - }, + xfail_streaming={}, xfail_nonstreaming={ "test_various_data_types": ( "Phi4MiniJsonToolParser regex has nesting limitations " @@ -108,3 +100,74 @@ def test_config(self) -> ToolParserTestConfig: ), }, ) + + +@pytest.fixture +def phi4mini_parser(default_tokenizer: TokenizerLike) -> ToolParser: + return ToolParserManager.get_tool_parser("phi4_mini_json")(default_tokenizer) + + +def test_streaming_split_marker_is_not_emitted_as_content( + phi4mini_parser: ToolParser, +) -> None: + """The `functools` marker spans several tokens and must not leak.""" + deltas = [ + "func", + "tools", + "[", + '{"name": "get_weather", ', + '"arguments": {"city": ', + '"Tokyo"}}]', + ] + + reconstructor = run_tool_extraction_streaming(phi4mini_parser, deltas) + + assert reconstructor.other_content == "", ( + f"Marker leaked into content: {reconstructor.other_content!r}" + ) + assert len(reconstructor.tool_calls) == 1, ( + f"Expected 1 tool call, got {len(reconstructor.tool_calls)}" + ) + assert reconstructor.tool_calls[0].function.name == "get_weather" + assert json.loads(reconstructor.tool_calls[0].function.arguments) == { + "city": "Tokyo" + } + + +def test_streaming_parameters_key_is_streamed_as_arguments( + phi4mini_parser: ToolParser, +) -> None: + """phi-4-mini templates emit either `arguments` or `parameters`.""" + deltas = [ + 'functools[{"name": "get_time", ', + '"parameters": {"timezone": "Asia/Tokyo"}}]', + ] + + reconstructor = run_tool_extraction_streaming(phi4mini_parser, deltas) + + assert len(reconstructor.tool_calls) == 1, ( + f"Expected 1 tool call, got {len(reconstructor.tool_calls)}" + ) + assert reconstructor.tool_calls[0].function.name == "get_time" + assert json.loads(reconstructor.tool_calls[0].function.arguments) == { + "timezone": "Asia/Tokyo" + } + + +def test_streaming_text_before_marker_is_preserved( + phi4mini_parser: ToolParser, +) -> None: + """Content preceding the marker must still reach the client.""" + deltas = [ + "Let me check that.\n", + 'functools[{"name": "get_weather", "arguments": {"city": "Tokyo"}}]', + ] + + reconstructor = run_tool_extraction_streaming(phi4mini_parser, deltas) + + assert reconstructor.other_content == "Let me check that.\n", ( + f"Expected leading content to be preserved, got {reconstructor.other_content!r}" + ) + assert len(reconstructor.tool_calls) == 1, ( + f"Expected 1 tool call, got {len(reconstructor.tool_calls)}" + ) diff --git a/vllm/tool_parsers/phi4mini_tool_parser.py b/vllm/tool_parsers/phi4mini_tool_parser.py index f2fa3ce9983c..5a4a3960c1a4 100644 --- a/vllm/tool_parsers/phi4mini_tool_parser.py +++ b/vllm/tool_parsers/phi4mini_tool_parser.py @@ -6,6 +6,8 @@ from typing import Any import regex as re +from partial_json_parser.core.exceptions import MalformedJSON +from partial_json_parser.core.options import Allow from transformers import PreTrainedTokenizerBase from vllm.entrypoints.chat_utils import make_tool_call_id @@ -13,7 +15,9 @@ ChatCompletionRequest, ) from vllm.entrypoints.openai.engine.protocol import ( + DeltaFunctionCall, DeltaMessage, + DeltaToolCall, ExtractedToolCallInformation, FunctionCall, ToolCall, @@ -23,10 +27,28 @@ Tool, ToolParser, ) +from vllm.tool_parsers.utils import ( + find_common_prefix, + partial_json_loads, +) logger = init_logger(__name__) +def _partial_bot_token_len(text: str, bot_token: str) -> int: + """Length of the longest suffix of ``text`` that starts ``bot_token``. + + The ``functools`` marker is not a special token, so it is split across + several ordinary tokens while streaming. Any trailing text that could still + grow into the marker is held back instead of being emitted as content. + """ + max_len = min(len(text), len(bot_token) - 1) + for length in range(max_len, 0, -1): + if bot_token.startswith(text[-length:]): + return length + return 0 + + class Phi4MiniJsonToolParser(ToolParser): """ Tool call parser for phi-4-mini models intended for use with the @@ -36,6 +58,8 @@ class Phi4MiniJsonToolParser(ToolParser): are all set """ + json_decoder: json.JSONDecoder = json.JSONDecoder() + def __init__( self, tokenizer: PreTrainedTokenizerBase, @@ -51,6 +75,8 @@ def __init__( str ] = [] # map what has been streamed for each tool so far to a list self.bot_token: str = "functools" + # number of leading characters already streamed back as content + self.streamed_content_len: int = 0 def extract_tool_calls( self, model_output: str, request: ChatCompletionRequest @@ -113,6 +139,98 @@ def extract_tool_calls( tools_called=False, tool_calls=[], content=model_output ) + def _stream_content(self, current_text: str, up_to: int) -> DeltaMessage | None: + """Emit the not-yet-streamed content in ``current_text[:up_to]``.""" + if up_to <= self.streamed_content_len: + return None + content = current_text[self.streamed_content_len : up_to] + self.streamed_content_len = up_to + return DeltaMessage(content=content) + + def _parse_partial_tool_calls( + self, tool_calls_text: str, flags: Allow + ) -> tuple[list[dict[str, Any]], list[bool]]: + """Parse the (possibly incomplete) ``[...]`` array after the marker. + + Returns the tool call objects decoded so far along with, for each of + them, whether its JSON object was already terminated. + """ + tool_call_arr: list[dict[str, Any]] = [] + is_complete: list[bool] = [] + + start_idx = 1 if tool_calls_text.startswith("[") else 0 + while start_idx < len(tool_calls_text): + while ( + start_idx < len(tool_calls_text) + and tool_calls_text[start_idx] in " \t\r\n," + ): + start_idx += 1 + if start_idx >= len(tool_calls_text) or tool_calls_text[start_idx] == "]": + break + + remainder = tool_calls_text[start_idx:] + try: + # A terminated object is decoded directly; this also keeps the + # array's closing bracket away from the partial JSON parser, + # which cannot cope with an unmatched closing bracket. + obj, end_idx = self.json_decoder.raw_decode(remainder) + complete = True + except json.JSONDecodeError: + try: + obj, end_idx = partial_json_loads(remainder, flags) + except (MalformedJSON, ValueError, IndexError): + # Not enough tokens to parse into JSON yet, or the model + # emitted something that is not a tool call at all. + break + complete = False + + if end_idx <= 0: + break + + # phi-4-mini templates emit either "arguments" or "parameters" + if "parameters" in obj and "arguments" not in obj: + obj["arguments"] = obj["parameters"] + + tool_call_arr.append(obj) + is_complete.append(complete) + start_idx += end_idx + + return tool_call_arr, is_complete + + def _argument_diff( + self, tool_call: dict[str, Any], is_complete: bool + ) -> str | None: + """Arguments of ``tool_call`` that have not been streamed yet.""" + cur_arguments = tool_call.get("arguments") + if not cur_arguments: + return None + + cur_args_json = json.dumps(cur_arguments, ensure_ascii=False) + sent = len(self.streamed_args_for_tool[self.current_tool_id]) + + if is_complete: + argument_diff = cur_args_json[sent:] + else: + prev_arguments = None + if self.current_tool_id < len(self.prev_tool_call_arr): + prev_arguments = self.prev_tool_call_arr[self.current_tool_id].get( + "arguments" + ) + if not prev_arguments: + return None + prev_args_json = json.dumps(prev_arguments, ensure_ascii=False) + if cur_args_json == prev_args_json: + return None + # Only stream the part both parses agree on, so that closing + # quotes and braces are not sent prematurely. + argument_diff = find_common_prefix(prev_args_json, cur_args_json)[sent:] + + if not argument_diff: + return None + + self.streamed_args_for_tool[self.current_tool_id] += argument_diff + return argument_diff + def extract_tool_calls_streaming( self, previous_text: str, @@ -123,4 +241,106 @@ def extract_tool_calls_streaming( delta_token_ids: Sequence[int], request: ChatCompletionRequest, ) -> DeltaMessage | None: - return None + if not previous_text: + # a new response is starting, drop any state left from the last one + self.prev_tool_call_arr = [] + self.current_tool_id = -1 + self.current_tool_name_sent = False + self.streamed_args_for_tool = [] + self.streamed_content_len = 0 + + bot_index = current_text.find(self.bot_token) + + # No tool call started yet: stream everything as content, holding back + # any trailing text that could still grow into the marker. + if bot_index == -1: + safe_end = len(current_text) - _partial_bot_token_len( + current_text, self.bot_token + ) + return self._stream_content(current_text, safe_end) + + # Any text before the marker is regular content and must be flushed + # before the tool call deltas are emitted. + content_delta = self._stream_content(current_text, bot_index) + if content_delta is not None: + return content_delta + + # 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 allows sending the entire tool/function name at once. + flags = Allow.ALL if self.current_tool_name_sent else Allow.ALL & ~Allow.STR + tool_calls_text = current_text[bot_index + len(self.bot_token) :] + + try: + tool_call_arr, is_complete = self._parse_partial_tool_calls( + tool_calls_text, flags + ) + + # case: nothing but the opening bracket has been streamed yet + if len(tool_call_arr) == 0: + return None + + # case: we are starting a new tool in the array. Any arguments of + # the previous tool that were completed by the same chunk must be + # flushed before the cursor moves on. + if len(tool_call_arr) > self.current_tool_id + 1: + pending = None + if self.current_tool_id >= 0 and self.current_tool_name_sent: + pending = self._argument_diff( + tool_call_arr[self.current_tool_id], is_complete=True + ) + + self.current_tool_id = len(tool_call_arr) - 1 + self.current_tool_name_sent = False + self.streamed_args_for_tool.append("") + logger.debug("starting on new tool %d", self.current_tool_id) + + if pending is not None: + self.prev_tool_call_arr = tool_call_arr + return pending + + current_tool_call: dict = tool_call_arr[self.current_tool_id] + + function_name = None + if not self.current_tool_name_sent: + function_name = current_tool_call.get("name") + if not function_name: + self.prev_tool_call_arr = tool_call_arr + return None + self.current_tool_name_sent = True + + # The name and the first arguments chunk may be emitted together, + # but only once the object is terminated: while it is still partial + # the arguments were parsed with string values suppressed, so the + # diff would not be a prefix of the final arguments. + argument_diff = None + if function_name is None or is_complete[self.current_tool_id]: + argument_diff = self._argument_diff( + current_tool_call, is_complete[self.current_tool_id] + ) + + self.prev_tool_call_arr = tool_call_arr + + if function_name is None and argument_diff is None: + return None + + return DeltaMessage( + tool_calls=[ + DeltaToolCall( + index=self.current_tool_id, + type="function" if function_name else None, + id=make_tool_call_id() if function_name else None, + function=DeltaFunctionCall( + name=function_name, + arguments=argument_diff, + ).model_dump(exclude_none=True), + ) + ] + ) + + except Exception: + logger.exception("Error trying to handle streaming tool call.") + logger.debug( + "Skipping chunk as a result of tool streaming extraction error" + ) + return None