inference-optimization/Kimi-K3-0.40B-MXFP4
73.1k
1"""Kimi K3 XTML encoding helpers.2 3This module keeps chat rendering in Python.4Callers that need token IDs should consume ``EncodeSegment`` objects directly:5structural markers may be encoded as tiktoken special tokens, while user/tool6text and attribute values are encoded as ordinary text.7"""8 9from __future__ import annotations10 11import json12from dataclasses import dataclass13from typing import Any, Iterable, Optional14 15OPEN_TOKEN = "<|open|>"16CLOSE_TOKEN = "<|close|>"17SEP_TOKEN = "<|sep|>"18END_OF_MSG_TOKEN = "<|end_of_msg|>"19IMAGE_PLACEHOLDER = "<|kimi_image_placeholder|>"20 21_VALID_THINKING_EFFORTS = {"low", "high", "max"}22 23 24@dataclass(frozen=True)25class EncodeSegment:26 text: str27 allow_special: bool = False28 29 30class _ImagePromptState:31 def __init__(self, image_prompts: Optional[list[str]] = None):32 self.image_prompts = image_prompts33 self.index = 034 35 def next_prompt(self) -> str:36 if self.image_prompts is None:37 return IMAGE_PLACEHOLDER38 if self.index >= len(self.image_prompts):39 raise ValueError("More image placeholders than image prompts.")40 prompt = self.image_prompts[self.index]41 self.index += 142 return prompt43 44 def assert_consumed(self) -> None:45 if self.image_prompts is None:46 return47 if self.index != len(self.image_prompts):48 raise ValueError(49 f"image prompt count {len(self.image_prompts)} != "50 f"consumed placeholder count {self.index}"51 )52 53 54def _segment(text: Any, *, allow_special: bool = False) -> list[EncodeSegment]:55 text = str(text)56 if not text:57 return []58 return [EncodeSegment(text, allow_special=allow_special)]59 60 61def _control(text: str) -> list[EncodeSegment]:62 return _segment(text, allow_special=True)63 64 65def _text(text: Any) -> list[EncodeSegment]:66 return _segment(text, allow_special=False)67 68 69def _append_text(70 segments: list[EncodeSegment],71 text: Any,72 image_state: _ImagePromptState,73) -> None:74 text = str(text)75 if text == "":76 return77 if image_state.image_prompts is None or IMAGE_PLACEHOLDER not in text:78 segments.extend(_text(text))79 return80 81 parts = text.split(IMAGE_PLACEHOLDER)82 for i, part in enumerate(parts):83 segments.extend(_text(part))84 if i < len(parts) - 1:85 segments.extend(_segment(image_state.next_prompt(),86 allow_special=True))87 88 89def _escape_attr_value(value: Any) -> str:90 return str(value).replace("&", "&").replace('"', """)91 92 93def _attr(key: str, value: Any) -> list[EncodeSegment]:94 return (95 _text(f" {key}")96 + _text('="')97 + _text(_escape_attr_value(value))98 + _text('"')99 )100 101 102def _open_tag(tag: str, attrs: Iterable[tuple[str, Any]] = ()) -> list[EncodeSegment]:103 segments: list[EncodeSegment] = []104 segments.extend(_control(OPEN_TOKEN))105 segments.extend(_text(tag))106 for key, value in attrs:107 segments.extend(_attr(key, value))108 segments.extend(_control(SEP_TOKEN))109 return segments110 111def _close_tag(tag: str) -> list[EncodeSegment]:112 segments: list[EncodeSegment] = []113 segments.extend(_control(CLOSE_TOKEN))114 segments.extend(_text(tag))115 segments.extend(_control(SEP_TOKEN))116 return segments117 118 119def _end_of_msg() -> list[EncodeSegment]:120 return _control(END_OF_MSG_TOKEN)121 122 123def _json_compact(value: Any) -> str:124 return json.dumps(value, ensure_ascii=False, separators=(",", ":"))125 126 127def _is_mapping(value: Any) -> bool:128 return isinstance(value, dict)129 130 131def _xtml_type(value: Any) -> str:132 if isinstance(value, bool):133 return "boolean"134 if value is None:135 return "null"136 if isinstance(value, (int, float)) and not isinstance(value, bool):137 return "number"138 if isinstance(value, str):139 return "string"140 if _is_mapping(value):141 return "object"142 return "array"143 144 145def _xtml_value(value: Any) -> str:146 if isinstance(value, str):147 return value148 return json.dumps(value, ensure_ascii=False)149 150 151def _get_value(obj: Any, key: str, default: Any = None) -> Any:152 if isinstance(obj, dict):153 return obj.get(key, default)154 return getattr(obj, key, default)155 156 157def extract_response_schema(response_format: Any) -> Any:158 if response_format is None:159 return None160 161 json_schema = _get_value(response_format, "json_schema")162 if json_schema is None:163 return None164 165 if isinstance(json_schema, dict):166 return json_schema.get(167 "schema",168 json_schema.get("json_schema", json_schema),169 )170 171 schema = _get_value(json_schema, "schema")172 if schema is not None:173 return schema174 175 schema = _get_value(json_schema, "json_schema")176 if schema is not None:177 return schema178 179 return json_schema180 181 182def deep_sort_dict(obj: Any) -> Any:183 if isinstance(obj, dict):184 return {k: deep_sort_dict(v) for k, v in sorted(obj.items())}185 if isinstance(obj, list):186 return [deep_sort_dict(item) for item in obj]187 return obj188 189 190# One normalized tool-call argument: (key, XTML type, rendered text). The191# text keeps the original JSON literal for non-string values and is the192# decoded string for string values.193XtmlArgument = tuple[str, str, str]194 195 196def _parse_arguments_object(s: str) -> list[XtmlArgument]:197 """Parse a JSON object string one level deep.198 199 Each top-level key-value pair yields one ``(key, type, text)`` triple. The200 text keeps the original JSON literal for non-string values (``1e2`` stays201 ``1e2``, ``[1,2]`` keeps its exact bytes) and is the decoded (unescaped)202 string for string values. Nested values are never re-serialized. Anything203 after the closing ``}`` is ignored. Raises ``ValueError`` on malformed204 input (including a valid non-object JSON document).205 """206 idx = 0207 208 def _skip_whitespaces() -> None:209 nonlocal idx210 while idx < len(s) and s[idx] in (" ", "\t", "\n", "\r"):211 idx += 1212 213 def _next_char() -> Optional[str]:214 nonlocal idx215 if idx < len(s):216 c = s[idx]217 idx += 1218 return c219 return None220 221 _skip_whitespaces()222 if _next_char() != "{":223 raise ValueError("JSON arguments must be an object")224 _skip_whitespaces()225 226 parsed: list[XtmlArgument] = []227 228 if idx >= len(s):229 raise ValueError("Unexpected end of JSON object")230 if s[idx] == "}":231 return parsed232 233 json_decoder = json.JSONDecoder(strict=False)234 235 def _raw_decode(start_idx: int) -> tuple[Any, int]:236 try:237 return json_decoder.raw_decode(s, idx=start_idx)238 except json.JSONDecodeError as e:239 raise ValueError(str(e)) from e240 241 while True:242 decoded_key, idx = _raw_decode(idx)243 if not isinstance(decoded_key, str):244 raise ValueError(f"JSON object key must be a string, got {decoded_key!r}")245 _skip_whitespaces()246 if _next_char() != ":":247 raise ValueError(f"Expects ':' after {decoded_key}")248 _skip_whitespaces()249 250 value_start_idx = idx251 decoded_value, idx = _raw_decode(idx)252 value_end_idx = idx253 254 text = (255 decoded_value256 if isinstance(decoded_value, str)257 else s[value_start_idx:value_end_idx]258 )259 parsed.append((decoded_key, _xtml_type(decoded_value), text))260 261 _skip_whitespaces()262 c = _next_char()263 _skip_whitespaces()264 265 if c == "}":266 break267 elif c == ",":268 continue269 else:270 raise ValueError(f"Expect '}}' or ',', got {c!r}")271 272 return parsed273 274 275def normalize_tool_arguments(276 arguments: Any,277) -> tuple[list[XtmlArgument], Optional[str]]:278 """Normalize tool call arguments for XTML rendering.279 280 Returns ``(argument_triples, raw_json_block)``. String arguments are parsed281 one level deep so non-string values keep their original JSON literal text,282 and anything after the closing ``}`` is discarded; dict arguments (already283 Python objects) are serialized on the spot. Any string that is not a284 well-formed JSON object -- unparseable, whitespace only, or valid285 non-object JSON -- falls back to the raw ``<json>`` block.286 """287 if arguments is None:288 return [], None289 if isinstance(arguments, dict):290 return [291 (str(key), _xtml_type(value), _xtml_value(value))292 for key, value in arguments.items()293 ], None294 if isinstance(arguments, str):295 if arguments == "":296 return [], None297 try:298 return _parse_arguments_object(arguments), None299 except ValueError:300 return [], arguments301 raise TypeError(302 "Kimi K3 tool call arguments must be a dict or a JSON object string."303 )304 305 306def normalize_message(message: Any) -> Any:307 if not isinstance(message, dict):308 return message309 310 normalized = dict(message)311 312 tools = normalized.get("tools")313 if tools is not None:314 normalized["tools"] = deep_sort_dict(tools)315 316 tool_calls = normalized.get("tool_calls")317 if not tool_calls:318 return normalized319 320 normalized_calls = []321 for tool_call in tool_calls:322 if not isinstance(tool_call, dict):323 normalized_calls.append(tool_call)324 continue325 326 tc = dict(tool_call)327 function = tc.get("function")328 if isinstance(function, dict):329 fn = dict(function)330 arguments, json_block = normalize_tool_arguments(fn.get("arguments"))331 fn["arguments"] = arguments332 if json_block is None:333 fn.pop("_xtml_json_block", None)334 else:335 fn["_xtml_json_block"] = json_block336 tc["function"] = fn337 else:338 arguments, json_block = normalize_tool_arguments(tc.get("arguments"))339 tc["arguments"] = arguments340 if json_block is None:341 tc.pop("_xtml_json_block", None)342 else:343 tc["_xtml_json_block"] = json_block344 normalized_calls.append(tc)345 346 normalized["tool_calls"] = normalized_calls347 return normalized348 349 350def normalize_conversation(conversation: Any) -> Any:351 if not isinstance(conversation, list):352 return conversation353 354 def normalize_messages(messages: list[Any]) -> list[Any]:355 return [normalize_message(message) for message in messages]356 357 if conversation and isinstance(conversation[0], list):358 return [normalize_messages(messages) for messages in conversation]359 return normalize_messages(conversation)360 361 362def _tool_call_id_index(tool_calls: Any) -> dict:363 """Map assistant ``tool_calls[].id`` to ``(1-based position, function name)``.364 365 The position mirrors the chat template's enumeration over ``tool_calls``366 (every entry advances the position, even an id-less one). Duplicate ids keep367 their first occurrence.368 """369 index: dict = {}370 if not isinstance(tool_calls, list):371 return index372 for position, tool_call in enumerate(tool_calls, start=1):373 if not isinstance(tool_call, dict):374 continue375 call_id = tool_call.get("id")376 if call_id is None:377 continue378 key = str(call_id)379 if key in index:380 continue381 function = tool_call.get("function")382 name = (383 function.get("name") if isinstance(function, dict) else tool_call.get("name")384 )385 index[key] = (position, name)386 return index387 388 389def normalize_xtml_tool_result_messages(messages: list[Any]) -> list[Any]:390 """Re-sort K3 XTML tool results into assistant ``tool_calls`` order.391 392 Serving frameworks generally deliver tool results already in call order. A393 direct Transformers caller, however, may pass OpenAI-style tool messages in any394 order, so each run of consecutive tool messages is matched against the most395 recent preceding assistant ``tool_calls`` by opaque ``tool_call_id`` ==396 ``tool_calls[].id`` (K3 drops the ``func:index`` format requirement) and397 sorted by the matched 1-based position. The matched call is authoritative,398 so each matched message's ``tool`` is set to that call's function name --399 this keeps an explicit (and possibly stale) ``tool``/``name`` from drifting400 out of sync with the reordered position. ``index`` is still derived from the401 rendered position by the chat template. A run that cannot be fully matched is402 left untouched. Re-running is idempotent.403 404 This function is side-effect free: matched tool messages are shallow-copied405 before their ``tool``/``name`` is rewritten, and every other message is406 appended to the output as-is. The input list and its message objects are407 never mutated.408 """409 if not isinstance(messages, list):410 return messages411 412 output: list[Any] = []413 current_index: dict = {}414 i = 0415 n = len(messages)416 417 while i < n:418 message = messages[i]419 420 if isinstance(message, dict) and message.get("role") == "assistant":421 tool_calls = message.get("tool_calls")422 current_index = _tool_call_id_index(tool_calls) if tool_calls else {}423 output.append(message)424 i += 1425 continue426 427 if not isinstance(message, dict) or message.get("role") != "tool":428 output.append(message)429 i += 1430 continue431 432 run: list[tuple] = [] # (position, original_offset, message, name)433 unresolved = False434 offset = 0435 while (436 i < n and isinstance(messages[i], dict) and messages[i].get("role") == "tool"437 ):438 tool_message = messages[i]439 call_id = tool_message.get("tool_call_id", tool_message.get("id"))440 matched = current_index.get(str(call_id)) if call_id is not None else None441 if matched is None:442 unresolved = True443 run.append((None, offset, tool_message, None))444 else:445 position, name = matched446 run.append((position, offset, tool_message, name))447 offset += 1448 i += 1449 450 if unresolved:451 output.extend(item[2] for item in run)452 else:453 run.sort(key=lambda item: (item[0], item[1]))454 for _, _, tool_message, name in run:455 if name is None:456 output.append(tool_message)457 continue458 # The id-matched call is authoritative: align tool (and any459 # explicit name) so the rendered XTML tool attribute cannot460 # disagree with the reordered position. Copy first so the461 # caller's message object is never mutated.462 resolved = dict(tool_message)463 resolved["tool"] = name464 if "name" in resolved:465 resolved["name"] = name466 output.append(resolved)467 468 return output469 470 471def is_batched_conversation(conversation: Any) -> bool:472 return (473 isinstance(conversation, list)474 and bool(conversation)475 and isinstance(conversation[0], list)476 )477 478 479def _render_content_segments(480 content: Any,481 image_state: _ImagePromptState,482) -> list[EncodeSegment]:483 segments: list[EncodeSegment] = []484 if isinstance(content, str):485 _append_text(segments, content, image_state)486 elif content is not None:487 for part in content:488 if part["type"] in ["image", "image_url"]:489 segments.extend(490 _segment(image_state.next_prompt(), allow_special=True))491 else:492 _append_text(segments, part["text"], image_state)493 return segments494 495 496def _internal_system_message(message_type: str, body: str) -> list[EncodeSegment]:497 segments: list[EncodeSegment] = []498 segments.extend(_open_tag("message", [("role", "system"), ("type", message_type)]))499 segments.extend(_text(body.strip()))500 segments.extend(_close_tag("message"))501 segments.extend(_end_of_msg())502 return segments503 504 505def _render_assistant_segments(506 message: dict[str, Any],507 image_state: _ImagePromptState,508 thinking: bool = True,509) -> list[EncodeSegment]:510 segments: list[EncodeSegment] = []511 # The <think> channel is structural: in thinking mode every assistant512 # message carries the open/close tags even when there is no reasoning513 # content to fill in. In non-thinking mode the channel is dropped514 # entirely.515 if thinking:516 reasoning_content = message.get("reasoning_content") or message.get(517 "reasoning"518 )519 segments.extend(_open_tag("think"))520 # Only an empty string counts as no reasoning, so whitespace-only521 # reasoning is still rendered into the think channel.522 if reasoning_content is not None and str(reasoning_content) != "":523 _append_text(segments, reasoning_content, image_state)524 segments.extend(_close_tag("think"))525 526 segments.extend(_open_tag("response"))527 segments.extend(_render_content_segments(message.get("content"), image_state))528 segments.extend(_close_tag("response"))529 530 tool_calls = message.get("tool_calls")531 if tool_calls:532 segments.extend(_open_tag("tools"))533 for index, tool_call in enumerate(tool_calls, start=1):534 fn = tool_call.get("function", tool_call)535 segments.extend(536 _open_tag("call", [("tool", fn["name"]), ("index", index)])537 )538 args = fn.get("arguments", [])539 json_block = fn.get("_xtml_json_block")540 if json_block is not None:541 segments.extend(_open_tag("json", [("type", "object")]))542 _append_text(segments, json_block, image_state)543 segments.extend(_close_tag("json"))544 else:545 for key, arg_type, arg_text in args:546 segments.extend(547 _open_tag("argument", [("key", key), ("type", arg_type)])548 )549 _append_text(segments, arg_text, image_state)550 segments.extend(_close_tag("argument"))551 segments.extend(_close_tag("call"))552 segments.extend(_close_tag("tools"))553 554 return segments555 556 557def _render_tool_declare(tools: Any, *, dynamic: bool = False) -> list[EncodeSegment]:558 if dynamic:559 body = (560 "## New Tools Available\n"561 "The system dynamically extends the toolset via lazy-loading.\n"562 "You have access to all existing and extended tools.\n"563 "Here are the specs for the extended tools.\n\n"564 "```json\n"565 f"{_json_compact(tools)}\n"566 "```"567 )568 else:569 body = (570 "# Tools\n"571 "Here are the available tools, described in JSONSchema.\n\n"572 "```json\n"573 f"{_json_compact(tools)}\n"574 "```"575 )576 segments: list[EncodeSegment] = []577 segments.extend(_open_tag("message", [("role", "system"), ("type", "tool-declare")]))578 segments.extend(_text(body))579 segments.extend(_close_tag("message"))580 segments.extend(_end_of_msg())581 return segments582 583 584def build_chat_segments(585 messages: list[Any],586 tools: Optional[list[dict]] = None,587 *,588 add_generation_prompt: bool = True,589 thinking: bool = True,590 image_prompts: Optional[list[str]] = None,591 **kwargs: Any,592) -> list[EncodeSegment]:593 # Re-sort tool results by tool_call_id at the lowest layer so every caller594 # (processor or direct tokenizer) gets correctly ordered XTML. The helper is595 # side-effect free, so the caller's message objects are left untouched.596 messages = normalize_xtml_tool_result_messages(messages)597 messages = normalize_conversation(messages)598 tools = deep_sort_dict(tools)599 600 kwargs = dict(kwargs)601 response_format = kwargs.get("response_format")602 if "response_schema" not in kwargs:603 response_schema = extract_response_schema(response_format)604 if response_schema is not None:605 kwargs["response_schema"] = response_schema606 if kwargs.get("response_schema") is not None:607 kwargs["response_schema"] = deep_sort_dict(kwargs["response_schema"])608 609 image_state = _ImagePromptState(image_prompts)610 segments: list[EncodeSegment] = []611 612 tool_calls = None613 tool_index = 0614 615 if tools:616 segments.extend(_render_tool_declare(tools))617 618 thinking_effort = kwargs.get("thinking_effort")619 if thinking and thinking_effort is not None:620 assert thinking_effort in _VALID_THINKING_EFFORTS, (621 f"Unsupported thinking_effort={thinking_effort!r}; "622 f"supported values are {sorted(_VALID_THINKING_EFFORTS)}."623 )624 if thinking and thinking_effort in _VALID_THINKING_EFFORTS:625 segments.extend(626 _internal_system_message(627 "thinking-effort",628 "`thinking_effort` guides on how much to think in your "629 "thinking channel (not including the response channel), "630 "supported values include `low`, `medium`, `high`, and `max`.\n"631 f"Now the system is invoked with `thinking_effort={thinking_effort}`.",632 )633 )634 635 for message_index, message in enumerate(messages):636 # Malformed messages are rejected instead of being silently skipped.637 if not isinstance(message, dict):638 raise ValueError(639 f"Kimi K3 messages must be dicts, got {type(message).__name__} "640 f"at index {message_index}."641 )642 643 role = message.get("role")644 if role == "user":645 attrs = [("role", "user")]646 if message.get("name"):647 attrs.append(("name", message["name"]))648 segments.extend(_open_tag("message", attrs))649 segments.extend(_render_content_segments(message.get("content"), image_state))650 segments.extend(_close_tag("message"))651 segments.extend(_end_of_msg())652 elif role == "system" and message.get("tools"):653 segments.extend(_render_tool_declare(message["tools"], dynamic=True))654 elif role == "system":655 attrs = [("role", "system")]656 if message.get("name"):657 attrs.append(("name", message["name"]))658 segments.extend(_open_tag("message", attrs))659 segments.extend(_render_content_segments(message.get("content"), image_state))660 segments.extend(_close_tag("message"))661 segments.extend(_end_of_msg())662 elif role == "tool":663 tool_index += 1664 tool_name = message.get("tool", message.get("name"))665 if (666 tool_name is None667 and tool_calls is not None668 and tool_index <= len(tool_calls)669 ):670 tc = tool_calls[tool_index - 1]671 fn = tc.get("function", tc)672 tool_name = fn["name"]673 if tool_name is None:674 raise ValueError(675 "Kimi K3 tool messages need a resolvable tool name: "676 "carry `tool`/`name`, or match a preceding assistant "677 "tool_call by order."678 )679 segments.extend(680 _open_tag(681 "message",682 [("role", "tool"), ("tool", tool_name), ("index", tool_index)],683 )684 )685 segments.extend(_render_content_segments(message.get("content"), image_state))686 segments.extend(_close_tag("message"))687 segments.extend(_end_of_msg())688 elif role == "assistant":689 tool_calls = message.get("tool_calls")690 tool_index = 0691 attrs = [("role", "assistant")]692 if message.get("name"):693 attrs.append(("name", message["name"]))694 segments.extend(_open_tag("message", attrs))695 segments.extend(_render_assistant_segments(message, image_state, thinking))696 segments.extend(_close_tag("message"))697 segments.extend(_end_of_msg())698 else:699 raise ValueError(700 f"Unknown message role {role!r} at index {message_index}."701 )702 703 tool_choice = kwargs.get("tool_choice")704 if tool_choice == "required":705 segments.extend(706 _internal_system_message(707 "tool-choice",708 "The system is invoked with `tool_choice=required`.\n"709 "You MUST call tools in the next message.",710 )711 )712 elif tool_choice == "none":713 segments.extend(714 _internal_system_message(715 "tool-choice",716 "The system is invoked with `tool_choice=none`.\n"717 "You MUST NOT call any tools in the next message.",718 )719 )720 721 rf = kwargs.get("response_format")722 rf_type = _get_value(rf, "type", rf) if isinstance(rf, dict) else rf723 if rf_type == "json_object":724 segments.extend(725 _internal_system_message(726 "response-format",727 "The system is invoked with `response_format=json_object`.\n"728 "Your response must be raw JSON data without markdown code "729 "blocks (```json) or any additional formatting.",730 )731 )732 elif rf_type == "json_schema":733 schema = _json_compact(kwargs.get("response_schema"))734 segments.extend(735 _internal_system_message(736 "response-format",737 "The system is invoked with `response_format=json_schema`.\n"738 "Your response must be raw JSON data without markdown code "739 "blocks (```json) or any additional formatting.\n"740 "The JSON data must match the following schema:\n"741 f"```json\n{schema}\n```",742 )743 )744 745 if add_generation_prompt:746 segments.extend(_open_tag("message", [("role", "assistant")]))747 segments.extend(_open_tag("think" if thinking else "response"))748 749 image_state.assert_consumed()750 return segments751