main
py 581 lines 20.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Gemma4 parser.
4
5 Handles channel-based reasoning plus custom tool call format in a single
6 state machine::
7
8 <|channel>thought
9 ...reasoning...<channel|>
10 <|tool_call>call:func_name{key:<|"|>value<|"|>,num:42}<tool_call|>
11 """
12
13 from __future__ import annotations
14
15 import functools
16 import json
17 from collections.abc import Sequence
18 from typing import TYPE_CHECKING
19
20 from vllm.entrypoints.generate.base.protocol import DeltaMessage
21 from vllm.logger import init_logger
22 from vllm.parser.engine.events import EventType, SemanticEvent
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
30 if TYPE_CHECKING:
31 from vllm.entrypoints.openai.chat_completion.protocol import (
32 ChatCompletionRequest,
33 )
34 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
35 from vllm.tokenizers import TokenizerLike
36 from vllm.tool_parsers.abstract_tool_parser import Tool
37
38 CHANNEL_START = "<|channel>"
39 CHANNEL_END = "<channel|>"
40 TOOL_CALL_START = "<|tool_call>"
41 TOOL_CALL_END = "<tool_call|>"
42 STRING_DELIM = '<|"|>'
43 _DELIM_LEN = len(STRING_DELIM)
44
45 logger = init_logger(__name__)
46
47
48 # ---------------------------------------------------------------------------
49 # Gemma4 argument parser
50 # ---------------------------------------------------------------------------
51
52 _PARTIAL_DELIM_SUFFIXES = tuple(
53 STRING_DELIM[:k] for k in range(len(STRING_DELIM), 0, -1)
54 )
55
56
57 def _strip_partial_delim(value: str) -> str:
58 """Strip a trailing partial ``STRING_DELIM`` prefix from *value*.
59
60 Prevents partial delimiters from leaking into the streamed JSON diff.
61 """
62 for suffix in _PARTIAL_DELIM_SUFFIXES:
63 if value.endswith(suffix):
64 return value[: -len(suffix)]
65 return value
66
67
68 def _parse_gemma4_args(args_str: str, *, partial: bool = False) -> dict:
69 """Parse Gemma4's custom key:value format into a Python dict.
70
71 Format examples::
72
73 location:<|"|>Tokyo<|"|>
74 location:<|"|>San Francisco<|"|>,unit:<|"|>celsius<|"|>
75 count:42,flag:true
76 nested:{inner_key:<|"|>val<|"|>}
77 items:[<|"|>a<|"|>,<|"|>b<|"|>]
78
79 Args:
80 args_str: The raw Gemma4 argument string.
81 partial: When True (streaming), bare values at end of string are
82 omitted because they may be incomplete and type-unstable
83 (e.g. partial boolean parsed as bare string).
84
85 Returns a dict ready for ``json.dumps()``.
86
87 """
88 if not args_str or not args_str.strip():
89 return {}
90
91 result: dict = {}
92 i = 0
93 n = len(args_str)
94
95 while i < n:
96 while i < n and args_str[i] in (" ", ",", "\n", "\t"):
97 i += 1
98 if i >= n:
99 break
100
101 key_start = i
102 while i < n and args_str[i] != ":":
103 i += 1
104 if i >= n:
105 break
106 key = args_str[key_start:i].strip()
107 if key.startswith(STRING_DELIM) and key.endswith(STRING_DELIM):
108 key = key[_DELIM_LEN:-_DELIM_LEN]
109 i += 1
110
111 if i >= n:
112 if not partial:
113 result[key] = ""
114 break
115
116 while i < n and args_str[i] in (" ", "\n", "\t"):
117 i += 1
118 if i >= n:
119 if not partial:
120 result[key] = ""
121 break
122
123 if args_str.startswith(STRING_DELIM, i):
124 i += _DELIM_LEN
125 val_start = i
126 end_pos = args_str.find(STRING_DELIM, i)
127 if end_pos == -1:
128 # Unterminated string — take rest, strip partial delimiter.
129 value = args_str[val_start:]
130 if partial:
131 value = _strip_partial_delim(value)
132 result[key] = value
133 break
134 result[key] = args_str[val_start:end_pos]
135 i = end_pos + _DELIM_LEN
136
137 elif args_str[i] == "{":
138 depth = 1
139 obj_start = i + 1
140 i += 1
141 while i < n and depth > 0:
142 if args_str.startswith(STRING_DELIM, i):
143 # Skip over string contents to avoid counting { inside strings
144 i += _DELIM_LEN
145 next_delim = args_str.find(STRING_DELIM, i)
146 i = n if next_delim == -1 else next_delim + _DELIM_LEN
147 continue
148 if args_str[i] == "{":
149 depth += 1
150 elif args_str[i] == "}":
151 depth -= 1
152 i += 1
153 if depth > 0:
154 # Incomplete nested object — use i (not i-1) to avoid
155 # dropping the last char, and recurse as partial.
156 result[key] = _parse_gemma4_args(args_str[obj_start:i], partial=True)
157 else:
158 result[key] = _parse_gemma4_args(args_str[obj_start : i - 1])
159
160 elif args_str[i] == "[":
161 depth = 1
162 arr_start = i + 1
163 i += 1
164 while i < n and depth > 0:
165 if args_str.startswith(STRING_DELIM, i):
166 i += _DELIM_LEN
167 next_delim = args_str.find(STRING_DELIM, i)
168 i = n if next_delim == -1 else next_delim + _DELIM_LEN
169 continue
170 if args_str[i] == "[":
171 depth += 1
172 elif args_str[i] == "]":
173 depth -= 1
174 i += 1
175 if depth > 0:
176 result[key] = _parse_gemma4_array(args_str[arr_start:i], partial=True)
177 else:
178 result[key] = _parse_gemma4_array(args_str[arr_start : i - 1])
179
180 else:
181 val_start = i
182 while i < n and args_str[i] not in (",", "}", "]"):
183 i += 1
184 if partial and i >= n:
185 # Value may be incomplete (e.g. partial boolean) —
186 # withhold to avoid type instability during streaming.
187 break
188 if i == val_start:
189 logger.warning(
190 "Gemma4 args parser made no progress at position %d; "
191 "aborting on malformed input.",
192 i,
193 )
194 break
195 raw_val = args_str[val_start:i].strip()
196 if partial and raw_val.endswith("."):
197 # Digits may still arrive (e.g. "108." -> "108.2");
198 # withhold to avoid corrupting the streaming diff.
199 break
200 result[key] = raw_val
201
202 return result
203
204
205 def _parse_gemma4_array(arr_str: str, *, partial: bool = False) -> list:
206 items: list = []
207 i = 0
208 n = len(arr_str)
209
210 while i < n:
211 while i < n and arr_str[i] in (" ", ",", "\n", "\t"):
212 i += 1
213 if i >= n:
214 break
215
216 if arr_str.startswith(STRING_DELIM, i):
217 i += _DELIM_LEN
218 end_pos = arr_str.find(STRING_DELIM, i)
219 if end_pos == -1:
220 items.append(arr_str[i:])
221 break
222 items.append(arr_str[i:end_pos])
223 i = end_pos + _DELIM_LEN
224
225 elif arr_str[i] == "{":
226 depth = 1
227 obj_start = i + 1
228 i += 1
229 while i < n and depth > 0:
230 if arr_str.startswith(STRING_DELIM, i):
231 i += _DELIM_LEN
232 nd = arr_str.find(STRING_DELIM, i)
233 i = nd + _DELIM_LEN if nd != -1 else n
234 continue
235 if arr_str[i] == "{":
236 depth += 1
237 elif arr_str[i] == "}":
238 depth -= 1
239 i += 1
240 if depth > 0:
241 items.append(_parse_gemma4_args(arr_str[obj_start:i], partial=True))
242 else:
243 items.append(_parse_gemma4_args(arr_str[obj_start : i - 1]))
244
245 elif arr_str[i] == "[":
246 depth = 1
247 sub_start = i + 1
248 i += 1
249 while i < n and depth > 0:
250 if arr_str.startswith(STRING_DELIM, i):
251 i += _DELIM_LEN
252 nd = arr_str.find(STRING_DELIM, i)
253 i = nd + _DELIM_LEN if nd != -1 else n
254 continue
255 if arr_str[i] == "[":
256 depth += 1
257 elif arr_str[i] == "]":
258 depth -= 1
259 i += 1
260 if depth > 0:
261 items.append(_parse_gemma4_array(arr_str[sub_start:i], partial=True))
262 else:
263 items.append(_parse_gemma4_array(arr_str[sub_start : i - 1]))
264
265 else:
266 val_start = i
267 while i < n and arr_str[i] not in (",", "]"):
268 i += 1
269 if partial and i >= n:
270 break
271 if i == val_start:
272 logger.warning(
273 "Gemma4 array parser made no progress at position %d; "
274 "aborting on malformed input.",
275 i,
276 )
277 break
278 raw_val = arr_str[val_start:i].strip()
279 if partial and raw_val.endswith("."):
280 break
281 items.append(raw_val)
282
283 return items
284
285
286 def _gemma4_arg_converter(raw_args: str, partial: bool) -> str:
287 """Convert Gemma4 custom arg format to a JSON string."""
288 text = raw_args.strip()
289 if text.endswith("}") or text.endswith(")") and text.count("(") < text.count(")"):
290 text = text[:-1]
291
292 parsed = _parse_gemma4_args(text, partial=partial)
293 return json.dumps(parsed, ensure_ascii=False)
294
295
296 @functools.cache
297 def gemma4_config(thinking: bool = False) -> ParserEngineConfig:
298 return ParserEngineConfig(
299 name="gemma4",
300 initial_state=ParserState.CONTENT,
301 # Decides a silent tail: the E-series HF templates write no marker
302 # with thinking off, and no template does after a tool response.
303 wait_for_reasoning=thinking,
304 turn_boundary_tokens=frozenset({"<|turn>", "<|tool_response>"}),
305 terminals={
306 "THINK_START": CHANNEL_START,
307 "THINK_END": CHANNEL_END,
308 "TOOL_START": TOOL_CALL_START,
309 "TOOL_END": TOOL_CALL_END,
310 "CALL_PREFIX": "call:",
311 "COLON": ":",
312 "OPEN_BRACE": "{",
313 "OPEN_PAREN": "(",
314 },
315 token_id_terminals={
316 "THINK_START": CHANNEL_START,
317 "THINK_END": CHANNEL_END,
318 "TOOL_START": TOOL_CALL_START,
319 "TOOL_END": TOOL_CALL_END,
320 },
321 transitions={
322 # -- Reasoning transitions --
323 (ParserState.CONTENT, "THINK_START"): Transition(
324 ParserState.REASONING,
325 (EventType.REASONING_START,),
326 ),
327 # No-op: if we pre-initialised the engine to REASONING from the
328 # prompt (see ``adjust_initial_state_from_prompt``) but the model
329 # still emits its own ``<|channel>`` opener, swallow it instead
330 # of leaking it as TEXT_CHUNK.
331 (ParserState.REASONING, "THINK_START"): Transition(
332 ParserState.REASONING,
333 (),
334 ),
335 (ParserState.REASONING, "THINK_END"): Transition(
336 ParserState.CONTENT,
337 (EventType.REASONING_END,),
338 ),
339 # Tool call directly from reasoning (no explicit <channel|>)
340 (ParserState.REASONING, "TOOL_START"): Transition(
341 ParserState.TOOL_PREAMBLE,
342 (EventType.REASONING_END, EventType.TOOL_CALL_START),
343 ),
344 # -- Tool call transitions --
345 (ParserState.CONTENT, "TOOL_START"): Transition(
346 ParserState.TOOL_PREAMBLE,
347 (EventType.REASONING_END, EventType.TOOL_CALL_START),
348 ),
349 (ParserState.TOOL_PREAMBLE, "TOOL_END"): Transition(
350 ParserState.CONTENT,
351 (EventType.TOOL_CALL_END,),
352 ),
353 (ParserState.TOOL_PREAMBLE, "CALL_PREFIX"): Transition(
354 ParserState.TOOL_NAME,
355 (),
356 ),
357 # Bare opener: some checkpoints emit "<|tool_call>:name{...}"
358 # without the "call" prefix. Match the colon directly so the
359 # tool call is not silently dropped.
360 (ParserState.TOOL_PREAMBLE, "COLON"): Transition(
361 ParserState.TOOL_NAME,
362 (),
363 ),
364 (ParserState.TOOL_NAME, "OPEN_BRACE"): Transition(
365 ParserState.TOOL_ARGS,
366 (),
367 ),
368 (ParserState.TOOL_NAME, "OPEN_PAREN"): Transition(
369 ParserState.TOOL_ARGS,
370 (),
371 ),
372 (ParserState.TOOL_NAME, "TOOL_END"): Transition(
373 ParserState.CONTENT,
374 (EventType.TOOL_CALL_END,),
375 ),
376 (ParserState.TOOL_ARGS, "TOOL_END"): Transition(
377 ParserState.CONTENT,
378 (EventType.TOOL_CALL_END,),
379 ),
380 # Back-to-back tool calls
381 (ParserState.CONTENT, "TOOL_END"): Transition(
382 ParserState.CONTENT,
383 (),
384 ),
385 # Absorb a bare <channel|> that arrives after we already
386 # returned to CONTENT; prevents leaking it as TEXT_CHUNK.
387 (ParserState.CONTENT, "THINK_END"): Transition(
388 ParserState.CONTENT,
389 (),
390 ),
391 },
392 content_events={
393 ParserState.CONTENT: EventType.TEXT_CHUNK,
394 ParserState.REASONING: EventType.REASONING_CHUNK,
395 ParserState.TOOL_NAME: EventType.TOOL_NAME,
396 ParserState.TOOL_ARGS: EventType.ARG_VALUE_CHUNK,
397 },
398 arg_converter=_gemma4_arg_converter,
399 tool_args_json=False,
400 arg_structural_chars=frozenset(",:{}[]<"),
401 preserve_tokens=frozenset({STRING_DELIM}),
402 )
403
404
405 _GEMMA4_THOUGHT_PREFIX = "thought\n"
406 _GEMMA4_THOUGHT_TOKEN = "thought"
407
408
409 class Gemma4Parser(ParserEngine):
410 """Gemma4 parser: ``<|channel>`` reasoning + ``<|tool_call>``
411 tool calls in a single engine.
412
413 - Strips the ``thought\\n`` prefix from reasoning content
414 - Sets ``skip_special_tokens=False`` so boundary tokens are visible
415 - Detects ``<|tool_call>`` token as implicit reasoning end
416 """
417
418 def __init__(
419 self,
420 tokenizer: TokenizerLike,
421 tools: list[Tool] | None = None,
422 **kwargs,
423 ) -> None:
424 chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
425 thinking = chat_kwargs.get("enable_thinking", False)
426 super().__init__(
427 tokenizer,
428 tools,
429 parser_engine_config=gemma4_config(thinking=thinking),
430 **kwargs,
431 )
432 vocab = self.vocab
433 self._tool_call_token_id: int | None = vocab.get("<|tool_call>")
434 self._new_turn_token_id: int | None = vocab.get("<|turn>")
435 self._tool_response_token_id: int | None = vocab.get("<|tool_response>")
436 self._reasoning_text: str = ""
437 self._prefix_stripped: bool = False
438 self._is_first_feed: bool = True
439
440 def _reset(self, initial_state=None) -> None:
441 super()._reset(initial_state=initial_state)
442 self._reasoning_text = ""
443 self._prefix_stripped = False
444 self._is_first_feed = True
445
446 def _preprocess_feed(
447 self,
448 delta_text: str,
449 delta_token_ids: Sequence[int],
450 ) -> tuple[str, Sequence[int]]:
451 if not self._is_first_feed:
452 return delta_text, delta_token_ids
453 self._is_first_feed = False
454
455 if (
456 not delta_text
457 or self._engine.state != ParserState.CONTENT
458 or self._reasoning_start_token_id is None
459 or self._reasoning_end_token_id is None
460 ):
461 return delta_text, delta_token_ids
462
463 if CHANNEL_START in delta_text:
464 return delta_text, delta_token_ids
465
466 needs_injection = (
467 CHANNEL_END in delta_text
468 or delta_text.startswith(_GEMMA4_THOUGHT_PREFIX)
469 or delta_text == _GEMMA4_THOUGHT_TOKEN
470 )
471 if not needs_injection:
472 return delta_text, delta_token_ids
473
474 delta_text = CHANNEL_START + delta_text
475 if delta_token_ids:
476 delta_token_ids = [self._reasoning_start_token_id, *delta_token_ids]
477
478 return delta_text, delta_token_ids
479
480 def _prompt_ends_in_open_reasoning(self, prompt_token_ids: Sequence[int]) -> bool:
481 """Whether the prompt tail is inside an open ``<|channel>`` block.
482
483 Scans backwards: a ``<|channel>`` start token seen before any
484 closing or turn-boundary token means the block is still open.
485 """
486 start_id = self._reasoning_start_token_id
487 if start_id is None:
488 return False
489 boundary_ids = {
490 tid
491 for tid in (
492 self._reasoning_end_token_id,
493 self._tool_call_token_id,
494 self._new_turn_token_id,
495 self._tool_response_token_id,
496 )
497 if tid is not None
498 }
499 for tid in reversed(prompt_token_ids):
500 if tid == start_id:
501 return True
502 if tid in boundary_ids:
503 return False
504 return False
505
506 def adjust_initial_state_from_prompt(self, prompt_token_ids: Sequence[int]) -> None:
507 """Pre-initialise the engine to ``REASONING`` when the prompt ends
508 inside an open ``<|channel>`` block.
509
510 This covers the post-tool-response continuation case where the chat
511 template leaves the prompt ending with ``<|channel>thought\n``
512 (issue #45834). A prompt that merely starts a new model turn must
513 not pre-initialise reasoning: the model may answer directly without
514 emitting any channel markers, and the non-streaming path classifies
515 such output as content (issue #48217). When the model does open its
516 own ``<|channel>``, the ``(CONTENT, THINK_START)`` transition
517 handles it, and the ``thought\n`` prefix in the first reasoning
518 chunk is stripped by ``_events_to_delta`` as in the default flow.
519 """
520 if not self._prompt_ends_in_open_reasoning(prompt_token_ids):
521 return
522 self._engine.reset(initial_state=ParserState.REASONING)
523 # Prevent a later default ``initialize_streaming()`` (e.g. from
524 # ``ParserEngineReasoningAdapter.extract_reasoning_streaming``) from
525 # clobbering this with ``CONTENT``.
526 self._streaming_initialized = True
527
528 def _events_to_delta(
529 self,
530 events: list[SemanticEvent],
531 finished: bool = False,
532 ) -> DeltaMessage | None:
533 delta = super()._events_to_delta(events, finished=finished)
534 if delta is None or delta.reasoning is None:
535 return delta
536
537 if self._prefix_stripped:
538 return delta
539 self._reasoning_text += delta.reasoning
540
541 if self._reasoning_text.startswith(_GEMMA4_THOUGHT_PREFIX):
542 prefix_len = len(_GEMMA4_THOUGHT_PREFIX)
543 prev_reasoning_len = len(self._reasoning_text) - len(delta.reasoning)
544 if prev_reasoning_len >= prefix_len:
545 self._prefix_stripped = True
546 return delta
547 chars_of_prefix_in_delta = prefix_len - prev_reasoning_len
548 stripped = delta.reasoning[chars_of_prefix_in_delta:]
549 if stripped:
550 self._prefix_stripped = True
551 delta.reasoning = stripped
552 return delta
553 if len(self._reasoning_text) >= prefix_len:
554 self._prefix_stripped = True
555 delta.reasoning = None
556 if delta.content is not None or delta.tool_calls:
557 return delta
558 return None
559 return None
560
561 if _GEMMA4_THOUGHT_PREFIX.startswith(self._reasoning_text):
562 if finished:
563 self._prefix_stripped = True
564 return None
565
566 self._prefix_stripped = True
567 delta.reasoning = self._reasoning_text
568 return delta
569
570 def extract_reasoning(
571 self,
572 model_output: str,
573 request: ChatCompletionRequest | ResponsesRequest,
574 ) -> tuple[str | None, str | None]:
575 reasoning, content = super().extract_reasoning(model_output, request)
576 if reasoning:
577 if reasoning.startswith(_GEMMA4_THOUGHT_PREFIX):
578 reasoning = reasoning[len(_GEMMA4_THOUGHT_PREFIX) :]
579 elif reasoning == _GEMMA4_THOUGHT_PREFIX.rstrip():
580 reasoning = None
581 return reasoning or None, content