diff --git a/tests/parser/engine/test_gemma4_streaming_reasoning.py b/tests/parser/engine/test_gemma4_streaming_reasoning.py index ae21d69fc3..c6158146fc 100644 --- a/tests/parser/engine/test_gemma4_streaming_reasoning.py +++ b/tests/parser/engine/test_gemma4_streaming_reasoning.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for the unified Gemma4 parser engine.""" +import dataclasses import json from unittest.mock import MagicMock @@ -17,8 +18,14 @@ from tests.parser.engine.streaming_helpers import ( from vllm.entrypoints.generate.base.protocol import DeltaMessage from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, + ChatCompletionToolsParam, +) +from vllm.parser.gemma4 import ( + STRING_DELIM, + Gemma4Parser, + _gemma4_arg_value_extender, + gemma4_config, ) -from vllm.parser.gemma4 import Gemma4Parser # ── Special token IDs (arbitrary but consistent) ───────────────────── CHANNEL_START_ID = 50 # <|channel> @@ -1649,3 +1656,219 @@ class TestGemma4IsReasoningEnd: ) def test_is_reasoning_end(self, ids, thinking, ended): assert self._parser(thinking).is_reasoning_end(ids) is ended + + +# ── Argument value extender ────────────────────────────────────────── + +_D = STRING_DELIM + +_EXTENDER_TOOL = ChatCompletionToolsParam( + type="function", + function={ + "name": "f", + "parameters": { + "type": "object", + "properties": { + "a": {"type": "string"}, + "b": {"type": "string"}, + "location": {"type": "string"}, + "n": {"type": "integer"}, + }, + }, + }, +) + +_EXTENDER_BODIES = [ + f"location:{_D}Paris, France{_D}", + f"a:{_D}x{_D},b:{_D}hello }} world{_D}", + f"a:{_D}smile :-) ok{_D},n:42", + f"a:{_D}trailing spaces {_D},n:42", + f'a:{_D}decoys <|x <|"y <|"|z{_D}', + f'a:{_D}quote " backslash \\ newline\nunicode 杭州{_D}', + f"a:{_D}x{_D},{_D}quoted key{_D}:{_D}val{_D}", + f"a:{_D}x{_D},nested:{{inner:{_D}deep value{_D},n:1}}", + f"items:[{_D}a{_D},{_D}b, c{_D}],a:{_D}last{_D}", + f"a:{_D}call:x{_D},b:{_D}brace {{ colon:{_D}", + f"a:{_D}{_D},b:{_D}}}}}{_D}", +] + +_TRUNCATED_BODIES = [ + f"a:{_D}hello }}", + f"a:{_D}hello )", + f"a:{_D}hello ", + f'a:{_D}hello<|"', + f"a:{_D}x{_D},b:{_D}open", + f"a:{_D}", +] + + +_TOOL_MODES = pytest.mark.parametrize( + "tools", [[], [_EXTENDER_TOOL]], ids=["no_schema", "schema"] +) + + +def _pieces(text: str, size: int) -> list[str]: + return [text[i : i + size] for i in range(0, len(text), size)] + + +def _stream_args(parser, request, chunks: list[str], finished: bool = True) -> str: + args = [] + for i, chunk in enumerate(chunks): + last = finished and i == len(chunks) - 1 + delta = parser.parse_delta(chunk, [], request, finished=last) + for tool_call in (delta and delta.tool_calls) or []: + if tool_call.function and tool_call.function.arguments: + args.append(tool_call.function.arguments) + return "".join(args) + + +class TestGemma4ArgValueExtender: + """The extender streams exactly what ``_gemma4_arg_converter`` would, so + the converter-only path, which only differs in when it emits, must agree + with it byte for byte on the accumulated arguments.""" + + @staticmethod + def _parser(tokenizer, tools, extender: bool): + config = gemma4_config() + if not extender: + config = dataclasses.replace(config, arg_value_extender=None) + return Gemma4Parser(tokenizer, tools, parser_engine_config=config) + + @_TOOL_MODES + @pytest.mark.parametrize("size", [1, 3, 7, 64]) + @pytest.mark.parametrize("body", _EXTENDER_BODIES) + def test_matches_converter_at_every_chunking( + self, tool_call_tokenizer, mock_request, body, size, tools + ): + mock_request.tools = tools + text = f"<|tool_call>call:f{{{body}}}" + chunks = _pieces(text, size) + + reference = self._parser(tool_call_tokenizer, tools, extender=True) + expected = reference.extract_tool_calls(text, mock_request) + expected = expected.tool_calls[0].function.arguments + + with_extender = _stream_args( + self._parser(tool_call_tokenizer, tools, extender=True), + mock_request, + chunks, + ) + converter_only = _stream_args( + self._parser(tool_call_tokenizer, tools, extender=False), + mock_request, + chunks, + ) + + assert with_extender == converter_only == expected + + @_TOOL_MODES + @pytest.mark.parametrize("size", [1, 3, 7]) + @pytest.mark.parametrize("body", _TRUNCATED_BODIES) + def test_truncated_stream_matches_converter( + self, tool_call_tokenizer, mock_request, body, size, tools + ): + """A stream cut inside the value ends with the converter's final + flush, so the extender must hold back the whitespace, closer or + partial delimiter that flush would drop.""" + mock_request.tools = tools + chunks = _pieces(f"<|tool_call>call:f{{{body}", size) + + with_extender = _stream_args( + self._parser(tool_call_tokenizer, tools, extender=True), + mock_request, + chunks, + ) + converter_only = _stream_args( + self._parser(tool_call_tokenizer, tools, extender=False), + mock_request, + chunks, + ) + + assert with_extender == converter_only + json.loads(with_extender) + + @pytest.mark.parametrize("offset", range(1, len(STRING_DELIM))) + def test_split_closing_delimiter(self, tool_call_parser, mock_request, offset): + """``STRING_DELIM`` is not a lexer terminal, so the extender may see + the close split across deltas and must window across them.""" + head = f"<|tool_call>call:f{{a:{_D}Paris" + tail = ",n:1}" + chunks = [head, _D[:offset], _D[offset:] + tail] + + streamed = _stream_args(tool_call_parser, mock_request, chunks) + + assert streamed == '{"a": "Paris", "n": "1"}' + + def test_string_value_streams_before_it_closes( + self, tool_call_tokenizer, mock_request + ): + """Body deltas carry no structural char, so without the extender the + gate holds the whole value until the close.""" + head = [f"<|tool_call>call:f{{location:{_D}", "Paris"] + + def streamed_before_close(extender: bool) -> str: + parser = self._parser(tool_call_tokenizer, [], extender) + return _stream_args(parser, mock_request, head, finished=False) + + assert streamed_before_close(extender=True) == '{"location": "Paris' + assert streamed_before_close(extender=False) == '{"location": "' + + def test_conversions_do_not_grow_with_value_length( + self, tool_call_tokenizer, mock_request + ): + def streaming_conversions(words: int) -> int: + calls = 0 + base = gemma4_config() + + def counting_converter(raw_args: str, partial: bool) -> str: + nonlocal calls + if partial: + calls += 1 + return base.arg_converter(raw_args, partial) + + parser = Gemma4Parser( + tool_call_tokenizer, + parser_engine_config=dataclasses.replace( + base, arg_converter=counting_converter + ), + ) + value = "word, " * words + text = f"<|tool_call>call:f{{a:{_D}{value}{_D},n:1}}" + _stream_args(parser, mock_request, _pieces(text, 4)) + return calls + + assert streaming_conversions(words=40) == streaming_conversions(words=400) + + @pytest.mark.parametrize( + "raw_args", + [ + f"a:{_D}closed{_D}", + f"a:{_D}x{_D},{_D}open ke", + f"nested:{{inner:{_D}deep", + f"a:{_D}x{_D},b:{_D}y{_D},a:{_D}dup", + "a:tr", + 'a:<|"', + ], + ids=["closed", "open_key", "nested", "duplicate_key", "bare", "opener"], + ) + def test_factory_declines_without_open_last_string_value(self, raw_args): + assert _gemma4_arg_value_extender(raw_args) is None + + def test_extender_holds_back_what_the_converter_hides(self): + extend = _gemma4_arg_value_extender(f"a:{_D}hello") + assert extend is not None + assert [extend(d) for d in [" }", "x", ")", "(", '<|"', "y", '<|"', "|>"]] == [ + " ", + "}x", + "", + ")(", + "", + '<|\\"y', + "", + None, + ] + + def test_factory_holds_back_the_tail_already_hidden(self): + extend = _gemma4_arg_value_extender(f"a:{_D}hello }} ") + assert extend is not None + assert extend("x") == "} x" diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index 25dfd546b0..811f9cb7b8 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -1685,6 +1685,57 @@ class TestArgDeltaWithConverter: ] +# ── TestArgValueExtender ────────────────────────────────────────── + + +def _extender_config(structural: frozenset[str] | None = None) -> ParserEngineConfig: + """``_kv_converter`` plus an extender for its trailing ``k=v`` record, + which is open until whitespace follows it; mirrors + ``_glm47_arg_value_extender``.""" + + def factory(raw_args: str): + if re.search(r"\w+=\S+$", raw_args) is None: + return None + + def extend(delta: str) -> str | None: + if re.search(r"\s", delta): + return None + return json.dumps(delta, ensure_ascii=False)[1:-1] + + return extend + + return dataclasses.replace( + _converter_config(), + arg_value_extender=factory, + arg_structural_chars=structural, + ) + + +class TestArgValueExtender: + """Engine contract for ``ParserEngineConfig.arg_value_extender``.""" + + def test_live_extender_sees_deltas_without_structural_chars(self): + """A delta without structural chars must still reach the live + extender: gating it out drops ``el`` from the streamed value and the + next accepted delta ``l=`` lands on the gap, so the converted JSON + no longer extends what was streamed.""" + engine = _make_engine(_extender_config(structural=frozenset("= "))) + assert _run_streaming_tool(engine, "f", ["a", "=h", "el", "l=", "o "]) == { + "a": "hell=o" + } + + @pytest.mark.parametrize( + "overrides", + [ + {"arg_converter": None}, + {"stream_arg_deltas": False}, + ], + ) + def test_extender_requires_streaming_conversion(self, overrides): + with pytest.raises(ValueError, match="arg_value_extender"): + dataclasses.replace(_extender_config(), **overrides) + + # ── TestSafeArgPrefix ──────────────────────────────────────────── diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index eaf9e4519e..95df03eab5 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -342,8 +342,10 @@ class ParserEngine(Parser): their unterminated content instead of buffering long arguments until the closing tag arrives. - Also returns whether the prefix ends inside that streamed trailing - string value. + Also returns whether the returned prefix stops inside that trailing + string, so streamed text may be appended to it. This holds even when + *json_str* closes the string: converters serialise an open raw value + as a closed JSON string. """ last_colon = -1 last_key: str | None = None @@ -1053,11 +1055,8 @@ class ParserEngine(Parser): if not self._stream_arg_deltas: return None - structural = self._arg_structural_chars - if structural is not None and structural.isdisjoint(raw_delta): - return None - slot = self._tool_slots[idx] + structural = self._arg_structural_chars if slot.extend_open_value is not None: diff = slot.extend_open_value(raw_delta) if diff is not None: @@ -1066,6 +1065,8 @@ class ParserEngine(Parser): slot.append_streamed(diff) return diff slot.extend_open_value = None + elif structural is not None and structural.isdisjoint(raw_delta): + return None try: current_json = converter(slot.args, True) diff --git a/vllm/parser/engine/parser_engine_config.py b/vllm/parser/engine/parser_engine_config.py index 2bb2ae267e..680023924c 100644 --- a/vllm/parser/engine/parser_engine_config.py +++ b/vllm/parser/engine/parser_engine_config.py @@ -94,9 +94,17 @@ class ParserEngineConfig: arg_structural_chars: frozenset[str] | None = None - # Called with the raw argument text when its converted JSON ends in an - # open value; returns an extender for that value, or None. Streaming - # skips ``arg_converter`` for every delta the extender accepts. + # Called with the raw argument text after each streaming conversion whose + # prefix-safe JSON extends the JSON streamed so far; returns an extender + # when the text ends in an open value that is the last value in the + # converted JSON, else None. While live, the extender receives every delta + # in place of ``arg_converter`` until it returns None, meaning the delta + # closed the value. Returned text is streamed as is, so it must match what + # ``arg_converter`` would serialise for those raw characters. Deltas are + # lexer output: a terminal is never split across deltas, but the lexer + # flushes a held partial terminal as text at the end of the stream or + # before a terminal matched by token ID, so a truncated value may end in + # a terminal prefix. Requires ``arg_converter`` and ``stream_arg_deltas``. arg_value_extender: Callable[[str], ArgValueExtender | None] | None = None # Special tokens exempt from auto-drop but not state-machine terminals. @@ -118,6 +126,15 @@ class ParserEngineConfig: # Reject tool calls whose names are absent from the request tools. validate_tool_names: bool = False + def __post_init__(self) -> None: + if self.arg_value_extender is not None and ( + self.arg_converter is None or not self.stream_arg_deltas + ): + raise ValueError( + f"{self.name}: arg_value_extender requires arg_converter " + "and stream_arg_deltas=True" + ) + def terminal_literal(self, name: str) -> str | None: """Canonical spelling of terminal *name*, or ``None`` if undeclared.""" value = self.terminals.get(name) diff --git a/vllm/parser/gemma4.py b/vllm/parser/gemma4.py index 2a6246041d..738ba58295 100644 --- a/vllm/parser/gemma4.py +++ b/vllm/parser/gemma4.py @@ -22,6 +22,7 @@ from vllm.logger import init_logger from vllm.parser.engine.events import EventType, SemanticEvent from vllm.parser.engine.parser_engine import ParserEngine from vllm.parser.engine.parser_engine_config import ( + ArgValueExtender, ParserEngineConfig, ParserState, Transition, @@ -66,6 +67,12 @@ def _strip_partial_delim(value: str) -> str: def _parse_gemma4_args(args_str: str, *, partial: bool = False) -> dict: + return _parse_gemma4_fields(args_str, partial=partial)[0] + + +def _parse_gemma4_fields( + args_str: str, *, partial: bool = False +) -> tuple[dict, str | None]: """Parse Gemma4's custom key:value format into a Python dict. Format examples:: @@ -82,13 +89,15 @@ def _parse_gemma4_args(args_str: str, *, partial: bool = False) -> dict: omitted because they may be incomplete and type-unstable (e.g. partial boolean parsed as bare string). - Returns a dict ready for ``json.dumps()``. + Returns a dict ready for ``json.dumps()`` and the key of the trailing + top-level ``STRING_DELIM`` value left unterminated, if any. """ if not args_str or not args_str.strip(): - return {} + return {}, None result: dict = {} + open_key: str | None = None i = 0 n = len(args_str) @@ -130,6 +139,7 @@ def _parse_gemma4_args(args_str: str, *, partial: bool = False) -> dict: if partial: value = _strip_partial_delim(value) result[key] = value + open_key = key break result[key] = args_str[val_start:end_pos] i = end_pos + _DELIM_LEN @@ -199,7 +209,7 @@ def _parse_gemma4_args(args_str: str, *, partial: bool = False) -> dict: break result[key] = raw_val - return result + return result, open_key def _parse_gemma4_array(arr_str: str, *, partial: bool = False) -> list: @@ -283,16 +293,64 @@ def _parse_gemma4_array(arr_str: str, *, partial: bool = False) -> list: return items +def _drop_call_closer(text: str, paren_excess: bool) -> str: + """Drop a trailing ``}`` or, with more ``)`` than ``(``, a trailing ``)``. + + The closer of the call body is dropped before parsing, so a value that + itself ends in one loses that character until more text arrives. + """ + if text.endswith("}") or (text.endswith(")") and paren_excess): + return text[:-1] + return text + + def _gemma4_arg_converter(raw_args: str, partial: bool) -> str: """Convert Gemma4 custom arg format to a JSON string.""" text = raw_args.strip() - if text.endswith("}") or text.endswith(")") and text.count("(") < text.count(")"): - text = text[:-1] - + text = _drop_call_closer(text, text.count("(") < text.count(")")) parsed = _parse_gemma4_args(text, partial=partial) return json.dumps(parsed, ensure_ascii=False) +def _gemma4_arg_value_extender(raw_args: str) -> ArgValueExtender | None: + """Extend the open top-level ``STRING_DELIM`` value at the end of + *raw_args*, if it is the last value in the converted JSON. + + The extender streams exactly what ``_gemma4_arg_converter`` would, so it + holds back the same tail: trailing whitespace, a dropped closer and a + partial ``STRING_DELIM``. ``STRING_DELIM`` is a special token but not a + lexer terminal, so unlike a tag it may reach the extender split across + deltas; the held tail doubles as the window that detects the close. + """ + text = raw_args.strip() + opens, closes = text.count("("), text.count(")") + body = _drop_call_closer(text, opens < closes) + params, open_key = _parse_gemma4_fields(body, partial=True) + if open_key is None or next(reversed(params)) != open_key: + return None + + held_len = ( + len(raw_args) + - len(raw_args.rstrip()) + + len(text) + - len(_strip_partial_delim(body)) + ) + held = raw_args[len(raw_args) - held_len :] if held_len else "" + + def extend(delta: str) -> str | None: + nonlocal held, opens, closes + text = held + delta + if STRING_DELIM in text: + return None + opens += delta.count("(") + closes += delta.count(")") + visible = _strip_partial_delim(_drop_call_closer(text.rstrip(), opens < closes)) + held = text[len(visible) :] + return json.dumps(visible, ensure_ascii=False)[1:-1] + + return extend + + @functools.cache def gemma4_config(thinking: bool = False) -> ParserEngineConfig: return ParserEngineConfig( @@ -396,6 +454,7 @@ def gemma4_config(thinking: bool = False) -> ParserEngineConfig: ParserState.TOOL_ARGS: EventType.ARG_VALUE_CHUNK, }, arg_converter=_gemma4_arg_converter, + arg_value_extender=_gemma4_arg_value_extender, tool_args_json=False, arg_structural_chars=frozenset(",:{}[]<"), preserve_tokens=frozenset({STRING_DELIM}), @@ -423,12 +482,8 @@ class Gemma4Parser(ParserEngine): ) -> None: chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} thinking = chat_kwargs.get("enable_thinking", False) - super().__init__( - tokenizer, - tools, - parser_engine_config=gemma4_config(thinking=thinking), - **kwargs, - ) + kwargs.setdefault("parser_engine_config", gemma4_config(thinking=thinking)) + super().__init__(tokenizer, tools, **kwargs) vocab = self.vocab self._tool_call_token_id: int | None = vocab.get("<|tool_call>") self._new_turn_token_id: int | None = vocab.get("<|turn>")