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)