| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | |
| 6 | from tests.reasoning.utils import run_reasoning_extraction |
| 7 | from vllm.parser.engine.registered_adapters import Step3p5ParserReasoningAdapter |
| 8 | from vllm.reasoning import ReasoningParserManager |
| 9 | from vllm.tokenizers import get_tokenizer |
| 10 | |
| 11 | parser_name = "step3p5" |
| 12 | start_token = "<think>" |
| 13 | end_token = "</think>" |
| 14 | |
| 15 | REASONING_MODEL_NAME = "stepfun-ai/Step-3.5-Flash" |
| 16 | |
| 17 | |
| 18 | @pytest.fixture(scope="module") |
| 19 | def step3p5_tokenizer(): |
| 20 | return get_tokenizer(tokenizer_name=REASONING_MODEL_NAME) |
| 21 | |
| 22 | |
| 23 | SIMPLE_REASONING = { |
| 24 | "output": "This is a reasoning section</think>This is the rest", |
| 25 | "reasoning": "This is a reasoning section", |
| 26 | "content": "This is the rest", |
| 27 | "is_reasoning_end": True, |
| 28 | } |
| 29 | COMPLETE_REASONING = { |
| 30 | "output": "This is a reasoning section</think>", |
| 31 | "reasoning": "This is a reasoning section", |
| 32 | "content": None, |
| 33 | "is_reasoning_end": True, |
| 34 | } |
| 35 | NO_CONTENT = { |
| 36 | "output": "This is content", |
| 37 | "reasoning": "This is content", |
| 38 | "content": None, |
| 39 | "is_reasoning_end": False, |
| 40 | } |
| 41 | NO_REASONING_STREAMING = { |
| 42 | "output": "This is a reasoning section", |
| 43 | "reasoning": "This is a reasoning section", |
| 44 | "content": None, |
| 45 | "is_reasoning_end": False, |
| 46 | } |
| 47 | MULTIPLE_LINES = { |
| 48 | "output": "This\nThat</think>This is the rest\nThat", |
| 49 | "reasoning": "This\nThat", |
| 50 | "content": "This is the rest\nThat", |
| 51 | "is_reasoning_end": True, |
| 52 | } |
| 53 | SHORTEST_REASONING_NO_STREAMING = { |
| 54 | "output": "</think>This is the rest", |
| 55 | "reasoning": None, |
| 56 | "content": "This is the rest", |
| 57 | "is_reasoning_end": True, |
| 58 | } |
| 59 | SHORTEST_REASONING = { |
| 60 | "output": "</think>This is the rest", |
| 61 | "reasoning": None, |
| 62 | "content": "This is the rest", |
| 63 | "is_reasoning_end": True, |
| 64 | } |
| 65 | REASONING_WITH_THINK = { |
| 66 | "output": "<think>This is a reasoning section</think>This is the rest", |
| 67 | "reasoning": "This is a reasoning section", |
| 68 | "content": "This is the rest", |
| 69 | "is_reasoning_end": True, |
| 70 | } |
| 71 | COMPLETE_REASONING_WITH_THINK = { |
| 72 | "output": "<think>This is a reasoning section</think>", |
| 73 | "reasoning": "This is a reasoning section", |
| 74 | "content": None, |
| 75 | "is_reasoning_end": True, |
| 76 | } |
| 77 | MULTIPLE_LINES_WITH_THINK = { |
| 78 | "output": "<think>This\nThat</think>This is the rest\nThat", |
| 79 | "reasoning": "This\nThat", |
| 80 | "content": "This is the rest\nThat", |
| 81 | "is_reasoning_end": True, |
| 82 | } |
| 83 | SHORTEST_REASONING_NO_STREAMING_WITH_THINK = { |
| 84 | "output": "</think>This is the rest", |
| 85 | "reasoning": None, |
| 86 | "content": "This is the rest", |
| 87 | "is_reasoning_end": True, |
| 88 | } |
| 89 | SHORTEST_REASONING_WITH_THINK = { |
| 90 | "output": "</think>This is the rest", |
| 91 | "reasoning": None, |
| 92 | "content": "This is the rest", |
| 93 | "is_reasoning_end": True, |
| 94 | } |
| 95 | THINK_NO_END = { |
| 96 | "output": "<think>This is a reasoning section", |
| 97 | "reasoning": "This is a reasoning section", |
| 98 | "content": None, |
| 99 | "is_reasoning_end": False, |
| 100 | } |
| 101 | EMPTY = { |
| 102 | "output": "", |
| 103 | "reasoning": None, |
| 104 | "content": None, |
| 105 | "is_reasoning_end": False, |
| 106 | } |
| 107 | EMPTY_STREAMING = { |
| 108 | "output": "", |
| 109 | "reasoning": None, |
| 110 | "content": None, |
| 111 | "is_reasoning_end": False, |
| 112 | } |
| 113 | NEW_LINE = { |
| 114 | "output": "\n<think>This is a reasoning section</think>\nThis is the rest", |
| 115 | "reasoning": "\nThis is a reasoning section", |
| 116 | "content": "\nThis is the rest", |
| 117 | "is_reasoning_end": True, |
| 118 | } |
| 119 | |
| 120 | NEW_LINE_STREAMING = { |
| 121 | "output": "\n<think>This is a reasoning section\n</think>\nThis is the rest", |
| 122 | "reasoning": "\nThis is a reasoning section", |
| 123 | "content": "\nThis is the rest", |
| 124 | "is_reasoning_end": True, |
| 125 | } |
| 126 | |
| 127 | NEW_LINE_STREAMING_COMPLEX_CONTENT = { |
| 128 | "output": "\n This is a \n reasoning section\n\n\n</think>\n\nThis is the rest", |
| 129 | "reasoning": "\n This is a \n reasoning section", |
| 130 | "content": "\n\nThis is the rest", |
| 131 | "is_reasoning_end": True, |
| 132 | } |
| 133 | |
| 134 | MULTI_TURN_PROMPT_CONTENT = { |
| 135 | "output": "<think> This is last turn's reasoning section </think> hello <think>", |
| 136 | "reasoning": "", |
| 137 | "content": "", |
| 138 | "is_reasoning_end": False, |
| 139 | } |
| 140 | |
| 141 | TEST_CASES = [ |
| 142 | pytest.param( |
| 143 | False, |
| 144 | SIMPLE_REASONING, |
| 145 | id="simple_reasoning", |
| 146 | ), |
| 147 | pytest.param( |
| 148 | True, |
| 149 | SIMPLE_REASONING, |
| 150 | id="simple_reasoning_streaming", |
| 151 | ), |
| 152 | pytest.param( |
| 153 | False, |
| 154 | COMPLETE_REASONING, |
| 155 | id="complete_reasoning", |
| 156 | ), |
| 157 | pytest.param( |
| 158 | True, |
| 159 | COMPLETE_REASONING, |
| 160 | id="complete_reasoning_streaming", |
| 161 | ), |
| 162 | pytest.param( |
| 163 | False, |
| 164 | NO_CONTENT, |
| 165 | id="no_content_token", |
| 166 | ), |
| 167 | pytest.param( |
| 168 | True, |
| 169 | NO_REASONING_STREAMING, |
| 170 | id="no_reasoning_token_streaming", |
| 171 | ), |
| 172 | pytest.param( |
| 173 | False, |
| 174 | MULTIPLE_LINES, |
| 175 | id="multiple_lines", |
| 176 | ), |
| 177 | pytest.param( |
| 178 | True, |
| 179 | MULTIPLE_LINES, |
| 180 | id="multiple_lines_streaming", |
| 181 | ), |
| 182 | pytest.param( |
| 183 | True, |
| 184 | SHORTEST_REASONING, |
| 185 | id="shortest", |
| 186 | ), |
| 187 | pytest.param( |
| 188 | False, |
| 189 | SHORTEST_REASONING_NO_STREAMING, |
| 190 | id="shortest_streaming", |
| 191 | ), |
| 192 | pytest.param( |
| 193 | False, |
| 194 | REASONING_WITH_THINK, |
| 195 | id="reasoning_with_think", |
| 196 | ), |
| 197 | pytest.param( |
| 198 | True, |
| 199 | REASONING_WITH_THINK, |
| 200 | id="reasoning_with_think_streaming", |
| 201 | ), |
| 202 | pytest.param( |
| 203 | False, |
| 204 | COMPLETE_REASONING_WITH_THINK, |
| 205 | id="complete_reasoning_with_think", |
| 206 | ), |
| 207 | pytest.param( |
| 208 | True, |
| 209 | COMPLETE_REASONING_WITH_THINK, |
| 210 | id="complete_reasoning_with_think_streaming", |
| 211 | ), |
| 212 | pytest.param( |
| 213 | False, |
| 214 | MULTIPLE_LINES_WITH_THINK, |
| 215 | id="multiple_lines_with_think", |
| 216 | ), |
| 217 | pytest.param( |
| 218 | True, |
| 219 | MULTIPLE_LINES_WITH_THINK, |
| 220 | id="multiple_lines_with_think_streaming", |
| 221 | ), |
| 222 | pytest.param( |
| 223 | False, |
| 224 | SHORTEST_REASONING_NO_STREAMING_WITH_THINK, |
| 225 | id="shortest_with_think", |
| 226 | ), |
| 227 | pytest.param( |
| 228 | True, |
| 229 | SHORTEST_REASONING_WITH_THINK, |
| 230 | id="shortest_with_think_streaming", |
| 231 | ), |
| 232 | pytest.param( |
| 233 | False, |
| 234 | THINK_NO_END, |
| 235 | id="think_no_end", |
| 236 | ), |
| 237 | pytest.param( |
| 238 | True, |
| 239 | THINK_NO_END, |
| 240 | id="think_no_end_streaming", |
| 241 | ), |
| 242 | pytest.param( |
| 243 | False, |
| 244 | EMPTY, |
| 245 | id="empty", |
| 246 | ), |
| 247 | pytest.param( |
| 248 | True, |
| 249 | EMPTY_STREAMING, |
| 250 | id="empty_streaming", |
| 251 | ), |
| 252 | pytest.param( |
| 253 | False, |
| 254 | NEW_LINE, |
| 255 | id="new_line", |
| 256 | ), |
| 257 | pytest.param( |
| 258 | True, |
| 259 | NEW_LINE_STREAMING, |
| 260 | id="new_line_streaming", |
| 261 | ), |
| 262 | pytest.param( |
| 263 | True, |
| 264 | NEW_LINE_STREAMING_COMPLEX_CONTENT, |
| 265 | id="new_line_streaming_complex_content", |
| 266 | ), |
| 267 | pytest.param( |
| 268 | True, |
| 269 | MULTI_TURN_PROMPT_CONTENT, |
| 270 | id="multi_turn_prompt_content", |
| 271 | ), |
| 272 | ] |
| 273 | |
| 274 | |
| 275 | @pytest.mark.parametrize("streaming, param_dict", TEST_CASES) |
| 276 | def test_reasoning( |
| 277 | streaming: bool, |
| 278 | param_dict: dict, |
| 279 | step3p5_tokenizer, |
| 280 | request, |
| 281 | ): |
| 282 | output = step3p5_tokenizer.tokenize(param_dict["output"]) |
| 283 | # decode everything to tokens |
| 284 | output_tokens: list[str] = [ |
| 285 | step3p5_tokenizer.convert_tokens_to_string([token]) for token in output |
| 286 | ] |
| 287 | parser = ReasoningParserManager.get_reasoning_parser(parser_name)(step3p5_tokenizer) |
| 288 | assert isinstance(parser, Step3p5ParserReasoningAdapter) |
| 289 | |
| 290 | reasoning, content = run_reasoning_extraction( |
| 291 | parser, output_tokens, streaming=streaming |
| 292 | ) |
| 293 | |
| 294 | print(f"reasoning: {reasoning}") |
| 295 | print(f"content: {content}") |
| 296 | test_id = request.node.callspec.id if hasattr(request.node, "callspec") else None |
| 297 | if request.node.callspec.id != "multi_turn_prompt_content": |
| 298 | assert reasoning == param_dict["reasoning"] |
| 299 | assert content == param_dict["content"] |
| 300 | |
| 301 | # Test is_reasoning_end |
| 302 | output_ids = step3p5_tokenizer.convert_tokens_to_ids(output) |
| 303 | is_reasoning_end = parser.is_reasoning_end(output_ids) |
| 304 | assert is_reasoning_end == param_dict["is_reasoning_end"] |
| 305 | |
| 306 | # Test is_reasoning_end_streaming |
| 307 | current_ids: list[int] = [] |
| 308 | for token_id in output_ids: |
| 309 | current_ids.append(token_id) |
| 310 | if token_id in parser.reasoning_end_token_ids: |
| 311 | assert parser.is_reasoning_end_streaming(current_ids, [token_id]) |
| 312 | if param_dict["is_reasoning_end"]: |
| 313 | assert parser.is_reasoning_end_streaming(output_ids, output_ids) |
| 314 | break |
| 315 | assert not parser.is_reasoning_end_streaming(current_ids, [token_id]) |
| 316 | else: |
| 317 | assert not parser.is_reasoning_end_streaming(output_ids, output_ids) |
| 318 | |
| 319 | # Test extract_content |
| 320 | if param_dict["content"] is not None: |
| 321 | content_ids = parser.extract_content_ids(output_ids) |
| 322 | # Fixed expected token ids for specific test cases |
| 323 | test_id = ( |
| 324 | request.node.callspec.id if hasattr(request.node, "callspec") else None |
| 325 | ) |
| 326 | # Match most specific first |
| 327 | if test_id not in [ |
| 328 | "new_line_streaming_complex_content", |
| 329 | "new_line_streaming", |
| 330 | "new_line", |
| 331 | "multi_turn_prompt_content", |
| 332 | ]: |
| 333 | expected_content_ids = step3p5_tokenizer.convert_tokens_to_ids( |
| 334 | step3p5_tokenizer.tokenize(param_dict["content"]) |
| 335 | ) |
| 336 | assert content_ids == expected_content_ids |
| 337 | elif param_dict["is_reasoning_end"]: |
| 338 | assert parser.extract_content_ids(output_ids) == [] |