From 44f0ddfd7b61e200aac323231dfc127d4cd53cfb Mon Sep 17 00:00:00 2001 From: Maroon Ayoub Date: Tue, 13 Oct 2026 00:28:35 +0300 Subject: [PATCH] [Bugfix][Parser] Stop re-parsing the whole GLM tool argument per chunk (#61253) Signed-off-by: Maroon Ayoub Co-authored-by: Flora Feng <4florafeng@gmail.com> --- tests/parser/engine/test_glm47_moe.py | 124 ++++++++++++++++++++- tests/parser/engine/test_parser_engine.py | 2 +- vllm/parser/engine/parser_engine.py | 81 ++++++++++---- vllm/parser/engine/parser_engine_config.py | 9 ++ vllm/parser/glm47_moe.py | 36 +++++- 5 files changed, 226 insertions(+), 26 deletions(-) diff --git a/tests/parser/engine/test_glm47_moe.py b/tests/parser/engine/test_glm47_moe.py index 58fe7aa28c..9f9d65ef20 100644 --- a/tests/parser/engine/test_glm47_moe.py +++ b/tests/parser/engine/test_glm47_moe.py @@ -5,18 +5,27 @@ GLM can transition straight from reasoning into a tool call without ever emitting ````, so ```` is an implicit reasoning terminator. These cover that path plus the multi-turn prompts where an -earlier turn's markers must not be read as the current turn's state. +earlier turn's markers must not be read as the current turn's state. The +streaming tests check that long tool arguments stream the same JSON as a +non-streaming parse, without converting the whole argument on every delta. """ +import dataclasses + import pytest from tests.parser.engine.conftest import make_mock_tokenizer +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionToolsParam, +) from vllm.parser.glm47_moe import ( THINK_END, THINK_START, TOOL_CALL_END, TOOL_CALL_START, Glm47MoeParser, + _glm47_arg_value_extender, + glm47_moe_config, ) THINK_S, THINK_E = 1000, 1001 @@ -106,3 +115,116 @@ class TestExtractContentIds: def test_thinking_disabled(self, no_thinking_parser): ids = [THINK_S, TEXT, TOOL_S] assert no_thinking_parser.extract_content_ids(ids) == ids + + +WRITE_TOOL = ChatCompletionToolsParam( + type="function", + function={ + "name": "Write", + "parameters": { + "type": "object", + "properties": { + "file_path": {"type": "string"}, + "line_count": {"type": "integer"}, + "content": {"type": "string"}, + }, + }, + }, +) +CALL_HEAD = ( + f"{THINK_END}{TOOL_CALL_START}Write\n" + "file_path/repo/app.py\n" + "line_count120\n" + "content" +) +CALL_TAIL = f"\n{TOOL_CALL_END}" +# Escapes, non-ASCII, and tag-like text that must stay inside the value. +CONTENT_LINE = 'if a < b > c: print("x\\y") # \t\u00e9\n' + + +def _pieces(text: str, size: int) -> list[str]: + return [text[i : i + size] for i in range(0, len(text), size)] + + +def _stream_arguments(parser, request, chunks: list[str]) -> str: + args = [] + for i, chunk in enumerate(chunks): + delta = parser.parse_delta(chunk, [], request, finished=i == len(chunks) - 1) + 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 TestStreamedToolArguments: + @pytest.mark.parametrize( + "tools, line_count", + [([], '"120"'), ([WRITE_TOOL], "120")], + ids=["no_schema", "schema"], + ) + def test_long_call_streams_non_streaming_arguments( + self, mock_request, tools, line_count + ): + mock_request.tools = tools + text = CALL_HEAD + CONTENT_LINE * 400 + CALL_TAIL + + streamed = _stream_arguments( + Glm47MoeParser(make_mock_tokenizer(VOCAB)), mock_request, _pieces(text, 7) + ) + result = Glm47MoeParser(make_mock_tokenizer(VOCAB)).extract_tool_calls( + text, mock_request + ) + + assert streamed == result.tool_calls[0].function.arguments + assert f'"line_count": {line_count},' in streamed + + def test_conversions_do_not_grow_with_argument_length(self, mock_request): + """The open value streams without re-running the converter.""" + mock_request.tools = [WRITE_TOOL] + + def streaming_conversions(lines: int) -> int: + calls = 0 + base = glm47_moe_config() + + def counting_converter(raw_args: str, partial: bool) -> str: + nonlocal calls + if partial: + calls += 1 + return base.arg_converter(raw_args, partial) + + parser = Glm47MoeParser( + make_mock_tokenizer(VOCAB), + parser_engine_config=dataclasses.replace( + base, arg_converter=counting_converter + ), + ) + chunks = [ + *_pieces(CALL_HEAD, 4), + *_pieces(CONTENT_LINE * lines, 4), + *_pieces(CALL_TAIL, 4), + ] + _stream_arguments(parser, mock_request, chunks) + return calls + + assert streaming_conversions(lines=40) == streaming_conversions(lines=400) + + def test_value_extender_stops_at_split_value_end(self): + extend = _glm47_arg_value_extender("kx") + assert [extend(d) for d in ['a"\n', ""]] == [ + 'a\\"\\n', + "kx", + "k1" + "j2" + "k3", + ], + ids=["closed_value", "repeated_key_not_last"], + ) + def test_no_value_extender_unless_open_value_is_last(self, raw_args): + assert _glm47_arg_value_extender(raw_args) is None diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index dc91cef740..25dfd546b0 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -1711,7 +1711,7 @@ class TestSafeArgPrefix: ], ) def test_safe_arg_prefix(self, json_str, expected): - assert ParserEngine._safe_arg_prefix(json_str) == expected + assert ParserEngine._safe_arg_prefix(json_str)[0] == expected # ── Coercion instability regression tests ──────────────────────── diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index c74be931d2..eaf9e4519e 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -24,7 +24,11 @@ from vllm.entrypoints.generate.base.protocol import ( from vllm.logger import init_logger from vllm.parser.abstract_parser import Parser, StreamState from vllm.parser.engine.events import EventType, SemanticEvent -from vllm.parser.engine.parser_engine_config import ParserEngineConfig, ParserState +from vllm.parser.engine.parser_engine_config import ( + ArgValueExtender, + ParserEngineConfig, + ParserState, +) from vllm.parser.engine.streaming_parser_engine import StreamingParserEngine from vllm.renderers.chat_utils import get_tool_call_id_type, make_tool_call_id from vllm.tool_parsers.utils import ( @@ -60,7 +64,10 @@ class ToolCallSlot: "_args_joined", "name_sent", "string_keys", - "streamed_json", + "_streamed_parts", + "_streamed_joined", + "extend_open_value", + "open_value_streams", ) def __init__(self) -> None: @@ -70,7 +77,10 @@ class ToolCallSlot: self._args_joined: str | None = "" self.name_sent: bool = False self.string_keys: set[str] | None = None - self.streamed_json: str = "" + self._streamed_parts: list[str] = [] + self._streamed_joined: str | None = "" + self.extend_open_value: ArgValueExtender | None = None + self.open_value_streams: bool = False @property def args(self) -> str: @@ -82,6 +92,21 @@ class ToolCallSlot: self._args_parts.append(value) self._args_joined = None + @property + def streamed_json(self) -> str: + if self._streamed_joined is None: + self._streamed_joined = "".join(self._streamed_parts) + return self._streamed_joined + + @streamed_json.setter + def streamed_json(self, value: str) -> None: + self._streamed_parts = [value] + self._streamed_joined = value + + def append_streamed(self, value: str) -> None: + self._streamed_parts.append(value) + self._streamed_joined = None + class ParserEngine(Parser): """A :class:`Parser` backed by a single declarative engine config. @@ -133,6 +158,7 @@ class ParserEngine(Parser): self._arg_converter = parser_engine_config.arg_converter self._arg_structural_chars = parser_engine_config.arg_structural_chars + self._arg_value_extender = parser_engine_config.arg_value_extender self._stream_arg_deltas = parser_engine_config.stream_arg_deltas self._strip_trailing_reasoning_ws = ( parser_engine_config.strip_trailing_reasoning_whitespace @@ -303,7 +329,9 @@ class ParserEngine(Parser): return args, changed @staticmethod - def _safe_arg_prefix(json_str: str, string_keys: set[str] | None = None) -> str: + def _safe_arg_prefix( + json_str: str, string_keys: set[str] | None = None + ) -> tuple[str, bool]: """Return the prefix of *json_str* up to the last top-level value. Middle values (followed by a comma) are stable across streaming @@ -313,6 +341,9 @@ class ParserEngine(Parser): String values for keys in ``string_keys`` are prefix-stable, so stream 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. """ last_colon = -1 last_key: str | None = None @@ -345,14 +376,14 @@ class ParserEngine(Parser): last_key = pending_key pending_key = None if last_colon < 0: - return "" + return "", False end = last_colon + 1 while end < len(json_str) and json_str[end] in (" ", "\t", "\n", "\r"): end += 1 if end >= len(json_str) or json_str[end] != '"': - return json_str[:end] + return json_str[:end], False if string_keys is not None and last_key not in string_keys: - return json_str[:end] + return json_str[:end], False escape = False for i in range(end + 1, len(json_str)): @@ -364,8 +395,8 @@ class ParserEngine(Parser): escape = True continue if c == '"': - return json_str[:i] - return json_str + return json_str[:i], True + return json_str, True @staticmethod def _streamable_string_keys(properties: dict) -> set[str] | None: @@ -1027,6 +1058,15 @@ class ParserEngine(Parser): return None slot = self._tool_slots[idx] + if slot.extend_open_value is not None: + diff = slot.extend_open_value(raw_delta) + if diff is not None: + if not diff or not slot.open_value_streams: + return None + slot.append_streamed(diff) + return diff + slot.extend_open_value = None + try: current_json = converter(slot.args, True) except (json.JSONDecodeError, ValueError, TypeError): @@ -1040,22 +1080,19 @@ class ParserEngine(Parser): current_json = self._fix_arg_types(current_json, slot.name) prev = slot.streamed_json - safe_json = self._safe_arg_prefix(current_json, slot.string_keys) + safe_json, value_streams = self._safe_arg_prefix(current_json, slot.string_keys) - if not safe_json or safe_json == prev: + if not safe_json: + return None + if prev and not safe_json.startswith(prev): return None - if prev: - if not safe_json.startswith(prev): - return None - diff = safe_json[len(prev) :] - else: - diff = safe_json - - if diff: - slot.streamed_json = safe_json - return diff - return None + slot.streamed_json = safe_json + if self._arg_value_extender is not None: + # Deltas that only extend the open value skip the conversion above. + slot.extend_open_value = self._arg_value_extender(slot.args) + slot.open_value_streams = value_streams + return safe_json[len(prev) :] or None def _flush_arg_converter(self, idx: int) -> str | None: converter = self._arg_converter diff --git a/vllm/parser/engine/parser_engine_config.py b/vllm/parser/engine/parser_engine_config.py index 215fa0d9bd..2bb2ae267e 100644 --- a/vllm/parser/engine/parser_engine_config.py +++ b/vllm/parser/engine/parser_engine_config.py @@ -24,6 +24,10 @@ from functools import cached_property from vllm.parser.engine.events import EventType +# Maps a raw argument delta to the JSON text it appends to the open trailing +# value, or returns None when the delta may change anything else. +ArgValueExtender = Callable[[str], str | None] + class ParserState(Enum): CONTENT = auto() @@ -90,6 +94,11 @@ 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. + arg_value_extender: Callable[[str], ArgValueExtender | None] | None = None + # Special tokens exempt from auto-drop but not state-machine terminals. preserve_tokens: frozenset[str] = field(default_factory=frozenset) diff --git a/vllm/parser/glm47_moe.py b/vllm/parser/glm47_moe.py index 9a18228fb0..44377e62b7 100644 --- a/vllm/parser/glm47_moe.py +++ b/vllm/parser/glm47_moe.py @@ -24,6 +24,7 @@ from vllm.logger import init_logger from vllm.parser.engine.events import EventType from vllm.parser.engine.parser_engine import ParserEngine from vllm.parser.engine.parser_engine_config import ( + ArgValueExtender, ParserEngineConfig, ParserState, Transition, @@ -63,12 +64,16 @@ _PARTIAL_ARG_RE = re.compile( ) -def _glm47_arg_converter(raw_args: str, partial: bool) -> str: +def _glm47_parse_args( + raw_args: str, partial: bool +) -> tuple[dict[str, object], str | None]: + """Return the arguments and the key of the open trailing value, if any.""" params: dict[str, object] = {} for match in _ARG_RE.finditer(raw_args): params[match.group("key").strip()] = match.group("value") + open_key = None if partial: remaining = _ARG_RE.sub("", raw_args) match = _PARTIAL_ARG_RE.search(remaining) @@ -76,8 +81,34 @@ def _glm47_arg_converter(raw_args: str, partial: bool) -> str: key = match.group("key").strip() if key: params[key] = match.group("value") + open_key = key - return json.dumps(params, ensure_ascii=False) + return params, open_key + + +def _glm47_arg_converter(raw_args: str, partial: bool) -> str: + return json.dumps(_glm47_parse_args(raw_args, partial)[0], ensure_ascii=False) + + +def _glm47_arg_value_extender(raw_args: str) -> ArgValueExtender | None: + # The open value runs to the end of raw_args and is the last JSON value, + # so a delta appends its JSON-escaped text unless it closes the value. + params, open_key = _glm47_parse_args(raw_args, partial=True) + if open_key is None or next(reversed(params)) != open_key: + return None + + keep = len(ARG_VALUE_END) - 1 + tail = raw_args[-keep:] + + def extend(delta: str) -> str | None: + nonlocal tail + text = tail + delta + if ARG_VALUE_END in text: + return None + tail = text[-keep:] + return json.dumps(delta, ensure_ascii=False)[1:-1] + + return extend def _glm53_always_thinks(tokenizer: TokenizerLike) -> bool: @@ -194,6 +225,7 @@ def glm47_moe_config(thinking: bool = True) -> ParserEngineConfig: }, turn_boundary_tokens=GLM_TURN_BOUNDARIES, arg_converter=_glm47_arg_converter, + arg_value_extender=_glm47_arg_value_extender, stream_arg_deltas=True, tool_args_json=False, validate_tool_names=True,