| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import dataclasses as dt |
| 5 | import enum |
| 6 | from collections.abc import Sequence |
| 7 | from typing import TYPE_CHECKING |
| 8 | |
| 9 | import regex as re |
| 10 | |
| 11 | from vllm.entrypoints.generate.base.protocol import DeltaMessage |
| 12 | from vllm.reasoning import ReasoningParser |
| 13 | |
| 14 | if TYPE_CHECKING: |
| 15 | from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest |
| 16 | from vllm.entrypoints.openai.responses.protocol import ResponsesRequest |
| 17 | from vllm.tokenizers import TokenizerLike |
| 18 | |
| 19 | |
| 20 | class Olmo3ReasoningState(enum.Enum): |
| 21 | REASONING = 1 |
| 22 | CONTENT = 2 |
| 23 | |
| 24 | |
| 25 | @dt.dataclass(frozen=True) |
| 26 | class Indices: |
| 27 | start: int |
| 28 | end: int |
| 29 | |
| 30 | def __len__(self): |
| 31 | return self.end - self.start |
| 32 | |
| 33 | |
| 34 | def string_overlap(a: str, b: str) -> tuple[Indices | None, Indices | None]: |
| 35 | """Find the longest overlap where the end of string a matches the start |
| 36 | of string b. |
| 37 | |
| 38 | Args: |
| 39 | a: First string |
| 40 | b: Second string |
| 41 | |
| 42 | Returns: |
| 43 | Tuple of IndicesTuples representing the overlapping portions in each |
| 44 | string, or a tuple of None if no overlap exists |
| 45 | |
| 46 | """ |
| 47 | # swap so a is always the shorter string |
| 48 | a, b, swap = (a, b, False) if len(a) < len(b) else (b, a, True) |
| 49 | |
| 50 | # first check: is a fully contained in b? |
| 51 | if a in b: |
| 52 | ind_a = Indices(0, len(a)) |
| 53 | ind_b = Indices(b.index(a), b.index(a) + len(a)) |
| 54 | return (ind_b, ind_a) if swap else (ind_a, ind_b) |
| 55 | |
| 56 | # second check: does the end of a overlap with the |
| 57 | # beginning of b? |
| 58 | for i in range(len(a) - 1, 0, -1): |
| 59 | if a[-i:] == b[:i]: |
| 60 | ind_a = Indices(len(a) - i, len(a)) |
| 61 | ind_b = Indices(0, i) |
| 62 | return (ind_b, ind_a) if swap else (ind_a, ind_b) |
| 63 | |
| 64 | # third check: does the beginning of a overlap with |
| 65 | # the end of b? |
| 66 | for i in range(len(a) - 1, 0, -1): |
| 67 | if b[-i:] == a[:i]: |
| 68 | ind_a = Indices(0, i) |
| 69 | ind_b = Indices(len(b) - i, len(b)) |
| 70 | return (ind_b, ind_a) if swap else (ind_a, ind_b) |
| 71 | |
| 72 | return None, None |
| 73 | |
| 74 | |
| 75 | @dt.dataclass |
| 76 | class Olmo3ReasoningBuffer: |
| 77 | think_start: str = "<think>" |
| 78 | think_end: str = "</think>" |
| 79 | buffer: str = "" |
| 80 | |
| 81 | # we start in reasoning state to support cases where we hardcode |
| 82 | # <think> as the start of the reasoning block. |
| 83 | # In those cases, the only token we will see is </think>, which |
| 84 | # is when we switch to content state. |
| 85 | state: Olmo3ReasoningState = Olmo3ReasoningState.REASONING |
| 86 | |
| 87 | def process_buffer(self) -> DeltaMessage | None: |
| 88 | start_think_idx = self.buffer.find(self.think_start) |
| 89 | |
| 90 | if start_think_idx >= 0: |
| 91 | self.state = Olmo3ReasoningState.REASONING |
| 92 | pretext, self.buffer = ( |
| 93 | self.buffer[:start_think_idx], |
| 94 | self.buffer[start_think_idx + len(self.think_start) :], |
| 95 | ) |
| 96 | if start_think_idx > 0: |
| 97 | # this covers the case there's content before |
| 98 | # the start of the reasoning block |
| 99 | return DeltaMessage(content=pretext) |
| 100 | |
| 101 | end_think_idx = self.buffer.rfind(self.think_end) |
| 102 | |
| 103 | if end_think_idx >= 0: |
| 104 | self.state = Olmo3ReasoningState.CONTENT |
| 105 | pretext, self.buffer = ( |
| 106 | self.buffer[:end_think_idx], |
| 107 | self.buffer[end_think_idx + len(self.think_end) :], |
| 108 | ) |
| 109 | if end_think_idx > 0: |
| 110 | # this covers the case there's content before |
| 111 | # the end of the reasoning block |
| 112 | return DeltaMessage(reasoning=pretext) |
| 113 | |
| 114 | if self.state == Olmo3ReasoningState.REASONING: |
| 115 | # we are inside reasoning block, return and empty |
| 116 | # the text buffer |
| 117 | ( |
| 118 | text_buffer, |
| 119 | self.buffer, |
| 120 | ) = self.buffer, "" |
| 121 | return DeltaMessage(reasoning=text_buffer) |
| 122 | |
| 123 | if self.state == Olmo3ReasoningState.CONTENT: |
| 124 | # we are outside reasoning block, return and empty |
| 125 | # the text buffer |
| 126 | ( |
| 127 | text_buffer, |
| 128 | self.buffer, |
| 129 | ) = self.buffer, "" |
| 130 | return DeltaMessage(content=text_buffer) |
| 131 | |
| 132 | # nothing to return unless we are in reasoning or content state |
| 133 | return None |
| 134 | |
| 135 | def __len__(self): |
| 136 | # is the length of the text buffer |
| 137 | return len(self.buffer) |
| 138 | |
| 139 | def add_text(self, delta_text: str) -> DeltaMessage | None: |
| 140 | # we start by adding the delta text to the buffer |
| 141 | self.buffer += delta_text |
| 142 | |
| 143 | # setting this to empty before starting |
| 144 | delta_message: DeltaMessage | None = None |
| 145 | |
| 146 | # we start by computing the overlap between the delta_text |
| 147 | # and start/end of think tokens. |
| 148 | _, overlap_think_start = string_overlap(delta_text, self.think_start) |
| 149 | _, overlap_think_end = string_overlap(delta_text, self.think_end) |
| 150 | |
| 151 | partial_overlap_start = overlap_think_start is not None and len( |
| 152 | overlap_think_start |
| 153 | ) < len(self.think_start) |
| 154 | partial_overlap_end = overlap_think_end is not None and len( |
| 155 | overlap_think_end |
| 156 | ) < len(self.think_end) |
| 157 | |
| 158 | if ( |
| 159 | partial_overlap_start |
| 160 | and self.think_start in self.buffer |
| 161 | and not partial_overlap_end |
| 162 | ): |
| 163 | # we can only process the buffer if partial overlap |
| 164 | # is the last part of think token (thus causing |
| 165 | # text_buffer to contain the start of think token) |
| 166 | # and there are no partial overlaps with end think |
| 167 | delta_message = self.process_buffer() |
| 168 | |
| 169 | elif partial_overlap_end and self.think_end in self.buffer: |
| 170 | # same as before (partial overlap only allowed) |
| 171 | # if the buffer contains the end think token, |
| 172 | # but we don't have to check for partial overlap |
| 173 | # with start think token because they are handled |
| 174 | # by the previous condition |
| 175 | delta_message = self.process_buffer() |
| 176 | |
| 177 | elif partial_overlap_start or partial_overlap_end: |
| 178 | # in general, if there are overlaps, we don't |
| 179 | # process the buffer because we want to wait until |
| 180 | # the think token is fully completed. |
| 181 | return None |
| 182 | else: |
| 183 | # we process the buffer as normal |
| 184 | delta_message = self.process_buffer() |
| 185 | |
| 186 | return delta_message |
| 187 | |
| 188 | |
| 189 | class Olmo3ReasoningParser(ReasoningParser): |
| 190 | """Reasoning parser for Olmo 3 model. |
| 191 | |
| 192 | Olmo3ReasoningParser |
| 193 | |
| 194 | This class implements a reasoning parser specifically designed for the |
| 195 | Olmo 3 family of models. Olmo 3 models do not use special tokens to |
| 196 | indicate reasoning; rather, reasoning trace is wrapped in `<think>` and |
| 197 | `</think>`, which are tokenized using standard vocabulary entries. |
| 198 | Because of this, the parser operates in string space, accumulating the |
| 199 | characters in a buffer until it sees `<think>` or `</think>`. tokens |
| 200 | to switch modes. |
| 201 | |
| 202 | Key Features: |
| 203 | - For non-stream output, Recognizes and extracts reasoning (text |
| 204 | bracketed by `<think>` and `</think>`) and content (everything |
| 205 | after the first `</think>`). |
| 206 | - For stream process, it uses a buffer to accumulate delta text, |
| 207 | and output progressive delta messages as soon as thinking starts |
| 208 | or ends. |
| 209 | - For reliability, some Olmo 3 models may hardcode the first |
| 210 | `<think>` token is the input text (similar to Deepseek R1, |
| 211 | or reasoning-only Qwen models). To support such variants, the |
| 212 | parser can optionally work in cases where the first `<think>` |
| 213 | token is missing from generation. |
| 214 | """ |
| 215 | |
| 216 | think_start: str = r"<think>" |
| 217 | think_end: str = r"</think>" |
| 218 | # </think> is split in 3 by the pre-tokenizer, first split can be tokenized |
| 219 | # with an optional leading space, so there are 2 possible tokenizations |
| 220 | think_end_first_split: list[str] = [r"Ġ</", r"</"] |
| 221 | think_end_rest_split: list[str] = [r"think", r">"] |
| 222 | # notice that the first think is optional; this allows template to |
| 223 | # work in cases when we hardcode a <think> at the beginning of the |
| 224 | # reasoning template. |
| 225 | reasoning_regex: re.Pattern = re.compile( |
| 226 | rf"^(?:{think_start})?(?P<reasoning>.*?)" |
| 227 | rf"{think_end}(?P<content>.*)$", |
| 228 | re.DOTALL, |
| 229 | ) |
| 230 | |
| 231 | def __init__(self, tokenizer: "TokenizerLike", *args, **kwargs): |
| 232 | super().__init__(tokenizer, *args, **kwargs) |
| 233 | self.buffer = Olmo3ReasoningBuffer( |
| 234 | think_start=self.think_start, think_end=self.think_end |
| 235 | ) |
| 236 | self.think_end_first_token_ids: list[int] = [ |
| 237 | self.vocab[token] for token in self.think_end_first_split |
| 238 | ] |
| 239 | self.think_end_rest_token_ids: list[int] = [ |
| 240 | self.vocab[token] for token in self.think_end_rest_split |
| 241 | ] |
| 242 | |
| 243 | @property |
| 244 | def reasoning_start_str(self) -> str: |
| 245 | return self.think_start |
| 246 | |
| 247 | @property |
| 248 | def reasoning_end_str(self) -> str: |
| 249 | return self.think_end |
| 250 | |
| 251 | def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: |
| 252 | rest_ids = self.think_end_rest_token_ids |
| 253 | rest_len = len(rest_ids) |
| 254 | for i in range(len(input_ids) - rest_len, -1, -1): |
| 255 | if ( |
| 256 | list(input_ids[i + 1 : i + 1 + rest_len]) == rest_ids |
| 257 | and input_ids[i] in self.think_end_first_token_ids |
| 258 | ): |
| 259 | return True |
| 260 | return False |
| 261 | |
| 262 | def extract_content_ids(self, input_ids: list[int]) -> list[int]: |
| 263 | # for Olmo 3 streaming reason parsing, the stream parse |
| 264 | # will call first, and the same token will be called in |
| 265 | # is_reasoning_end and extract_content_ids |
| 266 | # this id is not part of content, so just return [] here. |
| 267 | return [] |
| 268 | |
| 269 | def extract_reasoning( |
| 270 | self, |
| 271 | model_output: str, |
| 272 | request: "ChatCompletionRequest | ResponsesRequest", |
| 273 | ) -> tuple[str | None, str | None]: |
| 274 | """Extract the reasoning content & content sections, respectively. |
| 275 | If the sequence doesn't match what we expect, i.e., the model generates |
| 276 | something else, all content is considered non-reasoning content. |
| 277 | |
| 278 | Args: |
| 279 | model_output: Output of the model to be parsed. |
| 280 | request: Request being |
| 281 | processed. |
| 282 | |
| 283 | Returns: |
| 284 | tuple[Optional[str], Optional[str]]: Tuple pair containing the |
| 285 | reasoning content and non-reasoning content. |
| 286 | |
| 287 | """ |
| 288 | re_match = self.reasoning_regex.match(model_output) |
| 289 | if re_match: |
| 290 | reasoning = re_match.group("reasoning") or None |
| 291 | content = re_match.group("content") or None |
| 292 | return reasoning, content |
| 293 | |
| 294 | # no reasoning content |
| 295 | return None, model_output |
| 296 | |
| 297 | def extract_reasoning_streaming( |
| 298 | self, |
| 299 | previous_text: str, |
| 300 | current_text: str, |
| 301 | delta_text: str, |
| 302 | previous_token_ids: Sequence[int], |
| 303 | current_token_ids: Sequence[int], |
| 304 | delta_token_ids: Sequence[int], |
| 305 | ) -> DeltaMessage | None: |
| 306 | """Extract content using token ID sequence state machine.""" |
| 307 | delta_message = self.buffer.add_text(delta_text) |
| 308 | if delta_message is None and self.buffer.think_end in self.buffer.buffer: |
| 309 | # this is a bit hacky, but, because of how the buffer is |
| 310 | # constructed, if the last delta_text contains characters that |
| 311 | # marks the end of thinking tokens, then messages in the buffer |
| 312 | # would never be processed because we get no other turn. To get |
| 313 | # around that, we check if the text buffer contains the end of |
| 314 | # thinking tokens, and, if so, we reprocess the buffer again. |
| 315 | delta_message = self.buffer.process_buffer() |
| 316 | |
| 317 | return delta_message |