@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 15ed1262e7

[Bugfix][Tool Parser] Migrate Granite to the streaming Parser Engine (#49648)

Signed-off-by: Nikhil Kulkarni <nikhilkulkarni1755@gmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Chauncey <chaunceyjiang@gmail.com>

Nikhil Kulkarni committed Sep 23, 2026 at 03:37 UTC 15ed1262e70873a65582d82f7973784a0f11e5a2
14 files changed +631 -369
tests/parser/engine/test_delegating_replay.py
+4
@@ -99,6 +99,10 @@ def _discover_pairings() -> list[_PairingInfo]:
99 # terminal matching; combined-parser replay coverage lives in
100 # test_inkling.py.
101 continue
102 + if cfg.name == "granite":
103 + # Granite has a JSON-array tool body with no TOOL_END terminal;
104 + # its replay coverage lives in test_granite.py.
105 + continue
106
107 parser_cls = type(
108 f"_Delegating{engine_cls.__name__}",
tests/parser/engine/test_granite.py new
+234
@@ -0,0 +1,234 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the engine-based Granite parser.
4 +
5 +Granite emits tool calls as a JSON array following a marker
6 +(``<|tool_call|>`` for 3.0, ``<tool_call>`` for 3.1) with no closing
7 +terminal; the ``]`` ends the region and surrounding prose is content. The
8 +engine's ``tool_call_body_array`` mode splits the array into one call per
9 +element and ``_granite_arg_converter`` carves each element's ``arguments``.
10 +"""
11 +
12 +import json
13 +
14 +import pytest
15 +
16 +from tests.parser.engine.conftest import make_mock_tokenizer
17 +from tests.parser.engine.trace_builder import _GRANITE_SCENARIOS, _build_granite
18 +from vllm.parser.granite import GraniteParser, _granite_arg_converter
19 +
20 +TOOL_TOKEN = "<|tool_call|>"
21 +TOOL_STRING = "<tool_call>"
22 +
23 +_GRANITE_VOCAB = {TOOL_TOKEN: 49154}
24 +
25 +
26 +@pytest.fixture
27 +def mock_tokenizer():
28 + return make_mock_tokenizer(_GRANITE_VOCAB)
29 +
30 +
31 +@pytest.fixture
32 +def parser(mock_tokenizer):
33 + return GraniteParser(mock_tokenizer)
34 +
35 +
36 +def _tokenize(text: str) -> list[tuple[int, str]]:
37 + """Tokenize like the real stream: ``<|tool_call|>`` is one special token,
38 + plain text becomes one token per character."""
39 + tokens: list[tuple[int, str]] = []
40 + i = 0
41 + while i < len(text):
42 + if text.startswith(TOOL_TOKEN, i):
43 + tokens.append((_GRANITE_VOCAB[TOOL_TOKEN], TOOL_TOKEN))
44 + i += len(TOOL_TOKEN)
45 + else:
46 + tokens.append((ord(text[i]), text[i]))
47 + i += 1
48 + return tokens
49 +
50 +
51 +def _stream(parser, request, text: str, chunk_size: int):
52 + tokens = _tokenize(text)
53 + results = []
54 + previous_text = ""
55 + previous_token_ids: list[int] = []
56 + for start in range(0, len(tokens), chunk_size):
57 + batch = tokens[start : start + chunk_size]
58 + delta_text = "".join(t for _, t in batch)
59 + delta_token_ids = [tid for tid, _ in batch]
60 + current_text = previous_text + delta_text
61 + current_token_ids = previous_token_ids + delta_token_ids
62 + delta = parser.extract_tool_calls_streaming(
63 + previous_text=previous_text,
64 + current_text=current_text,
65 + delta_text=delta_text,
66 + previous_token_ids=tuple(previous_token_ids),
67 + current_token_ids=tuple(current_token_ids),
68 + delta_token_ids=tuple(delta_token_ids),
69 + request=request,
70 + )
71 + results.append(delta)
72 + previous_text = current_text
73 + previous_token_ids = current_token_ids
74 + finish = parser.finish_streaming()
75 + if finish is not None:
76 + results.append(finish)
77 + return results
78 +
79 +
80 +def _collect_content(results) -> str:
81 + return "".join(d.content for d in results if d and d.content)
82 +
83 +
84 +def _collect_names(results) -> list[str]:
85 + names: dict[int, str] = {}
86 + for d in results:
87 + if not (d and d.tool_calls):
88 + continue
89 + for tc in d.tool_calls:
90 + if tc.function and tc.function.name:
91 + names[tc.index] = names.get(tc.index, "") + tc.function.name
92 + return [names[i] for i in sorted(names)]
93 +
94 +
95 +def _collect_args(results) -> dict[int, str]:
96 + args: dict[int, str] = {}
97 + for d in results:
98 + if not (d and d.tool_calls):
99 + continue
100 + for tc in d.tool_calls:
101 + if tc.function and tc.function.arguments:
102 + args[tc.index] = args.get(tc.index, "") + tc.function.arguments
103 + return args
104 +
105 +
106 +class TestArgConverter:
107 + def test_complete_wrapper(self):
108 + raw = '{"name": "get_weather", "arguments": {"city": "SF"}}'
109 + assert _granite_arg_converter(raw, False) == '{"city": "SF"}'
110 +
111 + def test_arguments_before_name(self):
112 + raw = '{"arguments": {"city": "SF"}, "name": "get_weather"}'
113 + assert _granite_arg_converter(raw, False) == '{"city": "SF"}'
114 +
115 + def test_partial_before_arguments(self):
116 + assert _granite_arg_converter('{"name": "get_w', True) == ""
117 +
118 + def test_partial_inside_arguments(self):
119 + raw = '{"name": "x", "arguments": {"a": 1'
120 + assert _granite_arg_converter(raw, True) == '{"a": 1'
121 +
122 + def test_prefix_stability(self):
123 + # Each growing prefix must extend the previous converter output.
124 + full = '{"name": "x", "arguments": {"city": "Tokyo"}}'
125 + prev = ""
126 + for i in range(1, len(full) + 1):
127 + out = _granite_arg_converter(full[:i], True)
128 + assert out.startswith(prev) or prev.startswith(out) or out == ""
129 + if out:
130 + prev = out
131 +
132 + def test_missing_arguments_defaults(self):
133 + assert _granite_arg_converter('{"name": "x"}', False) == "{}"
134 + assert _granite_arg_converter('{"name": "x"}', True) == ""
135 +
136 + def test_non_object_arguments_rejected(self):
137 + with pytest.raises(ValueError, match="JSON object"):
138 + _granite_arg_converter('{"name": "x", "arguments": [1]}', False)
139 +
140 +
141 +class TestNonStreaming:
142 + def test_plain_text(self, parser, mock_request):
143 + reasoning, content, tools = parser.parse("Just a reply.", mock_request)
144 + assert reasoning is None
145 + assert content == "Just a reply."
146 + assert tools is None
147 +
148 + def test_single_tool_token_marker(self, parser, mock_request):
149 + text = (
150 + f'{TOOL_TOKEN} [{{"name": "get_weather", "arguments": {{"city": "SF"}}}}]'
151 + )
152 + _, content, tools = parser.parse(text, mock_request)
153 + assert content is None
154 + assert [t.name for t in tools] == ["get_weather"]
155 + assert json.loads(tools[0].arguments) == {"city": "SF"}
156 +
157 + def test_single_tool_string_marker(self, parser, mock_request):
158 + # Granite 3.1 uses the plain-text ``<tool_call>`` marker.
159 + text = f'{TOOL_STRING} [{{"name": "get_time", "arguments": {{}}}}]'
160 + _, content, tools = parser.parse(text, mock_request)
161 + assert [t.name for t in tools] == ["get_time"]
162 + assert json.loads(tools[0].arguments) == {}
163 +
164 + def test_parallel_calls_one_array(self, parser, mock_request):
165 + text = (
166 + f'{TOOL_TOKEN} [{{"name": "a", "arguments": {{"x": 1}}}}, '
167 + f'{{"name": "b", "arguments": {{"y": [1, 2]}}}}]'
168 + )
169 + _, _, tools = parser.parse(text, mock_request)
170 + assert [t.name for t in tools] == ["a", "b"]
171 + assert json.loads(tools[0].arguments) == {"x": 1}
172 + assert json.loads(tools[1].arguments) == {"y": [1, 2]}
173 +
174 + def test_surrounding_text_is_content(self, parser, mock_request):
175 + text = (
176 + "Let me check.\n"
177 + f'{TOOL_TOKEN} [{{"name": "get_weather", "arguments": {{"city": "SF"}}}}]'
178 + )
179 + _, content, tools = parser.parse(text, mock_request)
180 + assert content == "Let me check."
181 + assert [t.name for t in tools] == ["get_weather"]
182 +
183 + def test_arguments_before_name(self, parser, mock_request):
184 + text = (
185 + f'{TOOL_TOKEN} [{{"arguments": {{"city": "SF"}}, "name": "get_weather"}}]'
186 + )
187 + _, _, tools = parser.parse(text, mock_request)
188 + assert [t.name for t in tools] == ["get_weather"]
189 + assert json.loads(tools[0].arguments) == {"city": "SF"}
190 +
191 + def test_marker_without_array_is_not_a_tool_call(self, parser, mock_request):
192 + # A non-array body after the marker (malformed) yields no tool call.
193 + text = f'{TOOL_TOKEN} {{"name": "func", "arguments": {{}}}}'
194 + _, _, tools = parser.parse(text, mock_request)
195 + assert not tools
196 +
197 +
198 +class TestStreaming:
199 + @pytest.mark.parametrize("chunk_size", [1, 2, 3, 5, 11])
200 + def test_parallel_calls_chunk_invariance(self, parser, mock_request, chunk_size):
201 + text = (
202 + f'{TOOL_TOKEN} [{{"name": "get_weather", '
203 + f'"arguments": {{"city": "Tokyo"}}}}, '
204 + f'{{"name": "get_time", "arguments": {{"timezone": "Asia/Tokyo"}}}}]'
205 + )
206 + results = _stream(parser, mock_request, text, chunk_size)
207 + assert _collect_names(results) == ["get_weather", "get_time"]
208 + args = _collect_args(results)
209 + assert json.loads(args[0]) == {"city": "Tokyo"}
210 + assert json.loads(args[1]) == {"timezone": "Asia/Tokyo"}
211 +
212 + @pytest.mark.parametrize("chunk_size", [1, 3, 7])
213 + def test_surrounding_text_streaming(self, parser, mock_request, chunk_size):
214 + text = (
215 + "Let me check.\n"
216 + f'{TOOL_TOKEN} [{{"name": "get_weather", "arguments": {{"city": "SF"}}}}]'
217 + )
218 + results = _stream(parser, mock_request, text, chunk_size)
219 + assert _collect_names(results) == ["get_weather"]
220 + assert "Let me check." in _collect_content(results)
221 +
222 + def test_incomplete_array_at_eos(self, parser, mock_request):
223 + # Truncated mid-element must not raise and must not leak the marker.
224 + text = f'{TOOL_TOKEN} [{{"name": "func", "arguments": {{"a": 1'
225 + results = _stream(parser, mock_request, text, 3)
226 + assert TOOL_TOKEN not in _collect_content(results)
227 +
228 +
229 +class TestReplay:
230 + @pytest.mark.parametrize("scenario", _GRANITE_SCENARIOS, ids=lambda s: s.id)
231 + def test_replay_scenarios(self, scenario):
232 + # ``_build_granite`` replays the sample through GraniteParser at
233 + # chunk_size=1 and asserts reasoning/content/tool_calls (validate=True).
234 + _build_granite(scenario)
tests/parser/engine/test_replay.py
+5
@@ -78,6 +78,11 @@ def _discover_parsers() -> list[_ParserInfo]:
78 # Inkling opts out of token-id terminal matching and has typed
79 # structural blocks; its replay coverage lives in test_inkling.py.
80 continue
81 + if cfg.name == "granite":
82 + # Granite has a JSON-array tool body with no TOOL_END terminal, so
83 + # it does not fit this token-terminal harness; its replay coverage
84 + # lives in test_granite.py.
85 + continue
86 tool_end = cfg.token_id_terminals.get("TOOL_END")
87 if not tool_end:
88 raise RuntimeError(
tests/parser/engine/trace_builder.py
+85
@@ -34,6 +34,7 @@ from vllm.parser.engine.registered_adapters import (
34 DeepSeekV41Parser,
35 Gemma4Parser,
36 Glm47MoeParser,
37 + GraniteParser,
38 InklingParser,
39 KimiK2Parser,
40 MinimaxM2Parser,
@@ -1042,6 +1043,89 @@ def _build_inkling(scenario: Scenario, validate: bool = True) -> Sample:
1043 return sample
1044
1045
1046 +# ── Granite (JSON-array tool bodies, no reasoning) ───────────────────────
1047 +
1048 +_GRANITE_VOCAB: dict[str, int] = {
1049 + "<|tool_call|>": 49154,
1050 +}
1051 +
1052 +
1053 +def _granite_segments(scenario: Scenario) -> list[tuple[str, bool]]:
1054 + segs: list[tuple[str, bool]] = []
1055 + if scenario.content:
1056 + # Granite has no reasoning; prose is plain content preceding the marker.
1057 + segs.append((scenario.content, False))
1058 + if scenario.tool_calls:
1059 + segs.append(("<|tool_call|>", True))
1060 + payload = json.dumps(
1061 + [
1062 + {"name": tc.name, "arguments": tc.arguments}
1063 + for tc in scenario.tool_calls
1064 + ],
1065 + ensure_ascii=False,
1066 + separators=(", ", ": "),
1067 + )
1068 + segs.append((" " + payload, False))
1069 + return segs
1070 +
1071 +
1072 +def _granite_expected_content(scenario: Scenario) -> str | None:
1073 + if scenario.tool_calls:
1074 + if not scenario.content:
1075 + return None
1076 + return scenario.content.strip() or None
1077 + return scenario.content
1078 +
1079 +
1080 +def _build_granite(scenario: Scenario, validate: bool = True) -> Sample:
1081 + sample = _make_sample(
1082 + sample_id=f"granite-{scenario.id}",
1083 + description=scenario.description,
1084 + vocab=_GRANITE_VOCAB,
1085 + segments=_granite_segments(scenario),
1086 + expected_reasoning=None,
1087 + expected_content=_granite_expected_content(scenario),
1088 + expected_tool_calls=_expected_tc(scenario),
1089 + tools=_expected_tools(scenario),
1090 + )
1091 + if validate:
1092 + _validate_sample(sample, GraniteParser)
1093 + return sample
1094 +
1095 +
1096 +# Granite has no reasoning, so the shared reasoning-centric SCENARIOS do not
1097 +# apply; these exercise the JSON-array tool body (single, parallel, surrounding
1098 +# text) instead.
1099 +_GRANITE_SCENARIOS: list[Scenario] = [
1100 + Scenario(
1101 + id="single-tool",
1102 + description="Single tool call",
1103 + tool_calls=[_READ_TOOL],
1104 + ),
1105 + Scenario(
1106 + id="parallel-tools",
1107 + description="Parallel tool calls in one JSON array",
1108 + tool_calls=[_BASH_TOOL, _WEATHER_TOOL],
1109 + ),
1110 + Scenario(
1111 + id="complex-json-args",
1112 + description="Tool call with nested objects, arrays, numbers, booleans",
1113 + tool_calls=[_COMPLEX_TOOL],
1114 + ),
1115 + Scenario(
1116 + id="content-only",
1117 + description="Plain content response without tool calls",
1118 + content="Hello! How can I help you today?",
1119 + ),
1120 + Scenario(
1121 + id="surrounding-text",
1122 + description="Prose content preceding the tool call",
1123 + content="Let me check the weather.",
1124 + tool_calls=[_WEATHER_TOOL],
1125 + ),
1126 +]
1127 +
1128 +
1129 # ── Registry and public API ──────────────────────────────────────────
1130
1131 _BUILDERS: dict[str, Any] = {
@@ -1049,6 +1133,7 @@ _BUILDERS: dict[str, Any] = {
1133 "deepseek_v4": _build_deepseek_v4,
1134 "deepseek_v41": functools.partial(_build_deepseek_v4, v41=True),
1135 "gemma4": _build_gemma4,
1136 + "granite": _build_granite,
1137 "minimax_m2": _build_minimax_m2,
1138 "nemotron_v3": _build_nemotron_v3,
1139 "seed_oss": _build_seed_oss,
tests/tool_parsers/test_granite_tool_parser.py
+2 -19
@@ -16,7 +16,7 @@ from tests.tool_parsers.utils import (
16 split_string_into_token_deltas,
17 )
18 from vllm.tokenizers import get_tokenizer
19 -from vllm.tool_parsers.granite_tool_parser import GraniteToolParser
19 +from vllm.tool_parsers.granite_engine_tool_parser import GraniteEngineToolParser
20
21
22 class TestGraniteToolParser(ToolParserTests):
@@ -76,23 +76,6 @@ I'll get that information.""",
76 single_tool_call_expected_content=None,
77 parallel_tool_calls_count=2,
78 parallel_tool_calls_names=["get_weather", "get_time"],
79 - # xfail markers
80 - xfail_streaming={
81 - "test_malformed_input": (
82 - "Streaming mode incorrectly creates tool call from malformed JSON"
83 - ),
84 - "test_surrounding_text": (
85 - "Parser doesn't handle surrounding text correctly in streaming"
86 - ),
87 - "test_streaming_reconstruction": (
88 - "Streaming mode doesn't strip <|tool_call|> marker from content"
89 - ),
90 - },
91 - xfail_nonstreaming={
92 - "test_surrounding_text": (
93 - "Parser doesn't handle surrounding text correctly in non-streaming"
94 - ),
95 - },
79 )
80
81 # Granite-Specific Tests
@@ -140,7 +123,7 @@ def test_streaming_parallel_calls_batched_deltas(granite_tokenizer, chunk_size):
123 parallel calls must not drop the first call's name. granite streams
124 arguments before name, so the name only completes as the next call appears.
125 """
143 - parser = GraniteToolParser(granite_tokenizer)
126 + parser = GraniteEngineToolParser(granite_tokenizer)
127 model_output = (
128 '<|tool_call|> [{"arguments": {"city": "Tokyo"}, "name": "get_weather"}, '
129 '{"arguments": {"timezone": "Asia/Tokyo"}, "name": "get_time"}]'
vllm/parser/engine/parser_engine_config.py
+6
@@ -82,6 +82,12 @@ class ParserEngineConfig:
82
83 tool_args_json: bool = True
84
85 + # The tool region body is a single JSON array of ``{"name", "arguments"}``
86 + # objects (one marker → N calls, no closing terminal). The engine splits
87 + # the array into one tool call per element; the array's ``]`` ends the
88 + # region and any trailing text becomes content. Used by Granite.
89 + tool_call_body_array: bool = False
90 +
91 arg_structural_chars: frozenset[str] | None = None
92
93 # Special tokens exempt from auto-drop but not state-machine terminals.
vllm/parser/engine/registered_adapters.py
+6
@@ -13,6 +13,7 @@ from vllm.parser.deepseek_v41 import DeepSeekV41Parser
13 from vllm.parser.engine.adapters import make_adapters
14 from vllm.parser.gemma4 import Gemma4Parser
15 from vllm.parser.glm47_moe import Glm47MoeParser
16 +from vllm.parser.granite import GraniteParser
17 from vllm.parser.inkling import InklingParser
18 from vllm.parser.kimi_k2 import KimiK2Parser
19 from vllm.parser.minimax_m2 import MinimaxM2Parser
@@ -46,6 +47,11 @@ from vllm.parser.seed_oss import SeedOssParser
47 Gemma4ParserToolAdapter,
48 ) = make_adapters(Gemma4Parser)
49
50 +(
51 + GraniteParserReasoningAdapter,
52 + GraniteParserToolAdapter,
53 +) = make_adapters(GraniteParser)
54 +
55 (
56 NemotronV3ParserReasoningAdapter,
57 NemotronV3ParserToolAdapter,
vllm/parser/engine/streaming_parser_engine.py
+102 -1
@@ -190,6 +190,12 @@ class StreamingParserEngine:
190 self._args_in_string: bool = False
191 self._args_escape_next: bool = False
192
193 + def _reset_array_state(self) -> None:
194 + # State for ``tool_call_body_array``: whether the opening ``[`` has been
195 + # seen, and whether a ``{...}`` element is currently being streamed.
196 + self._array_started: bool = False
197 + self._array_in_element: bool = False
198 +
199 def reset(self, initial_state: ParserState | None = None) -> None:
200 """Reset mutable state for reuse across requests.
201
@@ -213,6 +219,7 @@ class StreamingParserEngine:
219 self._message_header_token_count = 0
220 self._in_skipped_tool_span = False
221 self._reset_args_state()
222 + self._reset_array_state()
223
224 def feed(
225 self,
@@ -306,7 +313,12 @@ class StreamingParserEngine:
313 ParserState.TOOL_NAME,
314 ParserState.TOOL_BETWEEN,
315 ):
309 - if self.tool_index >= 0:
316 + # In array mode a closed element has already emitted its
317 + # TOOL_CALL_END; only a truncated (still-open) element needs one.
318 + end_open_call = (
319 + self._array_in_element if self.config.tool_call_body_array else True
320 + )
321 + if end_open_call and self.tool_index >= 0:
322 events.append(
323 SemanticEvent(
324 EventType.TOOL_CALL_END,
@@ -474,6 +486,8 @@ class StreamingParserEngine:
486 self._message_header_token_count += token_count
487 return []
488 if self.state == ParserState.TOOL_ARGS:
489 + if self.config.tool_call_body_array:
490 + return self._feed_array_text(text)
491 if self.config.tool_args_json:
492 return self._feed_args_text(text)
493 return [
@@ -564,6 +578,8 @@ class StreamingParserEngine:
578 self._args_in_string = False
579 self._args_escape_next = False
580 self._args_safe_end = 0
581 + if self.config.tool_call_body_array:
582 + self._reset_array_state()
583
584 return events
585
@@ -634,3 +650,88 @@ class StreamingParserEngine:
650 tool_index=self.tool_index,
651 )
652 ]
653 +
654 + # ── JSON-array tool body (tool_call_body_array) ─────────────────────
655 +
656 + def _feed_array_text(self, text: str) -> list[SemanticEvent]:
657 + """Split a JSON-array tool body into one tool call per element.
658 +
659 + Granite emits ``[{"name":..,"arguments":{..}}, ...]`` after the tool
660 + marker. Element interiors are streamed through the string/brace-aware
661 + :meth:`_feed_args_char`, so each element behaves exactly like a
662 + single-object tool body; the array's ``]`` ends the region and any
663 + trailing text becomes content.
664 + """
665 + events: list[SemanticEvent] = []
666 + for i, ch in enumerate(text):
667 + if self.state != ParserState.TOOL_ARGS:
668 + # The array closed mid-chunk; the rest of the text is content.
669 + events.extend(self._emit_for_state(text[i:]))
670 + break
671 + events.extend(self._feed_array_char(ch))
672 + return events
673 +
674 + def _feed_array_char(self, ch: str) -> list[SemanticEvent]:
675 + if self._array_in_element:
676 + events = self._feed_args_char(ch)
677 + # Element object closed: top-level '}' with all braces balanced.
678 + if (
679 + ch == "}"
680 + and not self._args_in_string
681 + and not self._args_escape_next
682 + and self._args_brace_depth == 0
683 + ):
684 + events.extend(self._end_array_element())
685 + return events
686 +
687 + # Between elements only structural array characters are expected.
688 + if ch in " \t\r\n":
689 + return []
690 + if not self._array_started:
691 + if ch == "[":
692 + self._array_started = True
693 + return []
694 + return self._bail_array_to_content(ch)
695 + if ch == ",":
696 + return []
697 + if ch == "{":
698 + return self._start_array_element(ch)
699 + if ch == "]":
700 + self.state = ParserState.CONTENT
701 + self._reset_array_state()
702 + return []
703 + return self._bail_array_to_content(ch)
704 +
705 + def _start_array_element(self, ch: str) -> list[SemanticEvent]:
706 + self.tool_index += 1
707 + self._reset_args_state()
708 + self._array_in_element = True
709 + events = [SemanticEvent(EventType.TOOL_CALL_START, tool_index=self.tool_index)]
710 + events.extend(self._feed_args_char(ch))
711 + return events
712 +
713 + def _end_array_element(self) -> list[SemanticEvent]:
714 + self._array_in_element = False
715 + events: list[SemanticEvent] = []
716 + # Flush the held closing '}' so the slot body is the full element.
717 + if self._args_buffer:
718 + events.append(
719 + SemanticEvent(
720 + EventType.ARG_VALUE_CHUNK,
721 + value=self._args_buffer,
722 + tool_index=self.tool_index,
723 + )
724 + )
725 + self._args_buffer = ""
726 + self._args_safe_end = 0
727 + events.append(
728 + SemanticEvent(EventType.TOOL_CALL_END, tool_index=self.tool_index)
729 + )
730 + return events
731 +
732 + def _bail_array_to_content(self, ch: str) -> list[SemanticEvent]:
733 + # Marker not followed by a JSON array (malformed) — revert to content
734 + # so the remaining tokens are emitted as plain text, no tool call.
735 + self.state = ParserState.CONTENT
736 + self._reset_array_state()
737 + return self._emit_for_state(ch)
vllm/parser/granite.py new
+88
@@ -0,0 +1,88 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Granite parser for JSON-array tool calls.
4 +
5 +Granite 3.0/3.1 emit tool calls as a JSON array following a marker::
6 +
7 + <|tool_call|> [{"name": "get_weather", "arguments": {"city": "SF"}}] # 3.0
8 + <tool_call> [{"name": "get_weather", "arguments": {"city": "SF"}}] # 3.1
9 +
10 +``<|tool_call|>`` is a single special token; ``<tool_call>`` is plain text.
11 +The array may hold several calls, has no closing marker (the ``]`` ends it),
12 +and any surrounding prose is content. The engine's ``tool_call_body_array``
13 +mode splits the array into one call per element and extracts each name via the
14 +name-from-args path; ``_granite_arg_converter`` carves the ``arguments`` object
15 +out of each ``{"name":..,"arguments":{..}}`` element. Granite has no reasoning.
16 +"""
17 +
18 +from __future__ import annotations
19 +
20 +import functools
21 +from typing import TYPE_CHECKING
22 +
23 +from vllm.parser.engine.parser_engine import ParserEngine
24 +from vllm.parser.engine.parser_engine_config import (
25 + ParserEngineConfig,
26 + ParserState,
27 + Transition,
28 +)
29 +from vllm.parser.utils import extract_json_member_object
30 +
31 +if TYPE_CHECKING:
32 + from vllm.tokenizers import TokenizerLike
33 + from vllm.tool_parsers.abstract_tool_parser import Tool
34 +
35 +TOOL_TOKEN = "<|tool_call|>" # granite 3.0 (single special token)
36 +TOOL_STRING = "<tool_call>" # granite 3.1 (plain text)
37 +
38 +
39 +def _granite_arg_converter(raw_args: str, partial: bool) -> str:
40 + """Carve the ``arguments`` object out of a ``{"name":..,"arguments":{..}}``
41 + array element (see :func:`extract_json_member_object`)."""
42 + span = extract_json_member_object(raw_args, "arguments")
43 + if span is None:
44 + return "" if partial else "{}"
45 + return span
46 +
47 +
48 +@functools.cache
49 +def granite_config() -> ParserEngineConfig:
50 + return ParserEngineConfig(
51 + name="granite",
52 + initial_state=ParserState.CONTENT,
53 + terminals={
54 + "TOOL_TOKEN": TOOL_TOKEN,
55 + "TOOL_STRING": TOOL_STRING,
56 + },
57 + token_id_terminals={
58 + "TOOL_TOKEN": TOOL_TOKEN,
59 + },
60 + transitions={
61 + # The marker only switches into the tool region; the array splitter
62 + # emits TOOL_CALL_START per element (empty events here).
63 + (ParserState.CONTENT, "TOOL_TOKEN"): Transition(
64 + ParserState.TOOL_ARGS,
65 + (),
66 + ),
67 + (ParserState.CONTENT, "TOOL_STRING"): Transition(
68 + ParserState.TOOL_ARGS,
69 + (),
70 + ),
71 + },
72 + arg_converter=_granite_arg_converter,
73 + tool_call_body_array=True,
74 + stream_arg_deltas=True,
75 + )
76 +
77 +
78 +class GraniteParser(ParserEngine):
79 + CONFIG_NAME = "granite"
80 +
81 + def __init__(
82 + self,
83 + tokenizer: TokenizerLike,
84 + tools: list[Tool] | None = None,
85 + **kwargs,
86 + ) -> None:
87 + kwargs.setdefault("parser_engine_config", granite_config())
88 + super().__init__(tokenizer, tools, **kwargs)
vllm/parser/inkling.py
+4 -90
@@ -37,6 +37,7 @@ from vllm.parser.engine.parser_engine_config import (
37 ParserState,
38 Transition,
39 )
40 +from vllm.parser.utils import extract_json_member_object
41
42 if TYPE_CHECKING:
43 from vllm.tokenizers import TokenizerLike
@@ -62,84 +63,6 @@ INKLING_SPECIAL_TOKENS = (
63 END_MESSAGE,
64 )
65
65 -_WS = " \t\r\n"
66 -
67 -
68 -def _scan_json_value(raw: str, start: int) -> int | None:
69 - """Return the end index (exclusive) of the JSON object starting at
70 - ``raw[start]``, or ``None`` when the object is still unterminated."""
71 - depth = 0
72 - in_string = False
73 - escape = False
74 - for i in range(start, len(raw)):
75 - ch = raw[i]
76 - if escape:
77 - escape = False
78 - continue
79 - if in_string:
80 - if ch == "\\":
81 - escape = True
82 - elif ch == '"':
83 - in_string = False
84 - continue
85 - if ch == '"':
86 - in_string = True
87 - elif ch == "{":
88 - depth += 1
89 - elif ch == "}":
90 - depth -= 1
91 - if depth == 0:
92 - return i + 1
93 - return None
94 -
95 -
96 -def _args_value_span(raw: str) -> str | None:
97 - """Extract the raw text span of the top-level ``"args"`` value from a
98 - (possibly incomplete) ``{"name":...,"args":{...}}`` wrapper.
99 -
100 - Returns the verbatim substring (prefix-stable across growing input,
101 - which the engine's argument-delta diffing relies on), possibly an
102 - unterminated object prefix; ``None`` when the value has not started.
103 - Raises ``ValueError`` when the value is not a JSON object.
104 - """
105 - depth = 0
106 - in_string = False
107 - escape = False
108 - string_start = -1
109 - last_string: str | None = None
110 - for i, ch in enumerate(raw):
111 - if escape:
112 - escape = False
113 - continue
114 - if in_string:
115 - if ch == "\\":
116 - escape = True
117 - elif ch == '"':
118 - in_string = False
119 - if depth == 1:
120 - last_string = raw[string_start + 1 : i]
121 - continue
122 - if ch == '"':
123 - in_string = True
124 - string_start = i
125 - elif ch == ":" and depth == 1 and last_string == "args":
126 - value_start = i + 1
127 - while value_start < len(raw) and raw[value_start] in _WS:
128 - value_start += 1
129 - if value_start >= len(raw):
130 - return None
131 - if raw[value_start] != "{":
132 - raise ValueError("Inkling tool call args must be a JSON object")
133 - value_end = _scan_json_value(raw, value_start)
134 - if value_end is None:
135 - return raw[value_start:]
136 - return raw[value_start:value_end]
137 - elif ch in "{[":
138 - depth += 1
139 - elif ch in "}]":
140 - depth -= 1
141 - return None
142 -
66
67 def _inkling_arg_converter(raw_args: str, partial: bool) -> str:
68 """Carve the ``args`` object out of the tool-call JSON wrapper.
@@ -150,19 +73,10 @@ def _inkling_arg_converter(raw_args: str, partial: bool) -> str:
73 converter, ``_compute_arg_delta`` streams the wrapper verbatim into
74 the OpenAI ``arguments`` field (``converter is None -> raw delta``,
75 unconditionally; ``stream_arg_deltas=False`` does not stop it).
153 -
154 - Why a hand-rolled scanner instead of (partial) ``json.loads`` +
155 - ``json.dumps``: the engine diffs successive converter outputs and
156 - requires each to extend the previous one (``startswith``); a
157 - violation silently drops argument deltas. Re-serialization changes
158 - whitespace and closes unterminated structures differently across
159 - ticks, so the only prefix-stable output is a verbatim substring of
160 - the input. The scanner is also string/escape-aware so an ``"args"``
161 - literal inside the name or a string value cannot mislead it, and it
162 - still recovers a partial span when EOS truncates the wrapper (where
163 - ``json.loads`` would fail).
76 + :func:`extract_json_member_object` returns a prefix-stable verbatim
77 + substring, which the engine's argument-delta diffing requires.
78 """
165 - span = _args_value_span(raw_args)
79 + span = extract_json_member_object(raw_args, "args")
80 if span is None:
81 # No args value yet (streaming) or none at all (treat as empty).
82 return "" if partial else "{}"
vllm/parser/utils.py
+85
@@ -63,3 +63,88 @@ def count_history_tool_calls(
63 return 0
64
65 return count_response_history_tool_calls(request_input)
66 +
67 +
68 +_JSON_WS = " \t\r\n"
69 +
70 +
71 +def scan_json_object_end(raw: str, start: int) -> int | None:
72 + """Return the end index (exclusive) of the JSON object starting at
73 + ``raw[start]`` (which must be ``{``), or ``None`` when it is still
74 + unterminated. String- and escape-aware so braces inside strings do not
75 + mislead the depth count."""
76 + depth = 0
77 + in_string = False
78 + escape = False
79 + for i in range(start, len(raw)):
80 + ch = raw[i]
81 + if escape:
82 + escape = False
83 + continue
84 + if in_string:
85 + if ch == "\\":
86 + escape = True
87 + elif ch == '"':
88 + in_string = False
89 + continue
90 + if ch == '"':
91 + in_string = True
92 + elif ch == "{":
93 + depth += 1
94 + elif ch == "}":
95 + depth -= 1
96 + if depth == 0:
97 + return i + 1
98 + return None
99 +
100 +
101 +def extract_json_member_object(raw: str, key: str) -> str | None:
102 + """Extract the raw text span of the top-level object-valued member ``key``
103 + from a (possibly incomplete) ``{...}`` JSON wrapper.
104 +
105 + Returns the verbatim substring (prefix-stable as the input grows, which
106 + the engine's argument-delta diffing relies on), possibly an unterminated
107 + object prefix; ``None`` when the value has not started. Raises
108 + ``ValueError`` when the member is present but its value is not a JSON
109 + object. A verbatim substring (not re-serialized JSON) is required because
110 + the engine diffs successive converter outputs and expects each to extend
111 + the previous one; the scanner is string/escape-aware so a matching literal
112 + inside another value cannot mislead it, and it recovers a partial span when
113 + the wrapper is truncated (where ``json.loads`` would fail)."""
114 + depth = 0
115 + in_string = False
116 + escape = False
117 + string_start = -1
118 + last_string: str | None = None
119 + for i, ch in enumerate(raw):
120 + if escape:
121 + escape = False
122 + continue
123 + if in_string:
124 + if ch == "\\":
125 + escape = True
126 + elif ch == '"':
127 + in_string = False
128 + if depth == 1:
129 + last_string = raw[string_start + 1 : i]
130 + continue
131 + if ch == '"':
132 + in_string = True
133 + string_start = i
134 + elif ch == ":" and depth == 1 and last_string == key:
135 + value_start = i + 1
136 + while value_start < len(raw) and raw[value_start] in _JSON_WS:
137 + value_start += 1
138 + if value_start >= len(raw):
139 + return None
140 + if raw[value_start] != "{":
141 + raise ValueError(f"member {key!r} must be a JSON object")
142 + value_end = scan_json_object_end(raw, value_start)
143 + if value_end is None:
144 + return raw[value_start:]
145 + return raw[value_start:value_end]
146 + elif ch in "{[":
147 + depth += 1
148 + elif ch in "}]":
149 + depth -= 1
150 + return None
vllm/tool_parsers/__init__.py
+2 -2
@@ -75,8 +75,8 @@ _TOOL_PARSERS_TO_REGISTER = {
75 "Granite20bFCToolParser",
76 ),
77 "granite": (
78 - "granite_tool_parser",
79 - "GraniteToolParser",
78 + "granite_engine_tool_parser",
79 + "GraniteEngineToolParser",
80 ),
81 "granite4": (
82 "granite4_tool_parser",
vllm/tool_parsers/granite_engine_tool_parser.py new
+8
@@ -0,0 +1,8 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm.parser.engine.registered_adapters import GraniteParserToolAdapter
5 +
6 +
7 +class GraniteEngineToolParser(GraniteParserToolAdapter): # type: ignore[valid-type, misc]
8 + structural_tag_model = None
vllm/tool_parsers/granite_tool_parser.py deleted
-257
@@ -1,257 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import json
5 -from collections.abc import Sequence
6 -
7 -import partial_json_parser
8 -from partial_json_parser.core.options import Allow
9 -
10 -from vllm.entrypoints.chat_utils import make_tool_call_id
11 -from vllm.entrypoints.generate.base.protocol import (
12 - DeltaFunctionCall,
13 - DeltaMessage,
14 - DeltaToolCall,
15 - ExtractedToolCallInformation,
16 - FunctionCall,
17 - ToolCall,
18 -)
19 -from vllm.entrypoints.openai.chat_completion.protocol import (
20 - ChatCompletionRequest,
21 -)
22 -from vllm.logger import init_logger
23 -from vllm.tokenizers import TokenizerLike
24 -from vllm.tool_parsers.abstract_tool_parser import (
25 - Tool,
26 - ToolParser,
27 -)
28 -from vllm.tool_parsers.utils import (
29 - consume_space,
30 - find_common_prefix,
31 - is_complete_json,
32 - partial_json_loads,
33 -)
34 -
35 -logger = init_logger(__name__)
36 -
37 -
38 -class GraniteToolParser(ToolParser):
39 - """Tool call parser for the granite 3.0 models. Intended
40 - for use with the examples/tool_chat_template_granite.jinja
41 - template.
42 -
43 - Used when --enable-auto-tool-choice --tool-call-parser granite
44 - are all set
45 - """
46 -
47 - def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
48 - super().__init__(tokenizer, tools)
49 - # for granite 3.0, the token `<|tool_call|>`
50 - self.bot_token = "<|tool_call|>"
51 - # for granite 3.1, the string `<tool_call>`
52 - self.bot_string = "<tool_call>"
53 -
54 - def extract_tool_calls(
55 - self, model_output: str, request: ChatCompletionRequest
56 - ) -> ExtractedToolCallInformation:
57 - stripped = (
58 - model_output.strip()
59 - .removeprefix(self.bot_token)
60 - .removeprefix(self.bot_string)
61 - .lstrip()
62 - )
63 - if not stripped or stripped[0] != "[":
64 - return ExtractedToolCallInformation(
65 - tools_called=False, tool_calls=[], content=model_output
66 - )
67 - try:
68 - raw_function_calls = json.loads(stripped)
69 - if not isinstance(raw_function_calls, list):
70 - raise Exception(
71 - f"Expected dict or list, got {type(raw_function_calls)}"
72 - )
73 -
74 - logger.debug("Extracted %d tool calls", len(raw_function_calls))
75 - tool_calls = [
76 - ToolCall(
77 - type="function",
78 - function=FunctionCall(
79 - name=function_call["name"],
80 - # function call args are JSON but as a string
81 - arguments=json.dumps(
82 - function_call["arguments"], ensure_ascii=False
83 - ),
84 - ),
85 - )
86 - for function_call in raw_function_calls
87 - ]
88 -
89 - return ExtractedToolCallInformation(
90 - tools_called=True,
91 - tool_calls=tool_calls,
92 - content=None,
93 - )
94 -
95 - except Exception as e:
96 - logger.error("Error in extracting tool call from response %s", e)
97 - return ExtractedToolCallInformation(
98 - tools_called=False, tool_calls=[], content=model_output
99 - )
100 -
101 - def extract_tool_calls_streaming(
102 - self,
103 - previous_text: str,
104 - current_text: str,
105 - delta_text: str,
106 - previous_token_ids: Sequence[int],
107 - current_token_ids: Sequence[int],
108 - delta_token_ids: Sequence[int],
109 - request: ChatCompletionRequest,
110 - ) -> DeltaMessage | None:
111 - start_idx = consume_space(0, current_text)
112 - if current_text[start_idx:].startswith(self.bot_token):
113 - start_idx = consume_space(start_idx + len(self.bot_token), current_text)
114 - if current_text[start_idx:].startswith(self.bot_string):
115 - start_idx = consume_space(start_idx + len(self.bot_string), current_text)
116 - if (
117 - not current_text
118 - or start_idx >= len(current_text)
119 - or current_text[start_idx] != "["
120 - ):
121 - return DeltaMessage(content=delta_text)
122 -
123 - # bit mask flags for partial JSON parsing. If the name hasn't been
124 - # sent yet, don't allow sending
125 - # an incomplete string since OpenAI only ever (as far as I have
126 - # seen) allows sending the entire tool/ function name at once.
127 - flags = Allow.ALL if self.current_tool_name_sent else Allow.ALL & ~Allow.STR
128 - try:
129 - tool_call_arr = None
130 - is_complete = None
131 - try:
132 - tool_calls, end_idx = partial_json_loads(
133 - current_text[start_idx:], flags
134 - )
135 - if type(tool_calls) is list:
136 - tool_call_arr = tool_calls
137 - else:
138 - return DeltaMessage(content=delta_text)
139 -
140 - is_complete = [True] * len(tool_calls)
141 - if not is_complete_json(current_text[start_idx : start_idx + end_idx]):
142 - is_complete[-1] = False
143 - except partial_json_parser.core.exceptions.MalformedJSON:
144 - logger.debug("not enough tokens to parse into JSON yet")
145 - return None
146 -
147 - # case -- if no tokens have been streamed for the tool, e.g.
148 - # only the array brackets, stream nothing
149 - if not tool_call_arr:
150 - return None
151 -
152 - # select as the current tool call the one we're on the state at
153 - current_tool_call: dict = tool_call_arr[self.current_tool_id]
154 -
155 - delta = None
156 - # Only advance once the current tool name is streamed; granite
157 - # emits arguments before name, so advancing early would drop it.
158 - if len(tool_call_arr) > self.current_tool_id + 1 and (
159 - self.current_tool_id < 0 or self.current_tool_name_sent
160 - ):
161 - # if we're moving on to a new call, first make sure we
162 - # haven't missed anything in the previous one that was
163 - # auto-generated due to JSON completions, but wasn't
164 - # streamed to the client yet.
165 - if self.current_tool_id >= 0:
166 - cur_arguments = current_tool_call.get("arguments")
167 - if cur_arguments:
168 - cur_args_json = json.dumps(cur_arguments, ensure_ascii=False)
169 - sent = len(self.streamed_args_for_tool[self.current_tool_id])
170 - argument_diff = cur_args_json[sent:]
171 -
172 - logger.debug("got arguments diff: %s", argument_diff)
173 - delta = DeltaMessage(
174 - tool_calls=[
175 - DeltaToolCall(
176 - index=self.current_tool_id,
177 - function=DeltaFunctionCall(
178 - arguments=argument_diff
179 - ).model_dump(exclude_none=True),
180 - )
181 - ]
182 - )
183 - self.streamed_args_for_tool[self.current_tool_id] += (
184 - argument_diff
185 - )
186 -
187 - # re-set stuff pertaining to progress in the current tool
188 - self.current_tool_id += 1
189 - self.current_tool_name_sent = False
190 - self.streamed_args_for_tool.append("")
191 - logger.debug("starting on new tool %d", self.current_tool_id)
192 - return delta
193 -
194 - # if the current tool name hasn't been sent, send if available
195 - # - otherwise send nothing
196 - elif not self.current_tool_name_sent:
197 - function_name = current_tool_call.get("name")
198 - if function_name:
199 - delta = DeltaMessage(
200 - tool_calls=[
201 - DeltaToolCall(
202 - index=self.current_tool_id,
203 - type="function",
204 - id=make_tool_call_id(),
205 - function=DeltaFunctionCall(
206 - name=function_name
207 - ).model_dump(exclude_none=True),
208 - )
209 - ]
210 - )
211 - self.current_tool_name_sent = True
212 -
213 - # now we know we're on the same tool call and we're streaming
214 - # arguments
215 - else:
216 - cur_arguments = current_tool_call.get("arguments")
217 -
218 - if cur_arguments:
219 - sent = len(self.streamed_args_for_tool[self.current_tool_id])
220 - cur_args_json = json.dumps(cur_arguments, ensure_ascii=False)
221 - prev_arguments = self.prev_tool_call_arr[self.current_tool_id].get(
222 - "arguments"
223 - )
224 -
225 - argument_diff = None
226 - if is_complete[self.current_tool_id]:
227 - argument_diff = cur_args_json[sent:]
228 - elif prev_arguments:
229 - prev_args_json = json.dumps(prev_arguments, ensure_ascii=False)
230 - if cur_args_json != prev_args_json:
231 - prefix = find_common_prefix(prev_args_json, cur_args_json)
232 - argument_diff = prefix[sent:]
233 -
234 - if argument_diff is not None:
235 - delta = DeltaMessage(
236 - tool_calls=[
237 - DeltaToolCall(
238 - index=self.current_tool_id,
239 - function=DeltaFunctionCall(
240 - arguments=argument_diff
241 - ).model_dump(exclude_none=True),
242 - )
243 - ]
244 - )
245 - self.streamed_args_for_tool[self.current_tool_id] += (
246 - argument_diff
247 - )
248 -
249 - self.prev_tool_call_arr = tool_call_arr
250 - return delta
251 -
252 - except Exception as e:
253 - logger.error("Error trying to handle streaming tool call: %s", e)
254 - logger.debug(
255 - "Skipping chunk as a result of tool streaming extraction error"
256 - )
257 - return None