| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import Any, cast |
| 4 | |
| 5 | from vllm.config import VllmConfig |
| 6 | from vllm.entrypoints.chat_utils import ( |
| 7 | ChatCompletionMessageParam, |
| 8 | ConversationMessage, |
| 9 | parse_chat_messages, |
| 10 | parse_chat_messages_async, |
| 11 | ) |
| 12 | from vllm.exceptions import VLLMValidationError |
| 13 | from vllm.multimodal.media.connector import merge_media_io_kwargs |
| 14 | from vllm.tokenizers.hf import HfTokenizer |
| 15 | from vllm.utils.async_utils import make_async |
| 16 | |
| 17 | from .base import BaseRenderer |
| 18 | from .inputs import DictPrompt |
| 19 | from .inputs.preprocess import parse_dec_only_prompt |
| 20 | from .params import ChatParams |
| 21 | |
| 22 | # Keep the original image mode (including the alpha channel) for K3 instead of |
| 23 | # flattening images onto a background color. Server-level (--media-io-kwargs) |
| 24 | # and request-level media_io_kwargs still take precedence over this default. |
| 25 | _K3_MEDIA_IO_DEFAULTS: dict[str, dict[str, Any]] = {"image": {"image_mode": None}} |
| 26 | _K3_THINKING_EFFORTS = ("low", "high", "max") |
| 27 | |
| 28 | |
| 29 | def _merge_k3_media_io_kwargs( |
| 30 | media_io_kwargs: dict[str, dict[str, Any]] | None, |
| 31 | ) -> dict[str, dict[str, Any]] | None: |
| 32 | return merge_media_io_kwargs(_K3_MEDIA_IO_DEFAULTS, media_io_kwargs) |
| 33 | |
| 34 | |
| 35 | def _dump_k3_template_value(value: Any) -> Any: |
| 36 | model_dump = getattr(value, "model_dump", None) |
| 37 | if callable(model_dump): |
| 38 | return model_dump(mode="json", exclude_none=True) |
| 39 | |
| 40 | to_dict = getattr(value, "to_dict", None) |
| 41 | if callable(to_dict): |
| 42 | return to_dict() |
| 43 | |
| 44 | return value |
| 45 | |
| 46 | |
| 47 | def _apply_k3_thinking_kwargs(kwargs: dict[str, Any]) -> None: |
| 48 | if (enable_thinking := kwargs.pop("enable_thinking", None)) is not None: |
| 49 | kwargs.setdefault("thinking", enable_thinking) |
| 50 | |
| 51 | reasoning_effort = kwargs.pop("reasoning_effort", None) |
| 52 | if reasoning_effort == "none": |
| 53 | kwargs.setdefault("thinking", False) |
| 54 | elif reasoning_effort is not None: |
| 55 | kwargs.setdefault("thinking_effort", reasoning_effort) |
| 56 | |
| 57 | thinking_effort = kwargs.get("thinking_effort") |
| 58 | if thinking_effort is not None and thinking_effort not in _K3_THINKING_EFFORTS: |
| 59 | supported = ", ".join(_K3_THINKING_EFFORTS) |
| 60 | raise VLLMValidationError( |
| 61 | f"Kimi K3 supports thinking_effort values: {supported}", |
| 62 | parameter="thinking_effort", |
| 63 | value=thinking_effort, |
| 64 | ) |
| 65 | |
| 66 | |
| 67 | def _normalize_k3_tool_messages( |
| 68 | conversation: list[ConversationMessage], |
| 69 | ) -> list[dict[str, Any]]: |
| 70 | """Reorder tool-result messages to match assistant tool_call order. |
| 71 | |
| 72 | Supports matching by ``tool_call_id`` or by the synthetic |
| 73 | ``"{tool}:{zero_based_index}"`` alias. When any tool message in a |
| 74 | block cannot be resolved, the whole block is left in its original |
| 75 | order (graceful fallback). |
| 76 | |
| 77 | Returns a new list; caller-owned message dicts are not mutated. |
| 78 | """ |
| 79 | normalized: list[dict[str, Any]] = [] |
| 80 | i = 0 |
| 81 | |
| 82 | while i < len(conversation): |
| 83 | message = conversation[i] |
| 84 | normalized.append(dict(message)) |
| 85 | i += 1 |
| 86 | |
| 87 | if message.get("role") != "assistant": |
| 88 | continue |
| 89 | |
| 90 | tool_calls = message.get("tool_calls") |
| 91 | if not tool_calls: |
| 92 | continue |
| 93 | |
| 94 | # Build the lookup table from the assistant's tool_calls. |
| 95 | targets_by_id: dict[str, tuple[int, str]] = {} |
| 96 | for position, tool_call in enumerate(tool_calls): |
| 97 | function = tool_call.get("function") |
| 98 | if not isinstance(function, dict): |
| 99 | continue |
| 100 | name = function.get("name") |
| 101 | if not isinstance(name, str) or not name: |
| 102 | continue |
| 103 | |
| 104 | call_target = (position, name) |
| 105 | aliases = [f"{name}:{position}"] |
| 106 | if tool_call_id := tool_call.get("id"): |
| 107 | aliases.insert(0, str(tool_call_id)) |
| 108 | for alias in aliases: |
| 109 | targets_by_id.setdefault(alias, call_target) |
| 110 | |
| 111 | # Collect consecutive tool messages. |
| 112 | block_start = i |
| 113 | while i < len(conversation) and conversation[i].get("role") == "tool": |
| 114 | i += 1 |
| 115 | if i == block_start: |
| 116 | continue |
| 117 | |
| 118 | # Try to resolve every tool message in the block. |
| 119 | resolved: list[tuple[int, int, dict[str, Any]]] = [] |
| 120 | for order, tool_message in enumerate(conversation[block_start:i]): |
| 121 | tool_call_id = tool_message.get("tool_call_id") |
| 122 | resolved_target = ( |
| 123 | targets_by_id.get(str(tool_call_id)) |
| 124 | if tool_call_id is not None |
| 125 | else None |
| 126 | ) |
| 127 | if resolved_target is None: |
| 128 | normalized.extend(dict(item) for item in conversation[block_start:i]) |
| 129 | break |
| 130 | |
| 131 | position, name = resolved_target |
| 132 | enriched = dict(tool_message) |
| 133 | enriched["tool"] = name |
| 134 | enriched["index"] = position + 1 |
| 135 | resolved.append((position, order, enriched)) |
| 136 | else: |
| 137 | resolved.sort(key=lambda item: (item[0], item[1])) |
| 138 | normalized.extend(item for _, _, item in resolved) |
| 139 | |
| 140 | return normalized |
| 141 | |
| 142 | |
| 143 | class KimiK3Renderer(BaseRenderer[HfTokenizer]): |
| 144 | """Render chat prompts with Kimi K3's Python XTML encoding. |
| 145 | |
| 146 | K3 ships no Jinja chat template; its tokenizer renders messages through |
| 147 | ``encoding_k3`` instead. We tokenize eagerly so the structural markers keep |
| 148 | their special-token ids while user- and tool-supplied text stays ordinary. |
| 149 | """ |
| 150 | |
| 151 | def __init__(self, config: VllmConfig, tokenizer: HfTokenizer | None) -> None: |
| 152 | super().__init__(config, tokenizer) |
| 153 | |
| 154 | self._apply_chat_template_async = make_async( |
| 155 | self._apply_chat_template, executor=self._executor |
| 156 | ) |
| 157 | |
| 158 | def _apply_chat_template( |
| 159 | self, |
| 160 | conversation: list[dict[str, Any]], |
| 161 | params: ChatParams, |
| 162 | ) -> list[int]: |
| 163 | # Tokenize eagerly: K3 encodes structural markers as special tokens and |
| 164 | # user/tool text as ordinary tokens, so we cannot defer to a plain |
| 165 | # re-tokenization of the rendered string downstream. |
| 166 | kwargs = params.get_apply_chat_template_kwargs() |
| 167 | _apply_k3_thinking_kwargs(kwargs) |
| 168 | if params.tool_choice not in (None, "auto"): |
| 169 | kwargs["tool_choice"] = _dump_k3_template_value(params.tool_choice) |
| 170 | if params.response_format is not None: |
| 171 | kwargs["response_format"] = _dump_k3_template_value(params.response_format) |
| 172 | kwargs["tokenize"] = True |
| 173 | return self.get_tokenizer().apply_chat_template(conversation, **kwargs) |
| 174 | |
| 175 | def render_messages( |
| 176 | self, |
| 177 | messages: list[ChatCompletionMessageParam], |
| 178 | params: ChatParams, |
| 179 | ) -> tuple[list[ConversationMessage], DictPrompt]: |
| 180 | conversation, mm_data, mm_uuids = parse_chat_messages( |
| 181 | messages, |
| 182 | self.model_config, |
| 183 | content_format="string", |
| 184 | media_io_kwargs=_merge_k3_media_io_kwargs(params.media_io_kwargs), |
| 185 | mm_processor_kwargs=params.mm_processor_kwargs, |
| 186 | ) |
| 187 | |
| 188 | rendered_conversation = _normalize_k3_tool_messages(conversation) |
| 189 | prompt = parse_dec_only_prompt( |
| 190 | self._apply_chat_template(rendered_conversation, params) |
| 191 | ) |
| 192 | if mm_data is not None: |
| 193 | prompt["multi_modal_data"] = mm_data |
| 194 | if mm_uuids is not None: |
| 195 | prompt["multi_modal_uuids"] = mm_uuids |
| 196 | |
| 197 | return cast(list[ConversationMessage], rendered_conversation), prompt |
| 198 | |
| 199 | async def render_messages_async( |
| 200 | self, |
| 201 | messages: list[ChatCompletionMessageParam], |
| 202 | params: ChatParams, |
| 203 | ) -> tuple[list[ConversationMessage], DictPrompt]: |
| 204 | conversation, mm_data, mm_uuids = await parse_chat_messages_async( |
| 205 | messages, |
| 206 | self.model_config, |
| 207 | content_format="string", |
| 208 | media_io_kwargs=_merge_k3_media_io_kwargs(params.media_io_kwargs), |
| 209 | mm_processor_kwargs=params.mm_processor_kwargs, |
| 210 | ) |
| 211 | |
| 212 | rendered_conversation = _normalize_k3_tool_messages(conversation) |
| 213 | token_ids = await self._apply_chat_template_async(rendered_conversation, params) |
| 214 | prompt = parse_dec_only_prompt(token_ids) |
| 215 | if mm_data is not None: |
| 216 | prompt["multi_modal_data"] = mm_data |
| 217 | if mm_uuids is not None: |
| 218 | prompt["multi_modal_uuids"] = mm_uuids |
| 219 | |
| 220 | return cast(list[ConversationMessage], rendered_conversation), prompt |