| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tool-call parser for the Kimi K3 (XTML) chat format. |
| 4 | |
| 5 | This turns the generated XTML ``response`` and ``tools`` channels back into |
| 6 | OpenAI-compatible ``content`` and ``tool_calls``. |
| 7 | |
| 8 | K3 assistant tool calls live in a nested ``tools`` channel:: |
| 9 | |
| 10 | <|open|>tools<|sep|> |
| 11 | <|open|>call tool="python" index="1"<|sep|> |
| 12 | <|open|>argument key="code" type="string"<|sep|>print(1)<|close|>argument<|sep|> |
| 13 | <|open|>argument key="opts" type="object"<|sep|>{"a":1}<|close|>argument<|sep|> |
| 14 | <|close|>call<|sep|> |
| 15 | <|close|>tools<|sep|> |
| 16 | |
| 17 | where ``<|open|>``, ``<|close|>``, ``<|sep|>`` are dedicated special tokens. The plain |
| 18 | reply lives in a sibling ``response`` channel which is unwrapped into content. |
| 19 | ``response`` always precedes ``tools`` (see the protocol channel grammar). |
| 20 | |
| 21 | Argument decoding mirrors the template's type tagging (inverse encoding): |
| 22 | * type="string" -> value is the RAW text, no unescaping (template emits it raw) |
| 23 | * other types -> value is JSON-decoded (number/boolean/null/object/array) |
| 24 | |
| 25 | Attribute values (``tool=``, ``index=``, ``key=``, ``type=``) ARE escaped on the |
| 26 | encode side (``&`` -> ``&``, ``"`` -> ``"``), so ``_attrs`` reverses |
| 27 | them on decode (``"`` before ``&`` -- reverse of encode order). |
| 28 | |
| 29 | Known limitation: because string argument and response bodies are emitted raw, |
| 30 | a value that literally contains ``<|close|>argument<|sep|>`` or |
| 31 | ``<|close|>response<|sep|>`` is indistinguishable from a real closing marker. |
| 32 | """ |
| 33 | |
| 34 | import json |
| 35 | from collections.abc import Sequence |
| 36 | |
| 37 | import regex as re |
| 38 | from openai.types.responses import ToolChoiceFunction |
| 39 | |
| 40 | from vllm.entrypoints.generate.base.protocol import ( |
| 41 | DeltaFunctionCall, |
| 42 | DeltaMessage, |
| 43 | DeltaToolCall, |
| 44 | ExtractedToolCallInformation, |
| 45 | FunctionCall, |
| 46 | ToolCall, |
| 47 | ) |
| 48 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 49 | ChatCompletionNamedToolChoiceParam, |
| 50 | ChatCompletionRequest, |
| 51 | ) |
| 52 | from vllm.entrypoints.openai.responses.protocol import ResponsesRequest |
| 53 | from vllm.exceptions import VLLMValidationError |
| 54 | from vllm.logger import init_logger |
| 55 | from vllm.tokenizers import TokenizerLike |
| 56 | from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser |
| 57 | |
| 58 | logger = init_logger(__name__) |
| 59 | |
| 60 | _O, _C, _S = r"<\|open\|>", r"<\|close\|>", r"<\|sep\|>" |
| 61 | _TEXT_UNTIL_SEP = r"(?:(?!" + _S + r").)*?" |
| 62 | |
| 63 | |
| 64 | def _partial_tag_overlap(text: str, tag: str) -> int: |
| 65 | max_len = min(len(text), len(tag) - 1) |
| 66 | for n in range(max_len, 0, -1): |
| 67 | if text.endswith(tag[:n]): |
| 68 | return n |
| 69 | return 0 |
| 70 | |
| 71 | |
| 72 | class KimiK3ToolParser(ToolParser): |
| 73 | supports_required_and_named = False |
| 74 | # Enables the vLLM-side XTML structural tag builder |
| 75 | # (``get_kimi_k3_structural_tag`` in ``structural_tag_registry``). With |
| 76 | # ``VLLM_ENFORCE_STRICT_TOOL_CALLING`` on (default), ``_apply_structural_tag`` |
| 77 | # constrains generation to K3's ``<|open|>tools<|sep|>`` channel for |
| 78 | # ``required`` (and ``auto`` when a tool sets ``strict``) instead of the |
| 79 | # generic JSON guided decoding, which conflicts with the XTML format. |
| 80 | structural_tag_model = "kimi_k3" |
| 81 | |
| 82 | def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): |
| 83 | super().__init__(tokenizer, tools) |
| 84 | |
| 85 | self.tools_open = "<|open|>tools<|sep|>" |
| 86 | self.tools_close = "<|close|>tools<|sep|>" |
| 87 | self.response_open = "<|open|>response<|sep|>" |
| 88 | self.response_close = "<|close|>response<|sep|>" |
| 89 | |
| 90 | # Regexes operate on detokenized text. The XTML markers reach us as the |
| 91 | # literal strings <|open|>/<|close|>/<|sep|>. adjust_request keeps them |
| 92 | # from being stripped and suppresses spacing between adjacent token |
| 93 | # pieces. As DEFENSE-IN-DEPTH we still tolerate optional whitespace |
| 94 | # WITHIN a marker (e.g. "<|open|> tools <|sep|>") in case some serving |
| 95 | # path leaves vLLM's added-token spacing on -- this `\s*` is a no-op on |
| 96 | # clean input, so the normal byte-exact path is unaffected. We do NOT |
| 97 | # strip body whitespace, so clean content stays byte-exact. |
| 98 | # All bodies use non-greedy (.*?) so each block stops at its own first |
| 99 | # closing marker -- see the module-level KNOWN LIMITATION about |
| 100 | # literal-marker values. |
| 101 | self._tools_open_re = re.compile(_O + r"\s*tools\s*" + _S) |
| 102 | self._tools_close_re = re.compile(_C + r"\s*tools\s*" + _S) |
| 103 | self._response_open_re = re.compile(_O + r"\s*response\s*" + _S) |
| 104 | self._response_close_re = re.compile(_C + r"\s*response\s*" + _S) |
| 105 | self._message_close_re = re.compile(_C + r"\s*message\s*" + _S) |
| 106 | self._call_re = re.compile( |
| 107 | _O |
| 108 | + r"\s*call\s+(?P<attrs>" |
| 109 | + _TEXT_UNTIL_SEP |
| 110 | + r")" |
| 111 | + _S |
| 112 | + r"(?P<body>.*?)" |
| 113 | + _C |
| 114 | + r"\s*call\s*" |
| 115 | + _S, |
| 116 | re.DOTALL, |
| 117 | ) |
| 118 | self._arg_re = re.compile( |
| 119 | _O |
| 120 | + r"\s*argument\s+(?P<attrs>" |
| 121 | + _TEXT_UNTIL_SEP |
| 122 | + r")" |
| 123 | + _S |
| 124 | + r"(?P<val>.*?)" |
| 125 | + _C |
| 126 | + r"\s*argument\s*" |
| 127 | + _S, |
| 128 | re.DOTALL, |
| 129 | ) |
| 130 | # attr segment: key="value" (value already escaped on the encode side) |
| 131 | self._attr_re = re.compile(r'(?P<k>\w+)="(?P<v>[^"]*)"') |
| 132 | self._response_re = re.compile( |
| 133 | _O + r"\s*response\s*" + _S + r"(?P<c>.*?)" + _C + r"\s*response\s*" + _S, |
| 134 | re.DOTALL, |
| 135 | ) |
| 136 | |
| 137 | # streaming state |
| 138 | self._sent_content_idx = 0 |
| 139 | self._sent_tool_call_count = 0 |
| 140 | |
| 141 | if not self.model_tokenizer: |
| 142 | raise ValueError( |
| 143 | "The model tokenizer must be passed to the ToolParser " |
| 144 | "constructor during construction." |
| 145 | ) |
| 146 | |
| 147 | def adjust_request( |
| 148 | self, request: ChatCompletionRequest | ResponsesRequest |
| 149 | ) -> ChatCompletionRequest | ResponsesRequest: |
| 150 | named = isinstance( |
| 151 | request.tool_choice, |
| 152 | (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction), |
| 153 | ) |
| 154 | structured_outputs = getattr(request, "structured_outputs", None) |
| 155 | has_structural_tag = ( |
| 156 | structured_outputs is not None |
| 157 | and structured_outputs.structural_tag is not None |
| 158 | ) |
| 159 | if named and not has_structural_tag: |
| 160 | # Without the XTML structural tag there is no way to force the |
| 161 | # named call (the generic JSON guided-decoding path conflicts |
| 162 | # with the XTML channel format). |
| 163 | raise VLLMValidationError( |
| 164 | "Named tool choice for Kimi K3 requires strict tool calling " |
| 165 | "(VLLM_ENFORCE_STRICT_TOOL_CALLING) so the XTML structural " |
| 166 | "tag can force the call. Otherwise use `tool_choice` set to " |
| 167 | '"auto", "required", or "none".', |
| 168 | parameter="tool_choice", |
| 169 | value=request.tool_choice, |
| 170 | ) |
| 171 | |
| 172 | if request.tools and (request.tool_choice == "required" or named): |
| 173 | # K3 emits tool calls in XTML. When strict tool calling is enabled, |
| 174 | # DelegatingParser._apply_structural_tag has already attached the K3 |
| 175 | # XTML structural tag (structural_tag_model = "kimi_k3"). We return |
| 176 | # early to skip the generic parent path, which would otherwise attach |
| 177 | # JSON guided decoding when strict calling is off -- that JSON |
| 178 | # constraint conflicts with the <|open|>tools<|sep|> channel. |
| 179 | request.skip_special_tokens = False |
| 180 | if hasattr(request, "spaces_between_special_tokens"): |
| 181 | request.spaces_between_special_tokens = False |
| 182 | return request |
| 183 | |
| 184 | request = super().adjust_request(request) |
| 185 | # The XTML markers (<|open|>/<|close|>/<|sep|>, |
| 186 | # <|open|>response<|sep|> ...) must |
| 187 | # reach this parser as CONTIGUOUS literal text. Two request flags govern |
| 188 | # that, and vLLM's detokenizer couples them: |
| 189 | # effective_spaces_between = |
| 190 | # skip_special_tokens OR spaces_between_special_tokens |
| 191 | # (vllm/v1/engine/detokenizer.py). |
| 192 | # The detokenizer treats these control tokens as separate sub-texts and, |
| 193 | # when the effective flag is True, joins them with spaces -> |
| 194 | # "<|open|> response <|sep|>", which the regexes below would NOT match. |
| 195 | # Forcing BOTH flags off is the only way to suppress that spacing. We |
| 196 | # set both unconditionally: the response channel is unwrapped from these |
| 197 | # markers even with no tools. |
| 198 | request.skip_special_tokens = False |
| 199 | if hasattr(request, "spaces_between_special_tokens"): |
| 200 | request.spaces_between_special_tokens = False |
| 201 | return request |
| 202 | |
| 203 | def _attrs(self, s: str) -> dict[str, str]: |
| 204 | return { |
| 205 | m["k"]: m["v"].replace(""", '"').replace("&", "&") |
| 206 | for m in self._attr_re.finditer(s) |
| 207 | } |
| 208 | |
| 209 | def _decode_call(self, attrs: str, body: str) -> ToolCall | None: |
| 210 | """Decode one ``call`` block into a :class:`ToolCall`. |
| 211 | |
| 212 | ``attrs`` is the text between ``<|open|>call `` and ``<|sep|>`` (carries |
| 213 | ``tool=`` / ``index=``); ``body`` is the sequence of ``argument`` blocks. |
| 214 | Each argument is re-typed per its ``type=`` tag: strings pass through |
| 215 | raw, everything else is JSON-decoded (falling back to raw text if the |
| 216 | JSON is malformed, so a partial stream never raises). Returns ``None`` |
| 217 | when no tool name is present (an empty/garbage block is dropped). |
| 218 | """ |
| 219 | call_attrs = self._attrs(attrs) |
| 220 | tool_name = call_attrs.get("tool", "") |
| 221 | arguments: dict = {} |
| 222 | for arg_match in self._arg_re.finditer(body): |
| 223 | arg_attrs = self._attrs(arg_match["attrs"]) |
| 224 | key = arg_attrs.get("key", "") |
| 225 | arg_type = arg_attrs.get("type", "string") |
| 226 | raw_value = arg_match["val"] |
| 227 | if arg_type == "string": |
| 228 | arguments[key] = raw_value |
| 229 | else: |
| 230 | try: |
| 231 | arguments[key] = json.loads(raw_value) |
| 232 | except json.JSONDecodeError: |
| 233 | arguments[key] = raw_value |
| 234 | if not tool_name: |
| 235 | return None |
| 236 | return ToolCall( |
| 237 | type="function", |
| 238 | function=FunctionCall( |
| 239 | name=tool_name, |
| 240 | arguments=json.dumps(arguments, ensure_ascii=False), |
| 241 | ), |
| 242 | ) |
| 243 | |
| 244 | def _strip_response_content(self, text: str) -> str | None: |
| 245 | """Strip XTML response/message markers from generated response text. |
| 246 | |
| 247 | In chat serving, ``<|open|>response<|sep|>`` is often part of the prompt |
| 248 | generation prefix, so the model output may only contain the body plus |
| 249 | ``<|close|>response<|sep|>``. Handle both that consumed-prefix shape and a |
| 250 | complete ``<|open|>response<|sep|>... <|close|>response<|sep|>`` wrapper. |
| 251 | """ |
| 252 | m_open = self._response_open_re.search(text) |
| 253 | if m_open is not None: |
| 254 | m_close = self._response_close_re.search(text, m_open.end()) |
| 255 | if m_close is not None: |
| 256 | text = text[m_open.end() : m_close.start()] |
| 257 | else: |
| 258 | text = text[m_open.end() :] |
| 259 | else: |
| 260 | text = self._response_close_re.sub("", text) |
| 261 | text = self._message_close_re.sub("", text) |
| 262 | return text or None |
| 263 | |
| 264 | def _content(self, model_output: str, before: str) -> str | None: |
| 265 | # prefer the unwrapped response channel; else the text before the tools |
| 266 | m = self._response_re.search(model_output) |
| 267 | if m is not None: |
| 268 | return m["c"] or None |
| 269 | return self._strip_response_content(before) |
| 270 | |
| 271 | def _extract_response_content(self, current_text: str) -> str | None: |
| 272 | # Streaming response text is computed from the accumulated text. This is |
| 273 | # what keeps split markers from leaking: |
| 274 | # <|open|> / response / <|sep|>Hi -> emit only "Hi" after open closes |
| 275 | # Hi<|open|> / tools / <|sep|>... -> emit "Hi", hold the tools marker |
| 276 | # Hi<|close|> / response / <|sep|>... -> emit "Hi", hold the close marker |
| 277 | # Tool calls are even simpler: they are not emitted until a full |
| 278 | # <|close|>call<|sep|> is present in current_text. |
| 279 | m_open = self._response_open_re.search(current_text) |
| 280 | # In the normal chat path, the response-open marker may be consumed as |
| 281 | # the generation prefix. Then generated text starts directly with the |
| 282 | # response body or with <|close|>response<|sep|> before a tools channel. |
| 283 | body_start = m_open.end() if m_open is not None else 0 |
| 284 | m_tools = self._tools_open_re.search(current_text, body_start) |
| 285 | m_rclose = self._response_close_re.search(current_text, body_start) |
| 286 | tools_start = m_tools.start() if m_tools else -1 |
| 287 | response_end = m_rclose.start() if m_rclose else -1 |
| 288 | |
| 289 | candidates = [i for i in (tools_start, response_end) if i != -1] |
| 290 | if candidates: |
| 291 | sendable_idx = min(candidates) |
| 292 | else: |
| 293 | overlap = max( |
| 294 | _partial_tag_overlap(current_text, self.response_open), |
| 295 | _partial_tag_overlap(current_text, self.response_close), |
| 296 | _partial_tag_overlap(current_text, self.tools_open), |
| 297 | ) |
| 298 | sendable_idx = len(current_text) - overlap |
| 299 | |
| 300 | if sendable_idx <= body_start: |
| 301 | return None |
| 302 | if self._sent_content_idx < body_start: |
| 303 | self._sent_content_idx = body_start |
| 304 | if sendable_idx <= self._sent_content_idx: |
| 305 | return None |
| 306 | |
| 307 | content = current_text[self._sent_content_idx : sendable_idx] |
| 308 | self._sent_content_idx = sendable_idx |
| 309 | return content or None |
| 310 | |
| 311 | def extract_tool_calls( |
| 312 | self, model_output: str, request: ChatCompletionRequest |
| 313 | ) -> ExtractedToolCallInformation: |
| 314 | m_open = self._tools_open_re.search(model_output) |
| 315 | if m_open is None: |
| 316 | # no tools channel -> content is the response channel (unwrapped) |
| 317 | return ExtractedToolCallInformation( |
| 318 | tools_called=False, |
| 319 | tool_calls=[], |
| 320 | content=self._content(model_output, model_output), |
| 321 | ) |
| 322 | try: |
| 323 | before = model_output[: m_open.start()] |
| 324 | start = m_open.end() |
| 325 | m_close = self._tools_close_re.search(model_output, start) |
| 326 | section = ( |
| 327 | model_output[start:] |
| 328 | if m_close is None |
| 329 | else model_output[start : m_close.start()] |
| 330 | ) |
| 331 | |
| 332 | tool_calls = [ |
| 333 | tc |
| 334 | for m in self._call_re.finditer(section) |
| 335 | if (tc := self._decode_call(m["attrs"], m["body"])) is not None |
| 336 | ] |
| 337 | if not tool_calls: |
| 338 | return ExtractedToolCallInformation( |
| 339 | tools_called=False, |
| 340 | tool_calls=[], |
| 341 | content=self._content(model_output, before), |
| 342 | ) |
| 343 | return ExtractedToolCallInformation( |
| 344 | tools_called=True, |
| 345 | tool_calls=tool_calls, |
| 346 | content=self._content(model_output, before), |
| 347 | ) |
| 348 | except Exception: |
| 349 | logger.exception("Error extracting K3 tool calls.") |
| 350 | return ExtractedToolCallInformation( |
| 351 | tools_called=False, tool_calls=[], content=model_output |
| 352 | ) |
| 353 | |
| 354 | def extract_tool_calls_streaming( |
| 355 | self, |
| 356 | previous_text: str, |
| 357 | current_text: str, |
| 358 | delta_text: str, |
| 359 | previous_token_ids: Sequence[int], |
| 360 | current_token_ids: Sequence[int], |
| 361 | delta_token_ids: Sequence[int], |
| 362 | request: ChatCompletionRequest, |
| 363 | ) -> DeltaMessage | None: |
| 364 | # Conservative streaming: stream unwrapped response-channel text, then |
| 365 | # buffer tool calls and emit each call once its block closes. |
| 366 | content = self._extract_response_content(current_text) |
| 367 | |
| 368 | # tools channel is open: parse fully-closed calls we have not emitted yet |
| 369 | m_tools = self._tools_open_re.search(current_text) |
| 370 | if m_tools is None: |
| 371 | return DeltaMessage(content=content) if content else None |
| 372 | |
| 373 | section = current_text[m_tools.end() :] |
| 374 | calls = [ |
| 375 | tc |
| 376 | for m in self._call_re.finditer(section) |
| 377 | if (tc := self._decode_call(m["attrs"], m["body"])) is not None |
| 378 | ] |
| 379 | if len(calls) <= self._sent_tool_call_count: |
| 380 | return DeltaMessage(content=content) if content else None |
| 381 | new = calls[self._sent_tool_call_count :] |
| 382 | |
| 383 | deltas = [ |
| 384 | DeltaToolCall( |
| 385 | index=self._sent_tool_call_count + i, |
| 386 | id=tc.id, |
| 387 | type="function", |
| 388 | function=DeltaFunctionCall( |
| 389 | name=tc.function.name, arguments=tc.function.arguments |
| 390 | ).model_dump(exclude_none=True), |
| 391 | ) |
| 392 | for i, tc in enumerate(new) |
| 393 | ] |
| 394 | self._sent_tool_call_count = len(calls) |
| 395 | return DeltaMessage(content=content, tool_calls=deltas) |