| 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 run_reasoning_extraction |
| 8 | from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest |
| 9 | from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser |
| 10 | |
| 11 | |
| 12 | # Create a concrete test implementation of BaseThinkingReasoningParser |
| 13 | class TestThinkingReasoningParser(BaseThinkingReasoningParser): |
| 14 | """Test implementation of BaseThinkingReasoningParser.""" |
| 15 | |
| 16 | @property |
| 17 | def start_token(self) -> str: |
| 18 | return "<test:think>" |
| 19 | |
| 20 | @property |
| 21 | def end_token(self) -> str: |
| 22 | return "</test:think>" |
| 23 | |
| 24 | |
| 25 | class TestThinkingReasoningParserAlt(BaseThinkingReasoningParser): |
| 26 | """Alternative test implementation with different tokens.""" |
| 27 | |
| 28 | @property |
| 29 | def start_token(self) -> str: |
| 30 | return "<alt:start>" |
| 31 | |
| 32 | @property |
| 33 | def end_token(self) -> str: |
| 34 | return "<alt:end>" |
| 35 | |
| 36 | |
| 37 | # Use a test model |
| 38 | REASONING_MODEL_NAME = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" |
| 39 | |
| 40 | |
| 41 | @pytest.fixture(scope="module") |
| 42 | def test_tokenizer() -> PreTrainedTokenizerBase: |
| 43 | tokenizer: PreTrainedTokenizerBase = AutoTokenizer.from_pretrained( |
| 44 | REASONING_MODEL_NAME |
| 45 | ) |
| 46 | # Add custom test tokens |
| 47 | test_tokens = ["<test:think>", "</test:think>", "<alt:start>", "<alt:end>"] |
| 48 | existing_tokens = set(tokenizer.get_vocab().keys()) |
| 49 | new_tokens = [token for token in test_tokens if token not in existing_tokens] |
| 50 | if new_tokens: |
| 51 | tokenizer.add_tokens(new_tokens) |
| 52 | return tokenizer |
| 53 | |
| 54 | |
| 55 | class TestBaseThinkingReasoningParserInit: |
| 56 | """Test initialization and basic properties of |
| 57 | BaseThinkingReasoningParser. |
| 58 | """ |
| 59 | |
| 60 | def test_successful_initialization(self, test_tokenizer): |
| 61 | """Test successful initialization with valid tokens.""" |
| 62 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 63 | assert parser.start_token == "<test:think>" |
| 64 | assert parser.end_token == "</test:think>" |
| 65 | assert parser.start_token_id is not None |
| 66 | assert parser.end_token_id is not None |
| 67 | |
| 68 | def test_initialization_with_missing_tokenizer(self): |
| 69 | """Test that initialization fails without tokenizer.""" |
| 70 | with pytest.raises(ValueError, match="model tokenizer must be passed"): |
| 71 | TestThinkingReasoningParser(None) |
| 72 | |
| 73 | def test_initialization_with_missing_tokens(self, test_tokenizer): |
| 74 | """Test that initialization fails when tokens are not in vocabulary.""" |
| 75 | |
| 76 | # Create a parser with tokens not in vocabulary |
| 77 | class MissingTokenParser(BaseThinkingReasoningParser): |
| 78 | @property |
| 79 | def start_token(self) -> str: |
| 80 | return "<missing:start>" |
| 81 | |
| 82 | @property |
| 83 | def end_token(self) -> str: |
| 84 | return "<missing:end>" |
| 85 | |
| 86 | with pytest.raises( |
| 87 | RuntimeError, match="could not locate think start/end tokens" |
| 88 | ): |
| 89 | MissingTokenParser(test_tokenizer) |
| 90 | |
| 91 | def test_initialization_with_empty_tokens(self, test_tokenizer): |
| 92 | """Test that initialization fails with empty token strings.""" |
| 93 | |
| 94 | class EmptyTokenParser(BaseThinkingReasoningParser): |
| 95 | @property |
| 96 | def start_token(self) -> str: |
| 97 | return "" |
| 98 | |
| 99 | @property |
| 100 | def end_token(self) -> str: |
| 101 | return "" |
| 102 | |
| 103 | with pytest.raises( |
| 104 | ValueError, match="start_token and end_token must be defined" |
| 105 | ): |
| 106 | EmptyTokenParser(test_tokenizer) |
| 107 | |
| 108 | |
| 109 | class TestBaseThinkingReasoningParserMethods: |
| 110 | """Test the methods of BaseThinkingReasoningParser.""" |
| 111 | |
| 112 | def test_is_reasoning_end(self, test_tokenizer): |
| 113 | """Test the is_reasoning_end method.""" |
| 114 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 115 | end_token_id = parser.end_token_id |
| 116 | start_token_id = parser.start_token_id |
| 117 | # Test with end token present |
| 118 | assert parser.is_reasoning_end([1, 2, end_token_id, 4]) is True |
| 119 | |
| 120 | # Test without end token |
| 121 | assert parser.is_reasoning_end([1, 2, 3, 4]) is False |
| 122 | |
| 123 | # Test with empty list |
| 124 | assert parser.is_reasoning_end([]) is False |
| 125 | |
| 126 | # Test with interleaved thinking |
| 127 | assert parser.is_reasoning_end([1, start_token_id, 2, end_token_id]) is True |
| 128 | assert parser.is_reasoning_end([1, start_token_id, 2, 3]) is False |
| 129 | assert ( |
| 130 | parser.is_reasoning_end( |
| 131 | [1, start_token_id, 2, end_token_id, 2, 2, start_token_id] |
| 132 | ) |
| 133 | is False |
| 134 | ) |
| 135 | |
| 136 | def test_is_reasoning_end_streaming(self, test_tokenizer): |
| 137 | """Test the is_reasoning_end_streaming method.""" |
| 138 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 139 | end_token_id = parser.end_token_id |
| 140 | start_token_id = parser.start_token_id |
| 141 | |
| 142 | assert ( |
| 143 | parser.is_reasoning_end_streaming([1, 2, end_token_id], [end_token_id]) |
| 144 | is True |
| 145 | ) |
| 146 | assert parser.is_reasoning_end_streaming([1, 2, 3, 4], [4]) is False |
| 147 | assert parser.is_reasoning_end_streaming([], []) is False |
| 148 | assert ( |
| 149 | parser.is_reasoning_end_streaming( |
| 150 | [1, start_token_id, 2, end_token_id], [end_token_id] |
| 151 | ) |
| 152 | is True |
| 153 | ) |
| 154 | assert ( |
| 155 | parser.is_reasoning_end_streaming([1, start_token_id, 2, 3], [3]) is False |
| 156 | ) |
| 157 | assert ( |
| 158 | parser.is_reasoning_end_streaming( |
| 159 | [1, start_token_id, 2, end_token_id, 2, start_token_id, 2], |
| 160 | [2], |
| 161 | ) |
| 162 | is False |
| 163 | ) |
| 164 | assert ( |
| 165 | parser.is_reasoning_end_streaming( |
| 166 | [1, start_token_id, 2, end_token_id, 2, 2], [2] |
| 167 | ) |
| 168 | is False |
| 169 | ) |
| 170 | |
| 171 | def test_count_reasoning_tokens(self, test_tokenizer): |
| 172 | """Count tokens between start/end markers.""" |
| 173 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 174 | start = parser.start_token_id |
| 175 | end = parser.end_token_id |
| 176 | token_ids = [0, start, 11, 12, end, 99] |
| 177 | assert parser.count_reasoning_tokens(token_ids) == 2 |
| 178 | |
| 179 | def test_count_reasoning_tokens_nested(self, test_tokenizer): |
| 180 | """Ensure nested thinking spans count all inner tokens safely.""" |
| 181 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 182 | s = parser.start_token_id |
| 183 | e = parser.end_token_id |
| 184 | token_ids = [s, 1, s, 2, e, 3, e] |
| 185 | # Tokens 1,2,3 are inside reasoning (depth>0) => 3 tokens |
| 186 | assert parser.count_reasoning_tokens(token_ids) == 3 |
| 187 | |
| 188 | def test_extract_content_ids(self, test_tokenizer): |
| 189 | """Test the extract_content_ids method.""" |
| 190 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 191 | end_token_id = parser.end_token_id |
| 192 | |
| 193 | # Test with end token in the middle |
| 194 | input_ids = [1, 2, end_token_id, 4, 5] |
| 195 | content_ids = parser.extract_content_ids(input_ids) |
| 196 | assert content_ids == [4, 5] |
| 197 | |
| 198 | # Test with end token at the end |
| 199 | input_ids = [1, 2, 3, end_token_id] |
| 200 | content_ids = parser.extract_content_ids(input_ids) |
| 201 | assert content_ids == [] |
| 202 | |
| 203 | # Test without end token |
| 204 | input_ids = [1, 2, 3, 4] |
| 205 | content_ids = parser.extract_content_ids(input_ids) |
| 206 | assert content_ids == [] |
| 207 | |
| 208 | # Test with end token as last element (should not extract) |
| 209 | input_ids = [1, 2, 3, end_token_id] |
| 210 | content_ids = parser.extract_content_ids(input_ids) |
| 211 | assert content_ids == [] |
| 212 | |
| 213 | |
| 214 | class TestBaseThinkingReasoningParserExtraction: |
| 215 | """Test reasoning content extraction methods.""" |
| 216 | |
| 217 | def test_extract_reasoning_with_both_tokens(self, test_tokenizer): |
| 218 | """Test extraction when both start and end tokens are present.""" |
| 219 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 220 | request = ChatCompletionRequest(messages=[], model="test-model") |
| 221 | |
| 222 | model_output = "<test:think>This is reasoning</test:think>This is content" |
| 223 | reasoning, content = parser.extract_reasoning(model_output, request) |
| 224 | |
| 225 | assert reasoning == "This is reasoning" |
| 226 | assert content == "This is content" |
| 227 | |
| 228 | def test_extract_reasoning_only_end_token(self, test_tokenizer): |
| 229 | """Test extraction when only end token is present.""" |
| 230 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 231 | request = ChatCompletionRequest(messages=[], model="test-model") |
| 232 | |
| 233 | model_output = "This is reasoning</test:think>This is content" |
| 234 | reasoning, content = parser.extract_reasoning(model_output, request) |
| 235 | |
| 236 | assert reasoning == "This is reasoning" |
| 237 | assert content == "This is content" |
| 238 | |
| 239 | def test_extract_reasoning_no_end_token(self, test_tokenizer): |
| 240 | """Test extraction when no end token is present.""" |
| 241 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 242 | request = ChatCompletionRequest(messages=[], model="test-model") |
| 243 | |
| 244 | model_output = "This is just content" |
| 245 | reasoning, content = parser.extract_reasoning(model_output, request) |
| 246 | |
| 247 | assert reasoning == "This is just content" |
| 248 | assert content is None |
| 249 | |
| 250 | def test_extract_reasoning_empty_output(self, test_tokenizer): |
| 251 | """Test extraction with empty output.""" |
| 252 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 253 | request = ChatCompletionRequest(messages=[], model="test-model") |
| 254 | |
| 255 | model_output = "" |
| 256 | reasoning, content = parser.extract_reasoning(model_output, request) |
| 257 | |
| 258 | assert reasoning == "" |
| 259 | assert content is None |
| 260 | |
| 261 | def test_extract_reasoning_only_tokens(self, test_tokenizer): |
| 262 | """Test extraction with only tokens and no content.""" |
| 263 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 264 | request = ChatCompletionRequest(messages=[], model="test-model") |
| 265 | |
| 266 | model_output = "<test:think></test:think>" |
| 267 | reasoning, content = parser.extract_reasoning(model_output, request) |
| 268 | |
| 269 | assert reasoning == "" |
| 270 | assert content is None |
| 271 | |
| 272 | |
| 273 | class TestBaseThinkingReasoningParserStreaming: |
| 274 | """Test streaming functionality of BaseThinkingReasoningParser.""" |
| 275 | |
| 276 | @pytest.mark.parametrize("streaming", [True, False]) |
| 277 | def test_simple_reasoning_extraction(self, test_tokenizer, streaming): |
| 278 | """Test basic reasoning extraction in both |
| 279 | streaming and non-streaming modes. |
| 280 | """ |
| 281 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 282 | |
| 283 | model_output = [ |
| 284 | "<test:think>", |
| 285 | "Some ", |
| 286 | "reasoning ", |
| 287 | "content", |
| 288 | "</test:think>", |
| 289 | "Final ", |
| 290 | "answer", |
| 291 | ] |
| 292 | |
| 293 | reasoning, content = run_reasoning_extraction( |
| 294 | parser, model_output, streaming=streaming |
| 295 | ) |
| 296 | |
| 297 | assert reasoning == "Some reasoning content" |
| 298 | assert content == "Final answer" |
| 299 | |
| 300 | def test_streaming_with_incremental_deltas(self, test_tokenizer): |
| 301 | """Test streaming processing with small incremental deltas.""" |
| 302 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 303 | |
| 304 | deltas = [ |
| 305 | "<test:think>", |
| 306 | "Some ", |
| 307 | "reasoning ", |
| 308 | "content", |
| 309 | "</test:think>", |
| 310 | "Final ", |
| 311 | "answer", |
| 312 | ] |
| 313 | |
| 314 | reasoning, content = run_reasoning_extraction(parser, deltas, streaming=True) |
| 315 | |
| 316 | assert reasoning == "Some reasoning content" |
| 317 | assert content == "Final answer" |
| 318 | |
| 319 | def test_streaming_with_start_token(self, test_tokenizer): |
| 320 | """Test streaming with start token included.""" |
| 321 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 322 | |
| 323 | deltas = [ |
| 324 | "<test:think>", |
| 325 | "Some ", |
| 326 | "reasoning", |
| 327 | "</test:think>", |
| 328 | "Answer", |
| 329 | ] |
| 330 | |
| 331 | reasoning, content = run_reasoning_extraction(parser, deltas, streaming=True) |
| 332 | |
| 333 | assert reasoning == "Some reasoning" |
| 334 | assert content == "Answer" |
| 335 | |
| 336 | def test_streaming_no_end_token(self, test_tokenizer): |
| 337 | """Test streaming when no end token is encountered.""" |
| 338 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 339 | |
| 340 | deltas = [ |
| 341 | "<test:think>", |
| 342 | "Some ", |
| 343 | "reasoning ", |
| 344 | "without ", |
| 345 | "end", |
| 346 | ] |
| 347 | |
| 348 | reasoning, content = run_reasoning_extraction(parser, deltas, streaming=True) |
| 349 | |
| 350 | assert reasoning == "Some reasoning without end" |
| 351 | assert content is None |
| 352 | |
| 353 | def test_streaming_only_end_token(self, test_tokenizer): |
| 354 | """Test streaming when only end token appears.""" |
| 355 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 356 | |
| 357 | deltas = [ |
| 358 | "<test:think>", |
| 359 | "Reasoning ", |
| 360 | "content", |
| 361 | "</test:think>", |
| 362 | "Final", |
| 363 | ] |
| 364 | |
| 365 | reasoning, content = run_reasoning_extraction(parser, deltas, streaming=True) |
| 366 | |
| 367 | assert reasoning == "Reasoning content" |
| 368 | assert content == "Final" |
| 369 | |
| 370 | |
| 371 | class TestBaseThinkingReasoningParserMultipleImplementations: |
| 372 | """Test that multiple implementations of |
| 373 | BaseThinkingReasoningParser work correctly. |
| 374 | """ |
| 375 | |
| 376 | def test_different_token_implementations(self, test_tokenizer): |
| 377 | """Test that different implementations |
| 378 | with different tokens work independently. |
| 379 | """ |
| 380 | parser1 = TestThinkingReasoningParser(test_tokenizer) |
| 381 | parser2 = TestThinkingReasoningParserAlt(test_tokenizer) |
| 382 | |
| 383 | # Test parser1 |
| 384 | model_output1 = "Reasoning1</test:think>Content1" |
| 385 | reasoning1, content1 = run_reasoning_extraction(parser1, [model_output1]) |
| 386 | assert reasoning1 == "Reasoning1" |
| 387 | assert content1 == "Content1" |
| 388 | |
| 389 | # Test parser2 |
| 390 | model_output2 = "Reasoning2<alt:end>Content2" |
| 391 | reasoning2, content2 = run_reasoning_extraction(parser2, [model_output2]) |
| 392 | assert reasoning2 == "Reasoning2" |
| 393 | assert content2 == "Content2" |
| 394 | |
| 395 | # Verify tokens are different |
| 396 | assert parser1.start_token != parser2.start_token |
| 397 | assert parser1.end_token != parser2.end_token |
| 398 | assert parser1.start_token_id != parser2.start_token_id |
| 399 | assert parser1.end_token_id != parser2.end_token_id |
| 400 | |
| 401 | |
| 402 | class TestBaseThinkingReasoningParserEdgeCases: |
| 403 | """Test edge cases and error conditions.""" |
| 404 | |
| 405 | def test_multiple_end_tokens(self, test_tokenizer): |
| 406 | """Test behavior with multiple end tokens.""" |
| 407 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 408 | |
| 409 | model_output = "First</test:think>Middle</test:think>Last" |
| 410 | reasoning, content = run_reasoning_extraction(parser, [model_output]) |
| 411 | |
| 412 | # Should stop at first end token |
| 413 | assert reasoning == "First" |
| 414 | assert content == "Middle</test:think>Last" |
| 415 | |
| 416 | def test_nested_tokens(self, test_tokenizer): |
| 417 | """Test behavior with nested-like token patterns.""" |
| 418 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 419 | |
| 420 | model_output = "<test:think>Outer<test:think>Inner</test:think>Content" |
| 421 | reasoning, content = run_reasoning_extraction(parser, [model_output]) |
| 422 | |
| 423 | # Should process normally, start from first start token |
| 424 | assert reasoning == "Outer<test:think>Inner" |
| 425 | assert content == "Content" |
| 426 | |
| 427 | def test_malformed_tokens(self, test_tokenizer): |
| 428 | """Test behavior with malformed token-like strings.""" |
| 429 | parser = TestThinkingReasoningParser(test_tokenizer) |
| 430 | |
| 431 | model_output = "<test:thinking>Not a real token</test:thinking>Content" |
| 432 | reasoning, content = run_reasoning_extraction(parser, [model_output]) |
| 433 | |
| 434 | # Should treat as regular content since tokens don't match exactly |
| 435 | assert reasoning == ("<test:thinking>Not a real token</test:thinking>Content") |
| 436 | assert content is None |