mirror of
https://github.com/vllm-project/vllm.git
synced 2026-08-22 13:40:15 +00:00
[Frontend] Add Streaming Parser Engine and new Kimi k2.5/k2.6/k2.7 Parser (#46610)
Signed-off-by: chaunceyjiang <[email protected]> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Chauncey
co-authored by
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent
8cc242335d
commit
2bc20e8aba
@@ -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] = {
|
||||
"<think>": 50,
|
||||
"</think>": 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(("<think>", True))
|
||||
segs.append((scenario.reasoning, False))
|
||||
if scenario.content is not None or scenario.tool_calls is not None:
|
||||
segs.append(("</think>", 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:
|
||||
|
||||
@@ -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(
|
||||
"<think></think>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 </think>."""
|
||||
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, </think> 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: </think> 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", # </think> 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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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>"
|
||||
THINK_END = "</think>"
|
||||
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<id>.+:\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)
|
||||
@@ -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 <think>...</think> 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 </think> 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 = "<think>"
|
||||
self._end_token = "</think>"
|
||||
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 (</think>)
|
||||
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"]
|
||||
|
||||
@@ -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<tool_call_id>[^<]+:\d+)\s*"
|
||||
r"<\|tool_call_argument_begin\|>\s*"
|
||||
r"(?P<function_arguments>(?:(?!<\|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
|
||||
|
||||
Reference in New Issue
Block a user