# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Kimi K2 parser for reasoning and tool calls. Kimi K2 tool call format:: <|tool_calls_section_begin|> <|tool_call_begin|>functions.get_weather:0 <|tool_call_argument_begin|>{"city": "Tokyo"}<|tool_call_end|> <|tool_calls_section_end|> The header before ``<|tool_call_argument_begin|>`` is Kimi's native tool call id. The function name is the final component before ``:N``. """ from __future__ import annotations import functools from collections.abc import Sequence from typing import TYPE_CHECKING import regex as re from vllm.entrypoints.openai.engine.protocol import DeltaFunctionCall, DeltaToolCall from vllm.parser.engine.events import EventType from vllm.parser.engine.parser_engine import ParserEngine from vllm.parser.engine.parser_engine_config import ( ParserEngineConfig, ParserState, Transition, ) if TYPE_CHECKING: from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool THINK_START = "" THINK_END = "" TOOL_SECTION_START = "<|tool_calls_section_begin|>" TOOL_SECTION_END = "<|tool_calls_section_end|>" TOOL_CALL_START = "<|tool_call_begin|>" TOOL_CALL_END = "<|tool_call_end|>" TOOL_ARG_START = "<|tool_call_argument_begin|>" _TOOL_ID_RE = re.compile(r"(?P.+:\d+)") @functools.cache def kimi_k2_config(thinking: bool = True) -> ParserEngineConfig: reasoning_terminals = ( { "THINK_START": THINK_START, "THINK_END": THINK_END, } if thinking else {} ) reasoning_transitions = ( { (ParserState.REASONING, "THINK_START"): Transition( ParserState.REASONING, (), ), (ParserState.REASONING, "THINK_END"): Transition( ParserState.CONTENT, (EventType.REASONING_END,), ), (ParserState.CONTENT, "THINK_END"): Transition( ParserState.CONTENT, (), ), } if thinking else {} ) return ParserEngineConfig( name="kimi_k2", initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, terminals={ **reasoning_terminals, "TOOL_SECTION_START": TOOL_SECTION_START, "TOOL_SECTION_END": TOOL_SECTION_END, "TOOL_START": TOOL_CALL_START, "TOOL_END": TOOL_CALL_END, "ARG_START": TOOL_ARG_START, }, token_id_terminals={ **reasoning_terminals, "TOOL_SECTION_START": TOOL_SECTION_START, "TOOL_SECTION_END": TOOL_SECTION_END, "TOOL_START": TOOL_CALL_START, "TOOL_END": TOOL_CALL_END, "ARG_START": TOOL_ARG_START, }, transitions={ **reasoning_transitions, (ParserState.REASONING, "TOOL_SECTION_START"): Transition( ParserState.TOOL_PREAMBLE, (EventType.REASONING_END,), ), (ParserState.CONTENT, "TOOL_SECTION_START"): Transition( ParserState.TOOL_PREAMBLE, (), ), (ParserState.TOOL_PREAMBLE, "TOOL_START"): Transition( ParserState.TOOL_NAME, (EventType.TOOL_CALL_START,), ), (ParserState.TOOL_NAME, "ARG_START"): Transition( ParserState.TOOL_ARGS, (), ), (ParserState.TOOL_ARGS, "TOOL_END"): Transition( ParserState.TOOL_BETWEEN, (EventType.TOOL_CALL_END,), ), (ParserState.TOOL_ARGS, "TOOL_SECTION_END"): Transition( ParserState.TOOL_PREAMBLE, (EventType.TOOL_CALL_END,), ), (ParserState.TOOL_BETWEEN, "TOOL_START"): Transition( ParserState.TOOL_NAME, (EventType.TOOL_CALL_START,), ), # Keep the parser in a tool state after the section closes so # trailing model text after native tool calls is suppressed. (ParserState.TOOL_PREAMBLE, "TOOL_SECTION_END"): Transition( ParserState.TOOL_PREAMBLE, (), ), (ParserState.TOOL_BETWEEN, "TOOL_SECTION_END"): Transition( ParserState.TOOL_PREAMBLE, (), ), }, stream_arg_deltas=True, tool_args_json=True, strip_trailing_reasoning_whitespace=True, drop_whitespace_only_content_before_tools=True, strip_content_whitespace_with_tools=False, validate_tool_names=False, ) class KimiK2Parser(ParserEngine): """Kimi K2 parser backed by the declarative parser engine.""" def __init__( self, tokenizer: TokenizerLike, tools: list[Tool] | None = None, **kwargs, ) -> None: chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} thinking = chat_kwargs.get("thinking", None) enable_thinking = chat_kwargs.get("enable_thinking", None) self.thinking_enabled = ( True if thinking is None and enable_thinking is None else bool(thinking) or bool(enable_thinking) ) kwargs.setdefault( "parser_engine_config", kimi_k2_config(thinking=self.thinking_enabled), ) super().__init__(tokenizer, tools, **kwargs) vocab = self.vocab self._start_token_id = vocab.get(THINK_START) self._end_token_id = vocab.get(THINK_END) self._tool_section_start_token_id = vocab.get(TOOL_SECTION_START) @staticmethod def _extract_tool_id_and_name(header: str | None) -> tuple[str | None, str | None]: if header is None: return None, None match = _TOOL_ID_RE.match(header.strip()) if not match: return None, None tool_id = match.group("id").strip() tool_name = tool_id.split(":")[0].removeprefix("functions.") return tool_id, tool_name def _emit_name_delta( self, idx: int, deltas: list[DeltaToolCall], name: str | None, ) -> None: tool_id, tool_name = self._extract_tool_id_and_name(name) if not tool_name: if 0 <= idx < len(self._tool_slots): self._tool_slots[idx].name = "" return slot = self._tool_slots[idx] slot.id = tool_id or "" super()._emit_name_delta(idx, deltas, tool_name) def _handle_tool_end(self, event, deltas) -> None: idx = event.tool_index if 0 <= idx < len(self._tool_slots) and not self._tool_slots[idx].name_sent: tool_id, tool_name = self._extract_tool_id_and_name( self._tool_slots[idx].name ) if tool_name: self._tool_slots[idx].id = tool_id or "" self._tool_slots[idx].name = tool_name super()._handle_tool_end(event, deltas) def _handle_arg_chunk(self, event, deltas) -> None: idx = event.tool_index name_sent_before = ( 0 <= idx < len(self._tool_slots) and self._tool_slots[idx].name_sent ) super()._handle_arg_chunk(event, deltas) if ( event.value and not name_sent_before and 0 <= idx < len(self._tool_slots) and self._tool_slots[idx].name_sent ): deltas.append( DeltaToolCall( index=idx, function=DeltaFunctionCall(arguments=event.value), ) ) def _extract_args_json(self, raw_args: str, func_name: str) -> str: return raw_args.strip() or "{}" def is_reasoning_end(self, input_ids: list[int]) -> bool: if not self.thinking_enabled: return True start_id = self._start_token_id end_id = self._end_token_id tool_section_id = self._tool_section_start_token_id for i in range(len(input_ids) - 1, -1, -1): token_id = input_ids[i] if start_id is not None and token_id == start_id: return False if end_id is not None and token_id == end_id: return True if tool_section_id is not None and token_id == tool_section_id: return True return False def extract_content_ids(self, input_ids: list[int]) -> list[int]: if not self.thinking_enabled: return input_ids end_id = self._end_token_id if end_id is not None and end_id in input_ids: end_idx = len(input_ids) - 1 - input_ids[::-1].index(end_id) return input_ids[end_idx + 1 :] tool_section_id = self._tool_section_start_token_id if tool_section_id is not None and tool_section_id in input_ids: section_idx = len(input_ids) - 1 - input_ids[::-1].index(tool_section_id) return input_ids[section_idx:] return [] def extract_reasoning( self, model_output: str, request: ChatCompletionRequest | ResponsesRequest, ) -> tuple[str | None, str | None]: if not self.thinking_enabled: return None, model_output return super().extract_reasoning(model_output, request) def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: if not self.thinking_enabled: return 0 return super().count_reasoning_tokens(token_ids)