@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 150 lines 5.74 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 import pytest
6
7 from vllm.entrypoints.generate.base.protocol import DeltaMessage
8 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
9 from vllm.parser.kimi_k2 import KimiK2Parser
10 from vllm.reasoning.kimi_k2_reasoning_parser import KimiK2ReasoningParser
11 from vllm.tokenizers import get_tokenizer
12
13 REASONING_MODEL_NAME = "moonshotai/Kimi-K2.5"
14
15
16 @pytest.fixture(scope="module")
17 def kimi_k2_tokenizer():
18 return get_tokenizer(tokenizer_name=REASONING_MODEL_NAME, trust_remote_code=True)
19
20
21 def test_extract_reasoning_with_think_tags(kimi_k2_tokenizer):
22 parser = KimiK2ReasoningParser(kimi_k2_tokenizer)
23 request = ChatCompletionRequest(model="test-model", messages=[], temperature=1.0)
24
25 reasoning, content = parser.extract_reasoning(
26 "<think>step by step reasoning</think>final answer", request
27 )
28 assert reasoning == "step by step reasoning"
29 assert content == "final answer"
30
31
32 def test_extract_reasoning_empty_thinking(kimi_k2_tokenizer):
33 parser = KimiK2ReasoningParser(kimi_k2_tokenizer)
34 request = ChatCompletionRequest(model="test-model", messages=[], temperature=1.0)
35
36 reasoning, content = parser.extract_reasoning(
37 "<think></think>final answer", request
38 )
39 assert reasoning is None
40 assert content == "final answer"
41
42
43 def test_extract_reasoning_implicit_start(kimi_k2_tokenizer):
44 """When there's no <think> tag, everything is treated as reasoning."""
45 parser = KimiK2ReasoningParser(kimi_k2_tokenizer)
46 request = ChatCompletionRequest(model="test-model", messages=[], temperature=1.0)
47
48 reasoning, content = parser.extract_reasoning(
49 "implicit reasoning with no tags", request
50 )
51 assert reasoning == "implicit reasoning with no tags"
52 assert content is None
53
54
55 def test_extract_reasoning_tool_section_ends_reasoning(kimi_k2_tokenizer):
56 """<|tool_calls_section_begin|> implicitly ends reasoning."""
57 parser = KimiK2ReasoningParser(kimi_k2_tokenizer)
58 request = ChatCompletionRequest(model="test-model", messages=[], temperature=1.0)
59
60 text = "some reasoning<|tool_calls_section_begin|>tool call data"
61 reasoning, content = parser.extract_reasoning(text, request)
62 assert reasoning == "some reasoning"
63 assert content == "<|tool_calls_section_begin|>tool call data"
64
65
66 def test_streaming_reasoning_then_content(kimi_k2_tokenizer):
67 """Token-by-token streaming: reasoning tokens then content after </think>."""
68 parser = KimiK2ReasoningParser(kimi_k2_tokenizer)
69 parser_engine = parser._parser_engine
70 assert isinstance(parser_engine, KimiK2Parser)
71
72 think_id = parser_engine._start_token_id
73 end_think_id = parser_engine._end_token_id
74 assert think_id is not None
75 assert end_think_id is not None
76 # Use a real token ID from the tokenizer for regular content
77 regular_id = kimi_k2_tokenizer.encode("hello", add_special_tokens=False)[0]
78
79 # First token: <think> — single special token should be skipped
80 result = parser.extract_reasoning_streaming(
81 previous_text="",
82 current_text="<think>",
83 delta_text="<think>",
84 previous_token_ids=[],
85 current_token_ids=[think_id],
86 delta_token_ids=[think_id],
87 )
88 assert result is None
89
90 # Reasoning token
91 result = parser.extract_reasoning_streaming(
92 previous_text="<think>",
93 current_text="<think>step one",
94 delta_text="step one",
95 previous_token_ids=[think_id],
96 current_token_ids=[think_id, regular_id],
97 delta_token_ids=[regular_id],
98 )
99 assert isinstance(result, DeltaMessage)
100 assert result.reasoning == "step one"
101 assert result.content is None
102
103 # End token </think> as single token — should be skipped
104 result = parser.extract_reasoning_streaming(
105 previous_text="<think>step one",
106 current_text="<think>step one</think>",
107 delta_text="</think>",
108 previous_token_ids=[think_id, regular_id],
109 current_token_ids=[think_id, regular_id, end_think_id],
110 delta_token_ids=[end_think_id],
111 )
112 assert result is None
113
114 # Content after </think>
115 content_id = kimi_k2_tokenizer.encode("world", add_special_tokens=False)[0]
116 result = parser.extract_reasoning_streaming(
117 previous_text="<think>step one</think>",
118 current_text="<think>step one</think>answer",
119 delta_text="answer",
120 previous_token_ids=[think_id, regular_id, end_think_id],
121 current_token_ids=[think_id, regular_id, end_think_id, content_id],
122 delta_token_ids=[content_id],
123 )
124 assert isinstance(result, DeltaMessage)
125 assert result.content == "answer"
126
127
128 def test_streaming_tool_section_ends_reasoning(kimi_k2_tokenizer):
129 """<|tool_calls_section_begin|> in delta ends reasoning during streaming."""
130 parser = KimiK2ReasoningParser(kimi_k2_tokenizer)
131 parser_engine = parser._parser_engine
132 assert isinstance(parser_engine, KimiK2Parser)
133
134 think_id = parser_engine._start_token_id
135 tool_begin_id = parser_engine._tool_section_start_token_id
136 assert think_id is not None
137 assert tool_begin_id is not None
138 regular_id = kimi_k2_tokenizer.encode("hello", add_special_tokens=False)[0]
139
140 # Tool section token arrives — should transition from reasoning to content
141 result = parser.extract_reasoning_streaming(
142 previous_text="<think>thinking",
143 current_text="<think>thinking<|tool_calls_section_begin|>",
144 delta_text="<|tool_calls_section_begin|>",
145 previous_token_ids=[think_id, regular_id],
146 current_token_ids=[think_id, regular_id, tool_begin_id],
147 delta_token_ids=[tool_begin_id],
148 )
149 assert isinstance(result, DeltaMessage)
150 assert result.content == "<|tool_calls_section_begin|>"