| 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.entrypoints.openai.chat_completion.protocol import ( |
| 8 | ChatCompletionRequest, |
| 9 | ) |
| 10 | from vllm.reasoning import ReasoningParser, ReasoningParserManager |
| 11 | |
| 12 | # Using mistral tokenizer as a generic mock since the actual model is not on HF |
| 13 | from vllm.tokenizers.registry import get_tokenizer |
| 14 | |
| 15 | parser_name = "gemma4" |
| 16 | |
| 17 | |
| 18 | @pytest.fixture(scope="module") |
| 19 | def generic_tokenizer(): |
| 20 | return get_tokenizer("google/gemma-4-E2B-it") |
| 21 | |
| 22 | |
| 23 | INVALID_SIMPLE_NONSTREAMING = { |
| 24 | "output": "This is a reasoning section<channel|>This is the rest", |
| 25 | "reasoning": "This is a reasoning section", |
| 26 | "content": "This is the rest", |
| 27 | "is_reasoning_end": True, |
| 28 | } |
| 29 | INVALID_SIMPLE_STREAMING = { |
| 30 | "output": "This is a reasoning section<channel|>This is the rest", |
| 31 | "reasoning": None, |
| 32 | "content": "This is a reasoning sectionThis is the rest", |
| 33 | "is_reasoning_end": True, |
| 34 | } |
| 35 | INVALID_COMPLETE_NONSTREAMING = { |
| 36 | "output": "This is a reasoning section<channel|>", |
| 37 | "reasoning": "This is a reasoning section", |
| 38 | "content": None, |
| 39 | "is_reasoning_end": True, |
| 40 | } |
| 41 | INVALID_COMPLETE_STREAMING = { |
| 42 | "output": "This is a reasoning section<channel|>", |
| 43 | "reasoning": None, |
| 44 | "content": "This is a reasoning section", |
| 45 | "is_reasoning_end": True, |
| 46 | } |
| 47 | NO_CONTENT = { |
| 48 | "output": "<|channel>This is reasoning", |
| 49 | "reasoning": "This is reasoning", |
| 50 | "content": None, |
| 51 | "is_reasoning_end": False, |
| 52 | } |
| 53 | NO_REASONING = { |
| 54 | "output": "This is content", |
| 55 | "reasoning": None, |
| 56 | "content": "This is content", |
| 57 | "is_reasoning_end": True, |
| 58 | } |
| 59 | REASONING_WITH_CHANNEL = { |
| 60 | "output": "<|channel>This is a reasoning section<channel|>This is the rest", |
| 61 | "reasoning": "This is a reasoning section", |
| 62 | "content": "This is the rest", |
| 63 | "is_reasoning_end": True, |
| 64 | } |
| 65 | COMPLETE_REASONING_WITH_CHANNEL = { |
| 66 | "output": "<|channel>This is a reasoning section<channel|>", |
| 67 | "reasoning": "This is a reasoning section", |
| 68 | "content": None, |
| 69 | "is_reasoning_end": True, |
| 70 | } |
| 71 | MULTIPLE_LINES_WITH_CHANNEL = { |
| 72 | "output": "<|channel>This\nThat<channel|>This is the rest\nThat", |
| 73 | "reasoning": "This\nThat", |
| 74 | "content": "This is the rest\nThat", |
| 75 | "is_reasoning_end": True, |
| 76 | } |
| 77 | CHANNEL_NO_END = { |
| 78 | "output": "<|channel>This is a reasoning section", |
| 79 | "reasoning": "This is a reasoning section", |
| 80 | "content": None, |
| 81 | "is_reasoning_end": False, |
| 82 | } |
| 83 | EMPTY = { |
| 84 | "output": "", |
| 85 | "reasoning": None, |
| 86 | "content": None, |
| 87 | "is_reasoning_end": True, |
| 88 | } |
| 89 | NEW_LINE_NONSTREAMING = { |
| 90 | "output": ( |
| 91 | "Before\n<|channel>This is a reasoning section<channel|>\nThis is the rest" |
| 92 | ), |
| 93 | "reasoning": "This is a reasoning section", |
| 94 | "content": "Before\n\nThis is the rest", |
| 95 | "is_reasoning_end": True, |
| 96 | } |
| 97 | NEW_LINE_STREAMING = { |
| 98 | "output": ( |
| 99 | "Before\n<|channel>This is a reasoning section<channel|>\nThis is the rest" |
| 100 | ), |
| 101 | "reasoning": "This is a reasoning section", |
| 102 | "content": "Before\n\nThis is the rest", |
| 103 | "is_reasoning_end": True, |
| 104 | } |
| 105 | |
| 106 | THOUGHT_PREFIX = { |
| 107 | "output": "<|channel>thought\nActual reasoning here<channel|>Final answer", |
| 108 | "reasoning": "Actual reasoning here", |
| 109 | "content": "Final answer", |
| 110 | "is_reasoning_end": True, |
| 111 | } |
| 112 | THOUGHT_PREFIX_ONLY = { |
| 113 | "output": "<|channel>thought\n<channel|>", |
| 114 | "reasoning": None, |
| 115 | "content": None, |
| 116 | "is_reasoning_end": True, |
| 117 | } |
| 118 | THOUGHT_PREFIX_MULTILINE = { |
| 119 | "output": "<|channel>thought\nLine1\nLine2<channel|>Answer", |
| 120 | "reasoning": "Line1\nLine2", |
| 121 | "content": "Answer", |
| 122 | "is_reasoning_end": True, |
| 123 | } |
| 124 | # "thousand" starts like "thought" but diverges — exercises Case 2→3 in streaming. |
| 125 | THOUGHT_PREFIX_DIVERGE = { |
| 126 | "output": "<|channel>thousand reasons<channel|>Done", |
| 127 | "reasoning": "thousand reasons", |
| 128 | "content": "Done", |
| 129 | "is_reasoning_end": True, |
| 130 | } |
| 131 | # The model isn't reasoning if we're generating tool calls. |
| 132 | TOOL_CALL_STARTED = { |
| 133 | "output": "<|tool_call>", |
| 134 | "reasoning": None, |
| 135 | "content": "<|tool_call>", |
| 136 | "is_reasoning_end": True, |
| 137 | } |
| 138 | |
| 139 | TEST_CASES = [ |
| 140 | pytest.param(False, INVALID_SIMPLE_NONSTREAMING, id="invalid_simple"), |
| 141 | pytest.param(True, INVALID_SIMPLE_STREAMING, id="invalid_simple_streaming"), |
| 142 | pytest.param(False, INVALID_COMPLETE_NONSTREAMING, id="invalid_complete"), |
| 143 | pytest.param(True, INVALID_COMPLETE_STREAMING, id="invalid_complete_streaming"), |
| 144 | pytest.param(False, NO_CONTENT, id="no_content"), |
| 145 | pytest.param(False, NO_REASONING, id="no_reasoning"), |
| 146 | pytest.param(False, REASONING_WITH_CHANNEL, id="reasoning"), |
| 147 | pytest.param(True, REASONING_WITH_CHANNEL, id="reasoning_streaming"), |
| 148 | pytest.param(False, COMPLETE_REASONING_WITH_CHANNEL, id="complete_reasoning"), |
| 149 | pytest.param( |
| 150 | True, COMPLETE_REASONING_WITH_CHANNEL, id="complete_reasoning_streaming" |
| 151 | ), |
| 152 | pytest.param(False, MULTIPLE_LINES_WITH_CHANNEL, id="multiple_lines"), |
| 153 | pytest.param(True, MULTIPLE_LINES_WITH_CHANNEL, id="multiple_lines_streaming"), |
| 154 | pytest.param(False, CHANNEL_NO_END, id="no_end"), |
| 155 | pytest.param(True, CHANNEL_NO_END, id="no_end_streaming"), |
| 156 | pytest.param(False, EMPTY, id="empty"), |
| 157 | pytest.param(False, NEW_LINE_NONSTREAMING, id="new_line"), |
| 158 | pytest.param(True, NEW_LINE_STREAMING, id="new_line_streaming"), |
| 159 | pytest.param(False, THOUGHT_PREFIX, id="thought_prefix"), |
| 160 | pytest.param(True, THOUGHT_PREFIX, id="thought_prefix_streaming"), |
| 161 | pytest.param(False, THOUGHT_PREFIX_ONLY, id="thought_prefix_only"), |
| 162 | pytest.param(True, THOUGHT_PREFIX_ONLY, id="thought_prefix_only_streaming"), |
| 163 | pytest.param(False, THOUGHT_PREFIX_MULTILINE, id="thought_prefix_multiline"), |
| 164 | pytest.param( |
| 165 | True, THOUGHT_PREFIX_MULTILINE, id="thought_prefix_multiline_streaming" |
| 166 | ), |
| 167 | pytest.param(False, THOUGHT_PREFIX_DIVERGE, id="thought_prefix_diverge"), |
| 168 | pytest.param(True, THOUGHT_PREFIX_DIVERGE, id="thought_prefix_diverge_streaming"), |
| 169 | pytest.param(False, TOOL_CALL_STARTED, id="tool_call_started"), |
| 170 | pytest.param(True, TOOL_CALL_STARTED, id="tool_call_started_streaming"), |
| 171 | ] |
| 172 | |
| 173 | |
| 174 | def gemma4_encode_output(generic_tokenizer, output: str) -> list[int]: |
| 175 | # Resolve token IDs dynamically from the real tokenizer |
| 176 | vocab = generic_tokenizer.get_vocab() |
| 177 | start_token_id = vocab["<|channel>"] |
| 178 | end_token_id = vocab["<channel|>"] |
| 179 | |
| 180 | index_start = output.find("<|channel>") |
| 181 | len_start = len("<|channel>") |
| 182 | index_end = output.find("<channel|>") |
| 183 | len_end = len("<channel|>") |
| 184 | |
| 185 | output_tokens = [] |
| 186 | |
| 187 | def _encode(text: str) -> list[int]: |
| 188 | if not text: |
| 189 | return [] |
| 190 | # Handle both raw transformers and vLLM wrappers |
| 191 | enc = getattr(generic_tokenizer, "tokenizer", generic_tokenizer) |
| 192 | try: |
| 193 | return enc.encode(text, add_special_tokens=False) |
| 194 | except TypeError: |
| 195 | return enc.encode(text) |
| 196 | |
| 197 | if index_start != -1: |
| 198 | output_before = output[:index_start] |
| 199 | output_tokens += _encode(output_before) |
| 200 | output_tokens += [start_token_id] |
| 201 | |
| 202 | if index_end != -1: |
| 203 | output_middle = output[index_start + len_start : index_end] |
| 204 | output_after = output[index_end + len_end :] |
| 205 | output_tokens += _encode(output_middle) |
| 206 | output_tokens += [end_token_id] |
| 207 | output_tokens += _encode(output_after) |
| 208 | else: |
| 209 | output_middle = output[index_start + len_start :] |
| 210 | output_tokens += _encode(output_middle) |
| 211 | elif index_end != -1: |
| 212 | output_before = output[:index_end] |
| 213 | output_after = output[index_end + len_end :] |
| 214 | output_tokens += _encode(output_before) |
| 215 | output_tokens += [end_token_id] |
| 216 | output_tokens += _encode(output_after) |
| 217 | else: |
| 218 | output_tokens += _encode(output) |
| 219 | |
| 220 | return output_tokens |
| 221 | |
| 222 | |
| 223 | @pytest.mark.parametrize("streaming, param_dict", TEST_CASES) |
| 224 | def test_gemma4_reasoning( |
| 225 | streaming: bool, |
| 226 | param_dict: dict, |
| 227 | generic_tokenizer, |
| 228 | ): |
| 229 | output = param_dict["output"] |
| 230 | output_tokens = gemma4_encode_output(generic_tokenizer, output) |
| 231 | |
| 232 | parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)( |
| 233 | generic_tokenizer |
| 234 | ) |
| 235 | |
| 236 | # We use the generic run_reasoning_extraction from utils |
| 237 | # Use decode per token to get standard spaces instead of |
| 238 | # SentencePiece space characters |
| 239 | output_token_strings = [generic_tokenizer.decode([t]) for t in output_tokens] |
| 240 | reasoning, content = run_reasoning_extraction( |
| 241 | parser, output_token_strings, streaming=streaming |
| 242 | ) |
| 243 | |
| 244 | assert reasoning == param_dict["reasoning"] |
| 245 | assert content == param_dict["content"] |
| 246 | |
| 247 | # Test is_reasoning_end |
| 248 | is_reasoning_end = parser.is_reasoning_end(output_tokens) |
| 249 | assert is_reasoning_end == param_dict["is_reasoning_end"] |
| 250 | |
| 251 | |
| 252 | def test_gemma4_adjust_request(generic_tokenizer): |
| 253 | parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)( |
| 254 | generic_tokenizer |
| 255 | ) |
| 256 | |
| 257 | request = ChatCompletionRequest(messages=[], model="test-model") |
| 258 | assert request.skip_special_tokens is True |
| 259 | |
| 260 | result = parser.adjust_request(request) |
| 261 | assert result.skip_special_tokens is False |
| 262 | assert result is request |
| 263 | |
| 264 | |
| 265 | @pytest.mark.parametrize( |
| 266 | ("chat_template_kwargs", "expected_is_reasoning_end"), |
| 267 | [ |
| 268 | pytest.param({}, True, id="omitted_enable_thinking"), |
| 269 | pytest.param({"enable_thinking": False}, True, id="thinking_disabled"), |
| 270 | pytest.param({"enable_thinking": True}, False, id="thinking_enabled"), |
| 271 | ], |
| 272 | ) |
| 273 | def test_gemma4_new_turn_reasoning_end_matches_enable_thinking( |
| 274 | generic_tokenizer, |
| 275 | chat_template_kwargs, |
| 276 | expected_is_reasoning_end, |
| 277 | ): |
| 278 | output = ( |
| 279 | "<|channel>thought\n1st thought<channel|>1st content<turn|>\n" |
| 280 | "<|turn>user\nThanks<|turn>model\n" |
| 281 | ) |
| 282 | output_tokens = gemma4_encode_output(generic_tokenizer, output) |
| 283 | parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)( |
| 284 | generic_tokenizer, |
| 285 | chat_template_kwargs=chat_template_kwargs, |
| 286 | ) |
| 287 | is_reasoning_end = parser.is_reasoning_end(output_tokens) |
| 288 | assert is_reasoning_end is expected_is_reasoning_end |