From 2bc20e8abaf7a82ecc068d879ebb925c9317bd40 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Tue, 30 Jun 2026 15:53:17 +0800 Subject: [PATCH] [Frontend] Add Streaming Parser Engine and new Kimi k2.5/k2.6/k2.7 Parser (#46610) Signed-off-by: chaunceyjiang Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/parser/engine/trace_builder.py | 95 +++++- .../test_kimi_k2_reasoning_parser.py | 72 +---- vllm/parser/engine/registered_adapters.py | 6 + vllm/parser/kimi_k2.py | 285 ++++++++++++++++++ vllm/reasoning/kimi_k2_reasoning_parser.py | 243 +-------------- vllm/tool_parsers/kimi_k2_tool_parser.py | 266 +--------------- 6 files changed, 397 insertions(+), 570 deletions(-) create mode 100644 vllm/parser/kimi_k2.py diff --git a/tests/parser/engine/trace_builder.py b/tests/parser/engine/trace_builder.py index 7f41b2b951..93159a0ae2 100644 --- a/tests/parser/engine/trace_builder.py +++ b/tests/parser/engine/trace_builder.py @@ -31,6 +31,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( from vllm.parser.engine.registered_adapters import ( Gemma4Parser, Glm47MoeParser, + KimiK2Parser, MinimaxM2Parser, NemotronV3Parser, Qwen3Parser, @@ -717,6 +718,96 @@ def _build_glm47_moe(scenario: Scenario, validate: bool = True) -> Sample: return sample +# ── Kimi K2 (native tool-call section, starts in REASONING) ────────── + +_KIMI_K2_VOCAB: dict[str, int] = { + "": 50, + "": 51, + "<|tool_calls_section_begin|>": 60, + "<|tool_calls_section_end|>": 61, + "<|tool_call_begin|>": 62, + "<|tool_call_end|>": 63, + "<|tool_call_argument_begin|>": 64, +} + + +def _kimi_k2_tool_segments( + tool_calls: list[ToolCallSpec], +) -> list[tuple[str, bool]]: + segs: list[tuple[str, bool]] = [("<|tool_calls_section_begin|>", True)] + for index, tc in enumerate(tool_calls): + args = json.dumps(tc.arguments, ensure_ascii=False, separators=(",", ":")) + segs.extend( + [ + ("<|tool_call_begin|>", True), + (f"functions.{tc.name}:{index}\n", False), + ("<|tool_call_argument_begin|>", True), + (args, False), + ("<|tool_call_end|>", True), + ] + ) + segs.append(("<|tool_calls_section_end|>", True)) + return segs + + +def _kimi_k2_segments(scenario: Scenario) -> list[tuple[str, bool]]: + segs: list[tuple[str, bool]] = [] + if scenario.reasoning is not None: + segs.append(("", True)) + segs.append((scenario.reasoning, False)) + if scenario.content is not None or scenario.tool_calls is not None: + segs.append(("", True)) + if scenario.content is not None: + segs.append((scenario.content, False)) + if scenario.tool_calls is not None: + segs.extend(_kimi_k2_tool_segments(scenario.tool_calls)) + return segs + + +def _build_kimi_k2( + scenario: Scenario, + validate: bool = True, + thinking: bool = True, +) -> Sample: + expected_reasoning = ( + scenario.reasoning.rstrip() + if (thinking and scenario.reasoning is not None) + else None + ) + if thinking and scenario.reasoning is None: + expected_reasoning = "" + + sample = _make_sample( + sample_id=f"kimi_k2-{scenario.id}", + description=scenario.description, + vocab=_KIMI_K2_VOCAB, + segments=_kimi_k2_segments(scenario), + expected_reasoning=expected_reasoning, + expected_content=_qwen3_expected_content(scenario), + expected_tool_calls=_expected_tc(scenario), + tools=_expected_tools(scenario), + chat_template_kwargs=None if thinking else {"thinking": False}, + ) + if validate: + _validate_sample( + sample, + KimiK2Parser, + chat_template_kwargs=sample.chat_template_kwargs, + ) + return sample + + +_KIMI_K2_SCENARIOS = [ + *SCENARIOS, + Scenario( + id="trailing-reasoning-whitespace", + description="Reasoning trailing whitespace is stripped", + reasoning="Reasoning with trailing whitespace. \n\t", + content="Done.", + ), +] + + # ── Registry and public API ────────────────────────────────────────── _BUILDERS: dict[str, Any] = { @@ -726,6 +817,7 @@ _BUILDERS: dict[str, Any] = { "nemotron_v3": _build_nemotron_v3, "seed_oss": _build_seed_oss, "glm47_moe": _build_glm47_moe, + "kimi_k2": _build_kimi_k2, } @@ -733,7 +825,8 @@ _BUILDERS: dict[str, Any] = { def build_samples(model: str) -> tuple[Sample, ...]: """Build all scenario samples for a model, self-validated.""" builder = _BUILDERS[model] - return tuple(builder(s) for s in SCENARIOS) + scenarios = _KIMI_K2_SCENARIOS if model == "kimi_k2" else SCENARIOS + return tuple(builder(s) for s in scenarios) def build_sample(model: str, scenario: Scenario) -> Sample: diff --git a/tests/reasoning/test_kimi_k2_reasoning_parser.py b/tests/reasoning/test_kimi_k2_reasoning_parser.py index dfce2075c6..e39f8cee74 100644 --- a/tests/reasoning/test_kimi_k2_reasoning_parser.py +++ b/tests/reasoning/test_kimi_k2_reasoning_parser.py @@ -7,7 +7,6 @@ import pytest from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.engine.protocol import DeltaMessage -from vllm.reasoning.identity_reasoning_parser import IdentityReasoningParser from vllm.reasoning.kimi_k2_reasoning_parser import KimiK2ReasoningParser from vllm.tokenizers import get_tokenizer @@ -33,20 +32,6 @@ def kimi_k2_tokenizer(): return get_tokenizer(tokenizer_name=REASONING_MODEL_NAME, trust_remote_code=True) -def test_parser_selection_thinking_enabled(kimi_k2_tokenizer): - parser = KimiK2ReasoningParser( - kimi_k2_tokenizer, chat_template_kwargs={"thinking": True} - ) - assert parser._identity_parser is None - - -def test_parser_selection_thinking_disabled(kimi_k2_tokenizer): - parser = KimiK2ReasoningParser( - kimi_k2_tokenizer, chat_template_kwargs={"thinking": False} - ) - assert isinstance(parser._identity_parser, IdentityReasoningParser) - - def test_extract_reasoning_with_think_tags(kimi_k2_tokenizer): parser = KimiK2ReasoningParser(kimi_k2_tokenizer) request = ChatCompletionRequest(model="test-model", messages=[], temperature=1.0) @@ -65,7 +50,7 @@ def test_extract_reasoning_empty_thinking(kimi_k2_tokenizer): reasoning, content = parser.extract_reasoning( "final answer", request ) - assert reasoning == "" + assert reasoning is None assert content == "final answer" @@ -96,8 +81,8 @@ def test_streaming_reasoning_then_content(kimi_k2_tokenizer): """Token-by-token streaming: reasoning tokens then content after .""" parser = KimiK2ReasoningParser(kimi_k2_tokenizer) - think_id = parser._start_token_id - end_think_id = parser._end_token_id + think_id = parser._parser_engine._start_token_id + end_think_id = parser._parser_engine._end_token_id # Use a real token ID from the tokenizer for regular content regular_id = kimi_k2_tokenizer.encode("hello", add_special_tokens=False)[0] @@ -154,8 +139,8 @@ def test_streaming_tool_section_ends_reasoning(kimi_k2_tokenizer): """<|tool_calls_section_begin|> in delta ends reasoning during streaming.""" parser = KimiK2ReasoningParser(kimi_k2_tokenizer) - think_id = parser._start_token_id - tool_begin_id = parser._tool_section_start_token_id + think_id = parser._parser_engine._start_token_id + tool_begin_id = parser._parser_engine._tool_section_start_token_id regular_id = kimi_k2_tokenizer.encode("hello", add_special_tokens=False)[0] # Tool section token arrives — should transition from reasoning to content @@ -169,50 +154,3 @@ def test_streaming_tool_section_ends_reasoning(kimi_k2_tokenizer): ) assert isinstance(result, DeltaMessage) assert result.content == "<|tool_calls_section_begin|>" - - -def test_streaming_end_token_id_buffered(mock_kimi_k2_tokenizer): - """When stop sequences buffer text, ID arrives before its text. - - The token ID is present in delta_token_ids but the actual string is not - yet in delta_text (still buffered). The parser must return None to wait - for the next delta, instead of calling find() which returns -1 and - silently corrupting the text split. - """ - parser = KimiK2ReasoningParser(mock_kimi_k2_tokenizer) - think_id = parser._start_token_id - end_think_id = parser._end_token_id - - # Simulate: ID arrived but text not yet flushed. - # Two token IDs in delta to bypass the single-special-token guard. - result = parser.extract_reasoning_streaming( - previous_text="some reasoning", - current_text="some reasoning extra", - delta_text="extra", # text not yet flushed - previous_token_ids=[think_id], - current_token_ids=[think_id, end_think_id, 999], - delta_token_ids=[end_think_id, 999], - ) - assert result is None - - -def test_streaming_tool_section_id_buffered(mock_kimi_k2_tokenizer): - """When stop sequences buffer text, tool section start ID arrives before its text. - - Same buffering scenario as above but for <|tool_calls_section_begin|>. - Without the guard, find() returns -1 and delta_text[:tool_index] silently - drops the last character of reasoning. - """ - parser = KimiK2ReasoningParser(mock_kimi_k2_tokenizer) - think_id = parser._start_token_id - tool_begin_id = parser._tool_section_start_token_id - - result = parser.extract_reasoning_streaming( - previous_text="some reasoning", - current_text="some reasoning extra", - delta_text="extra", # tool section text not yet flushed - previous_token_ids=[think_id], - current_token_ids=[think_id, tool_begin_id, 999], - delta_token_ids=[tool_begin_id, 999], - ) - assert result is None diff --git a/vllm/parser/engine/registered_adapters.py b/vllm/parser/engine/registered_adapters.py index c5250abf82..fdcbf81ee2 100644 --- a/vllm/parser/engine/registered_adapters.py +++ b/vllm/parser/engine/registered_adapters.py @@ -10,6 +10,7 @@ names so that :class:`ReasoningParserManager` and from vllm.parser.engine.adapters import make_adapters from vllm.parser.gemma4 import Gemma4Parser from vllm.parser.glm47_moe import Glm47MoeParser +from vllm.parser.kimi_k2 import KimiK2Parser from vllm.parser.minimax_m2 import MinimaxM2Parser from vllm.parser.nemotron_v3 import NemotronV3Parser from vllm.parser.qwen3 import Qwen3Parser @@ -44,3 +45,8 @@ from vllm.parser.seed_oss import SeedOssParser Glm47MoeParserReasoningAdapter, Glm47MoeParserToolAdapter, ) = make_adapters(Glm47MoeParser) + +( + KimiK2ParserReasoningAdapter, + KimiK2ParserToolAdapter, +) = make_adapters(KimiK2Parser) diff --git a/vllm/parser/kimi_k2.py b/vllm/parser/kimi_k2.py new file mode 100644 index 0000000000..03a0aa6cec --- /dev/null +++ b/vllm/parser/kimi_k2.py @@ -0,0 +1,285 @@ +# 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) diff --git a/vllm/reasoning/kimi_k2_reasoning_parser.py b/vllm/reasoning/kimi_k2_reasoning_parser.py index 0b64c5c62e..45f99965bc 100644 --- a/vllm/reasoning/kimi_k2_reasoning_parser.py +++ b/vllm/reasoning/kimi_k2_reasoning_parser.py @@ -1,245 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Iterable, Sequence -from typing import TYPE_CHECKING +from vllm.parser.engine.registered_adapters import KimiK2ParserReasoningAdapter -from transformers import PreTrainedTokenizerBase +KimiK2ReasoningParser = KimiK2ParserReasoningAdapter -from vllm.entrypoints.openai.engine.protocol import DeltaMessage -from vllm.reasoning.abs_reasoning_parsers import ReasoningParser -from vllm.reasoning.identity_reasoning_parser import IdentityReasoningParser - -if TYPE_CHECKING: - from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest - from vllm.entrypoints.openai.responses.protocol import ResponsesRequest - - -class KimiK2ReasoningParser(ReasoningParser): - """ - Reasoning parser for Kimi K2 model. - - The Kimi K2 model uses ... tokens to denote reasoning text, - and may implicitly end reasoning by starting a tool call section using - <|tool_calls_section_begin|>. - Thinking may also begin without a token. - - Kimi's thinking mode can be disabled via chat_template_kwargs. - """ - - def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs): - super().__init__(tokenizer, *args, **kwargs) - - if not self.model_tokenizer: - raise ValueError( - "The model tokenizer must be passed to the ReasoningParser " - "constructor during construction." - ) - - # Check if thinking is disabled via chat_template_kwargs - chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} - thinking = bool(chat_kwargs.get("thinking", True)) - - # If thinking is not enabled, use identity parser to fall through - self._identity_parser: IdentityReasoningParser | None - if not thinking: - self._identity_parser = IdentityReasoningParser(tokenizer, *args, **kwargs) - else: - self._identity_parser = None - - # Token definitions - self._start_token = "" - self._end_token = "" - self._tool_section_start_token = "<|tool_calls_section_begin|>" - - # Get token IDs - self._start_token_id = self.vocab.get(self._start_token) - self._end_token_id = self.vocab.get(self._end_token) - self._tool_section_start_token_id = self.vocab.get( - self._tool_section_start_token - ) - - if self._start_token_id is None or self._end_token_id is None: - raise RuntimeError( - "KimiK2ReasoningParser could not locate think start/end " - "tokens in the tokenizer!" - ) - - @property - def reasoning_start_str(self) -> str | None: - return self._start_token - - @property - def reasoning_end_str(self) -> str | None: - return self._end_token - - def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: - """ - Check if the reasoning content ends in the input_ids. - - Reasoning ends when we see either: - 1. The end token () - 2. The tool section start token (<|tool_calls_section_begin|>) - """ - if self._identity_parser is not None: - return self._identity_parser.is_reasoning_end(input_ids) - - start_token_id = self._start_token_id - end_token_id = self._end_token_id - tool_section_start_token_id = self._tool_section_start_token_id - - for i in range(len(input_ids) - 1, -1, -1): - if input_ids[i] == start_token_id: - return False - if input_ids[i] == end_token_id: - return True - # Implicit reasoning end via tool call section - if ( - tool_section_start_token_id is not None - and input_ids[i] == tool_section_start_token_id - ): - return True - return False - - def is_reasoning_end_streaming( - self, input_ids: Sequence[int], delta_ids: Iterable[int] - ) -> bool: - """ - Check if the reasoning content ends in the input_ids on a decode step. - """ - if self._identity_parser is not None: - return self._identity_parser.is_reasoning_end_streaming( - input_ids, delta_ids - ) - - # Materialize iterable for membership checks - delta_ids_set = set(delta_ids) - - # Check for explicit end token or implicit tool section start in delta - if self._end_token_id in delta_ids_set: - return True - return ( - self._tool_section_start_token_id is not None - and self._tool_section_start_token_id in delta_ids_set - ) - - def extract_content_ids(self, input_ids: list[int]) -> list[int]: - """ - Extract content token ids from the input_ids. - """ - if self._identity_parser is not None: - return self._identity_parser.extract_content_ids(input_ids) - - if self._end_token_id in input_ids: - end_token_index = ( - len(input_ids) - 1 - input_ids[::-1].index(self._end_token_id) - ) - - if end_token_index != -1: - return input_ids[end_token_index + 1 :] - - if ( - self._tool_section_start_token_id is not None - and self._tool_section_start_token_id in input_ids - ): - tool_section_index = ( - len(input_ids) - - 1 - - input_ids[::-1].index(self._tool_section_start_token_id) - ) - - if tool_section_index != -1: - return input_ids[tool_section_index:] - - # still reasoning (no content) - return [] - - def extract_reasoning( - self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest" - ) -> tuple[str | None, str | None]: - """ - Extract reasoning content from the model output. - """ - if self._identity_parser is not None: - return self._identity_parser.extract_reasoning(model_output, request) - - # thinking does not require a think start token but consume it if present - start_token_index = model_output.find(self._start_token) - start_token_index = 0 if start_token_index != 0 else len(self._start_token) - end_token_index = model_output.find(self._end_token) - - if end_token_index != -1: - return ( - model_output[start_token_index:end_token_index], - model_output[end_token_index + len(self._end_token) :] or None, - ) - - tool_section_index = model_output.find(self._tool_section_start_token) - if tool_section_index != -1: - return ( - model_output[start_token_index:tool_section_index], - model_output[tool_section_index:] or None, - ) - - # still reasoning (no content) - return ( - model_output[start_token_index:], - None, - ) - - def extract_reasoning_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - ) -> DeltaMessage | None: - """ - Extract reasoning content from a delta message during streaming. - """ - if self._identity_parser is not None: - return self._identity_parser.extract_reasoning_streaming( - previous_text, - current_text, - delta_text, - previous_token_ids, - current_token_ids, - delta_token_ids, - ) - - # If reasoning has already ended in previous tokens, this is content - if self.is_reasoning_end(previous_token_ids): - return DeltaMessage(content=delta_text) - - # Skip single special tokens - if len(delta_token_ids) == 1 and delta_token_ids[0] in [ - self._start_token_id, - self._end_token_id, - ]: - return None - - if self._end_token_id in delta_token_ids: - if self._end_token not in delta_text: - # Token ID arrived before text was flushed (stop-sequence buffering). - # Wait for the next delta when the text becomes visible. - return None - end_index = delta_text.find(self._end_token) - reasoning = delta_text[:end_index] - content = delta_text[end_index + len(self._end_token) :] - return DeltaMessage( - reasoning=reasoning, content=content if content else None - ) - - if self._tool_section_start_token_id in delta_token_ids: - if self._tool_section_start_token not in delta_text: - # Token ID arrived before text was flushed (stop-sequence buffering). - return None - tool_index = delta_text.find(self._tool_section_start_token) - reasoning = delta_text[:tool_index] - content = delta_text[tool_index:] - return DeltaMessage(reasoning=reasoning, content=content) - - # still reasoning (no end token) - return DeltaMessage(reasoning=delta_text) +__all__ = ["KimiK2ReasoningParser"] diff --git a/vllm/tool_parsers/kimi_k2_tool_parser.py b/vllm/tool_parsers/kimi_k2_tool_parser.py index 18f242fffe..d1cc483018 100644 --- a/vllm/tool_parsers/kimi_k2_tool_parser.py +++ b/vllm/tool_parsers/kimi_k2_tool_parser.py @@ -1,278 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Sequence - -import regex as re - from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) -from vllm.entrypoints.openai.engine.protocol import ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, - ExtractedToolCallInformation, - FunctionCall, - ToolCall, -) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -from vllm.logger import init_logger -from vllm.tokenizers import TokenizerLike -from vllm.tool_parsers.abstract_tool_parser import ( - Tool, - ToolParser, -) -from vllm.tool_parsers.utils import partial_tag_overlap - -logger = init_logger(__name__) +from vllm.parser.engine.registered_adapters import KimiK2ParserToolAdapter -class KimiK2ToolParser(ToolParser): +class KimiK2ToolParser(KimiK2ParserToolAdapter): # type: ignore[valid-type, misc] structural_tag_model = "kimi" - def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): - super().__init__(tokenizer, tools) - - # Streaming state - self._sent_content_idx: int = 0 - self.prev_tool_call_arr: list[dict] = [] - self.streamed_args_for_tool: list[str] = [] - - # Section marker - self.tool_calls_start_token: str = "<|tool_calls_section_begin|>" - - # Individual tool call markers - self.tool_call_start_token: str = "<|tool_call_begin|>" - self.tool_call_end_token: str = "<|tool_call_end|>" - self.tool_call_arg_token: str = "<|tool_call_argument_begin|>" - - # Regex for non-streaming extraction - self.tool_call_regex = re.compile( - r"<\|tool_call_begin\|>\s*(?P[^<]+:\d+)\s*" - r"<\|tool_call_argument_begin\|>\s*" - r"(?P(?:(?!<\|tool_call_begin\|>).)*?)\s*" - r"<\|tool_call_end\|>", - re.DOTALL, - ) - - if not self.model_tokenizer: - raise ValueError( - "The model tokenizer must be passed to the ToolParser " - "constructor during construction." - ) - def adjust_request( - self, request: ChatCompletionRequest | ResponsesRequest + self, + request: ChatCompletionRequest | ResponsesRequest, ) -> ChatCompletionRequest | ResponsesRequest: - request = super().adjust_request(request) if request.tools and request.tool_choice != "none": - # Ensure special-token markers appear as literal text in - # current_text so we can do pure text-based parsing. request.skip_special_tokens = False return request - - def extract_tool_calls( - self, - model_output: str, - request: ChatCompletionRequest, - ) -> ExtractedToolCallInformation: - # sanity check; avoid unnecessary processing - if self.tool_calls_start_token not in model_output: - return ExtractedToolCallInformation( - tools_called=False, tool_calls=[], content=model_output - ) - - else: - try: - # there are two possible captures - between tags, or between a - # tag and end-of-string so the result of - # findall is an array of tuples where one is a function call and - # the other is None - function_call_tuples = self.tool_call_regex.findall(model_output) - - logger.debug("function_call_tuples: %s", function_call_tuples) - - tool_calls = [] - for match in function_call_tuples: - function_id, function_args = match - # function_id: functions.get_weather:0 or get_weather:0 - function_name = function_id.split(":")[0].split(".")[-1] - tool_calls.append( - ToolCall( - id=function_id, - type="function", - function=FunctionCall( - name=function_name, arguments=function_args - ), - ) - ) - - content = model_output[: model_output.find(self.tool_calls_start_token)] - return ExtractedToolCallInformation( - tools_called=True, - tool_calls=tool_calls, - content=content if content else None, - ) - - except Exception: - logger.exception("Error in extracting tool call from response.") - return ExtractedToolCallInformation( - tools_called=False, tool_calls=[], content=model_output - ) - - def _extract_content(self, current_text: str) -> str | None: - """Return unsent content before the tool-calls section, or None. - - Holds back any trailing suffix that partially matches - ``<|tool_calls_section_begin|>`` to avoid leaking marker bytes. - """ - if self.tool_calls_start_token not in current_text: - overlap = partial_tag_overlap(current_text, self.tool_calls_start_token) - sendable_idx = len(current_text) - overlap - else: - sendable_idx = current_text.index(self.tool_calls_start_token) - - if sendable_idx > self._sent_content_idx: - content = current_text[self._sent_content_idx : sendable_idx] - self._sent_content_idx = sendable_idx - return content - return None - - def _extract_tool_calls(self, current_text: str) -> list[str]: - """Extract raw bodies from ``<|tool_call_begin|>…<|tool_call_end|>`` blocks.""" - if self.tool_calls_start_token not in current_text: - return [] - - results: list[str] = [] - pos = current_text.index(self.tool_calls_start_token) - while True: - start = current_text.find(self.tool_call_start_token, pos) - if start == -1: - break - tc_start = start + len(self.tool_call_start_token) - end = current_text.find(self.tool_call_end_token, tc_start) - - if end != -1: - tool_call = current_text[tc_start:end] - pos = end + len(self.tool_call_end_token) - else: - tool_call = current_text[tc_start:] - overlap = partial_tag_overlap(tool_call, self.tool_call_end_token) - if overlap: - tool_call = tool_call[:-overlap] - - results.append(tool_call) - - if end == -1: - break - return results - - @staticmethod - def _extract_tool_id_and_name( - header: str | None, - ) -> tuple[str | None, str | None]: - """Parse ``(tool_id, tool_name)`` from a header - like ``"functions.get_weather:0"``.""" - if header is None: - return None, None - match = re.match(r"(.+:\d+)", header) - if not match: - return None, None - - tool_id = match.group(1).strip() - tool_name = tool_id.split(":")[0].split(".")[-1] - return tool_id, tool_name - - def _split_tool_call(self, tool_call: str) -> tuple[str | None, str | None]: - """Split a tool-call body into ``(header, arguments)`` at the argument marker. - - Example:: - 'get_weather:0 <|tool_call_argument_begin|>{"c' - -> ("get_weather:0", '{"c') - """ - arg_pos = tool_call.find(self.tool_call_arg_token) - if arg_pos == -1: - return None, None - header = tool_call[:arg_pos].strip() - tool_args = tool_call[arg_pos + len(self.tool_call_arg_token) :] - return header, tool_args - - def _compute_args_diff(self, index: int, tool_args: str | None) -> str | None: - """Return new argument text not yet sent for tool `index`, or None.""" - if tool_args is None: - return None - prev = self.streamed_args_for_tool[index] - if len(tool_args) <= len(prev): - return None - diff = tool_args[len(prev) :] - self.streamed_args_for_tool[index] = tool_args - self.prev_tool_call_arr[index]["arguments"] = tool_args - return diff - - def extract_tool_calls_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - request: ChatCompletionRequest, - ) -> DeltaMessage | None: - try: - # Extract any content before tool calls. - content = self._extract_content(current_text) - tool_calls = self._extract_tool_calls(current_text) - tool_call_deltas: list[DeltaToolCall] = [] - - for i, tool_call in enumerate(tool_calls): - # First time seeing tool call at index i. - if i >= len(self.prev_tool_call_arr): - # Initialize streaming state. - self.prev_tool_call_arr.append({}) - self.streamed_args_for_tool.append("") - - header, tool_args = self._split_tool_call(tool_call) - - # Stream back tool name. - if "name" not in self.prev_tool_call_arr[i]: - tool_id, tool_name = self._extract_tool_id_and_name(header) - if not tool_name: - # Can't skip to tool i+1 if i isn't ready - break - self.prev_tool_call_arr[i]["name"] = tool_name - self.prev_tool_call_arr[i]["id"] = tool_id - tool_call_deltas.append( - DeltaToolCall( - index=i, - type="function", - id=tool_id, - function=DeltaFunctionCall(name=tool_name).model_dump( - exclude_none=True - ), - ) - ) - - # Stream back new tool args by diffing against what was sent. - args_diff = self._compute_args_diff(i, tool_args) - if args_diff: - tool_call_deltas.append( - DeltaToolCall( - index=i, - function=DeltaFunctionCall(arguments=args_diff).model_dump( - exclude_none=True - ), - ) - ) - - if content or tool_call_deltas: - return DeltaMessage( - content=content, - tool_calls=tool_call_deltas, - ) - return None - - except Exception: - logger.exception("Error trying to handle streaming tool call.") - return None