| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Kimi K2 parser for reasoning and tool calls. |
| 4 | |
| 5 | Kimi K2 tool call format:: |
| 6 | |
| 7 | <|tool_calls_section_begin|> |
| 8 | <|tool_call_begin|>functions.get_weather:0 |
| 9 | <|tool_call_argument_begin|>{"city": "Tokyo"}<|tool_call_end|> |
| 10 | <|tool_calls_section_end|> |
| 11 | |
| 12 | The header before ``<|tool_call_argument_begin|>`` is Kimi's native tool |
| 13 | call id. The function name is the final component before ``:N``. |
| 14 | """ |
| 15 | |
| 16 | from __future__ import annotations |
| 17 | |
| 18 | import functools |
| 19 | from collections.abc import Sequence |
| 20 | from typing import TYPE_CHECKING |
| 21 | |
| 22 | import regex as re |
| 23 | |
| 24 | from vllm.entrypoints.generate.base.protocol import DeltaFunctionCall, DeltaToolCall |
| 25 | from vllm.parser.engine.events import EventType |
| 26 | from vllm.parser.engine.parser_engine import ParserEngine |
| 27 | from vllm.parser.engine.parser_engine_config import ( |
| 28 | ParserEngineConfig, |
| 29 | ParserState, |
| 30 | Transition, |
| 31 | ) |
| 32 | |
| 33 | if TYPE_CHECKING: |
| 34 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 35 | ChatCompletionRequest, |
| 36 | ) |
| 37 | from vllm.entrypoints.openai.responses.protocol import ResponsesRequest |
| 38 | from vllm.tokenizers import TokenizerLike |
| 39 | from vllm.tool_parsers.abstract_tool_parser import Tool |
| 40 | |
| 41 | THINK_START = "<think>" |
| 42 | THINK_END = "</think>" |
| 43 | TOOL_SECTION_START = "<|tool_calls_section_begin|>" |
| 44 | TOOL_SECTION_END = "<|tool_calls_section_end|>" |
| 45 | TOOL_CALL_START = "<|tool_call_begin|>" |
| 46 | TOOL_CALL_END = "<|tool_call_end|>" |
| 47 | TOOL_ARG_START = "<|tool_call_argument_begin|>" |
| 48 | |
| 49 | _TOOL_ID_RE = re.compile(r"(?P<id>.+:\d+)") |
| 50 | |
| 51 | |
| 52 | @functools.cache |
| 53 | def kimi_k2_config(thinking: bool = True) -> ParserEngineConfig: |
| 54 | reasoning_terminals = ( |
| 55 | { |
| 56 | "THINK_START": THINK_START, |
| 57 | "THINK_END": THINK_END, |
| 58 | } |
| 59 | if thinking |
| 60 | else {} |
| 61 | ) |
| 62 | reasoning_transitions = ( |
| 63 | { |
| 64 | (ParserState.REASONING, "THINK_START"): Transition( |
| 65 | ParserState.REASONING, |
| 66 | (), |
| 67 | ), |
| 68 | (ParserState.REASONING, "THINK_END"): Transition( |
| 69 | ParserState.CONTENT, |
| 70 | (EventType.REASONING_END,), |
| 71 | ), |
| 72 | (ParserState.CONTENT, "THINK_END"): Transition( |
| 73 | ParserState.CONTENT, |
| 74 | (), |
| 75 | ), |
| 76 | } |
| 77 | if thinking |
| 78 | else {} |
| 79 | ) |
| 80 | |
| 81 | return ParserEngineConfig( |
| 82 | name="kimi_k2", |
| 83 | initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, |
| 84 | wait_for_reasoning=thinking, |
| 85 | terminals={ |
| 86 | **reasoning_terminals, |
| 87 | "TOOL_SECTION_START": TOOL_SECTION_START, |
| 88 | "TOOL_SECTION_END": TOOL_SECTION_END, |
| 89 | "TOOL_START": TOOL_CALL_START, |
| 90 | "TOOL_END": TOOL_CALL_END, |
| 91 | "ARG_START": TOOL_ARG_START, |
| 92 | }, |
| 93 | token_id_terminals={ |
| 94 | **reasoning_terminals, |
| 95 | "TOOL_SECTION_START": TOOL_SECTION_START, |
| 96 | "TOOL_SECTION_END": TOOL_SECTION_END, |
| 97 | "TOOL_START": TOOL_CALL_START, |
| 98 | "TOOL_END": TOOL_CALL_END, |
| 99 | "ARG_START": TOOL_ARG_START, |
| 100 | }, |
| 101 | transitions={ |
| 102 | **reasoning_transitions, |
| 103 | (ParserState.REASONING, "TOOL_SECTION_START"): Transition( |
| 104 | ParserState.TOOL_PREAMBLE, |
| 105 | (EventType.REASONING_END,), |
| 106 | ), |
| 107 | (ParserState.CONTENT, "TOOL_SECTION_START"): Transition( |
| 108 | ParserState.TOOL_PREAMBLE, |
| 109 | (), |
| 110 | ), |
| 111 | (ParserState.TOOL_PREAMBLE, "TOOL_START"): Transition( |
| 112 | ParserState.TOOL_NAME, |
| 113 | (EventType.TOOL_CALL_START,), |
| 114 | ), |
| 115 | (ParserState.TOOL_NAME, "ARG_START"): Transition( |
| 116 | ParserState.TOOL_ARGS, |
| 117 | (), |
| 118 | ), |
| 119 | (ParserState.TOOL_ARGS, "TOOL_END"): Transition( |
| 120 | ParserState.TOOL_BETWEEN, |
| 121 | (EventType.TOOL_CALL_END,), |
| 122 | ), |
| 123 | (ParserState.TOOL_ARGS, "TOOL_SECTION_END"): Transition( |
| 124 | ParserState.TOOL_PREAMBLE, |
| 125 | (EventType.TOOL_CALL_END,), |
| 126 | ), |
| 127 | (ParserState.TOOL_BETWEEN, "TOOL_START"): Transition( |
| 128 | ParserState.TOOL_NAME, |
| 129 | (EventType.TOOL_CALL_START,), |
| 130 | ), |
| 131 | # Keep the parser in a tool state after the section closes so |
| 132 | # trailing model text after native tool calls is suppressed. |
| 133 | (ParserState.TOOL_PREAMBLE, "TOOL_SECTION_END"): Transition( |
| 134 | ParserState.TOOL_PREAMBLE, |
| 135 | (), |
| 136 | ), |
| 137 | (ParserState.TOOL_BETWEEN, "TOOL_SECTION_END"): Transition( |
| 138 | ParserState.TOOL_PREAMBLE, |
| 139 | (), |
| 140 | ), |
| 141 | }, |
| 142 | stream_arg_deltas=True, |
| 143 | tool_args_json=True, |
| 144 | strip_trailing_reasoning_whitespace=True, |
| 145 | drop_whitespace_only_content_before_tools=True, |
| 146 | strip_content_whitespace_with_tools=False, |
| 147 | validate_tool_names=False, |
| 148 | ) |
| 149 | |
| 150 | |
| 151 | class KimiK2Parser(ParserEngine): |
| 152 | """Kimi K2 parser backed by the declarative parser engine.""" |
| 153 | |
| 154 | def __init__( |
| 155 | self, |
| 156 | tokenizer: TokenizerLike, |
| 157 | tools: list[Tool] | None = None, |
| 158 | **kwargs, |
| 159 | ) -> None: |
| 160 | chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} |
| 161 | thinking = chat_kwargs.get("thinking", None) |
| 162 | enable_thinking = chat_kwargs.get("enable_thinking", None) |
| 163 | self.thinking_enabled = ( |
| 164 | True |
| 165 | if thinking is None and enable_thinking is None |
| 166 | else bool(thinking) or bool(enable_thinking) |
| 167 | ) |
| 168 | kwargs.setdefault( |
| 169 | "parser_engine_config", |
| 170 | kimi_k2_config(thinking=self.thinking_enabled), |
| 171 | ) |
| 172 | super().__init__(tokenizer, tools, **kwargs) |
| 173 | |
| 174 | vocab = self.vocab |
| 175 | self._start_token_id = vocab.get(THINK_START) |
| 176 | self._end_token_id = vocab.get(THINK_END) |
| 177 | self._tool_section_start_token_id = vocab.get(TOOL_SECTION_START) |
| 178 | |
| 179 | @staticmethod |
| 180 | def _extract_tool_id_and_name(header: str | None) -> tuple[str | None, str | None]: |
| 181 | if header is None: |
| 182 | return None, None |
| 183 | match = _TOOL_ID_RE.match(header.strip()) |
| 184 | if not match: |
| 185 | return None, None |
| 186 | |
| 187 | tool_id = match.group("id").strip() |
| 188 | tool_name = tool_id.split(":")[0].removeprefix("functions.") |
| 189 | return tool_id, tool_name |
| 190 | |
| 191 | def _emit_name_delta( |
| 192 | self, |
| 193 | idx: int, |
| 194 | deltas: list[DeltaToolCall], |
| 195 | name: str | None, |
| 196 | ) -> None: |
| 197 | tool_id, tool_name = self._extract_tool_id_and_name(name) |
| 198 | if not tool_name: |
| 199 | if 0 <= idx < len(self._tool_slots): |
| 200 | self._tool_slots[idx].name = "" |
| 201 | return |
| 202 | |
| 203 | slot = self._tool_slots[idx] |
| 204 | slot.id = tool_id or "" |
| 205 | super()._emit_name_delta(idx, deltas, tool_name) |
| 206 | |
| 207 | def _handle_tool_end(self, event, deltas) -> None: |
| 208 | idx = event.tool_index |
| 209 | if 0 <= idx < len(self._tool_slots) and not self._tool_slots[idx].name_sent: |
| 210 | tool_id, tool_name = self._extract_tool_id_and_name( |
| 211 | self._tool_slots[idx].name |
| 212 | ) |
| 213 | if tool_name: |
| 214 | self._tool_slots[idx].id = tool_id or "" |
| 215 | self._tool_slots[idx].name = tool_name |
| 216 | super()._handle_tool_end(event, deltas) |
| 217 | |
| 218 | def _handle_arg_chunk(self, event, deltas) -> None: |
| 219 | idx = event.tool_index |
| 220 | name_sent_before = ( |
| 221 | 0 <= idx < len(self._tool_slots) and self._tool_slots[idx].name_sent |
| 222 | ) |
| 223 | super()._handle_arg_chunk(event, deltas) |
| 224 | if ( |
| 225 | event.value |
| 226 | and not name_sent_before |
| 227 | and 0 <= idx < len(self._tool_slots) |
| 228 | and self._tool_slots[idx].name_sent |
| 229 | ): |
| 230 | deltas.append( |
| 231 | DeltaToolCall( |
| 232 | index=idx, |
| 233 | function=DeltaFunctionCall(arguments=event.value), |
| 234 | ) |
| 235 | ) |
| 236 | |
| 237 | def _extract_args_json(self, raw_args: str, func_name: str) -> str: |
| 238 | return raw_args.strip() or "{}" |
| 239 | |
| 240 | def extract_content_ids(self, input_ids: list[int]) -> list[int]: |
| 241 | if not self.thinking_enabled: |
| 242 | return input_ids |
| 243 | |
| 244 | end_id = self._end_token_id |
| 245 | if end_id is not None and end_id in input_ids: |
| 246 | end_idx = len(input_ids) - 1 - input_ids[::-1].index(end_id) |
| 247 | return input_ids[end_idx + 1 :] |
| 248 | |
| 249 | tool_section_id = self._tool_section_start_token_id |
| 250 | if tool_section_id is not None and tool_section_id in input_ids: |
| 251 | section_idx = len(input_ids) - 1 - input_ids[::-1].index(tool_section_id) |
| 252 | return input_ids[section_idx:] |
| 253 | |
| 254 | return [] |
| 255 | |
| 256 | def extract_reasoning( |
| 257 | self, |
| 258 | model_output: str, |
| 259 | request: ChatCompletionRequest | ResponsesRequest, |
| 260 | ) -> tuple[str | None, str | None]: |
| 261 | if not self.thinking_enabled: |
| 262 | return None, model_output |
| 263 | return super().extract_reasoning(model_output, request) |
| 264 | |
| 265 | def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: |
| 266 | if not self.thinking_enabled: |
| 267 | return 0 |
| 268 | return super().count_reasoning_tokens(token_ids) |