| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import string |
| 5 | from collections.abc import Sequence |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest |
| 10 | from vllm.reasoning import ReasoningParserManager |
| 11 | from vllm.reasoning.minimax_m3_reasoning_parser import MiniMaxM3ReasoningParser |
| 12 | |
| 13 | pytestmark = pytest.mark.skip_global_cleanup |
| 14 | |
| 15 | |
| 16 | class MiniMaxM3Tokenizer: |
| 17 | """Small tokenizer with MiniMax M3 reasoning tags as special tokens.""" |
| 18 | |
| 19 | special_tokens = ("<mm:think>", "</mm:think>") |
| 20 | |
| 21 | def __init__(self): |
| 22 | self._token_to_id: dict[str, int] = {} |
| 23 | self._id_to_token: dict[int, str] = {} |
| 24 | for token in self.special_tokens: |
| 25 | self._add_token(token) |
| 26 | for char in string.printable: |
| 27 | self._add_token(char) |
| 28 | |
| 29 | def _add_token(self, token: str) -> int: |
| 30 | token_id = self._token_to_id.get(token) |
| 31 | if token_id is None: |
| 32 | token_id = len(self._token_to_id) + 1 |
| 33 | self._token_to_id[token] = token_id |
| 34 | self._id_to_token[token_id] = token |
| 35 | return token_id |
| 36 | |
| 37 | def get_vocab(self) -> dict[str, int]: |
| 38 | return dict(self._token_to_id) |
| 39 | |
| 40 | def encode( |
| 41 | self, |
| 42 | text: str, |
| 43 | truncation: bool | None = None, |
| 44 | max_length: int | None = None, |
| 45 | add_special_tokens: bool = True, |
| 46 | ) -> list[int]: |
| 47 | return [self._add_token(token) for token in self.tokenize(text)] |
| 48 | |
| 49 | def decode( |
| 50 | self, ids: Sequence[int] | int, skip_special_tokens: bool = False |
| 51 | ) -> str: |
| 52 | if isinstance(ids, int): |
| 53 | ids = [ids] |
| 54 | return "".join(self._id_to_token[token_id] for token_id in ids) |
| 55 | |
| 56 | def tokenize(self, text: str) -> list[str]: |
| 57 | tokens: list[str] = [] |
| 58 | pos = 0 |
| 59 | while pos < len(text): |
| 60 | for special_token in self.special_tokens: |
| 61 | if text.startswith(special_token, pos): |
| 62 | tokens.append(special_token) |
| 63 | pos += len(special_token) |
| 64 | break |
| 65 | else: |
| 66 | tokens.append(text[pos]) |
| 67 | pos += 1 |
| 68 | return tokens |
| 69 | |
| 70 | def convert_ids_to_tokens( |
| 71 | self, |
| 72 | ids: Sequence[int], |
| 73 | skip_special_tokens: bool = False, |
| 74 | ) -> list[str]: |
| 75 | return [self._id_to_token[token_id] for token_id in ids] |
| 76 | |
| 77 | def convert_tokens_to_ids(self, tokens: str | list[str]) -> int | list[int]: |
| 78 | if isinstance(tokens, str): |
| 79 | return self._add_token(tokens) |
| 80 | return [self._add_token(token) for token in tokens] |
| 81 | |
| 82 | def convert_tokens_to_string(self, tokens: list[str]) -> str: |
| 83 | return "".join(tokens) |
| 84 | |
| 85 | |
| 86 | class SplitMiniMaxM3Tokenizer(MiniMaxM3Tokenizer): |
| 87 | """Tokenizer that exposes marker vocab entries but encodes them as text.""" |
| 88 | |
| 89 | def tokenize(self, text: str) -> list[str]: |
| 90 | return list(text) |
| 91 | |
| 92 | |
| 93 | class RuntimeSplitMiniMaxM3Tokenizer(MiniMaxM3Tokenizer): |
| 94 | """Tokenizer whose runtime output splits markers despite atomic encodes.""" |
| 95 | |
| 96 | def encode_runtime(self, text: str) -> list[int]: |
| 97 | return [self._add_token(token) for token in list(text)] |
| 98 | |
| 99 | |
| 100 | def make_parser( |
| 101 | chat_template_kwargs: dict[str, str] | None = None, |
| 102 | ) -> tuple[MiniMaxM3ReasoningParser, MiniMaxM3Tokenizer]: |
| 103 | tokenizer = MiniMaxM3Tokenizer() |
| 104 | return ( |
| 105 | MiniMaxM3ReasoningParser(tokenizer, chat_template_kwargs=chat_template_kwargs), |
| 106 | tokenizer, |
| 107 | ) |
| 108 | |
| 109 | |
| 110 | def run_streaming( |
| 111 | parser: MiniMaxM3ReasoningParser, |
| 112 | tokenizer: MiniMaxM3Tokenizer, |
| 113 | chunks: list[str], |
| 114 | ) -> tuple[str | None, str | None, list[bool]]: |
| 115 | previous_text = "" |
| 116 | previous_token_ids: list[int] = [] |
| 117 | reasoning_parts: list[str] = [] |
| 118 | content_parts: list[str] = [] |
| 119 | reasoning_end_states: list[bool] = [] |
| 120 | |
| 121 | for chunk in chunks: |
| 122 | encode_runtime = getattr(tokenizer, "encode_runtime", tokenizer.encode) |
| 123 | delta_token_ids = encode_runtime(chunk) |
| 124 | current_text = previous_text + chunk |
| 125 | current_token_ids = previous_token_ids + delta_token_ids |
| 126 | delta = parser.extract_reasoning_streaming( |
| 127 | previous_text=previous_text, |
| 128 | current_text=current_text, |
| 129 | delta_text=chunk, |
| 130 | previous_token_ids=previous_token_ids, |
| 131 | current_token_ids=current_token_ids, |
| 132 | delta_token_ids=delta_token_ids, |
| 133 | ) |
| 134 | reasoning_end_states.append( |
| 135 | parser.is_reasoning_end_streaming(current_token_ids, delta_token_ids) |
| 136 | ) |
| 137 | |
| 138 | if delta is not None: |
| 139 | if delta.reasoning is not None: |
| 140 | reasoning_parts.append(delta.reasoning) |
| 141 | if delta.content is not None: |
| 142 | content_parts.append(delta.content) |
| 143 | |
| 144 | previous_text = current_text |
| 145 | previous_token_ids = current_token_ids |
| 146 | |
| 147 | return ( |
| 148 | "".join(reasoning_parts) or None, |
| 149 | "".join(content_parts) or None, |
| 150 | reasoning_end_states, |
| 151 | ) |
| 152 | |
| 153 | |
| 154 | def test_parser_registration(): |
| 155 | parser_cls = ReasoningParserManager.get_reasoning_parser("minimax_m3") |
| 156 | |
| 157 | assert parser_cls is MiniMaxM3ReasoningParser |
| 158 | |
| 159 | |
| 160 | def test_nonstreaming_extracts_explicit_reasoning_block(): |
| 161 | parser, _ = make_parser() |
| 162 | request = ChatCompletionRequest(messages=[], model="test-model") |
| 163 | |
| 164 | reasoning, content = parser.extract_reasoning( |
| 165 | "<mm:think>plan</mm:think>answer", request |
| 166 | ) |
| 167 | |
| 168 | assert reasoning == "plan" |
| 169 | assert content == "answer" |
| 170 | |
| 171 | |
| 172 | def test_nonstreaming_without_start_tag_is_content(): |
| 173 | parser, _ = make_parser() |
| 174 | request = ChatCompletionRequest(messages=[], model="test-model") |
| 175 | |
| 176 | reasoning, content = parser.extract_reasoning("plain answer", request) |
| 177 | |
| 178 | assert reasoning is None |
| 179 | assert content == "plain answer" |
| 180 | |
| 181 | |
| 182 | def test_nonstreaming_drops_leading_end_tag(): |
| 183 | parser, _ = make_parser() |
| 184 | request = ChatCompletionRequest(messages=[], model="test-model") |
| 185 | |
| 186 | reasoning, content = parser.extract_reasoning("</mm:think>answer", request) |
| 187 | |
| 188 | assert reasoning is None |
| 189 | assert content == "answer" |
| 190 | |
| 191 | |
| 192 | def test_nonstreaming_non_leading_end_tag_is_content(): |
| 193 | parser, _ = make_parser() |
| 194 | request = ChatCompletionRequest(messages=[], model="test-model") |
| 195 | |
| 196 | reasoning, content = parser.extract_reasoning("XXX</mm:think>YYY", request) |
| 197 | |
| 198 | assert reasoning is None |
| 199 | assert content == "XXX</mm:think>YYY" |
| 200 | |
| 201 | |
| 202 | def test_nonstreaming_enabled_mode_starts_in_reasoning(): |
| 203 | parser, _ = make_parser(chat_template_kwargs={"thinking_mode": "enabled"}) |
| 204 | request = ChatCompletionRequest(messages=[], model="test-model") |
| 205 | |
| 206 | reasoning, content = parser.extract_reasoning("plan</mm:think>answer", request) |
| 207 | |
| 208 | assert reasoning == "plan" |
| 209 | assert content == "answer" |
| 210 | |
| 211 | |
| 212 | def test_nonstreaming_open_reasoning_block(): |
| 213 | parser, _ = make_parser() |
| 214 | request = ChatCompletionRequest(messages=[], model="test-model") |
| 215 | |
| 216 | reasoning, content = parser.extract_reasoning("<mm:think>still thinking", request) |
| 217 | |
| 218 | assert reasoning == "still thinking" |
| 219 | assert content is None |
| 220 | |
| 221 | |
| 222 | def test_streaming_reasoning_tags_are_not_returned(): |
| 223 | parser, tokenizer = make_parser() |
| 224 | |
| 225 | reasoning, content, end_states = run_streaming( |
| 226 | parser, |
| 227 | tokenizer, |
| 228 | ["<mm:think>", "plan", "</mm:think>", "answer"], |
| 229 | ) |
| 230 | |
| 231 | assert reasoning == "plan" |
| 232 | assert content == "answer" |
| 233 | assert end_states == [False, False, True, True] |
| 234 | |
| 235 | |
| 236 | def test_streaming_boundary_can_emit_reasoning_and_content(): |
| 237 | parser, tokenizer = make_parser() |
| 238 | |
| 239 | reasoning, content, end_states = run_streaming( |
| 240 | parser, |
| 241 | tokenizer, |
| 242 | ["<mm:think>plan</mm:think>answer"], |
| 243 | ) |
| 244 | |
| 245 | assert reasoning == "plan" |
| 246 | assert content == "answer" |
| 247 | assert end_states == [True] |
| 248 | |
| 249 | |
| 250 | def test_streaming_drops_leading_end_tag(): |
| 251 | parser, tokenizer = make_parser() |
| 252 | |
| 253 | reasoning, content, end_states = run_streaming( |
| 254 | parser, |
| 255 | tokenizer, |
| 256 | ["</mm:think>", "answer"], |
| 257 | ) |
| 258 | |
| 259 | assert reasoning is None |
| 260 | assert content == "answer" |
| 261 | assert end_states == [True, True] |
| 262 | |
| 263 | |
| 264 | def test_streaming_non_leading_end_tag_is_content(): |
| 265 | parser, tokenizer = make_parser() |
| 266 | |
| 267 | reasoning, content, end_states = run_streaming( |
| 268 | parser, |
| 269 | tokenizer, |
| 270 | ["XXX</mm:think>YYY"], |
| 271 | ) |
| 272 | |
| 273 | assert reasoning is None |
| 274 | assert content == "XXX</mm:think>YYY" |
| 275 | assert end_states == [True] |
| 276 | |
| 277 | |
| 278 | def test_streaming_enabled_mode_starts_in_reasoning(): |
| 279 | parser, tokenizer = make_parser(chat_template_kwargs={"thinking_mode": "enabled"}) |
| 280 | |
| 281 | reasoning, content, end_states = run_streaming( |
| 282 | parser, |
| 283 | tokenizer, |
| 284 | ["plan", "</mm:think>", "answer"], |
| 285 | ) |
| 286 | |
| 287 | assert reasoning == "plan" |
| 288 | assert content == "answer" |
| 289 | assert end_states == [False, True, True] |
| 290 | |
| 291 | |
| 292 | def test_streaming_plain_content_ends_reasoning_phase(): |
| 293 | parser, tokenizer = make_parser() |
| 294 | |
| 295 | reasoning, content, end_states = run_streaming( |
| 296 | parser, |
| 297 | tokenizer, |
| 298 | ["plain ", "answer"], |
| 299 | ) |
| 300 | |
| 301 | assert reasoning is None |
| 302 | assert content == "plain answer" |
| 303 | assert end_states == [True, True] |
| 304 | |
| 305 | |
| 306 | def test_streaming_split_marker_tokens_are_not_returned(): |
| 307 | tokenizer = RuntimeSplitMiniMaxM3Tokenizer() |
| 308 | parser = MiniMaxM3ReasoningParser(tokenizer) |
| 309 | |
| 310 | reasoning, content, end_states = run_streaming( |
| 311 | parser, |
| 312 | tokenizer, |
| 313 | ["<mm:think>", "Reasoning", " content", "</mm:think>", "content"], |
| 314 | ) |
| 315 | |
| 316 | assert reasoning == "Reasoning content" |
| 317 | assert content == "content" |
| 318 | assert end_states == [False, False, False, True, True] |
| 319 | |
| 320 | |
| 321 | def test_streaming_split_marker_text_drives_end_state(): |
| 322 | tokenizer = RuntimeSplitMiniMaxM3Tokenizer() |
| 323 | parser = MiniMaxM3ReasoningParser(tokenizer) |
| 324 | previous_text = "" |
| 325 | previous_token_ids: list[int] = [] |
| 326 | |
| 327 | for chunk in ["<mm:think>", "Reasoning", " content", "</mm:think>"]: |
| 328 | delta_token_ids = tokenizer.encode_runtime(chunk) |
| 329 | current_text = previous_text + chunk |
| 330 | current_token_ids = previous_token_ids + delta_token_ids |
| 331 | parser.extract_reasoning_streaming( |
| 332 | previous_text=previous_text, |
| 333 | current_text=current_text, |
| 334 | delta_text=chunk, |
| 335 | previous_token_ids=previous_token_ids, |
| 336 | current_token_ids=current_token_ids, |
| 337 | delta_token_ids=delta_token_ids, |
| 338 | ) |
| 339 | previous_text = current_text |
| 340 | previous_token_ids = current_token_ids |
| 341 | |
| 342 | assert parser.is_reasoning_end_streaming(previous_token_ids, []) is True |
| 343 | |
| 344 | |
| 345 | def test_streaming_split_end_marker_content_ids_are_stripped(): |
| 346 | tokenizer = RuntimeSplitMiniMaxM3Tokenizer() |
| 347 | parser = MiniMaxM3ReasoningParser(tokenizer) |
| 348 | previous_text = "<mm:think>Reasoning" |
| 349 | previous_token_ids = tokenizer.encode_runtime(previous_text) |
| 350 | delta_text = "</mm:think>content" |
| 351 | delta_token_ids = tokenizer.encode_runtime(delta_text) |
| 352 | current_token_ids = previous_token_ids + delta_token_ids |
| 353 | |
| 354 | parser.extract_reasoning_streaming( |
| 355 | previous_text=previous_text, |
| 356 | current_text=previous_text + delta_text, |
| 357 | delta_text=delta_text, |
| 358 | previous_token_ids=previous_token_ids, |
| 359 | current_token_ids=current_token_ids, |
| 360 | delta_token_ids=delta_token_ids, |
| 361 | ) |
| 362 | |
| 363 | assert parser.is_reasoning_end_streaming(current_token_ids, delta_token_ids) |
| 364 | assert tokenizer.decode(parser.extract_content_ids(delta_token_ids)) == "content" |
| 365 | |
| 366 | |
| 367 | def test_streaming_split_marker_tokens_enabled_mode(): |
| 368 | tokenizer = RuntimeSplitMiniMaxM3Tokenizer() |
| 369 | parser = MiniMaxM3ReasoningParser( |
| 370 | tokenizer, chat_template_kwargs={"thinking_mode": "enabled"} |
| 371 | ) |
| 372 | |
| 373 | reasoning, content, end_states = run_streaming( |
| 374 | parser, |
| 375 | tokenizer, |
| 376 | ["Reasoning", " content", "</mm:think>", "content"], |
| 377 | ) |
| 378 | |
| 379 | assert reasoning == "Reasoning content" |
| 380 | assert content == "content" |
| 381 | assert end_states == [False, False, True, True] |
| 382 | |
| 383 | |
| 384 | def test_streaming_split_marker_text_across_deltas(): |
| 385 | tokenizer = RuntimeSplitMiniMaxM3Tokenizer() |
| 386 | parser = MiniMaxM3ReasoningParser(tokenizer) |
| 387 | |
| 388 | reasoning, content, end_states = run_streaming( |
| 389 | parser, |
| 390 | tokenizer, |
| 391 | ["<mm:", "think>", "Reasoning", " content", "</mm:", "think>", "content"], |
| 392 | ) |
| 393 | |
| 394 | assert reasoning == "Reasoning content" |
| 395 | assert content == "content" |
| 396 | assert end_states == [False, False, False, False, False, True, True] |
| 397 | |
| 398 | |
| 399 | def test_streaming_split_leading_end_marker_text_across_deltas(): |
| 400 | tokenizer = RuntimeSplitMiniMaxM3Tokenizer() |
| 401 | parser = MiniMaxM3ReasoningParser(tokenizer) |
| 402 | |
| 403 | reasoning, content, end_states = run_streaming( |
| 404 | parser, |
| 405 | tokenizer, |
| 406 | ["</mm:", "think>", "content"], |
| 407 | ) |
| 408 | |
| 409 | assert reasoning is None |
| 410 | assert content == "content" |
| 411 | assert end_states == [False, True, True] |
| 412 | |
| 413 | |
| 414 | def test_token_id_helpers_with_split_marker_tokens(): |
| 415 | tokenizer = SplitMiniMaxM3Tokenizer() |
| 416 | parser = MiniMaxM3ReasoningParser(tokenizer) |
| 417 | output_ids = tokenizer.encode( |
| 418 | "<mm:think>abc</mm:think>def", add_special_tokens=False |
| 419 | ) |
| 420 | open_reasoning_ids = tokenizer.encode("<mm:think>abc", add_special_tokens=False) |
| 421 | content_ids = tokenizer.encode("plain", add_special_tokens=False) |
| 422 | |
| 423 | assert parser.is_reasoning_end(output_ids) |
| 424 | assert not parser.is_reasoning_end(open_reasoning_ids) |
| 425 | assert not parser.is_reasoning_end(content_ids) |
| 426 | assert tokenizer.decode(parser.extract_content_ids(output_ids)) == "def" |
| 427 | assert parser.extract_content_ids(open_reasoning_ids) == [] |
| 428 | assert parser.extract_content_ids(content_ids) == content_ids |
| 429 | assert parser.count_reasoning_tokens(output_ids) == len(tokenizer.encode("abc")) |
| 430 | |
| 431 | |
| 432 | def test_token_id_helpers(): |
| 433 | parser, tokenizer = make_parser() |
| 434 | output_ids = tokenizer.encode( |
| 435 | "<mm:think>abc</mm:think>def", add_special_tokens=False |
| 436 | ) |
| 437 | open_reasoning_ids = tokenizer.encode("<mm:think>abc", add_special_tokens=False) |
| 438 | content_ids = tokenizer.encode("plain", add_special_tokens=False) |
| 439 | |
| 440 | assert parser.is_reasoning_end(output_ids) |
| 441 | assert not parser.is_reasoning_end(open_reasoning_ids) |
| 442 | assert not parser.is_reasoning_end(content_ids) |
| 443 | assert tokenizer.decode(parser.extract_content_ids(output_ids)) == "def" |
| 444 | assert parser.extract_content_ids(open_reasoning_ids) == [] |
| 445 | assert parser.extract_content_ids(content_ids) == content_ids |
| 446 | assert parser.count_reasoning_tokens(output_ids) == len(tokenizer.encode("abc")) |
| 447 | |
| 448 | |
| 449 | def test_token_id_helpers_enabled_mode(): |
| 450 | parser, tokenizer = make_parser(chat_template_kwargs={"thinking_mode": "enabled"}) |
| 451 | output_ids = tokenizer.encode("abc</mm:think>def", add_special_tokens=False) |
| 452 | open_reasoning_ids = tokenizer.encode("abc", add_special_tokens=False) |
| 453 | |
| 454 | assert parser.is_reasoning_end(output_ids) |
| 455 | assert not parser.is_reasoning_end(open_reasoning_ids) |
| 456 | assert tokenizer.decode(parser.extract_content_ids(output_ids)) == "def" |
| 457 | assert parser.extract_content_ids(open_reasoning_ids) == [] |
| 458 | assert parser.count_reasoning_tokens(output_ids) == len(tokenizer.encode("abc")) |
| 459 | assert parser.count_reasoning_tokens(open_reasoning_ids) == len( |
| 460 | tokenizer.encode("abc") |
| 461 | ) |