| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | from transformers import AutoTokenizer, PreTrainedTokenizerBase |
| 6 | |
| 7 | from tests.reasoning.utils import ( |
| 8 | StreamingReasoningReconstructor, |
| 9 | run_reasoning_extraction, |
| 10 | run_reasoning_extraction_streaming, |
| 11 | ) |
| 12 | from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest |
| 13 | from vllm.reasoning import ReasoningParser, ReasoningParserManager |
| 14 | |
| 15 | parser_name = "qwen3" |
| 16 | start_token = "<think>" |
| 17 | end_token = "</think>" |
| 18 | |
| 19 | REASONING_MODEL_NAMES = [ |
| 20 | "Qwen/Qwen3-0.6B", |
| 21 | "Qwen/Qwen3.5-397B-A17B", |
| 22 | "Qwen/Qwen3-4B-Thinking-2507", |
| 23 | ] |
| 24 | |
| 25 | |
| 26 | @pytest.fixture(scope="module", params=REASONING_MODEL_NAMES) |
| 27 | def qwen3_tokenizer(request) -> PreTrainedTokenizerBase: |
| 28 | return AutoTokenizer.from_pretrained(request.param) |
| 29 | |
| 30 | |
| 31 | # --- <think> in prompt, only </think> in output (typical) --- |
| 32 | |
| 33 | WITHOUT_START_TOKEN = { |
| 34 | "output": "This is a reasoning section</think>This is the rest", |
| 35 | "reasoning": "This is a reasoning section", |
| 36 | "content": "This is the rest", |
| 37 | } |
| 38 | WITHOUT_START_TOKEN_STREAM = { |
| 39 | "output": "This is a reasoning section</think>This is the rest", |
| 40 | "reasoning": "This is a reasoning section", |
| 41 | "content": "This is the rest", |
| 42 | } |
| 43 | WITHOUT_START_TOKEN_COMPLETE_REASONING = { |
| 44 | "output": "This is a reasoning section</think>", |
| 45 | "reasoning": "This is a reasoning section", |
| 46 | "content": None, |
| 47 | } |
| 48 | |
| 49 | # --- <think> present in output (old template / edge case) --- |
| 50 | |
| 51 | WITH_THINK = { |
| 52 | "output": "<think>This is a reasoning section</think>This is the rest", |
| 53 | "reasoning": "This is a reasoning section", |
| 54 | "content": "This is the rest", |
| 55 | } |
| 56 | WITH_THINK_STREAM = { |
| 57 | "output": "<think>This is a reasoning section</think>This is the rest", |
| 58 | "reasoning": "This is a reasoning section", |
| 59 | "content": "This is the rest", |
| 60 | } |
| 61 | |
| 62 | # --- No think tokens at all (thinking enabled, truncated) --- |
| 63 | |
| 64 | # With thinking enabled (default), no think tokens means the output was |
| 65 | # truncated before </think> could be generated. All output is reasoning. |
| 66 | WITHOUT_THINK = { |
| 67 | "output": "This is the rest", |
| 68 | "reasoning": "This is the rest", |
| 69 | "content": None, |
| 70 | } |
| 71 | # In streaming, the parser cannot distinguish "thinking disabled" from |
| 72 | # "reasoning in progress" when no think tokens have appeared yet. |
| 73 | # It assumes reasoning. The serving layer handles the "thinking disabled" |
| 74 | # case by checking prompt_is_reasoning_end_arr before calling the parser. |
| 75 | WITHOUT_THINK_STREAM = { |
| 76 | "output": "This is the rest", |
| 77 | "reasoning": "This is the rest", |
| 78 | "content": None, |
| 79 | } |
| 80 | |
| 81 | # --- <tool_call> without </think> (implicit reasoning end) --- |
| 82 | |
| 83 | TOOL_CALL_BODY = ( |
| 84 | "<tool_call>\n<function=bash>\n<parameter=command>" |
| 85 | "\ncat /etc/hosts\n</parameter>\n</function>\n</tool_call>" |
| 86 | ) |
| 87 | |
| 88 | TOOL_CALL_NO_THINK_END = { |
| 89 | "output": "I need to read the file.\n\n" + TOOL_CALL_BODY, |
| 90 | "reasoning": "I need to read the file.\n\n", |
| 91 | "content": TOOL_CALL_BODY, |
| 92 | } |
| 93 | |
| 94 | TOOL_CALL_WITH_THINK_NO_END = { |
| 95 | "output": "<think>I need to read the file.\n\n" + TOOL_CALL_BODY, |
| 96 | "reasoning": "I need to read the file.\n\n", |
| 97 | "content": TOOL_CALL_BODY, |
| 98 | } |
| 99 | |
| 100 | # --- Edge cases --- |
| 101 | |
| 102 | COMPLETE_REASONING = { |
| 103 | "output": "<think>This is a reasoning section</think>", |
| 104 | "reasoning": "This is a reasoning section", |
| 105 | "content": None, |
| 106 | } |
| 107 | MULTILINE_REASONING = { |
| 108 | "output": "<think>This is a reasoning\nsection</think>This is the rest\nThat", |
| 109 | "reasoning": "This is a reasoning\nsection", |
| 110 | "content": "This is the rest\nThat", |
| 111 | } |
| 112 | # Truncated output: <think> present but no </think> (thinking enabled). |
| 113 | # Everything is reasoning because the output was cut off mid-thought. |
| 114 | ONLY_OPEN_TAG = { |
| 115 | "output": "<think>This is a reasoning section", |
| 116 | "reasoning": "This is a reasoning section", |
| 117 | "content": None, |
| 118 | } |
| 119 | |
| 120 | ONLY_OPEN_TAG_STREAM = { |
| 121 | "output": "<think>This is a reasoning section", |
| 122 | "reasoning": "This is a reasoning section", |
| 123 | "content": None, |
| 124 | } |
| 125 | |
| 126 | # Truncated output without <think> prefix (Qwen3.5 style where <think> |
| 127 | # is in the prompt). No </think> means truncation — all is reasoning. |
| 128 | TRUNCATED_NO_START_TOKEN = { |
| 129 | "output": "This is a reasoning section", |
| 130 | "reasoning": "This is a reasoning section", |
| 131 | "content": None, |
| 132 | } |
| 133 | |
| 134 | TRUNCATED_NO_START_TOKEN_STREAM = { |
| 135 | "output": "This is a reasoning section", |
| 136 | "reasoning": "This is a reasoning section", |
| 137 | "content": None, |
| 138 | } |
| 139 | |
| 140 | TEST_CASES = [ |
| 141 | pytest.param( |
| 142 | False, |
| 143 | WITHOUT_START_TOKEN, |
| 144 | id="without_start_token", |
| 145 | ), |
| 146 | pytest.param( |
| 147 | True, |
| 148 | WITHOUT_START_TOKEN_STREAM, |
| 149 | id="without_start_token_stream", |
| 150 | ), |
| 151 | pytest.param( |
| 152 | False, |
| 153 | WITHOUT_START_TOKEN_COMPLETE_REASONING, |
| 154 | id="without_start_token_complete_reasoning", |
| 155 | ), |
| 156 | pytest.param( |
| 157 | True, |
| 158 | WITHOUT_START_TOKEN_COMPLETE_REASONING, |
| 159 | id="without_start_token_complete_reasoning_stream", |
| 160 | ), |
| 161 | pytest.param( |
| 162 | False, |
| 163 | WITH_THINK, |
| 164 | id="with_think", |
| 165 | ), |
| 166 | pytest.param( |
| 167 | True, |
| 168 | WITH_THINK_STREAM, |
| 169 | id="with_think_stream", |
| 170 | ), |
| 171 | pytest.param( |
| 172 | False, |
| 173 | WITHOUT_THINK, |
| 174 | id="without_think", |
| 175 | ), |
| 176 | pytest.param( |
| 177 | True, |
| 178 | WITHOUT_THINK_STREAM, |
| 179 | id="without_think_stream", |
| 180 | ), |
| 181 | pytest.param( |
| 182 | False, |
| 183 | COMPLETE_REASONING, |
| 184 | id="complete_reasoning", |
| 185 | ), |
| 186 | pytest.param( |
| 187 | True, |
| 188 | COMPLETE_REASONING, |
| 189 | id="complete_reasoning_stream", |
| 190 | ), |
| 191 | pytest.param( |
| 192 | False, |
| 193 | MULTILINE_REASONING, |
| 194 | id="multiline_reasoning", |
| 195 | ), |
| 196 | pytest.param( |
| 197 | True, |
| 198 | MULTILINE_REASONING, |
| 199 | id="multiline_reasoning_stream", |
| 200 | ), |
| 201 | pytest.param( |
| 202 | False, |
| 203 | ONLY_OPEN_TAG, |
| 204 | id="only_open_tag", |
| 205 | ), |
| 206 | pytest.param( |
| 207 | True, |
| 208 | ONLY_OPEN_TAG_STREAM, |
| 209 | id="only_open_tag_stream", |
| 210 | ), |
| 211 | pytest.param( |
| 212 | False, |
| 213 | TRUNCATED_NO_START_TOKEN, |
| 214 | id="truncated_no_start_token", |
| 215 | ), |
| 216 | pytest.param( |
| 217 | True, |
| 218 | TRUNCATED_NO_START_TOKEN_STREAM, |
| 219 | id="truncated_no_start_token_stream", |
| 220 | ), |
| 221 | pytest.param( |
| 222 | False, |
| 223 | TOOL_CALL_NO_THINK_END, |
| 224 | id="tool_call_no_think_end", |
| 225 | ), |
| 226 | pytest.param( |
| 227 | True, |
| 228 | TOOL_CALL_NO_THINK_END, |
| 229 | id="tool_call_no_think_end_stream", |
| 230 | ), |
| 231 | pytest.param( |
| 232 | False, |
| 233 | TOOL_CALL_WITH_THINK_NO_END, |
| 234 | id="tool_call_with_think_no_end", |
| 235 | ), |
| 236 | pytest.param( |
| 237 | True, |
| 238 | TOOL_CALL_WITH_THINK_NO_END, |
| 239 | id="tool_call_with_think_no_end_stream", |
| 240 | ), |
| 241 | ] |
| 242 | |
| 243 | |
| 244 | @pytest.mark.parametrize("streaming, param_dict", TEST_CASES) |
| 245 | def test_reasoning( |
| 246 | streaming: bool, |
| 247 | param_dict: dict, |
| 248 | qwen3_tokenizer, |
| 249 | ): |
| 250 | output = qwen3_tokenizer.tokenize(param_dict["output"]) |
| 251 | output_tokens: list[str] = [ |
| 252 | qwen3_tokenizer.convert_tokens_to_string([token]) for token in output |
| 253 | ] |
| 254 | parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)( |
| 255 | qwen3_tokenizer |
| 256 | ) |
| 257 | |
| 258 | reasoning, content = run_reasoning_extraction( |
| 259 | parser, output_tokens, streaming=streaming |
| 260 | ) |
| 261 | |
| 262 | assert reasoning == param_dict["reasoning"] |
| 263 | assert content == param_dict["content"] |
| 264 | |
| 265 | |
| 266 | # Multi-token delta tests: simulate real-world streaming where a single |
| 267 | # delta can contain multiple tokens (e.g., speculative decoding). |
| 268 | MULTI_TOKEN_DELTA_CASES = [ |
| 269 | pytest.param( |
| 270 | # <think> grouped with following text in one delta |
| 271 | ["<think>This is a reasoning section", "</think>", "This is the rest"], |
| 272 | "This is a reasoning section", |
| 273 | "This is the rest", |
| 274 | id="start_token_grouped_with_text", |
| 275 | ), |
| 276 | pytest.param( |
| 277 | # </think> grouped with following content in one delta |
| 278 | ["reasoning section", "</think>This is the rest"], |
| 279 | "reasoning section", |
| 280 | "This is the rest", |
| 281 | id="end_token_grouped_with_content", |
| 282 | ), |
| 283 | pytest.param( |
| 284 | # <think> and </think> in the same delta, no content after |
| 285 | ["<think>reasoning</think>"], |
| 286 | "reasoning", |
| 287 | None, |
| 288 | id="start_and_end_in_one_delta_no_content", |
| 289 | ), |
| 290 | pytest.param( |
| 291 | # No start token, end grouped with content (Qwen3.5 style) |
| 292 | ["reasoning section", "</think>content"], |
| 293 | "reasoning section", |
| 294 | "content", |
| 295 | id="no_start_end_grouped_with_content", |
| 296 | ), |
| 297 | pytest.param( |
| 298 | # <tool_call> arrives in a separate delta after reasoning text |
| 299 | ["I need to read the file.\n\n", "<tool_call>\n<function=bash>"], |
| 300 | "I need to read the file.\n\n", |
| 301 | "<tool_call>\n<function=bash>", |
| 302 | id="tool_call_implicit_reasoning_end", |
| 303 | ), |
| 304 | ] |
| 305 | |
| 306 | |
| 307 | @pytest.mark.parametrize( |
| 308 | "deltas, expected_reasoning, expected_content", MULTI_TOKEN_DELTA_CASES |
| 309 | ) |
| 310 | def test_reasoning_streaming_multi_token_deltas( |
| 311 | deltas: list[str], |
| 312 | expected_reasoning: str | None, |
| 313 | expected_content: str | None, |
| 314 | qwen3_tokenizer, |
| 315 | ): |
| 316 | """Test that multi-token deltas don't leak <think> into reasoning.""" |
| 317 | parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)( |
| 318 | qwen3_tokenizer |
| 319 | ) |
| 320 | |
| 321 | reconstructor: StreamingReasoningReconstructor = run_reasoning_extraction_streaming( |
| 322 | parser, deltas |
| 323 | ) |
| 324 | |
| 325 | assert reconstructor.reasoning == expected_reasoning |
| 326 | assert (reconstructor.other_content or None) == expected_content |
| 327 | |
| 328 | |
| 329 | # --- Tests for enable_thinking=False (thinking explicitly disabled) --- |
| 330 | |
| 331 | |
| 332 | THINKING_DISABLED_CASES = [ |
| 333 | pytest.param( |
| 334 | "This is plain content", |
| 335 | None, |
| 336 | "This is plain content", |
| 337 | id="thinking_disabled_plain_content", |
| 338 | ), |
| 339 | pytest.param( |
| 340 | "Some output without think tokens", |
| 341 | None, |
| 342 | "Some output without think tokens", |
| 343 | id="thinking_disabled_no_think_tokens", |
| 344 | ), |
| 345 | pytest.param( |
| 346 | "I need to read the file.\n\n" + TOOL_CALL_BODY, |
| 347 | None, |
| 348 | "I need to read the file.\n\n" + TOOL_CALL_BODY, |
| 349 | id="thinking_disabled_with_tool_call", |
| 350 | ), |
| 351 | ] |
| 352 | |
| 353 | |
| 354 | @pytest.mark.parametrize( |
| 355 | "output, expected_reasoning, expected_content", THINKING_DISABLED_CASES |
| 356 | ) |
| 357 | def test_reasoning_thinking_disabled( |
| 358 | output: str, |
| 359 | expected_reasoning: str | None, |
| 360 | expected_content: str | None, |
| 361 | qwen3_tokenizer, |
| 362 | ): |
| 363 | """When enable_thinking=False, output without </think> is all content.""" |
| 364 | parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)( |
| 365 | qwen3_tokenizer, |
| 366 | chat_template_kwargs={"enable_thinking": False}, |
| 367 | ) |
| 368 | |
| 369 | reasoning, content = parser.extract_reasoning( |
| 370 | model_output=output, |
| 371 | request=ChatCompletionRequest(messages=[], model="test-model"), |
| 372 | ) |
| 373 | |
| 374 | assert reasoning == expected_reasoning |
| 375 | assert content == expected_content |