[Frontend] Add Streaming Parser Engine and new Kimi k2.5/k2.6/k2.7 Parser (#46610)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Chauncey committed
Jun 30, 2026 at 15:53 UTC
2bc20e8abaf7a82ecc068d879ebb925c9317bd40
6 files changed
+397
-570
tests/parser/engine/trace_builder.py
+94
-1
@@ -31,6 +31,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
31
from vllm.parser.engine.registered_adapters import (
32
Gemma4Parser,
33
Glm47MoeParser,
34
+ KimiK2Parser,
35
MinimaxM2Parser,
36
NemotronV3Parser,
37
Qwen3Parser,
@@ -717,6 +718,96 @@ def _build_glm47_moe(scenario: Scenario, validate: bool = True) -> Sample:
718
return sample
719
720
721
+# ── Kimi K2 (native tool-call section, starts in REASONING) ──────────
722
+
723
+_KIMI_K2_VOCAB: dict[str, int] = {
724
+ "<think>": 50,
725
+ "</think>": 51,
726
+ "<|tool_calls_section_begin|>": 60,
727
+ "<|tool_calls_section_end|>": 61,
728
+ "<|tool_call_begin|>": 62,
729
+ "<|tool_call_end|>": 63,
730
+ "<|tool_call_argument_begin|>": 64,
731
+}
732
+
733
+
734
+def _kimi_k2_tool_segments(
735
+ tool_calls: list[ToolCallSpec],
736
+) -> list[tuple[str, bool]]:
737
+ segs: list[tuple[str, bool]] = [("<|tool_calls_section_begin|>", True)]
738
+ for index, tc in enumerate(tool_calls):
739
+ args = json.dumps(tc.arguments, ensure_ascii=False, separators=(",", ":"))
740
+ segs.extend(
741
+ [
742
+ ("<|tool_call_begin|>", True),
743
+ (f"functions.{tc.name}:{index}\n", False),
744
+ ("<|tool_call_argument_begin|>", True),
745
+ (args, False),
746
+ ("<|tool_call_end|>", True),
747
+ ]
748
+ )
749
+ segs.append(("<|tool_calls_section_end|>", True))
750
+ return segs
751
+
752
+
753
+def _kimi_k2_segments(scenario: Scenario) -> list[tuple[str, bool]]:
754
+ segs: list[tuple[str, bool]] = []
755
+ if scenario.reasoning is not None:
756
+ segs.append(("<think>", True))
757
+ segs.append((scenario.reasoning, False))
758
+ if scenario.content is not None or scenario.tool_calls is not None:
759
+ segs.append(("</think>", True))
760
+ if scenario.content is not None:
761
+ segs.append((scenario.content, False))
762
+ if scenario.tool_calls is not None:
763
+ segs.extend(_kimi_k2_tool_segments(scenario.tool_calls))
764
+ return segs
765
+
766
+
767
+def _build_kimi_k2(
768
+ scenario: Scenario,
769
+ validate: bool = True,
770
+ thinking: bool = True,
771
+) -> Sample:
772
+ expected_reasoning = (
773
+ scenario.reasoning.rstrip()
774
+ if (thinking and scenario.reasoning is not None)
775
+ else None
776
+ )
777
+ if thinking and scenario.reasoning is None:
778
+ expected_reasoning = ""
779
+
780
+ sample = _make_sample(
781
+ sample_id=f"kimi_k2-{scenario.id}",
782
+ description=scenario.description,
783
+ vocab=_KIMI_K2_VOCAB,
784
+ segments=_kimi_k2_segments(scenario),
785
+ expected_reasoning=expected_reasoning,
786
+ expected_content=_qwen3_expected_content(scenario),
787
+ expected_tool_calls=_expected_tc(scenario),
788
+ tools=_expected_tools(scenario),
789
+ chat_template_kwargs=None if thinking else {"thinking": False},
790
+ )
791
+ if validate:
792
+ _validate_sample(
793
+ sample,
794
+ KimiK2Parser,
795
+ chat_template_kwargs=sample.chat_template_kwargs,
796
+ )
797
+ return sample
798
+
799
+
800
+_KIMI_K2_SCENARIOS = [
801
+ *SCENARIOS,
802
+ Scenario(
803
+ id="trailing-reasoning-whitespace",
804
+ description="Reasoning trailing whitespace is stripped",
805
+ reasoning="Reasoning with trailing whitespace. \n\t",
806
+ content="Done.",
807
+ ),
808
+]
809
+
810
+
811
# ── Registry and public API ──────────────────────────────────────────
812
813
_BUILDERS: dict[str, Any] = {
@@ -726,6 +817,7 @@ _BUILDERS: dict[str, Any] = {
817
"nemotron_v3": _build_nemotron_v3,
818
"seed_oss": _build_seed_oss,
819
"glm47_moe": _build_glm47_moe,
820
+ "kimi_k2": _build_kimi_k2,
821
}
822
823
@@ -733,7 +825,8 @@ _BUILDERS: dict[str, Any] = {
825
def build_samples(model: str) -> tuple[Sample, ...]:
826
"""Build all scenario samples for a model, self-validated."""
827
builder = _BUILDERS[model]
736
- return tuple(builder(s) for s in SCENARIOS)
828
+ scenarios = _KIMI_K2_SCENARIOS if model == "kimi_k2" else SCENARIOS
829
+ return tuple(builder(s) for s in scenarios)
830
831
832
def build_sample(model: str, scenario: Scenario) -> Sample:
tests/reasoning/test_kimi_k2_reasoning_parser.py
+5
-67
@@ -7,7 +7,6 @@ import pytest
7
8
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
9
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
10
-from vllm.reasoning.identity_reasoning_parser import IdentityReasoningParser
10
from vllm.reasoning.kimi_k2_reasoning_parser import KimiK2ReasoningParser
11
from vllm.tokenizers import get_tokenizer
12
@@ -33,20 +32,6 @@ def kimi_k2_tokenizer():
32
return get_tokenizer(tokenizer_name=REASONING_MODEL_NAME, trust_remote_code=True)
33
34
36
-def test_parser_selection_thinking_enabled(kimi_k2_tokenizer):
37
- parser = KimiK2ReasoningParser(
38
- kimi_k2_tokenizer, chat_template_kwargs={"thinking": True}
39
- )
40
- assert parser._identity_parser is None
41
-
42
-
43
-def test_parser_selection_thinking_disabled(kimi_k2_tokenizer):
44
- parser = KimiK2ReasoningParser(
45
- kimi_k2_tokenizer, chat_template_kwargs={"thinking": False}
46
- )
47
- assert isinstance(parser._identity_parser, IdentityReasoningParser)
48
-
49
-
35
def test_extract_reasoning_with_think_tags(kimi_k2_tokenizer):
36
parser = KimiK2ReasoningParser(kimi_k2_tokenizer)
37
request = ChatCompletionRequest(model="test-model", messages=[], temperature=1.0)
@@ -65,7 +50,7 @@ def test_extract_reasoning_empty_thinking(kimi_k2_tokenizer):
50
reasoning, content = parser.extract_reasoning(
51
"<think></think>final answer", request
52
)
68
- assert reasoning == ""
53
+ assert reasoning is None
54
assert content == "final answer"
55
56
@@ -96,8 +81,8 @@ def test_streaming_reasoning_then_content(kimi_k2_tokenizer):
81
"""Token-by-token streaming: reasoning tokens then content after </think>."""
82
parser = KimiK2ReasoningParser(kimi_k2_tokenizer)
83
99
- think_id = parser._start_token_id
100
- end_think_id = parser._end_token_id
84
+ think_id = parser._parser_engine._start_token_id
85
+ end_think_id = parser._parser_engine._end_token_id
86
# Use a real token ID from the tokenizer for regular content
87
regular_id = kimi_k2_tokenizer.encode("hello", add_special_tokens=False)[0]
88
@@ -154,8 +139,8 @@ def test_streaming_tool_section_ends_reasoning(kimi_k2_tokenizer):
139
"""<|tool_calls_section_begin|> in delta ends reasoning during streaming."""
140
parser = KimiK2ReasoningParser(kimi_k2_tokenizer)
141
157
- think_id = parser._start_token_id
158
- tool_begin_id = parser._tool_section_start_token_id
142
+ think_id = parser._parser_engine._start_token_id
143
+ tool_begin_id = parser._parser_engine._tool_section_start_token_id
144
regular_id = kimi_k2_tokenizer.encode("hello", add_special_tokens=False)[0]
145
146
# Tool section token arrives — should transition from reasoning to content
@@ -169,50 +154,3 @@ def test_streaming_tool_section_ends_reasoning(kimi_k2_tokenizer):
154
)
155
assert isinstance(result, DeltaMessage)
156
assert result.content == "<|tool_calls_section_begin|>"
172
-
173
-
174
-def test_streaming_end_token_id_buffered(mock_kimi_k2_tokenizer):
175
- """When stop sequences buffer text, </think> ID arrives before its text.
176
-
177
- The token ID is present in delta_token_ids but the actual string is not
178
- yet in delta_text (still buffered). The parser must return None to wait
179
- for the next delta, instead of calling find() which returns -1 and
180
- silently corrupting the text split.
181
- """
182
- parser = KimiK2ReasoningParser(mock_kimi_k2_tokenizer)
183
- think_id = parser._start_token_id
184
- end_think_id = parser._end_token_id
185
-
186
- # Simulate: </think> ID arrived but text not yet flushed.
187
- # Two token IDs in delta to bypass the single-special-token guard.
188
- result = parser.extract_reasoning_streaming(
189
- previous_text="some reasoning",
190
- current_text="some reasoning extra",
191
- delta_text="extra", # </think> text not yet flushed
192
- previous_token_ids=[think_id],
193
- current_token_ids=[think_id, end_think_id, 999],
194
- delta_token_ids=[end_think_id, 999],
195
- )
196
- assert result is None
197
-
198
-
199
-def test_streaming_tool_section_id_buffered(mock_kimi_k2_tokenizer):
200
- """When stop sequences buffer text, tool section start ID arrives before its text.
201
-
202
- Same buffering scenario as above but for <|tool_calls_section_begin|>.
203
- Without the guard, find() returns -1 and delta_text[:tool_index] silently
204
- drops the last character of reasoning.
205
- """
206
- parser = KimiK2ReasoningParser(mock_kimi_k2_tokenizer)
207
- think_id = parser._start_token_id
208
- tool_begin_id = parser._tool_section_start_token_id
209
-
210
- result = parser.extract_reasoning_streaming(
211
- previous_text="some reasoning",
212
- current_text="some reasoning extra",
213
- delta_text="extra", # tool section text not yet flushed
214
- previous_token_ids=[think_id],
215
- current_token_ids=[think_id, tool_begin_id, 999],
216
- delta_token_ids=[tool_begin_id, 999],
217
- )
218
- assert result is None
vllm/parser/engine/registered_adapters.py
+6
@@ -10,6 +10,7 @@ names so that :class:`ReasoningParserManager` and
10
from vllm.parser.engine.adapters import make_adapters
11
from vllm.parser.gemma4 import Gemma4Parser
12
from vllm.parser.glm47_moe import Glm47MoeParser
13
+from vllm.parser.kimi_k2 import KimiK2Parser
14
from vllm.parser.minimax_m2 import MinimaxM2Parser
15
from vllm.parser.nemotron_v3 import NemotronV3Parser
16
from vllm.parser.qwen3 import Qwen3Parser
@@ -44,3 +45,8 @@ from vllm.parser.seed_oss import SeedOssParser
45
Glm47MoeParserReasoningAdapter,
46
Glm47MoeParserToolAdapter,
47
) = make_adapters(Glm47MoeParser)
48
+
49
+(
50
+ KimiK2ParserReasoningAdapter,
51
+ KimiK2ParserToolAdapter,
52
+) = make_adapters(KimiK2Parser)
vllm/parser/kimi_k2.py
new
+285
@@ -0,0 +1,285 @@
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.openai.engine.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
+ terminals={
85
+ **reasoning_terminals,
86
+ "TOOL_SECTION_START": TOOL_SECTION_START,
87
+ "TOOL_SECTION_END": TOOL_SECTION_END,
88
+ "TOOL_START": TOOL_CALL_START,
89
+ "TOOL_END": TOOL_CALL_END,
90
+ "ARG_START": TOOL_ARG_START,
91
+ },
92
+ token_id_terminals={
93
+ **reasoning_terminals,
94
+ "TOOL_SECTION_START": TOOL_SECTION_START,
95
+ "TOOL_SECTION_END": TOOL_SECTION_END,
96
+ "TOOL_START": TOOL_CALL_START,
97
+ "TOOL_END": TOOL_CALL_END,
98
+ "ARG_START": TOOL_ARG_START,
99
+ },
100
+ transitions={
101
+ **reasoning_transitions,
102
+ (ParserState.REASONING, "TOOL_SECTION_START"): Transition(
103
+ ParserState.TOOL_PREAMBLE,
104
+ (EventType.REASONING_END,),
105
+ ),
106
+ (ParserState.CONTENT, "TOOL_SECTION_START"): Transition(
107
+ ParserState.TOOL_PREAMBLE,
108
+ (),
109
+ ),
110
+ (ParserState.TOOL_PREAMBLE, "TOOL_START"): Transition(
111
+ ParserState.TOOL_NAME,
112
+ (EventType.TOOL_CALL_START,),
113
+ ),
114
+ (ParserState.TOOL_NAME, "ARG_START"): Transition(
115
+ ParserState.TOOL_ARGS,
116
+ (),
117
+ ),
118
+ (ParserState.TOOL_ARGS, "TOOL_END"): Transition(
119
+ ParserState.TOOL_BETWEEN,
120
+ (EventType.TOOL_CALL_END,),
121
+ ),
122
+ (ParserState.TOOL_ARGS, "TOOL_SECTION_END"): Transition(
123
+ ParserState.TOOL_PREAMBLE,
124
+ (EventType.TOOL_CALL_END,),
125
+ ),
126
+ (ParserState.TOOL_BETWEEN, "TOOL_START"): Transition(
127
+ ParserState.TOOL_NAME,
128
+ (EventType.TOOL_CALL_START,),
129
+ ),
130
+ # Keep the parser in a tool state after the section closes so
131
+ # trailing model text after native tool calls is suppressed.
132
+ (ParserState.TOOL_PREAMBLE, "TOOL_SECTION_END"): Transition(
133
+ ParserState.TOOL_PREAMBLE,
134
+ (),
135
+ ),
136
+ (ParserState.TOOL_BETWEEN, "TOOL_SECTION_END"): Transition(
137
+ ParserState.TOOL_PREAMBLE,
138
+ (),
139
+ ),
140
+ },
141
+ stream_arg_deltas=True,
142
+ tool_args_json=True,
143
+ strip_trailing_reasoning_whitespace=True,
144
+ drop_whitespace_only_content_before_tools=True,
145
+ strip_content_whitespace_with_tools=False,
146
+ validate_tool_names=False,
147
+ )
148
+
149
+
150
+class KimiK2Parser(ParserEngine):
151
+ """Kimi K2 parser backed by the declarative parser engine."""
152
+
153
+ def __init__(
154
+ self,
155
+ tokenizer: TokenizerLike,
156
+ tools: list[Tool] | None = None,
157
+ **kwargs,
158
+ ) -> None:
159
+ chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
160
+ thinking = chat_kwargs.get("thinking", None)
161
+ enable_thinking = chat_kwargs.get("enable_thinking", None)
162
+ self.thinking_enabled = (
163
+ True
164
+ if thinking is None and enable_thinking is None
165
+ else bool(thinking) or bool(enable_thinking)
166
+ )
167
+ kwargs.setdefault(
168
+ "parser_engine_config",
169
+ kimi_k2_config(thinking=self.thinking_enabled),
170
+ )
171
+ super().__init__(tokenizer, tools, **kwargs)
172
+
173
+ vocab = self.vocab
174
+ self._start_token_id = vocab.get(THINK_START)
175
+ self._end_token_id = vocab.get(THINK_END)
176
+ self._tool_section_start_token_id = vocab.get(TOOL_SECTION_START)
177
+
178
+ @staticmethod
179
+ def _extract_tool_id_and_name(header: str | None) -> tuple[str | None, str | None]:
180
+ if header is None:
181
+ return None, None
182
+ match = _TOOL_ID_RE.match(header.strip())
183
+ if not match:
184
+ return None, None
185
+
186
+ tool_id = match.group("id").strip()
187
+ tool_name = tool_id.split(":")[0].removeprefix("functions.")
188
+ return tool_id, tool_name
189
+
190
+ def _emit_name_delta(
191
+ self,
192
+ idx: int,
193
+ deltas: list[DeltaToolCall],
194
+ name: str | None,
195
+ ) -> None:
196
+ tool_id, tool_name = self._extract_tool_id_and_name(name)
197
+ if not tool_name:
198
+ if 0 <= idx < len(self._tool_slots):
199
+ self._tool_slots[idx].name = ""
200
+ return
201
+
202
+ slot = self._tool_slots[idx]
203
+ slot.id = tool_id or ""
204
+ super()._emit_name_delta(idx, deltas, tool_name)
205
+
206
+ def _handle_tool_end(self, event, deltas) -> None:
207
+ idx = event.tool_index
208
+ if 0 <= idx < len(self._tool_slots) and not self._tool_slots[idx].name_sent:
209
+ tool_id, tool_name = self._extract_tool_id_and_name(
210
+ self._tool_slots[idx].name
211
+ )
212
+ if tool_name:
213
+ self._tool_slots[idx].id = tool_id or ""
214
+ self._tool_slots[idx].name = tool_name
215
+ super()._handle_tool_end(event, deltas)
216
+
217
+ def _handle_arg_chunk(self, event, deltas) -> None:
218
+ idx = event.tool_index
219
+ name_sent_before = (
220
+ 0 <= idx < len(self._tool_slots) and self._tool_slots[idx].name_sent
221
+ )
222
+ super()._handle_arg_chunk(event, deltas)
223
+ if (
224
+ event.value
225
+ and not name_sent_before
226
+ and 0 <= idx < len(self._tool_slots)
227
+ and self._tool_slots[idx].name_sent
228
+ ):
229
+ deltas.append(
230
+ DeltaToolCall(
231
+ index=idx,
232
+ function=DeltaFunctionCall(arguments=event.value),
233
+ )
234
+ )
235
+
236
+ def _extract_args_json(self, raw_args: str, func_name: str) -> str:
237
+ return raw_args.strip() or "{}"
238
+
239
+ def is_reasoning_end(self, input_ids: list[int]) -> bool:
240
+ if not self.thinking_enabled:
241
+ return True
242
+
243
+ start_id = self._start_token_id
244
+ end_id = self._end_token_id
245
+ tool_section_id = self._tool_section_start_token_id
246
+
247
+ for i in range(len(input_ids) - 1, -1, -1):
248
+ token_id = input_ids[i]
249
+ if start_id is not None and token_id == start_id:
250
+ return False
251
+ if end_id is not None and token_id == end_id:
252
+ return True
253
+ if tool_section_id is not None and token_id == tool_section_id:
254
+ return True
255
+ return False
256
+
257
+ def extract_content_ids(self, input_ids: list[int]) -> list[int]:
258
+ if not self.thinking_enabled:
259
+ return input_ids
260
+
261
+ end_id = self._end_token_id
262
+ if end_id is not None and end_id in input_ids:
263
+ end_idx = len(input_ids) - 1 - input_ids[::-1].index(end_id)
264
+ return input_ids[end_idx + 1 :]
265
+
266
+ tool_section_id = self._tool_section_start_token_id
267
+ if tool_section_id is not None and tool_section_id in input_ids:
268
+ section_idx = len(input_ids) - 1 - input_ids[::-1].index(tool_section_id)
269
+ return input_ids[section_idx:]
270
+
271
+ return []
272
+
273
+ def extract_reasoning(
274
+ self,
275
+ model_output: str,
276
+ request: ChatCompletionRequest | ResponsesRequest,
277
+ ) -> tuple[str | None, str | None]:
278
+ if not self.thinking_enabled:
279
+ return None, model_output
280
+ return super().extract_reasoning(model_output, request)
281
+
282
+ def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
283
+ if not self.thinking_enabled:
284
+ return 0
285
+ return super().count_reasoning_tokens(token_ids)
vllm/reasoning/kimi_k2_reasoning_parser.py
+3
-240
@@ -1,245 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Iterable, Sequence
5
-from typing import TYPE_CHECKING
4
+from vllm.parser.engine.registered_adapters import KimiK2ParserReasoningAdapter
5
7
-from transformers import PreTrainedTokenizerBase
6
+KimiK2ReasoningParser = KimiK2ParserReasoningAdapter
7
9
-from vllm.entrypoints.openai.engine.protocol import DeltaMessage
10
-from vllm.reasoning.abs_reasoning_parsers import ReasoningParser
11
-from vllm.reasoning.identity_reasoning_parser import IdentityReasoningParser
12
-
13
-if TYPE_CHECKING:
14
- from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
15
- from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
16
-
17
-
18
-class KimiK2ReasoningParser(ReasoningParser):
19
- """
20
- Reasoning parser for Kimi K2 model.
21
-
22
- The Kimi K2 model uses <think>...</think> tokens to denote reasoning text,
23
- and may implicitly end reasoning by starting a tool call section using
24
- <|tool_calls_section_begin|>.
25
- Thinking may also begin without a </think> token.
26
-
27
- Kimi's thinking mode can be disabled via chat_template_kwargs.
28
- """
29
-
30
- def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs):
31
- super().__init__(tokenizer, *args, **kwargs)
32
-
33
- if not self.model_tokenizer:
34
- raise ValueError(
35
- "The model tokenizer must be passed to the ReasoningParser "
36
- "constructor during construction."
37
- )
38
-
39
- # Check if thinking is disabled via chat_template_kwargs
40
- chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
41
- thinking = bool(chat_kwargs.get("thinking", True))
42
-
43
- # If thinking is not enabled, use identity parser to fall through
44
- self._identity_parser: IdentityReasoningParser | None
45
- if not thinking:
46
- self._identity_parser = IdentityReasoningParser(tokenizer, *args, **kwargs)
47
- else:
48
- self._identity_parser = None
49
-
50
- # Token definitions
51
- self._start_token = "<think>"
52
- self._end_token = "</think>"
53
- self._tool_section_start_token = "<|tool_calls_section_begin|>"
54
-
55
- # Get token IDs
56
- self._start_token_id = self.vocab.get(self._start_token)
57
- self._end_token_id = self.vocab.get(self._end_token)
58
- self._tool_section_start_token_id = self.vocab.get(
59
- self._tool_section_start_token
60
- )
61
-
62
- if self._start_token_id is None or self._end_token_id is None:
63
- raise RuntimeError(
64
- "KimiK2ReasoningParser could not locate think start/end "
65
- "tokens in the tokenizer!"
66
- )
67
-
68
- @property
69
- def reasoning_start_str(self) -> str | None:
70
- return self._start_token
71
-
72
- @property
73
- def reasoning_end_str(self) -> str | None:
74
- return self._end_token
75
-
76
- def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
77
- """
78
- Check if the reasoning content ends in the input_ids.
79
-
80
- Reasoning ends when we see either:
81
- 1. The end token (</think>)
82
- 2. The tool section start token (<|tool_calls_section_begin|>)
83
- """
84
- if self._identity_parser is not None:
85
- return self._identity_parser.is_reasoning_end(input_ids)
86
-
87
- start_token_id = self._start_token_id
88
- end_token_id = self._end_token_id
89
- tool_section_start_token_id = self._tool_section_start_token_id
90
-
91
- for i in range(len(input_ids) - 1, -1, -1):
92
- if input_ids[i] == start_token_id:
93
- return False
94
- if input_ids[i] == end_token_id:
95
- return True
96
- # Implicit reasoning end via tool call section
97
- if (
98
- tool_section_start_token_id is not None
99
- and input_ids[i] == tool_section_start_token_id
100
- ):
101
- return True
102
- return False
103
-
104
- def is_reasoning_end_streaming(
105
- self, input_ids: Sequence[int], delta_ids: Iterable[int]
106
- ) -> bool:
107
- """
108
- Check if the reasoning content ends in the input_ids on a decode step.
109
- """
110
- if self._identity_parser is not None:
111
- return self._identity_parser.is_reasoning_end_streaming(
112
- input_ids, delta_ids
113
- )
114
-
115
- # Materialize iterable for membership checks
116
- delta_ids_set = set(delta_ids)
117
-
118
- # Check for explicit end token or implicit tool section start in delta
119
- if self._end_token_id in delta_ids_set:
120
- return True
121
- return (
122
- self._tool_section_start_token_id is not None
123
- and self._tool_section_start_token_id in delta_ids_set
124
- )
125
-
126
- def extract_content_ids(self, input_ids: list[int]) -> list[int]:
127
- """
128
- Extract content token ids from the input_ids.
129
- """
130
- if self._identity_parser is not None:
131
- return self._identity_parser.extract_content_ids(input_ids)
132
-
133
- if self._end_token_id in input_ids:
134
- end_token_index = (
135
- len(input_ids) - 1 - input_ids[::-1].index(self._end_token_id)
136
- )
137
-
138
- if end_token_index != -1:
139
- return input_ids[end_token_index + 1 :]
140
-
141
- if (
142
- self._tool_section_start_token_id is not None
143
- and self._tool_section_start_token_id in input_ids
144
- ):
145
- tool_section_index = (
146
- len(input_ids)
147
- - 1
148
- - input_ids[::-1].index(self._tool_section_start_token_id)
149
- )
150
-
151
- if tool_section_index != -1:
152
- return input_ids[tool_section_index:]
153
-
154
- # still reasoning (no content)
155
- return []
156
-
157
- def extract_reasoning(
158
- self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest"
159
- ) -> tuple[str | None, str | None]:
160
- """
161
- Extract reasoning content from the model output.
162
- """
163
- if self._identity_parser is not None:
164
- return self._identity_parser.extract_reasoning(model_output, request)
165
-
166
- # thinking does not require a think start token but consume it if present
167
- start_token_index = model_output.find(self._start_token)
168
- start_token_index = 0 if start_token_index != 0 else len(self._start_token)
169
- end_token_index = model_output.find(self._end_token)
170
-
171
- if end_token_index != -1:
172
- return (
173
- model_output[start_token_index:end_token_index],
174
- model_output[end_token_index + len(self._end_token) :] or None,
175
- )
176
-
177
- tool_section_index = model_output.find(self._tool_section_start_token)
178
- if tool_section_index != -1:
179
- return (
180
- model_output[start_token_index:tool_section_index],
181
- model_output[tool_section_index:] or None,
182
- )
183
-
184
- # still reasoning (no content)
185
- return (
186
- model_output[start_token_index:],
187
- None,
188
- )
189
-
190
- def extract_reasoning_streaming(
191
- self,
192
- previous_text: str,
193
- current_text: str,
194
- delta_text: str,
195
- previous_token_ids: Sequence[int],
196
- current_token_ids: Sequence[int],
197
- delta_token_ids: Sequence[int],
198
- ) -> DeltaMessage | None:
199
- """
200
- Extract reasoning content from a delta message during streaming.
201
- """
202
- if self._identity_parser is not None:
203
- return self._identity_parser.extract_reasoning_streaming(
204
- previous_text,
205
- current_text,
206
- delta_text,
207
- previous_token_ids,
208
- current_token_ids,
209
- delta_token_ids,
210
- )
211
-
212
- # If reasoning has already ended in previous tokens, this is content
213
- if self.is_reasoning_end(previous_token_ids):
214
- return DeltaMessage(content=delta_text)
215
-
216
- # Skip single special tokens
217
- if len(delta_token_ids) == 1 and delta_token_ids[0] in [
218
- self._start_token_id,
219
- self._end_token_id,
220
- ]:
221
- return None
222
-
223
- if self._end_token_id in delta_token_ids:
224
- if self._end_token not in delta_text:
225
- # Token ID arrived before text was flushed (stop-sequence buffering).
226
- # Wait for the next delta when the text becomes visible.
227
- return None
228
- end_index = delta_text.find(self._end_token)
229
- reasoning = delta_text[:end_index]
230
- content = delta_text[end_index + len(self._end_token) :]
231
- return DeltaMessage(
232
- reasoning=reasoning, content=content if content else None
233
- )
234
-
235
- if self._tool_section_start_token_id in delta_token_ids:
236
- if self._tool_section_start_token not in delta_text:
237
- # Token ID arrived before text was flushed (stop-sequence buffering).
238
- return None
239
- tool_index = delta_text.find(self._tool_section_start_token)
240
- reasoning = delta_text[:tool_index]
241
- content = delta_text[tool_index:]
242
- return DeltaMessage(reasoning=reasoning, content=content)
243
-
244
- # still reasoning (no end token)
245
- return DeltaMessage(reasoning=delta_text)
8
+__all__ = ["KimiK2ReasoningParser"]
vllm/tool_parsers/kimi_k2_tool_parser.py
+4
-262
@@ -1,278 +1,20 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Sequence
5
-
6
-import regex as re
7
-
4
from vllm.entrypoints.openai.chat_completion.protocol import (
5
ChatCompletionRequest,
6
)
11
-from vllm.entrypoints.openai.engine.protocol import (
12
- DeltaFunctionCall,
13
- DeltaMessage,
14
- DeltaToolCall,
15
- ExtractedToolCallInformation,
16
- FunctionCall,
17
- ToolCall,
18
-)
7
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
20
-from vllm.logger import init_logger
21
-from vllm.tokenizers import TokenizerLike
22
-from vllm.tool_parsers.abstract_tool_parser import (
23
- Tool,
24
- ToolParser,
25
-)
26
-from vllm.tool_parsers.utils import partial_tag_overlap
27
-
28
-logger = init_logger(__name__)
8
+from vllm.parser.engine.registered_adapters import KimiK2ParserToolAdapter
9
10
31
-class KimiK2ToolParser(ToolParser):
11
+class KimiK2ToolParser(KimiK2ParserToolAdapter): # type: ignore[valid-type, misc]
12
structural_tag_model = "kimi"
13
34
- def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
35
- super().__init__(tokenizer, tools)
36
-
37
- # Streaming state
38
- self._sent_content_idx: int = 0
39
- self.prev_tool_call_arr: list[dict] = []
40
- self.streamed_args_for_tool: list[str] = []
41
-
42
- # Section marker
43
- self.tool_calls_start_token: str = "<|tool_calls_section_begin|>"
44
-
45
- # Individual tool call markers
46
- self.tool_call_start_token: str = "<|tool_call_begin|>"
47
- self.tool_call_end_token: str = "<|tool_call_end|>"
48
- self.tool_call_arg_token: str = "<|tool_call_argument_begin|>"
49
-
50
- # Regex for non-streaming extraction
51
- self.tool_call_regex = re.compile(
52
- r"<\|tool_call_begin\|>\s*(?P<tool_call_id>[^<]+:\d+)\s*"
53
- r"<\|tool_call_argument_begin\|>\s*"
54
- r"(?P<function_arguments>(?:(?!<\|tool_call_begin\|>).)*?)\s*"
55
- r"<\|tool_call_end\|>",
56
- re.DOTALL,
57
- )
58
-
59
- if not self.model_tokenizer:
60
- raise ValueError(
61
- "The model tokenizer must be passed to the ToolParser "
62
- "constructor during construction."
63
- )
64
-
14
def adjust_request(
66
- self, request: ChatCompletionRequest | ResponsesRequest
15
+ self,
16
+ request: ChatCompletionRequest | ResponsesRequest,
17
) -> ChatCompletionRequest | ResponsesRequest:
68
- request = super().adjust_request(request)
18
if request.tools and request.tool_choice != "none":
70
- # Ensure special-token markers appear as literal text in
71
- # current_text so we can do pure text-based parsing.
19
request.skip_special_tokens = False
20
return request
74
-
75
- def extract_tool_calls(
76
- self,
77
- model_output: str,
78
- request: ChatCompletionRequest,
79
- ) -> ExtractedToolCallInformation:
80
- # sanity check; avoid unnecessary processing
81
- if self.tool_calls_start_token not in model_output:
82
- return ExtractedToolCallInformation(
83
- tools_called=False, tool_calls=[], content=model_output
84
- )
85
-
86
- else:
87
- try:
88
- # there are two possible captures - between tags, or between a
89
- # tag and end-of-string so the result of
90
- # findall is an array of tuples where one is a function call and
91
- # the other is None
92
- function_call_tuples = self.tool_call_regex.findall(model_output)
93
-
94
- logger.debug("function_call_tuples: %s", function_call_tuples)
95
-
96
- tool_calls = []
97
- for match in function_call_tuples:
98
- function_id, function_args = match
99
- # function_id: functions.get_weather:0 or get_weather:0
100
- function_name = function_id.split(":")[0].split(".")[-1]
101
- tool_calls.append(
102
- ToolCall(
103
- id=function_id,
104
- type="function",
105
- function=FunctionCall(
106
- name=function_name, arguments=function_args
107
- ),
108
- )
109
- )
110
-
111
- content = model_output[: model_output.find(self.tool_calls_start_token)]
112
- return ExtractedToolCallInformation(
113
- tools_called=True,
114
- tool_calls=tool_calls,
115
- content=content if content else None,
116
- )
117
-
118
- except Exception:
119
- logger.exception("Error in extracting tool call from response.")
120
- return ExtractedToolCallInformation(
121
- tools_called=False, tool_calls=[], content=model_output
122
- )
123
-
124
- def _extract_content(self, current_text: str) -> str | None:
125
- """Return unsent content before the tool-calls section, or None.
126
-
127
- Holds back any trailing suffix that partially matches
128
- ``<|tool_calls_section_begin|>`` to avoid leaking marker bytes.
129
- """
130
- if self.tool_calls_start_token not in current_text:
131
- overlap = partial_tag_overlap(current_text, self.tool_calls_start_token)
132
- sendable_idx = len(current_text) - overlap
133
- else:
134
- sendable_idx = current_text.index(self.tool_calls_start_token)
135
-
136
- if sendable_idx > self._sent_content_idx:
137
- content = current_text[self._sent_content_idx : sendable_idx]
138
- self._sent_content_idx = sendable_idx
139
- return content
140
- return None
141
-
142
- def _extract_tool_calls(self, current_text: str) -> list[str]:
143
- """Extract raw bodies from ``<|tool_call_begin|>…<|tool_call_end|>`` blocks."""
144
- if self.tool_calls_start_token not in current_text:
145
- return []
146
-
147
- results: list[str] = []
148
- pos = current_text.index(self.tool_calls_start_token)
149
- while True:
150
- start = current_text.find(self.tool_call_start_token, pos)
151
- if start == -1:
152
- break
153
- tc_start = start + len(self.tool_call_start_token)
154
- end = current_text.find(self.tool_call_end_token, tc_start)
155
-
156
- if end != -1:
157
- tool_call = current_text[tc_start:end]
158
- pos = end + len(self.tool_call_end_token)
159
- else:
160
- tool_call = current_text[tc_start:]
161
- overlap = partial_tag_overlap(tool_call, self.tool_call_end_token)
162
- if overlap:
163
- tool_call = tool_call[:-overlap]
164
-
165
- results.append(tool_call)
166
-
167
- if end == -1:
168
- break
169
- return results
170
-
171
- @staticmethod
172
- def _extract_tool_id_and_name(
173
- header: str | None,
174
- ) -> tuple[str | None, str | None]:
175
- """Parse ``(tool_id, tool_name)`` from a header
176
- like ``"functions.get_weather:0"``."""
177
- if header is None:
178
- return None, None
179
- match = re.match(r"(.+:\d+)", header)
180
- if not match:
181
- return None, None
182
-
183
- tool_id = match.group(1).strip()
184
- tool_name = tool_id.split(":")[0].split(".")[-1]
185
- return tool_id, tool_name
186
-
187
- def _split_tool_call(self, tool_call: str) -> tuple[str | None, str | None]:
188
- """Split a tool-call body into ``(header, arguments)`` at the argument marker.
189
-
190
- Example::
191
- 'get_weather:0 <|tool_call_argument_begin|>{"c'
192
- -> ("get_weather:0", '{"c')
193
- """
194
- arg_pos = tool_call.find(self.tool_call_arg_token)
195
- if arg_pos == -1:
196
- return None, None
197
- header = tool_call[:arg_pos].strip()
198
- tool_args = tool_call[arg_pos + len(self.tool_call_arg_token) :]
199
- return header, tool_args
200
-
201
- def _compute_args_diff(self, index: int, tool_args: str | None) -> str | None:
202
- """Return new argument text not yet sent for tool `index`, or None."""
203
- if tool_args is None:
204
- return None
205
- prev = self.streamed_args_for_tool[index]
206
- if len(tool_args) <= len(prev):
207
- return None
208
- diff = tool_args[len(prev) :]
209
- self.streamed_args_for_tool[index] = tool_args
210
- self.prev_tool_call_arr[index]["arguments"] = tool_args
211
- return diff
212
-
213
- def extract_tool_calls_streaming(
214
- self,
215
- previous_text: str,
216
- current_text: str,
217
- delta_text: str,
218
- previous_token_ids: Sequence[int],
219
- current_token_ids: Sequence[int],
220
- delta_token_ids: Sequence[int],
221
- request: ChatCompletionRequest,
222
- ) -> DeltaMessage | None:
223
- try:
224
- # Extract any content before tool calls.
225
- content = self._extract_content(current_text)
226
- tool_calls = self._extract_tool_calls(current_text)
227
- tool_call_deltas: list[DeltaToolCall] = []
228
-
229
- for i, tool_call in enumerate(tool_calls):
230
- # First time seeing tool call at index i.
231
- if i >= len(self.prev_tool_call_arr):
232
- # Initialize streaming state.
233
- self.prev_tool_call_arr.append({})
234
- self.streamed_args_for_tool.append("")
235
-
236
- header, tool_args = self._split_tool_call(tool_call)
237
-
238
- # Stream back tool name.
239
- if "name" not in self.prev_tool_call_arr[i]:
240
- tool_id, tool_name = self._extract_tool_id_and_name(header)
241
- if not tool_name:
242
- # Can't skip to tool i+1 if i isn't ready
243
- break
244
- self.prev_tool_call_arr[i]["name"] = tool_name
245
- self.prev_tool_call_arr[i]["id"] = tool_id
246
- tool_call_deltas.append(
247
- DeltaToolCall(
248
- index=i,
249
- type="function",
250
- id=tool_id,
251
- function=DeltaFunctionCall(name=tool_name).model_dump(
252
- exclude_none=True
253
- ),
254
- )
255
- )
256
-
257
- # Stream back new tool args by diffing against what was sent.
258
- args_diff = self._compute_args_diff(i, tool_args)
259
- if args_diff:
260
- tool_call_deltas.append(
261
- DeltaToolCall(
262
- index=i,
263
- function=DeltaFunctionCall(arguments=args_diff).model_dump(
264
- exclude_none=True
265
- ),
266
- )
267
- )
268
-
269
- if content or tool_call_deltas:
270
- return DeltaMessage(
271
- content=content,
272
- tool_calls=tool_call_deltas,
273
- )
274
- return None
275
-
276
- except Exception:
277
- logger.exception("Error trying to handle streaming tool call.")
278
- return None