| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import cast |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | from vllm.entrypoints.generate.base.protocol import DeltaMessage |
| 9 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 10 | ChatCompletionRequest, |
| 11 | ) |
| 12 | from vllm.parser.kimi_k3 import KimiK3Parser |
| 13 | from vllm.parser.parser_manager import ParserManager |
| 14 | from vllm.reasoning.kimi_k3_reasoning_parser import KimiK3ReasoningParser |
| 15 | from vllm.tokenizers import TokenizerLike |
| 16 | |
| 17 | pytestmark = pytest.mark.skip_global_cleanup |
| 18 | |
| 19 | OPEN = "<|open|>" |
| 20 | CLOSE = "<|close|>" |
| 21 | SEP = "<|sep|>" |
| 22 | THINK_OPEN = f"{OPEN}think{SEP}" |
| 23 | THINK_CLOSE = f"{CLOSE}think{SEP}" |
| 24 | RESPONSE_OPEN = f"{OPEN}response{SEP}" |
| 25 | OPEN_IDS = [1, 2, 3] |
| 26 | CLOSE_IDS = [4, 2, 3] |
| 27 | RESPONSE_OPEN_IDS = [ord(ch) for ch in RESPONSE_OPEN] |
| 28 | |
| 29 | |
| 30 | class DummyTokenizer: |
| 31 | def __len__(self) -> int: |
| 32 | return 1 |
| 33 | |
| 34 | def get_vocab(self) -> dict[str, int]: |
| 35 | return {} |
| 36 | |
| 37 | def encode(self, text: str, *args, **kwargs) -> list[int]: |
| 38 | if text == THINK_OPEN: |
| 39 | return [1, 2, 3] |
| 40 | if text == THINK_CLOSE: |
| 41 | return [4, 2, 3] |
| 42 | return [ord(ch) for ch in text] |
| 43 | |
| 44 | |
| 45 | def _dummy_tokenizer() -> TokenizerLike: |
| 46 | return cast(TokenizerLike, DummyTokenizer()) |
| 47 | |
| 48 | |
| 49 | class ReasoningOnlyParser(KimiK3Parser): |
| 50 | reasoning_parser_cls = KimiK3ReasoningParser |
| 51 | |
| 52 | |
| 53 | def test_parser_manager_selects_kimi_k3_parser_for_reasoning_only(): |
| 54 | parser_cls = ParserManager.get_parser(reasoning_parser_name="kimi_k3") |
| 55 | |
| 56 | assert parser_cls is not None |
| 57 | assert issubclass(parser_cls, KimiK3Parser) |
| 58 | assert parser_cls.reasoning_parser_cls is KimiK3ReasoningParser |
| 59 | assert parser_cls.tool_parser_cls is None |
| 60 | |
| 61 | |
| 62 | def test_parser_selection_thinking_disabled(): |
| 63 | parser = KimiK3ReasoningParser( |
| 64 | DummyTokenizer(), chat_template_kwargs={"thinking": False} |
| 65 | ) |
| 66 | |
| 67 | assert parser._thinking_enabled is False |
| 68 | |
| 69 | |
| 70 | def test_extract_reasoning_with_xtml_tags(): |
| 71 | parser = KimiK3ReasoningParser(DummyTokenizer()) |
| 72 | request = ChatCompletionRequest(model="test-model", messages=[]) |
| 73 | |
| 74 | reasoning, content = parser.extract_reasoning_content( |
| 75 | f"{THINK_OPEN}step{THINK_CLOSE}{RESPONSE_OPEN}answer", |
| 76 | request, |
| 77 | ) |
| 78 | |
| 79 | assert reasoning == "step" |
| 80 | assert content == "answer" |
| 81 | |
| 82 | |
| 83 | def test_extract_reasoning_with_generation_prefix_consumed(): |
| 84 | parser = KimiK3ReasoningParser(DummyTokenizer()) |
| 85 | request = ChatCompletionRequest(model="test-model", messages=[]) |
| 86 | |
| 87 | reasoning, content = parser.extract_reasoning_content( |
| 88 | f"step{THINK_CLOSE}{RESPONSE_OPEN}answer", |
| 89 | request, |
| 90 | ) |
| 91 | |
| 92 | assert reasoning == "step" |
| 93 | assert content == "answer" |
| 94 | |
| 95 | |
| 96 | def test_delegating_parser_strips_response_wrapper_without_tool_parser(): |
| 97 | parser = ReasoningOnlyParser(_dummy_tokenizer()) |
| 98 | request = ChatCompletionRequest(model="test-model", messages=[]) |
| 99 | |
| 100 | reasoning, content, tool_calls = parser.parse( |
| 101 | f"{THINK_OPEN}step{THINK_CLOSE}{RESPONSE_OPEN}answer", |
| 102 | request, |
| 103 | ) |
| 104 | |
| 105 | assert reasoning == "step" |
| 106 | assert content == "answer" |
| 107 | assert tool_calls == [] |
| 108 | |
| 109 | |
| 110 | def test_is_reasoning_end_uses_full_input_ids(): |
| 111 | parser = KimiK3ReasoningParser(DummyTokenizer()) |
| 112 | |
| 113 | assert not parser.is_reasoning_end([4, 2]) |
| 114 | assert parser.is_reasoning_end([4, 2, 3]) |
| 115 | |
| 116 | |
| 117 | def test_is_reasoning_end_ignores_stale_close_from_prior_turn(): |
| 118 | # DummyTokenizer: THINK_OPEN -> [1, 2, 3], THINK_CLOSE -> [4, 2, 3]. |
| 119 | # Multi-turn / agent continuation: a prior turn's think channel (its close |
| 120 | # marker) is kept in the prompt, then the current turn opens a new think |
| 121 | # block that has not closed yet. Reasoning must read as NOT ended, otherwise |
| 122 | # the structured-output gate constrains the current turn's reasoning. |
| 123 | parser = KimiK3ReasoningParser(DummyTokenizer()) |
| 124 | |
| 125 | stale_close = [4, 2, 3] |
| 126 | new_open = [1, 2, 3] |
| 127 | # prior close, then current-turn open still unclosed -> not ended |
| 128 | assert not parser.is_reasoning_end([*stale_close, *new_open]) |
| 129 | # ...then the current turn emits its own close -> ended |
| 130 | assert parser.is_reasoning_end([*stale_close, *new_open, *stale_close]) |
| 131 | # open with no close yet -> not ended |
| 132 | assert not parser.is_reasoning_end([*new_open]) |
| 133 | |
| 134 | |
| 135 | @pytest.mark.parametrize( |
| 136 | ("token_ids", "expected"), |
| 137 | [ |
| 138 | pytest.param( |
| 139 | [*OPEN_IDS, 9, 10, *CLOSE_IDS, *RESPONSE_OPEN_IDS, 11], |
| 140 | 2, |
| 141 | id="open_and_close_markers", |
| 142 | ), |
| 143 | pytest.param( |
| 144 | [9, 10, *CLOSE_IDS, *RESPONSE_OPEN_IDS, 11], |
| 145 | 2, |
| 146 | id="open_marker_consumed_as_generation_prefix", |
| 147 | ), |
| 148 | pytest.param([*RESPONSE_OPEN_IDS, 11, 12], 0, id="response_only"), |
| 149 | pytest.param([*OPEN_IDS, 9, 10], 2, id="unterminated_after_open"), |
| 150 | pytest.param([9, 10, 11], 3, id="unterminated_without_markers"), |
| 151 | pytest.param([], 0, id="empty"), |
| 152 | pytest.param([9, 4, 2, 10], 4, id="partial_close_marker_is_reasoning"), |
| 153 | pytest.param( |
| 154 | [*OPEN_IDS, 9, *RESPONSE_OPEN_IDS, 11], |
| 155 | 2 + len(RESPONSE_OPEN_IDS), |
| 156 | id="response_open_inside_open_think_is_reasoning", |
| 157 | ), |
| 158 | ], |
| 159 | ) |
| 160 | def test_count_reasoning_tokens_matches_think_channel(token_ids, expected): |
| 161 | """reasoning_tokens must cover exactly what extract_reasoning labels as |
| 162 | reasoning: marker tokens are excluded and a consumed generation prefix |
| 163 | means the output starts inside the think channel.""" |
| 164 | parser = KimiK3ReasoningParser(DummyTokenizer()) |
| 165 | |
| 166 | assert parser.count_reasoning_tokens(token_ids) == expected |
| 167 | |
| 168 | |
| 169 | def test_count_reasoning_tokens_is_zero_when_thinking_disabled(): |
| 170 | parser = KimiK3ReasoningParser( |
| 171 | DummyTokenizer(), chat_template_kwargs={"thinking": False} |
| 172 | ) |
| 173 | |
| 174 | assert parser.count_reasoning_tokens([*OPEN_IDS, 9, 10, *CLOSE_IDS]) == 0 |
| 175 | |
| 176 | |
| 177 | def test_count_reasoning_tokens_through_delegating_parser(): |
| 178 | parser = ReasoningOnlyParser(_dummy_tokenizer()) |
| 179 | |
| 180 | assert parser.count_reasoning_tokens([*OPEN_IDS, 9, *CLOSE_IDS, 11]) == 1 |
| 181 | |
| 182 | |
| 183 | def test_streaming_split_open_marker_is_held_back(): |
| 184 | parser = KimiK3ReasoningParser(DummyTokenizer()) |
| 185 | |
| 186 | first = parser.extract_reasoning_content_streaming( |
| 187 | previous_text="", |
| 188 | current_text=OPEN, |
| 189 | delta_text=OPEN, |
| 190 | previous_token_ids=[], |
| 191 | current_token_ids=[1], |
| 192 | delta_token_ids=[1], |
| 193 | ) |
| 194 | second = parser.extract_reasoning_content_streaming( |
| 195 | previous_text=OPEN, |
| 196 | current_text=f"{OPEN}think", |
| 197 | delta_text="think", |
| 198 | previous_token_ids=[1], |
| 199 | current_token_ids=[1, 2], |
| 200 | delta_token_ids=[2], |
| 201 | ) |
| 202 | third = parser.extract_reasoning_content_streaming( |
| 203 | previous_text=f"{OPEN}think", |
| 204 | current_text=THINK_OPEN + "step", |
| 205 | delta_text=f"{SEP}step", |
| 206 | previous_token_ids=[1, 2], |
| 207 | current_token_ids=[1, 2, 3, 9], |
| 208 | delta_token_ids=[3, 9], |
| 209 | ) |
| 210 | |
| 211 | assert first is None |
| 212 | assert second is None |
| 213 | assert isinstance(third, DeltaMessage) |
| 214 | assert third.reasoning == "step" |
| 215 | |
| 216 | |
| 217 | def test_streaming_split_close_marker_hands_content_downstream(): |
| 218 | parser = KimiK3ReasoningParser(DummyTokenizer()) |
| 219 | |
| 220 | previous_text = f"{THINK_OPEN}step" |
| 221 | partial_close = parser.extract_reasoning_content_streaming( |
| 222 | previous_text=previous_text, |
| 223 | current_text=previous_text + CLOSE, |
| 224 | delta_text=CLOSE, |
| 225 | previous_token_ids=[1, 2, 3, 9], |
| 226 | current_token_ids=[1, 2, 3, 9, 4], |
| 227 | delta_token_ids=[4], |
| 228 | ) |
| 229 | closed = parser.extract_reasoning_content_streaming( |
| 230 | previous_text=previous_text + CLOSE, |
| 231 | current_text=previous_text + f"{THINK_CLOSE}{RESPONSE_OPEN}answer", |
| 232 | delta_text=f"think{SEP}{RESPONSE_OPEN}answer", |
| 233 | previous_token_ids=[1, 2, 3, 9, 4], |
| 234 | current_token_ids=[1, 2, 3, 9, 4, 2, 3, 10], |
| 235 | delta_token_ids=[2, 3, 10], |
| 236 | ) |
| 237 | |
| 238 | assert partial_close is None |
| 239 | assert isinstance(closed, DeltaMessage) |
| 240 | assert closed.reasoning is None |
| 241 | assert closed.content == f"{RESPONSE_OPEN}answer" |
| 242 | assert parser.extract_content_ids([2, 3, 10]) == [10] |
| 243 | |
| 244 | |
| 245 | def test_thinking_disabled_streams_content(): |
| 246 | parser = KimiK3ReasoningParser( |
| 247 | DummyTokenizer(), chat_template_kwargs={"enable_thinking": False} |
| 248 | ) |
| 249 | |
| 250 | delta = parser.extract_reasoning_content_streaming( |
| 251 | previous_text="", |
| 252 | current_text=f"{RESPONSE_OPEN}answer", |
| 253 | delta_text=f"{RESPONSE_OPEN}answer", |
| 254 | previous_token_ids=[], |
| 255 | current_token_ids=[1], |
| 256 | delta_token_ids=[1], |
| 257 | ) |
| 258 | |
| 259 | assert isinstance(delta, DeltaMessage) |
| 260 | assert delta.content == f"{RESPONSE_OPEN}answer" |
| 261 | assert delta.reasoning is None |
| 262 | |
| 263 | |
| 264 | def test_delegating_parser_thinking_false_streams_response_content(): |
| 265 | parser = ReasoningOnlyParser( |
| 266 | _dummy_tokenizer(), |
| 267 | chat_template_kwargs={"thinking": False}, |
| 268 | ) |
| 269 | request = ChatCompletionRequest( |
| 270 | model="test-model", |
| 271 | messages=[], |
| 272 | chat_template_kwargs={"thinking": False}, |
| 273 | ) |
| 274 | |
| 275 | first = parser.parse_delta( |
| 276 | delta_text="OK", |
| 277 | delta_token_ids=[10], |
| 278 | request=request, |
| 279 | prompt_token_ids=[1], |
| 280 | finished=False, |
| 281 | ) |
| 282 | partial_close = parser.parse_delta( |
| 283 | delta_text=CLOSE, |
| 284 | delta_token_ids=[2], |
| 285 | request=request, |
| 286 | prompt_token_ids=[1], |
| 287 | finished=False, |
| 288 | ) |
| 289 | closed = parser.parse_delta( |
| 290 | delta_text=f"response{SEP}", |
| 291 | delta_token_ids=[3, 4], |
| 292 | request=request, |
| 293 | prompt_token_ids=[1], |
| 294 | finished=False, |
| 295 | ) |
| 296 | |
| 297 | assert first is not None |
| 298 | assert first.content == "OK" |
| 299 | assert first.reasoning is None |
| 300 | assert partial_close is None |
| 301 | assert closed is None |
| 302 | |
| 303 | |
| 304 | def test_adjust_request_keeps_xtml_markers_contiguous(): |
| 305 | parser = KimiK3ReasoningParser(DummyTokenizer()) |
| 306 | request = ChatCompletionRequest(model="test-model", messages=[]) |
| 307 | |
| 308 | adjusted = parser.adjust_request(request) |
| 309 | |
| 310 | assert adjusted.skip_special_tokens is False |
| 311 | if hasattr(adjusted, "spaces_between_special_tokens"): |
| 312 | assert adjusted.spaces_between_special_tokens is False |
| 313 | |
| 314 | |
| 315 | def _reference_is_reasoning_end(input_ids: list[int]) -> bool: |
| 316 | """Full-sequence reference for the streaming check to be measured against. |
| 317 | |
| 318 | Two independent last-occurrence scans, i.e. the straightforward reading of |
| 319 | "reasoning ended iff the newest think marker is a close marker". |
| 320 | """ |
| 321 | |
| 322 | def last(needle: list[int]) -> int: |
| 323 | for i in range(len(input_ids) - len(needle), -1, -1): |
| 324 | if input_ids[i : i + len(needle)] == needle: |
| 325 | return i |
| 326 | return -1 |
| 327 | |
| 328 | last_close, last_open = last(CLOSE_IDS), last(OPEN_IDS) |
| 329 | if last_open == -1: |
| 330 | return last_close != -1 |
| 331 | return last_close > last_open |
| 332 | |
| 333 | |
| 334 | def test_is_reasoning_end_streaming_only_scans_the_step_window(): |
| 335 | """The decode-step check must not re-derive the answer from the whole |
| 336 | sequence: an already-closed think block earlier in the sequence is the |
| 337 | engine's business (it latches `reasoning_ended`), not this call's.""" |
| 338 | parser = KimiK3ReasoningParser(DummyTokenizer()) |
| 339 | prompt = [*CLOSE_IDS, *OPEN_IDS, 9, 9, 9] |
| 340 | |
| 341 | assert not parser.is_reasoning_end_streaming([*prompt, 7], [7]) |
| 342 | assert parser.is_reasoning_end_streaming([*prompt, *CLOSE_IDS], CLOSE_IDS) |
| 343 | |
| 344 | |
| 345 | def test_is_reasoning_end_streaming_detects_marker_across_step_boundary(): |
| 346 | """A 3-token marker can be split over decode steps; the check carries the |
| 347 | preceding len(marker)-1 tokens so the final token still completes it.""" |
| 348 | parser = KimiK3ReasoningParser(DummyTokenizer()) |
| 349 | history = [*OPEN_IDS, 5, *CLOSE_IDS[:2]] |
| 350 | |
| 351 | assert not parser.is_reasoning_end_streaming(history, [CLOSE_IDS[1]]) |
| 352 | assert parser.is_reasoning_end_streaming([*history, CLOSE_IDS[2]], [CLOSE_IDS[2]]) |
| 353 | |
| 354 | |
| 355 | def test_is_reasoning_end_streaming_reopened_block_is_not_ended(): |
| 356 | """Kimi K3 may close and immediately reopen the think channel inside one |
| 357 | window (speculative decoding lands several tokens per step). The newest |
| 358 | marker wins, matching is_reasoning_end.""" |
| 359 | parser = KimiK3ReasoningParser(DummyTokenizer()) |
| 360 | window = [*CLOSE_IDS, *OPEN_IDS] |
| 361 | |
| 362 | assert not parser.is_reasoning_end_streaming([*OPEN_IDS, 5, *window], window) |
| 363 | |
| 364 | |
| 365 | def test_is_reasoning_end_streaming_accepts_an_iterator_delta(): |
| 366 | """`should_advance` may hand over an islice rather than a list.""" |
| 367 | parser = KimiK3ReasoningParser(DummyTokenizer()) |
| 368 | full = [*OPEN_IDS, 5, *CLOSE_IDS] |
| 369 | |
| 370 | assert parser.is_reasoning_end_streaming(full, iter(CLOSE_IDS)) |
| 371 | |
| 372 | |
| 373 | def test_is_reasoning_end_streaming_thinking_disabled(): |
| 374 | parser = KimiK3ReasoningParser( |
| 375 | DummyTokenizer(), chat_template_kwargs={"thinking": False} |
| 376 | ) |
| 377 | |
| 378 | assert parser.is_reasoning_end_streaming([1], [1]) |
| 379 | |
| 380 | |
| 381 | @pytest.mark.parametrize("seed", range(24)) |
| 382 | def test_reasoning_end_matches_reference_over_marker_dense_sequences(seed): |
| 383 | """Both the full-sequence check and the per-step check must agree with the |
| 384 | two-scan reference, on sequences where the markers (which share the suffix |
| 385 | [2, 3]) collide and overlap constantly.""" |
| 386 | import random |
| 387 | |
| 388 | rnd = random.Random(seed) |
| 389 | parser = KimiK3ReasoningParser(DummyTokenizer()) |
| 390 | |
| 391 | for _ in range(200): |
| 392 | head = [rnd.choice([1, 2, 3, 4]) for _ in range(rnd.randrange(0, 14))] |
| 393 | delta = [rnd.choice([1, 2, 3, 4]) for _ in range(rnd.randrange(1, 5))] |
| 394 | full = [*head, *delta] |
| 395 | |
| 396 | assert parser.is_reasoning_end(full) == _reference_is_reasoning_end(full) |
| 397 | |
| 398 | # The engine only calls the streaming check while reasoning is still |
| 399 | # open, so that is the only case it has to agree on. |
| 400 | if not _reference_is_reasoning_end(head): |
| 401 | assert parser.is_reasoning_end_streaming( |
| 402 | full, delta |
| 403 | ) == _reference_is_reasoning_end(full), (head, delta) |