test_base_thinking_reasoning_parser.py 15.4 KB
Newer Older
1
2
3
4
5
6
7
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import pytest
from transformers import AutoTokenizer

from tests.reasoning.utils import run_reasoning_extraction
8
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser


# Create a concrete test implementation of BaseThinkingReasoningParser
class TestThinkingReasoningParser(BaseThinkingReasoningParser):
    """Test implementation of BaseThinkingReasoningParser."""

    @property
    def start_token(self) -> str:
        return "<test:think>"

    @property
    def end_token(self) -> str:
        return "</test:think>"


class TestThinkingReasoningParserAlt(BaseThinkingReasoningParser):
    """Alternative test implementation with different tokens."""

    @property
    def start_token(self) -> str:
        return "<alt:start>"

    @property
    def end_token(self) -> str:
        return "<alt:end>"


# Use a test model
REASONING_MODEL_NAME = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"


@pytest.fixture(scope="module")
def test_tokenizer():
    tokenizer = AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)
    # Add custom test tokens
    test_tokens = ["<test:think>", "</test:think>", "<alt:start>", "<alt:end>"]
    existing_tokens = set(tokenizer.get_vocab().keys())
47
    new_tokens = [token for token in test_tokens if token not in existing_tokens]
48
49
50
51
52
53
54
    if new_tokens:
        tokenizer.add_tokens(new_tokens)
    return tokenizer


class TestBaseThinkingReasoningParserInit:
    """
55
56
    Test initialization and basic properties of
    BaseThinkingReasoningParser.
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
    """

    def test_successful_initialization(self, test_tokenizer):
        """Test successful initialization with valid tokens."""
        parser = TestThinkingReasoningParser(test_tokenizer)
        assert parser.start_token == "<test:think>"
        assert parser.end_token == "</test:think>"
        assert parser.start_token_id is not None
        assert parser.end_token_id is not None

    def test_initialization_with_missing_tokenizer(self):
        """Test that initialization fails without tokenizer."""
        with pytest.raises(ValueError, match="model tokenizer must be passed"):
            TestThinkingReasoningParser(None)

    def test_initialization_with_missing_tokens(self, test_tokenizer):
        """Test that initialization fails when tokens are not in vocabulary."""

        # Create a parser with tokens not in vocabulary
        class MissingTokenParser(BaseThinkingReasoningParser):
            @property
            def start_token(self) -> str:
                return "<missing:start>"

            @property
            def end_token(self) -> str:
                return "<missing:end>"

85
86
87
        with pytest.raises(
            RuntimeError, match="could not locate think start/end tokens"
        ):
88
89
90
91
92
93
94
95
96
97
98
99
100
101
            MissingTokenParser(test_tokenizer)

    def test_initialization_with_empty_tokens(self, test_tokenizer):
        """Test that initialization fails with empty token strings."""

        class EmptyTokenParser(BaseThinkingReasoningParser):
            @property
            def start_token(self) -> str:
                return ""

            @property
            def end_token(self) -> str:
                return ""

102
103
104
        with pytest.raises(
            ValueError, match="start_token and end_token must be defined"
        ):
105
106
107
108
109
110
111
112
113
114
            EmptyTokenParser(test_tokenizer)


class TestBaseThinkingReasoningParserMethods:
    """Test the methods of BaseThinkingReasoningParser."""

    def test_is_reasoning_end(self, test_tokenizer):
        """Test the is_reasoning_end method."""
        parser = TestThinkingReasoningParser(test_tokenizer)
        end_token_id = parser.end_token_id
115
        start_token_id = parser.start_token_id
116
117
118
119
120
121
122
123
124
        # Test with end token present
        assert parser.is_reasoning_end([1, 2, end_token_id, 4]) is True

        # Test without end token
        assert parser.is_reasoning_end([1, 2, 3, 4]) is False

        # Test with empty list
        assert parser.is_reasoning_end([]) is False

125
126
127
128
129
130
131
132
        # Test with interleaved thinking
        assert parser.is_reasoning_end([1, start_token_id, 2, end_token_id]) is True
        assert parser.is_reasoning_end([1, start_token_id, 2, 3]) is False
        assert (
            parser.is_reasoning_end(
                [1, start_token_id, 2, end_token_id, 2, 2, start_token_id]
            )
            is False
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
        )

    def test_is_reasoning_end_streaming(self, test_tokenizer):
        """Test the is_reasoning_end_streaming method."""
        parser = TestThinkingReasoningParser(test_tokenizer)
        end_token_id = parser.end_token_id
        start_token_id = parser.start_token_id

        assert (
            parser.is_reasoning_end_streaming([1, 2, end_token_id], [end_token_id])
            is True
        )
        assert parser.is_reasoning_end_streaming([1, 2, 3, 4], [4]) is False
        assert parser.is_reasoning_end_streaming([], []) is False
        assert (
            parser.is_reasoning_end_streaming(
                [1, start_token_id, 2, end_token_id], [end_token_id]
            )
            is True
        )
        assert (
            parser.is_reasoning_end_streaming([1, start_token_id, 2, 3], [3]) is False
        )
        assert (
            parser.is_reasoning_end_streaming(
                [1, start_token_id, 2, end_token_id, 2, start_token_id, 2],
                [2],
            )
            is False
        )
        assert (
            parser.is_reasoning_end_streaming(
                [1, start_token_id, 2, end_token_id, 2, 2], [2]
            )
            is False
168
169
        )

170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
    def test_count_reasoning_tokens(self, test_tokenizer):
        """Count tokens between start/end markers."""
        parser = TestThinkingReasoningParser(test_tokenizer)
        start = parser.start_token_id
        end = parser.end_token_id
        token_ids = [0, start, 11, 12, end, 99]
        assert parser.count_reasoning_tokens(token_ids) == 2

    def test_count_reasoning_tokens_nested(self, test_tokenizer):
        """Ensure nested thinking spans count all inner tokens safely."""
        parser = TestThinkingReasoningParser(test_tokenizer)
        s = parser.start_token_id
        e = parser.end_token_id
        token_ids = [s, 1, s, 2, e, 3, e]
        # Tokens 1,2,3 are inside reasoning (depth>0) => 3 tokens
        assert parser.count_reasoning_tokens(token_ids) == 3

187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
    def test_extract_content_ids(self, test_tokenizer):
        """Test the extract_content_ids method."""
        parser = TestThinkingReasoningParser(test_tokenizer)
        end_token_id = parser.end_token_id

        # Test with end token in the middle
        input_ids = [1, 2, end_token_id, 4, 5]
        content_ids = parser.extract_content_ids(input_ids)
        assert content_ids == [4, 5]

        # Test with end token at the end
        input_ids = [1, 2, 3, end_token_id]
        content_ids = parser.extract_content_ids(input_ids)
        assert content_ids == []

        # Test without end token
        input_ids = [1, 2, 3, 4]
        content_ids = parser.extract_content_ids(input_ids)
        assert content_ids == []

        # Test with end token as last element (should not extract)
        input_ids = [1, 2, 3, end_token_id]
        content_ids = parser.extract_content_ids(input_ids)
        assert content_ids == []


class TestBaseThinkingReasoningParserExtraction:
    """Test reasoning content extraction methods."""

216
    def test_extract_reasoning_with_both_tokens(self, test_tokenizer):
217
218
219
220
        """Test extraction when both start and end tokens are present."""
        parser = TestThinkingReasoningParser(test_tokenizer)
        request = ChatCompletionRequest(messages=[], model="test-model")

221
        model_output = "<test:think>This is reasoning</test:think>This is content"
222
        reasoning, content = parser.extract_reasoning(model_output, request)
223
224
225
226

        assert reasoning == "This is reasoning"
        assert content == "This is content"

227
    def test_extract_reasoning_only_end_token(self, test_tokenizer):
228
229
230
231
        """Test extraction when only end token is present."""
        parser = TestThinkingReasoningParser(test_tokenizer)
        request = ChatCompletionRequest(messages=[], model="test-model")

232
        model_output = "This is reasoning</test:think>This is content"
233
        reasoning, content = parser.extract_reasoning(model_output, request)
234
235
236
237

        assert reasoning == "This is reasoning"
        assert content == "This is content"

238
    def test_extract_reasoning_no_end_token(self, test_tokenizer):
239
240
241
242
243
        """Test extraction when no end token is present."""
        parser = TestThinkingReasoningParser(test_tokenizer)
        request = ChatCompletionRequest(messages=[], model="test-model")

        model_output = "This is just content"
244
        reasoning, content = parser.extract_reasoning(model_output, request)
245
246
247
248

        assert reasoning == "This is just content"
        assert content is None

249
    def test_extract_reasoning_empty_output(self, test_tokenizer):
250
251
252
253
254
        """Test extraction with empty output."""
        parser = TestThinkingReasoningParser(test_tokenizer)
        request = ChatCompletionRequest(messages=[], model="test-model")

        model_output = ""
255
        reasoning, content = parser.extract_reasoning(model_output, request)
256
257
258
259

        assert reasoning == ""
        assert content is None

260
    def test_extract_reasoning_only_tokens(self, test_tokenizer):
261
262
263
264
        """Test extraction with only tokens and no content."""
        parser = TestThinkingReasoningParser(test_tokenizer)
        request = ChatCompletionRequest(messages=[], model="test-model")

265
        model_output = "<test:think></test:think>"
266
        reasoning, content = parser.extract_reasoning(model_output, request)
267
268
269
270
271
272
273
274
275
276
277

        assert reasoning == ""
        assert content is None


class TestBaseThinkingReasoningParserStreaming:
    """Test streaming functionality of BaseThinkingReasoningParser."""

    @pytest.mark.parametrize("streaming", [True, False])
    def test_simple_reasoning_extraction(self, test_tokenizer, streaming):
        """
278
279
        Test basic reasoning extraction in both
        streaming and non-streaming modes.
280
281
282
283
        """
        parser = TestThinkingReasoningParser(test_tokenizer)

        model_output = [
284
285
286
287
288
289
290
            "<test:think>",
            "Some ",
            "reasoning ",
            "content",
            "</test:think>",
            "Final ",
            "answer",
291
292
        ]

293
294
295
        reasoning, content = run_reasoning_extraction(
            parser, model_output, streaming=streaming
        )
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313

        assert reasoning == "Some reasoning content"
        assert content == "Final answer"

    def test_streaming_with_incremental_deltas(self, test_tokenizer):
        """Test streaming processing with small incremental deltas."""
        parser = TestThinkingReasoningParser(test_tokenizer)

        deltas = [
            "<test:think>",
            "Some ",
            "reasoning ",
            "content",
            "</test:think>",
            "Final ",
            "answer",
        ]

314
        reasoning, content = run_reasoning_extraction(parser, deltas, streaming=True)
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330

        assert reasoning == "Some reasoning content"
        assert content == "Final answer"

    def test_streaming_with_start_token(self, test_tokenizer):
        """Test streaming with start token included."""
        parser = TestThinkingReasoningParser(test_tokenizer)

        deltas = [
            "<test:think>",
            "Some ",
            "reasoning",
            "</test:think>",
            "Answer",
        ]

331
        reasoning, content = run_reasoning_extraction(parser, deltas, streaming=True)
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347

        assert reasoning == "Some reasoning"
        assert content == "Answer"

    def test_streaming_no_end_token(self, test_tokenizer):
        """Test streaming when no end token is encountered."""
        parser = TestThinkingReasoningParser(test_tokenizer)

        deltas = [
            "<test:think>",
            "Some ",
            "reasoning ",
            "without ",
            "end",
        ]

348
        reasoning, content = run_reasoning_extraction(parser, deltas, streaming=True)
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364

        assert reasoning == "Some reasoning without end"
        assert content is None

    def test_streaming_only_end_token(self, test_tokenizer):
        """Test streaming when only end token appears."""
        parser = TestThinkingReasoningParser(test_tokenizer)

        deltas = [
            "<test:think>",
            "Reasoning ",
            "content",
            "</test:think>",
            "Final",
        ]

365
        reasoning, content = run_reasoning_extraction(parser, deltas, streaming=True)
366
367
368
369
370
371
372

        assert reasoning == "Reasoning content"
        assert content == "Final"


class TestBaseThinkingReasoningParserMultipleImplementations:
    """
373
374
    Test that multiple implementations of
    BaseThinkingReasoningParser work correctly.
375
376
377
378
    """

    def test_different_token_implementations(self, test_tokenizer):
        """
379
380
        Test that different implementations
        with different tokens work independently.
381
382
383
384
385
        """
        parser1 = TestThinkingReasoningParser(test_tokenizer)
        parser2 = TestThinkingReasoningParserAlt(test_tokenizer)

        # Test parser1
386
387
        model_output1 = "Reasoning1</test:think>Content1"
        reasoning1, content1 = run_reasoning_extraction(parser1, [model_output1])
388
389
390
391
392
        assert reasoning1 == "Reasoning1"
        assert content1 == "Content1"

        # Test parser2
        model_output2 = "Reasoning2<alt:end>Content2"
393
        reasoning2, content2 = run_reasoning_extraction(parser2, [model_output2])
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
        assert reasoning2 == "Reasoning2"
        assert content2 == "Content2"

        # Verify tokens are different
        assert parser1.start_token != parser2.start_token
        assert parser1.end_token != parser2.end_token
        assert parser1.start_token_id != parser2.start_token_id
        assert parser1.end_token_id != parser2.end_token_id


class TestBaseThinkingReasoningParserEdgeCases:
    """Test edge cases and error conditions."""

    def test_multiple_end_tokens(self, test_tokenizer):
        """Test behavior with multiple end tokens."""
        parser = TestThinkingReasoningParser(test_tokenizer)

411
        model_output = "First</test:think>Middle</test:think>Last"
412
413
414
415
416
417
418
419
420
421
        reasoning, content = run_reasoning_extraction(parser, [model_output])

        # Should stop at first end token
        assert reasoning == "First"
        assert content == "Middle</test:think>Last"

    def test_nested_tokens(self, test_tokenizer):
        """Test behavior with nested-like token patterns."""
        parser = TestThinkingReasoningParser(test_tokenizer)

422
        model_output = "<test:think>Outer<test:think>Inner</test:think>Content"
423
424
425
426
427
428
429
430
431
432
        reasoning, content = run_reasoning_extraction(parser, [model_output])

        # Should process normally, start from first start token
        assert reasoning == "Outer<test:think>Inner"
        assert content == "Content"

    def test_malformed_tokens(self, test_tokenizer):
        """Test behavior with malformed token-like strings."""
        parser = TestThinkingReasoningParser(test_tokenizer)

433
        model_output = "<test:thinking>Not a real token</test:thinking>Content"
434
435
436
        reasoning, content = run_reasoning_extraction(parser, [model_output])

        # Should treat as regular content since tokens don't match exactly
437
        assert reasoning == ("<test:thinking>Not a real token</test:thinking>Content")
438
        assert content is None