| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Reasoning-content parser for MuseGlimmer. |
| 4 | Port of the ``reasoning_content`` rule from the HuggingFace MuseGlimmer |
| 5 | ``MUSE_GLIMMER_RESPONSE_SCHEMA`` (synced with internal master). MuseGlimmer emits |
| 6 | chain-of-thought in ``to=self`` channels delimited by ``<|message|>`` ... ``<|eom|>``: |
| 7 | to=self<|message|>...reasoning...<|eom|> |
| 8 | A turn may contain several ``to=self`` blocks interleaved with tool calls, and a |
| 9 | tool call or final answer follows in its own channel. |
| 10 | Because MuseGlimmer's framing markers (``<|message|>``, ``<|eom|>``) are not guaranteed |
| 11 | to be single vocab tokens across every checkpoint's tokenizer, this parser works |
| 12 | on the decoded text with regexes rather than the single start/end-token base class. |
| 13 | Usage: ``--reasoning-parser muse_glimmer`` |
| 14 | """ |
| 15 | |
| 16 | from __future__ import annotations |
| 17 | |
| 18 | from collections.abc import Iterable, Sequence |
| 19 | |
| 20 | import regex as re |
| 21 | |
| 22 | from vllm.entrypoints.generate.base.protocol import DeltaMessage |
| 23 | from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest |
| 24 | from vllm.entrypoints.openai.responses.protocol import ResponsesRequest |
| 25 | from vllm.reasoning.abs_reasoning_parsers import ReasoningParser |
| 26 | |
| 27 | _EOM = "<|eom|>" |
| 28 | _EOT = "<|eot|>" |
| 29 | _FUNCTION_CALLS_OPEN = "<atem:function_calls>" |
| 30 | _REASONING_OPEN = "to=self<|message|>" |
| 31 | _ASSISTANT_TURN_OPEN = "<|start|>assistant" |
| 32 | # A channel header: ``to=<recipient><|message|>`` where recipient is ``self`` |
| 33 | # (reasoning), ``user`` (final answer) or ``<tool>[.<fn>]`` (tool call). |
| 34 | _CHANNEL_HEADER_RE = re.compile(r"to=(?P<recipient>[^\s<]+)<\|message\|>") |
| 35 | _HEADER_PAT = r"to=[^\s<]+<\|message\|>" |
| 36 | # Collapse the gap between reasoning blocks so multiple to=self spans join. |
| 37 | _COLLAPSE_RE = re.compile( |
| 38 | r"<\|eom\|>(?:(?!to=self<\|message\|>).)*?to=self<\|message\|>", re.DOTALL |
| 39 | ) |
| 40 | _REASONING_RE = re.compile(r"to=self<\|message\|>(.*?)<\|eom\|>", re.DOTALL) |
| 41 | _CONTENT_RE = re.compile( |
| 42 | r"to=user<\|message\|>(.*?)(?=<\|eot\|>|<\|eom\|>|$)", re.DOTALL |
| 43 | ) |
| 44 | # Strip a CLOSED reasoning span (header .. <|eom|>). |
| 45 | _STRIP_REASONING_RE = re.compile( |
| 46 | r"(?:<\|start\|>assistant\s*)?to=self<\|message\|>.*?<\|eom\|>", re.DOTALL |
| 47 | ) |
| 48 | # An UNTERMINATED trailing reasoning span. The model sometimes leaves the |
| 49 | # analysis channel WITHOUT emitting <|eom|>, writing a bare |
| 50 | # ``to=<tool><|message|>`` header instead (observed deterministically for a call |
| 51 | # with EMPTY arguments on a tool that has optional parameters; reproduced on |
| 52 | # other engines too, so it is a model-side defect, not engine-specific). |
| 53 | # |
| 54 | # These two patterns MUST therefore stop at the next channel header rather than |
| 55 | # running to end-of-text. An unbounded ``...$`` version consumes the real tool |
| 56 | # call along with the reasoning: `is_reasoning_end` then never fires, the parser |
| 57 | # never leaves the reasoning phase, the tool parser is never invoked, and the |
| 58 | # entire generation is dropped (empty reasoning, empty content, no tool call). |
| 59 | _STRIP_OPEN_REASONING_RE = re.compile( |
| 60 | r"(?:<\|start\|>assistant\s*)?to=self<\|message\|>" |
| 61 | r"(?:(?!<\|eom\|>)(?!" + _HEADER_PAT + r").)*" |
| 62 | r"(?=" + _HEADER_PAT + r"|$)", |
| 63 | re.DOTALL, |
| 64 | ) |
| 65 | _OPEN_REASONING_RE = re.compile( |
| 66 | r"to=self<\|message\|>((?:(?!<\|eom\|>)(?!" + _HEADER_PAT + r").)*)" |
| 67 | r"(?=" + _HEADER_PAT + r"|$)", |
| 68 | re.DOTALL, |
| 69 | ) |
| 70 | # Markers whose PREFIX could appear at the tail of an OPEN (still-streaming) body. |
| 71 | _HOLDBACK_MARKERS = (_EOM, _EOT, "<|start|>", "<|message|>") |
| 72 | # A trailing fragment that could still grow into a channel header (" t", " to", |
| 73 | # " to=", " to=skill"). Without this the recipient name leaks into reasoning and |
| 74 | # then has to be un-emitted once ``<|message|>`` arrives. |
| 75 | _OPEN_TAIL_HEADER_RE = re.compile(r"[\s](?:t|to|to=[^\s<]*)$") |
| 76 | |
| 77 | |
| 78 | def _current_assistant_turn(text: str) -> str: |
| 79 | """Return only the text generated in the current assistant turn. |
| 80 | ``is_reasoning_end`` is evaluated on the PROMPT token-ids at stream start, |
| 81 | and an MuseGlimmer prompt legitimately contains ATEM markers (``render_tool_defs`` |
| 82 | writes a literal ``<atem:function_calls>`` example into the system message, |
| 83 | and prior assistant turns may carry real tool calls). Anchoring on the last |
| 84 | channel-open keeps prompt text from deciding the phase. |
| 85 | """ |
| 86 | idx = text.rfind(_ASSISTANT_TURN_OPEN) |
| 87 | return text[idx + len(_ASSISTANT_TURN_OPEN) :] if idx != -1 else text |
| 88 | |
| 89 | |
| 90 | def _trim_open_body(body: str) -> str: |
| 91 | """Hold back any tail of a still-growing body that could still be framing. |
| 92 | Iterated to a fixpoint because the two cases compose: `" to=skill<"` needs |
| 93 | the partial-marker trim (``<``) before the partial-header trim can see |
| 94 | `" to=skill"`. Trimming only once leaks the recipient name as reasoning. |
| 95 | """ |
| 96 | while True: |
| 97 | trimmed = body |
| 98 | for marker in _HOLDBACK_MARKERS: |
| 99 | for k in range(min(len(marker) - 1, len(trimmed)), 0, -1): |
| 100 | if trimmed.endswith(marker[:k]): |
| 101 | trimmed = trimmed[:-k] |
| 102 | break |
| 103 | else: |
| 104 | continue |
| 105 | break |
| 106 | header_tail = _OPEN_TAIL_HEADER_RE.search(trimmed) |
| 107 | if header_tail is not None: |
| 108 | trimmed = trimmed[: header_tail.start()] |
| 109 | if trimmed == body: |
| 110 | return body |
| 111 | body = trimmed |
| 112 | |
| 113 | |
| 114 | class MuseGlimmerReasoningParser(ReasoningParser): |
| 115 | def __init__(self, tokenizer, *args, **kwargs) -> None: |
| 116 | super().__init__(tokenizer, *args, **kwargs) |
| 117 | # Cursors over what was ACTUALLY emitted. Diffing a freshly reclassified |
| 118 | # `previous_text` is unsafe: a classified body legitimately shrinks when |
| 119 | # a partial header becomes recognisable, and diffing against the shrunken |
| 120 | # value re-emits text that already went out. |
| 121 | self._emitted_reasoning: str = "" |
| 122 | self._emitted_content: str = "" |
| 123 | self._tool_handoff_done: bool = False |
| 124 | |
| 125 | def adjust_request( |
| 126 | self, request: ChatCompletionRequest | ResponsesRequest |
| 127 | ) -> ChatCompletionRequest | ResponsesRequest: |
| 128 | """Preserve MuseGlimmer's ATEM framing tokens in the decoded output. |
| 129 | vLLM's serving default is ``skip_special_tokens=True``, which strips |
| 130 | ``<|start|>`` / ``<|message|>`` / ``<|eom|>`` / ``<|eot|>`` before the |
| 131 | parsers run, collapsing reasoning into content and breaking channel |
| 132 | scoping. Unlike the base tool-parser hook we do NOT touch |
| 133 | ``structured_outputs`` -- MuseGlimmer emits native ATEM, not JSON. |
| 134 | """ |
| 135 | request.skip_special_tokens = False |
| 136 | return request |
| 137 | |
| 138 | def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: |
| 139 | """Whether the model has left reasoning and opened a TOOL channel. |
| 140 | A ``to=user`` answer is NOT a reason to leave the reasoning phase -- this |
| 141 | parser surfaces that content itself. Only a real tool channel switches |
| 142 | the ``DelegatingParser`` phase machine over to the tool parser. |
| 143 | Both closed and unterminated reasoning spans are stripped before the |
| 144 | check, so an ``<atem:invoke>`` the model merely echoes inside its CoT |
| 145 | never flips the phase. |
| 146 | """ |
| 147 | try: |
| 148 | text = self.model_tokenizer.decode(input_ids) |
| 149 | except Exception: |
| 150 | return False |
| 151 | remainder = self._tool_channel_remainder(text) |
| 152 | return _FUNCTION_CALLS_OPEN in remainder or "<atem:invoke" in remainder |
| 153 | |
| 154 | def is_reasoning_end_streaming( |
| 155 | self, input_ids: Sequence[int], delta_ids: Iterable[int] |
| 156 | ) -> bool: |
| 157 | return self.is_reasoning_end(list(input_ids)) |
| 158 | |
| 159 | def extract_content_ids(self, input_ids: list[int]) -> list[int]: |
| 160 | # Content-id slicing is unreliable for multi-token markers; the serving |
| 161 | # path uses extract_reasoning() for the final split. |
| 162 | return [] |
| 163 | |
| 164 | @classmethod |
| 165 | def _scoped_turn(cls, text: str) -> str: |
| 166 | """Current assistant turn with reasoning spans removed.""" |
| 167 | scoped = _current_assistant_turn(text) |
| 168 | scoped = _STRIP_REASONING_RE.sub("", scoped) |
| 169 | return _STRIP_OPEN_REASONING_RE.sub("", scoped) |
| 170 | |
| 171 | @classmethod |
| 172 | def _tool_channel_remainder(cls, text: str) -> str: |
| 173 | """Text from the first tool-channel header onward, framing INCLUDED. |
| 174 | ``DelegatingParser.parse_delta`` rebuilds ``current_text`` from whatever |
| 175 | this parser returns as ``.content`` on the transition delta and commits |
| 176 | it; anything not returned is destroyed. It must start AT the |
| 177 | ``to=<name><|message|>`` header -- handing over the text after the header |
| 178 | loses the recipient, and the tool parser then sees a bare ``<|message|>``, |
| 179 | classifies it as the content channel, and leaks the ATEM markup. |
| 180 | """ |
| 181 | scoped = cls._scoped_turn(text) |
| 182 | for match in _CHANNEL_HEADER_RE.finditer(scoped): |
| 183 | if match.group("recipient") not in ("self", "user"): |
| 184 | return scoped[match.start() :] |
| 185 | return "" |
| 186 | |
| 187 | @staticmethod |
| 188 | def _classify_bodies(text: str) -> tuple[str, str]: |
| 189 | """Split ``text`` into (reasoning_body, content_body), channel-aware. |
| 190 | Framing markers and tool channels contribute nothing -- the tool parser |
| 191 | owns those. A body ends at ``<|eom|>`` / ``<|eot|>``, at the next channel |
| 192 | header, or at end-of-text (an OPEN body, which is held back). |
| 193 | """ |
| 194 | reasoning_parts: list[str] = [] |
| 195 | content_parts: list[str] = [] |
| 196 | pos = 0 |
| 197 | n = len(text) |
| 198 | while pos < n: |
| 199 | match = _CHANNEL_HEADER_RE.search(text, pos) |
| 200 | if not match: |
| 201 | break |
| 202 | recipient = match.group("recipient") |
| 203 | body_start = match.end() |
| 204 | eom = text.find(_EOM, body_start) |
| 205 | eot = text.find(_EOT, body_start) |
| 206 | terminators = [p for p in (eom, eot) if p != -1] |
| 207 | next_header = _CHANNEL_HEADER_RE.search(text, body_start) |
| 208 | if next_header is not None: |
| 209 | terminators.append(next_header.start()) |
| 210 | body_end = min(terminators) if terminators else n |
| 211 | body = text[body_start:body_end] |
| 212 | if not terminators: |
| 213 | body = _trim_open_body(body) |
| 214 | if recipient == "self": |
| 215 | reasoning_parts.append(body) |
| 216 | elif ( |
| 217 | # Never surface tool XML echoed into a user channel. |
| 218 | recipient == "user" |
| 219 | and _FUNCTION_CALLS_OPEN not in body |
| 220 | and "<atem:invoke" not in body |
| 221 | ): |
| 222 | content_parts.append(body) |
| 223 | if terminators and body_end in (eom, eot): |
| 224 | pos = body_end + len(_EOM if body_end == eom else _EOT) |
| 225 | else: |
| 226 | pos = body_end |
| 227 | return "".join(reasoning_parts), "".join(content_parts) |
| 228 | |
| 229 | def get_streaming_fallback_content( |
| 230 | self, |
| 231 | previous_text: str, |
| 232 | request: ChatCompletionRequest | ResponsesRequest, |
| 233 | ) -> str | None: |
| 234 | """Promote un-surfaced content when the stream ends mid-reasoning. |
| 235 | ``DelegatingParser.finalize_generation`` calls this when |
| 236 | ``reasoning_ended`` is still False. Returns only the channel-classified |
| 237 | ``to=user`` body, and only the portion not already streamed. |
| 238 | """ |
| 239 | _, content_body = self._classify_bodies(previous_text) |
| 240 | remainder = content_body[len(self._emitted_content) :] |
| 241 | return remainder or None |
| 242 | |
| 243 | def extract_reasoning( |
| 244 | self, |
| 245 | model_output: str, |
| 246 | request: ChatCompletionRequest | ResponsesRequest, |
| 247 | ) -> tuple[str | None, str | None]: |
| 248 | collapsed = _COLLAPSE_RE.sub("\n", model_output) |
| 249 | matches = _REASONING_RE.findall(collapsed) |
| 250 | reasoning = "\n".join(matches) if matches else None |
| 251 | # Truncation fallback: generation stopped inside a to=self block, so |
| 252 | # there is no closing <|eom|>. Bounded at the next channel header so a |
| 253 | # real tool call that follows a header-less channel switch is not |
| 254 | # absorbed into the reasoning field. |
| 255 | open_match = _OPEN_REASONING_RE.search(model_output) |
| 256 | if open_match and open_match.group(1): |
| 257 | partial = open_match.group(1) |
| 258 | reasoning = f"{reasoning}\n{partial}" if reasoning else partial |
| 259 | # Content is everything that is not a reasoning block. In a |
| 260 | # reasoning+tool-call turn there is no to=user answer, but the tool |
| 261 | # channels MUST be forwarded -- the unified parser runs the tool parser |
| 262 | # on this returned `content`, not on the original model_output. |
| 263 | remainder = _STRIP_REASONING_RE.sub("", model_output) |
| 264 | remainder = _STRIP_OPEN_REASONING_RE.sub("", remainder) |
| 265 | if "<atem:invoke" in remainder or _FUNCTION_CALLS_OPEN in remainder: |
| 266 | return reasoning, (remainder or None) |
| 267 | content_match = _CONTENT_RE.search(model_output) |
| 268 | if content_match: |
| 269 | content = content_match.group(1) or None |
| 270 | elif _REASONING_OPEN in model_output: |
| 271 | content = None |
| 272 | else: |
| 273 | content = model_output or None |
| 274 | reasoning = None |
| 275 | return reasoning, content |
| 276 | |
| 277 | def extract_reasoning_streaming( |
| 278 | self, |
| 279 | previous_text: str, |
| 280 | current_text: str, |
| 281 | delta_text: str, |
| 282 | previous_token_ids: Sequence[int], |
| 283 | current_token_ids: Sequence[int], |
| 284 | delta_token_ids: Sequence[int], |
| 285 | ) -> DeltaMessage | None: |
| 286 | """Channel-aware streaming split of reasoning vs content. |
| 287 | Classifies the full ``current_text`` and emits only what has not been |
| 288 | emitted yet, so no framing token is ever surfaced and a delta straddling |
| 289 | a channel boundary only contributes the portion inside a real body. |
| 290 | """ |
| 291 | curr_reason, curr_content = self._classify_bodies(current_text) |
| 292 | reasoning_delta = "" |
| 293 | if curr_reason.startswith(self._emitted_reasoning) and len(curr_reason) > len( |
| 294 | self._emitted_reasoning |
| 295 | ): |
| 296 | reasoning_delta = curr_reason[len(self._emitted_reasoning) :] |
| 297 | self._emitted_reasoning = curr_reason |
| 298 | content_delta = "" |
| 299 | if curr_content.startswith(self._emitted_content) and len(curr_content) > len( |
| 300 | self._emitted_content |
| 301 | ): |
| 302 | content_delta = curr_content[len(self._emitted_content) :] |
| 303 | self._emitted_content = curr_content |
| 304 | # Hand the tool channel to the tool parser exactly once, starting at its |
| 305 | # header. parse_delta discards anything not returned here. |
| 306 | # |
| 307 | # This MUST fire on the same delta where is_reasoning_end() flips, i.e. |
| 308 | # only once the tool channel actually contains ATEM. Emitting it earlier |
| 309 | # -- when only the bare `to=<name><|message|>` header has arrived -- keeps |
| 310 | # the parser in the reasoning phase, so parse_delta never replaces this |
| 311 | # DeltaMessage with the tool parser's and the header is delivered to the |
| 312 | # client as visible content. |
| 313 | handoff = "" |
| 314 | if not self._tool_handoff_done: |
| 315 | remainder = self._tool_channel_remainder(current_text) |
| 316 | if _FUNCTION_CALLS_OPEN in remainder or "<atem:invoke" in remainder: |
| 317 | handoff = remainder |
| 318 | self._tool_handoff_done = True |
| 319 | if handoff: |
| 320 | return DeltaMessage(reasoning=reasoning_delta or None, content=handoff) |
| 321 | if reasoning_delta and content_delta: |
| 322 | return DeltaMessage(reasoning=reasoning_delta, content=content_delta) |
| 323 | if reasoning_delta: |
| 324 | return DeltaMessage(reasoning=reasoning_delta) |
| 325 | if content_delta: |
| 326 | return DeltaMessage(content=content_delta) |
| 327 | return None |