| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass, field |
| 4 | from typing import Any |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | from vllm.exceptions import VLLMValidationError |
| 9 | from vllm.renderers import ChatParams |
| 10 | from vllm.renderers.kimi_k3 import KimiK3Renderer, _merge_k3_media_io_kwargs |
| 11 | from vllm.renderers.registry import RENDERER_REGISTRY |
| 12 | from vllm.tokenizers.registry import TokenizerRegistry |
| 13 | |
| 14 | |
| 15 | class StubTokenizer: |
| 16 | """Stands in for the model's TikTokenTokenizer. |
| 17 | |
| 18 | Records the kwargs it is called with and returns fixed token ids, so tests |
| 19 | can assert how the renderer drives ``apply_chat_template`` without the real |
| 20 | (git-LFS) tiktoken vocabulary. |
| 21 | """ |
| 22 | |
| 23 | def __init__(self, token_ids: list[int]) -> None: |
| 24 | self.token_ids = token_ids |
| 25 | self.calls: list[dict[str, Any]] = [] |
| 26 | self.conversations: list[list[dict[str, Any]]] = [] |
| 27 | |
| 28 | def apply_chat_template(self, conversation, **kwargs) -> list[int]: |
| 29 | self.conversations.append(conversation) |
| 30 | self.calls.append(kwargs) |
| 31 | return list(self.token_ids) |
| 32 | |
| 33 | |
| 34 | @dataclass |
| 35 | class MockHFConfig: |
| 36 | model_type: str = "kimi_k3" |
| 37 | |
| 38 | |
| 39 | @dataclass |
| 40 | class MockModelConfig: |
| 41 | runner_type: str = "generate" |
| 42 | is_multimodal_model: bool = False |
| 43 | supports_multimodal_inputs: bool = False |
| 44 | multimodal_config: Any = None |
| 45 | hf_config: MockHFConfig = field(default_factory=MockHFConfig) |
| 46 | allowed_local_media_path: str = "" |
| 47 | allowed_media_domains: Any = None |
| 48 | enable_prompt_embeds: bool = False |
| 49 | renderer_num_workers: int = 1 |
| 50 | |
| 51 | |
| 52 | @dataclass |
| 53 | class MockParallelConfig: |
| 54 | _api_process_rank: int = 0 |
| 55 | |
| 56 | |
| 57 | @dataclass |
| 58 | class MockVllmConfig: |
| 59 | model_config: MockModelConfig |
| 60 | parallel_config: MockParallelConfig |
| 61 | |
| 62 | |
| 63 | def _make_renderer(tokenizer: StubTokenizer) -> KimiK3Renderer: |
| 64 | config = MockVllmConfig(MockModelConfig(), MockParallelConfig()) |
| 65 | return KimiK3Renderer(config, tokenizer) |
| 66 | |
| 67 | |
| 68 | def test_kimi_k3_registered(): |
| 69 | assert RENDERER_REGISTRY.load_renderer_cls("kimi_k3").__name__ == "KimiK3Renderer" |
| 70 | assert ( |
| 71 | TokenizerRegistry.load_tokenizer_cls("kimi_k3").__name__ == "CachedHfTokenizer" |
| 72 | ) |
| 73 | |
| 74 | |
| 75 | def test_k3_media_io_defaults_preserve_original_mode(): |
| 76 | # Default: K3 keeps the original image mode (no background flattening). |
| 77 | assert _merge_k3_media_io_kwargs(None) == {"image": {"image_mode": None}} |
| 78 | |
| 79 | # Server-/request-level values take precedence over the K3 default. |
| 80 | assert _merge_k3_media_io_kwargs({"image": {"image_mode": "RGB"}}) == { |
| 81 | "image": {"image_mode": "RGB"} |
| 82 | } |
| 83 | |
| 84 | # Unrelated image kwargs are merged with the default. |
| 85 | assert _merge_k3_media_io_kwargs( |
| 86 | {"image": {"rgba_background_color": (0, 0, 0)}} |
| 87 | ) == {"image": {"image_mode": None, "rgba_background_color": (0, 0, 0)}} |
| 88 | |
| 89 | |
| 90 | def test_apply_chat_template_forces_tokenize_and_pins_return_dict(): |
| 91 | tokenizer = StubTokenizer([7, 8, 9]) |
| 92 | renderer = _make_renderer(tokenizer) |
| 93 | tools = [{"type": "function", "function": {"name": "search"}}] |
| 94 | params = ChatParams( |
| 95 | chat_template_kwargs={"tools": tools, "tokenize": False, "thinking": True} |
| 96 | ) |
| 97 | |
| 98 | token_ids = renderer._apply_chat_template( |
| 99 | [{"role": "user", "content": "hi"}], params |
| 100 | ) |
| 101 | |
| 102 | assert token_ids == [7, 8, 9] |
| 103 | kwargs = tokenizer.calls[-1] |
| 104 | # tokenize is forced on even though the request asked for False, so K3 keeps |
| 105 | # the special-vs-ordinary token distinction instead of re-tokenizing a string. |
| 106 | assert kwargs["tokenize"] is True |
| 107 | # return_dict is pinned False so we always get a flat list of ids. |
| 108 | assert kwargs["return_dict"] is False |
| 109 | assert kwargs["tools"] == tools |
| 110 | assert kwargs["thinking"] is True |
| 111 | |
| 112 | |
| 113 | def test_apply_chat_template_translates_standard_thinking_kwargs(): |
| 114 | # Standard enable_thinking/reasoning_effort kwargs must be translated |
| 115 | # to K3's native thinking/thinking_effort. |
| 116 | tokenizer = StubTokenizer([7, 8, 9]) |
| 117 | renderer = _make_renderer(tokenizer) |
| 118 | params = ChatParams( |
| 119 | chat_template_kwargs={"enable_thinking": False, "reasoning_effort": "none"} |
| 120 | ) |
| 121 | |
| 122 | renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) |
| 123 | |
| 124 | kwargs = tokenizer.calls[-1] |
| 125 | assert kwargs["thinking"] is False |
| 126 | assert "thinking_effort" not in kwargs |
| 127 | assert "enable_thinking" not in kwargs |
| 128 | assert "reasoning_effort" not in kwargs |
| 129 | |
| 130 | |
| 131 | @pytest.mark.parametrize("reasoning_effort", ["low", "high", "max"]) |
| 132 | def test_apply_chat_template_translates_supported_reasoning_effort( |
| 133 | reasoning_effort: str, |
| 134 | ): |
| 135 | tokenizer = StubTokenizer([7, 8, 9]) |
| 136 | renderer = _make_renderer(tokenizer) |
| 137 | params = ChatParams(chat_template_kwargs={"reasoning_effort": reasoning_effort}) |
| 138 | |
| 139 | renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) |
| 140 | |
| 141 | kwargs = tokenizer.calls[-1] |
| 142 | assert kwargs["thinking_effort"] == reasoning_effort |
| 143 | assert "reasoning_effort" not in kwargs |
| 144 | |
| 145 | |
| 146 | @pytest.mark.parametrize("reasoning_effort", ["minimal", "medium", "xhigh"]) |
| 147 | def test_apply_chat_template_rejects_unsupported_reasoning_effort( |
| 148 | reasoning_effort: str, |
| 149 | ): |
| 150 | tokenizer = StubTokenizer([7, 8, 9]) |
| 151 | renderer = _make_renderer(tokenizer) |
| 152 | params = ChatParams(chat_template_kwargs={"reasoning_effort": reasoning_effort}) |
| 153 | |
| 154 | with pytest.raises(VLLMValidationError, match="thinking_effort") as exc_info: |
| 155 | renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) |
| 156 | |
| 157 | assert exc_info.value.parameter == "thinking_effort" |
| 158 | assert exc_info.value.value == reasoning_effort |
| 159 | assert tokenizer.calls == [] |
| 160 | |
| 161 | |
| 162 | def test_apply_chat_template_validates_canonical_native_thinking_effort(): |
| 163 | tokenizer = StubTokenizer([7, 8, 9]) |
| 164 | renderer = _make_renderer(tokenizer) |
| 165 | params = ChatParams( |
| 166 | chat_template_kwargs={ |
| 167 | "thinking_effort": "low", |
| 168 | "reasoning_effort": "medium", |
| 169 | } |
| 170 | ) |
| 171 | |
| 172 | renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) |
| 173 | |
| 174 | assert tokenizer.calls[-1]["thinking_effort"] == "low" |
| 175 | |
| 176 | |
| 177 | @pytest.mark.parametrize("thinking_effort", ["none", "minimal", "medium", "xhigh"]) |
| 178 | def test_apply_chat_template_rejects_unsupported_native_thinking_effort( |
| 179 | thinking_effort: str, |
| 180 | ): |
| 181 | tokenizer = StubTokenizer([7, 8, 9]) |
| 182 | renderer = _make_renderer(tokenizer) |
| 183 | params = ChatParams(chat_template_kwargs={"thinking_effort": thinking_effort}) |
| 184 | |
| 185 | with pytest.raises(VLLMValidationError, match="thinking_effort") as exc_info: |
| 186 | renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) |
| 187 | |
| 188 | assert exc_info.value.parameter == "thinking_effort" |
| 189 | assert exc_info.value.value == thinking_effort |
| 190 | assert tokenizer.calls == [] |
| 191 | |
| 192 | |
| 193 | def test_apply_chat_template_native_k3_kwargs_take_precedence(): |
| 194 | tokenizer = StubTokenizer([7, 8, 9]) |
| 195 | renderer = _make_renderer(tokenizer) |
| 196 | params = ChatParams( |
| 197 | chat_template_kwargs={ |
| 198 | "thinking": True, |
| 199 | "enable_thinking": False, |
| 200 | "thinking_effort": "low", |
| 201 | "reasoning_effort": "high", |
| 202 | } |
| 203 | ) |
| 204 | |
| 205 | renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) |
| 206 | |
| 207 | kwargs = tokenizer.calls[-1] |
| 208 | assert kwargs["thinking"] is True |
| 209 | assert kwargs["thinking_effort"] == "low" |
| 210 | |
| 211 | |
| 212 | def test_apply_chat_template_adds_k3_api_metadata(): |
| 213 | tokenizer = StubTokenizer([7, 8, 9]) |
| 214 | renderer = _make_renderer(tokenizer) |
| 215 | response_format = {"type": "json_object"} |
| 216 | params = ChatParams( |
| 217 | tool_choice="required", |
| 218 | response_format=response_format, |
| 219 | ) |
| 220 | |
| 221 | renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) |
| 222 | |
| 223 | kwargs = tokenizer.calls[-1] |
| 224 | assert kwargs["tool_choice"] == "required" |
| 225 | assert kwargs["response_format"] == response_format |
| 226 | |
| 227 | |
| 228 | def test_apply_chat_template_auto_tool_choice_keeps_template_kwarg(): |
| 229 | tokenizer = StubTokenizer([7, 8, 9]) |
| 230 | renderer = _make_renderer(tokenizer) |
| 231 | params = ChatParams( |
| 232 | chat_template_kwargs={"tool_choice": "required"}, |
| 233 | tool_choice="auto", |
| 234 | ) |
| 235 | |
| 236 | renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) |
| 237 | |
| 238 | assert tokenizer.calls[-1]["tool_choice"] == "required" |
| 239 | |
| 240 | |
| 241 | def test_apply_chat_template_omits_tool_choice_without_tools(): |
| 242 | tokenizer = StubTokenizer([7, 8, 9]) |
| 243 | renderer = _make_renderer(tokenizer) |
| 244 | |
| 245 | renderer._apply_chat_template( |
| 246 | [{"role": "user", "content": "hi"}], ChatParams(tool_choice=None) |
| 247 | ) |
| 248 | |
| 249 | assert "tool_choice" not in tokenizer.calls[-1] |
| 250 | |
| 251 | |
| 252 | def test_render_messages_returns_token_prompt(): |
| 253 | renderer = _make_renderer(StubTokenizer([1, 2, 3])) |
| 254 | |
| 255 | conversation, prompt = renderer.render_messages( |
| 256 | [{"role": "user", "content": "hi"}], ChatParams() |
| 257 | ) |
| 258 | |
| 259 | assert prompt == {"prompt_token_ids": [1, 2, 3]} |
| 260 | assert "multi_modal_data" not in prompt |
| 261 | assert conversation[0]["role"] == "user" |
| 262 | |
| 263 | |
| 264 | def test_render_messages_derives_private_xtml_tool_attrs(): |
| 265 | tokenizer = StubTokenizer([1, 2, 3]) |
| 266 | renderer = _make_renderer(tokenizer) |
| 267 | |
| 268 | conversation, _ = renderer.render_messages( |
| 269 | [ |
| 270 | { |
| 271 | "role": "assistant", |
| 272 | "content": "", |
| 273 | "tool_calls": [ |
| 274 | { |
| 275 | "id": "lookup:0", |
| 276 | "type": "function", |
| 277 | "function": {"name": "lookup", "arguments": "{}"}, |
| 278 | }, |
| 279 | { |
| 280 | "id": "lookup:1", |
| 281 | "type": "function", |
| 282 | "function": {"name": "lookup", "arguments": "{}"}, |
| 283 | }, |
| 284 | ], |
| 285 | }, |
| 286 | { |
| 287 | "role": "tool", |
| 288 | "tool_call_id": "lookup:1", |
| 289 | "tool": "client-supplied-name", |
| 290 | "index": 99, |
| 291 | "content": "second", |
| 292 | }, |
| 293 | { |
| 294 | "role": "tool", |
| 295 | "tool_call_id": "lookup:0", |
| 296 | "content": "first", |
| 297 | }, |
| 298 | ], |
| 299 | ChatParams(), |
| 300 | ) |
| 301 | |
| 302 | assert [message["content"] for message in conversation[1:]] == [ |
| 303 | "first", |
| 304 | "second", |
| 305 | ] |
| 306 | assert conversation[1]["tool"] == "lookup" |
| 307 | assert conversation[1]["index"] == 1 |
| 308 | assert conversation[2]["tool"] == "lookup" |
| 309 | assert conversation[2]["index"] == 2 |
| 310 | assert tokenizer.conversations[-1] == conversation |
| 311 | |
| 312 | |
| 313 | def test_render_messages_ignores_client_supplied_xtml_tool_attrs(): |
| 314 | tokenizer = StubTokenizer([1, 2, 3]) |
| 315 | renderer = _make_renderer(tokenizer) |
| 316 | |
| 317 | conversation, _ = renderer.render_messages( |
| 318 | [ |
| 319 | { |
| 320 | "role": "assistant", |
| 321 | "content": "", |
| 322 | "tool_calls": [ |
| 323 | { |
| 324 | "id": "lookup:0", |
| 325 | "type": "function", |
| 326 | "function": {"name": "lookup", "arguments": "{}"}, |
| 327 | } |
| 328 | ], |
| 329 | }, |
| 330 | { |
| 331 | "role": "tool", |
| 332 | "tool_call_id": "unknown", |
| 333 | "tool": "lookup", |
| 334 | "index": 1, |
| 335 | "content": "result", |
| 336 | }, |
| 337 | ], |
| 338 | ChatParams(), |
| 339 | ) |
| 340 | |
| 341 | assert "tool" not in conversation[1] |
| 342 | assert "index" not in conversation[1] |
| 343 | |
| 344 | |
| 345 | @pytest.mark.asyncio |
| 346 | async def test_render_messages_async_returns_token_prompt(): |
| 347 | renderer = _make_renderer(StubTokenizer([4, 5])) |
| 348 | |
| 349 | conversation, prompt = await renderer.render_messages_async( |
| 350 | [{"role": "user", "content": "hi"}], ChatParams() |
| 351 | ) |
| 352 | |
| 353 | assert prompt == {"prompt_token_ids": [4, 5]} |
| 354 | assert conversation[0]["role"] == "user" |