diff --git a/tests/tool_parsers/test_functiongemma_tool_parser.py b/tests/tool_parsers/test_functiongemma_tool_parser.py index d32aba3085e5..12941ae7dd90 100644 --- a/tests/tool_parsers/test_functiongemma_tool_parser.py +++ b/tests/tool_parsers/test_functiongemma_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 @@ -120,6 +121,61 @@ def test_argument_with_spaces(self, parser): assert result == {"message": "Hello World"} +class TestExtractToolCallsStreaming: + def _stream(self, parser, chunks): + current_text = "" + collected_args = "" + for chunk in chunks: + previous_text = current_text + current_text += chunk + result = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=chunk, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=None, + ) + if result and result.tool_calls: + function = result.tool_calls[0].function + arguments = getattr(function, "arguments", None) + if arguments: + collected_args += arguments + return collected_args + + def test_streamed_arguments_form_valid_json_with_two_keys(self, parser): + # Regression test: each streaming delta previously diffed + # json.dumps(current_args) by raw string length against the last-sent + # string. Since json.dumps() closes the object with '}' on every call, + # the delta after the first key arrived re-sent a stale closing brace + # mid-object, corrupting the client-side reconstructed JSON. + chunks = [ + "call:get_weather{", + "ci", + "ty:NYC", + "unit:celsius", + "}", + ] + reconstructed = self._stream(parser, chunks) + assert json.loads(reconstructed) == {"city": "NYC", "unit": "celsius"} + + def test_streamed_arguments_form_valid_json_with_three_keys(self, parser): + chunks = [ + "call:search{", + "query:python", + "limit:10", + "sort:relevance", + "}", + ] + reconstructed = self._stream(parser, chunks) + assert json.loads(reconstructed) == { + "query": "python", + "limit": 10, + "sort": "relevance", + } + + class TestAdjustRequest: def test_skip_special_tokens_disabled(self, parser, mock_request): mock_request.tools = [{"type": "function", "function": {"name": "test"}}] diff --git a/vllm/tool_parsers/functiongemma_tool_parser.py b/vllm/tool_parsers/functiongemma_tool_parser.py index 776792ea1d6d..3c145e31332b 100644 --- a/vllm/tool_parsers/functiongemma_tool_parser.py +++ b/vllm/tool_parsers/functiongemma_tool_parser.py @@ -22,6 +22,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 find_common_prefix logger = init_logger(__name__) @@ -238,15 +239,23 @@ def extract_tool_calls_streaming( if self.current_tool_name_sent and args_part: current_args = self._parse_arguments(args_part) if current_args: + # Mid-stream, more keys may still arrive, so the + # object isn't closed yet -- drop the trailing '}' + # that json.dumps() always adds for a complete + # dict. The real closing brace is only sent once, + # from the "function call just ended" branch below. current_args_json = json.dumps( current_args, ensure_ascii=False - ) + )[:-1] prev_streamed = self.streamed_args_for_tool[ self.current_tool_id ] - if len(current_args_json) > len(prev_streamed): - diff = current_args_json[len(prev_streamed) :] + if current_args_json != prev_streamed: + common = find_common_prefix( + current_args_json, prev_streamed + ) + diff = current_args_json[len(common) :] self.streamed_args_for_tool[ self.current_tool_id ] = current_args_json @@ -288,8 +297,9 @@ def extract_tool_calls_streaming( prev_streamed = self.streamed_args_for_tool[ self.current_tool_id ] - if len(args_json) > len(prev_streamed): - diff = args_json[len(prev_streamed) :] + if args_json != prev_streamed: + common = find_common_prefix(args_json, prev_streamed) + diff = args_json[len(common) :] self.streamed_args_for_tool[self.current_tool_id] = ( args_json )