abacusai/Smaug-Agentic
25200
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 190def normalize_tool_arguments(arguments: Any) -> tuple[dict[str, Any], Optional[str]]:191 if arguments is None:192 return {}, None193 if isinstance(arguments, dict):194 return arguments, None195 if isinstance(arguments, str):196 if not arguments.strip():197 return {}, None198 try:199 parsed = json.loads(arguments)200 except json.JSONDecodeError:201 return {}, arguments202 if not isinstance(parsed, dict):203 raise ValueError("Kimi K3 tool call arguments must be a JSON object.")204 return parsed, None205 raise TypeError(206 "Kimi K3 tool call arguments must be a dict or a JSON object string."207 )208 209 210def normalize_message(message: Any) -> Any:211 if not isinstance(message, dict):212 return message213 214 normalized = dict(message)215 216 tools = normalized.get("tools")217 if tools is not None:218 normalized["tools"] = deep_sort_dict(tools)219 220 tool_calls = normalized.get("tool_calls")221 if not tool_calls:222 return normalized223 224 normalized_calls = []225 for tool_call in tool_calls:226 if not isinstance(tool_call, dict):227 normalized_calls.append(tool_call)228 continue229 230 tc = dict(tool_call)231 function = tc.get("function")232 if isinstance(function, dict):233 fn = dict(function)234 arguments, json_block = normalize_tool_arguments(fn.get("arguments"))235 fn["arguments"] = arguments236 if json_block is None:237 fn.pop("_xtml_json_block", None)238 else:239 fn["_xtml_json_block"] = json_block240 tc["function"] = fn241 else:242 arguments, json_block = normalize_tool_arguments(tc.get("arguments"))243 tc["arguments"] = arguments244 if json_block is None:245 tc.pop("_xtml_json_block", None)246 else:247 tc["_xtml_json_block"] = json_block248 normalized_calls.append(tc)249 250 normalized["tool_calls"] = normalized_calls251 return normalized252 253 254def normalize_conversation(conversation: Any) -> Any:255 if not isinstance(conversation, list):256 return conversation257 258 def normalize_messages(messages: list[Any]) -> list[Any]:259 return [normalize_message(message) for message in messages]260 261 if conversation and isinstance(conversation[0], list):262 return [normalize_messages(messages) for messages in conversation]263 return normalize_messages(conversation)264 265 266def _tool_call_id_index(tool_calls: Any) -> dict:267 """Map assistant ``tool_calls[].id`` to ``(1-based position, function name)``.268 269 The position mirrors the chat template's enumeration over ``tool_calls``270 (every entry advances the position, even an id-less one). Duplicate ids keep271 their first occurrence.272 """273 index: dict = {}274 if not isinstance(tool_calls, list):275 return index276 for position, tool_call in enumerate(tool_calls, start=1):277 if not isinstance(tool_call, dict):278 continue279 call_id = tool_call.get("id")280 if call_id is None:281 continue282 key = str(call_id)283 if key in index:284 continue285 function = tool_call.get("function")286 name = (287 function.get("name") if isinstance(function, dict) else tool_call.get("name")288 )289 index[key] = (position, name)290 return index291 292 293def normalize_xtml_tool_result_messages(messages: list[Any]) -> list[Any]:294 """Re-sort K3 XTML tool results into assistant ``tool_calls`` order.295 296 Serving frameworks generally deliver tool results already in call order. A297 direct Transformers caller, however, may pass OpenAI-style tool messages in any298 order, so each run of consecutive tool messages is matched against the most299 recent preceding assistant ``tool_calls`` by opaque ``tool_call_id`` ==300 ``tool_calls[].id`` (K3 drops the ``func:index`` format requirement) and301 sorted by the matched 1-based position. The matched call is authoritative,302 so each matched message's ``tool`` is set to that call's function name --303 this keeps an explicit (and possibly stale) ``tool``/``name`` from drifting304 out of sync with the reordered position. ``index`` is still derived from the305 rendered position by the chat template. A run that cannot be fully matched is306 left untouched. Re-running is idempotent.307 308 This function is side-effect free: matched tool messages are shallow-copied309 before their ``tool``/``name`` is rewritten, and every other message is310 appended to the output as-is. The input list and its message objects are311 never mutated.312 """313 if not isinstance(messages, list):314 return messages315 316 output: list[Any] = []317 current_index: dict = {}318 i = 0319 n = len(messages)320 321 while i < n:322 message = messages[i]323 324 if isinstance(message, dict) and message.get("role") == "assistant":325 tool_calls = message.get("tool_calls")326 current_index = _tool_call_id_index(tool_calls) if tool_calls else {}327 output.append(message)328 i += 1329 continue330 331 if not isinstance(message, dict) or message.get("role") != "tool":332 output.append(message)333 i += 1334 continue335 336 run: list[tuple] = [] # (position, original_offset, message, name)337 unresolved = False338 offset = 0339 while (340 i < n and isinstance(messages[i], dict) and messages[i].get("role") == "tool"341 ):342 tool_message = messages[i]343 call_id = tool_message.get("tool_call_id", tool_message.get("id"))344 matched = current_index.get(str(call_id)) if call_id is not None else None345 if matched is None:346 unresolved = True347 run.append((None, offset, tool_message, None))348 else:349 position, name = matched350 run.append((position, offset, tool_message, name))351 offset += 1352 i += 1353 354 if unresolved:355 output.extend(item[2] for item in run)356 else:357 run.sort(key=lambda item: (item[0], item[1]))358 for _, _, tool_message, name in run:359 if name is None:360 output.append(tool_message)361 continue362 # The id-matched call is authoritative: align tool (and any363 # explicit name) so the rendered XTML tool attribute cannot364 # disagree with the reordered position. Copy first so the365 # caller's message object is never mutated.366 resolved = dict(tool_message)367 resolved["tool"] = name368 if "name" in resolved:369 resolved["name"] = name370 output.append(resolved)371 372 return output373 374 375def is_batched_conversation(conversation: Any) -> bool:376 return (377 isinstance(conversation, list)378 and bool(conversation)379 and isinstance(conversation[0], list)380 )381 382 383def _render_content_segments(384 content: Any,385 image_state: _ImagePromptState,386) -> list[EncodeSegment]:387 segments: list[EncodeSegment] = []388 if isinstance(content, str):389 _append_text(segments, content, image_state)390 elif content is not None:391 for part in content:392 if part["type"] in ["image", "image_url"]:393 segments.extend(394 _segment(image_state.next_prompt(), allow_special=True))395 else:396 _append_text(segments, part["text"], image_state)397 return segments398 399 400def _internal_system_message(message_type: str, body: str) -> list[EncodeSegment]:401 segments: list[EncodeSegment] = []402 segments.extend(_open_tag("message", [("role", "system"), ("type", message_type)]))403 segments.extend(_text(body.strip()))404 segments.extend(_close_tag("message"))405 segments.extend(_end_of_msg())406 return segments407 408 409def _render_assistant_segments(410 message: dict[str, Any],411 image_state: _ImagePromptState,412 thinking: bool = True,413) -> list[EncodeSegment]:414 segments: list[EncodeSegment] = []415 # The <think> channel is structural: in thinking mode every assistant416 # message carries the open/close tags even when there is no reasoning417 # content to fill in. In non-thinking mode the channel is dropped418 # entirely.419 if thinking:420 reasoning_content = message.get("reasoning_content") or message.get(421 "reasoning"422 )423 segments.extend(_open_tag("think"))424 if reasoning_content is not None and str(reasoning_content).strip():425 _append_text(segments, reasoning_content, image_state)426 segments.extend(_close_tag("think"))427 428 segments.extend(_open_tag("response"))429 segments.extend(_render_content_segments(message.get("content"), image_state))430 segments.extend(_close_tag("response"))431 432 tool_calls = message.get("tool_calls")433 if tool_calls:434 segments.extend(_open_tag("tools"))435 for index, tool_call in enumerate(tool_calls, start=1):436 fn = tool_call.get("function", tool_call)437 segments.extend(438 _open_tag("call", [("tool", fn["name"]), ("index", index)])439 )440 args = fn.get("arguments", {})441 json_block = fn.get("_xtml_json_block")442 if json_block is not None:443 segments.extend(_open_tag("json", [("type", "object")]))444 _append_text(segments, json_block, image_state)445 segments.extend(_close_tag("json"))446 elif _is_mapping(args):447 for key, value in args.items():448 segments.extend(449 _open_tag(450 "argument",451 [("key", key), ("type", _xtml_type(value))],452 )453 )454 _append_text(segments, _xtml_value(value), image_state)455 segments.extend(_close_tag("argument"))456 segments.extend(_close_tag("call"))457 segments.extend(_close_tag("tools"))458 459 return segments460 461 462def _render_tool_declare(tools: Any, *, dynamic: bool = False) -> list[EncodeSegment]:463 if dynamic:464 body = (465 "## New Tools Available\n"466 "The system dynamically extends the toolset via lazy-loading.\n"467 "You have access to all existing and extended tools.\n"468 "Here are the specs for the extended tools.\n\n"469 "```json\n"470 f"{_json_compact(tools)}\n"471 "```"472 )473 else:474 body = (475 "# Tools\n"476 "Here are the available tools, described in JSONSchema.\n\n"477 "```json\n"478 f"{_json_compact(tools)}\n"479 "```"480 )481 segments: list[EncodeSegment] = []482 segments.extend(_open_tag("message", [("role", "system"), ("type", "tool-declare")]))483 segments.extend(_text(body))484 segments.extend(_close_tag("message"))485 segments.extend(_end_of_msg())486 return segments487 488 489def build_chat_segments(490 messages: list[Any],491 tools: Optional[list[dict]] = None,492 *,493 add_generation_prompt: bool = True,494 thinking: bool = True,495 image_prompts: Optional[list[str]] = None,496 **kwargs: Any,497) -> list[EncodeSegment]:498 # Re-sort tool results by tool_call_id at the lowest layer so every caller499 # (processor or direct tokenizer) gets correctly ordered XTML. The helper is500 # side-effect free, so the caller's message objects are left untouched.501 messages = normalize_xtml_tool_result_messages(messages)502 messages = normalize_conversation(messages)503 tools = deep_sort_dict(tools)504 505 kwargs = dict(kwargs)506 response_format = kwargs.get("response_format")507 if "response_schema" not in kwargs:508 response_schema = extract_response_schema(response_format)509 if response_schema is not None:510 kwargs["response_schema"] = response_schema511 if kwargs.get("response_schema") is not None:512 kwargs["response_schema"] = deep_sort_dict(kwargs["response_schema"])513 514 image_state = _ImagePromptState(image_prompts)515 segments: list[EncodeSegment] = []516 517 tool_calls = None518 tool_index = 0519 520 if tools:521 segments.extend(_render_tool_declare(tools))522 523 thinking_effort = kwargs.get("thinking_effort")524 if thinking and thinking_effort is not None:525 assert thinking_effort in _VALID_THINKING_EFFORTS, (526 f"Unsupported thinking_effort={thinking_effort!r}; "527 f"supported values are {sorted(_VALID_THINKING_EFFORTS)}."528 )529 if thinking and thinking_effort in _VALID_THINKING_EFFORTS:530 segments.extend(531 _internal_system_message(532 "thinking-effort",533 "`thinking_effort` guides on how much to think in your "534 "thinking channel (not including the response channel), "535 "supported values include `low`, `medium`, `high`, and `max`.\n"536 f"Now the system is invoked with `thinking_effort={thinking_effort}`.",537 )538 )539 540 for message_index, message in enumerate(messages):541 if not isinstance(message, dict):542 continue543 544 role = message["role"]545 if role == "user":546 attrs = [("role", "user")]547 if message.get("name"):548 attrs.append(("name", message["name"]))549 segments.extend(_open_tag("message", attrs))550 segments.extend(_render_content_segments(message.get("content"), image_state))551 segments.extend(_close_tag("message"))552 segments.extend(_end_of_msg())553 elif role == "system" and message.get("tools"):554 segments.extend(_render_tool_declare(message["tools"], dynamic=True))555 elif role == "system":556 attrs = [("role", "system")]557 if message.get("name"):558 attrs.append(("name", message["name"]))559 segments.extend(_open_tag("message", attrs))560 segments.extend(_render_content_segments(message.get("content"), image_state))561 segments.extend(_close_tag("message"))562 segments.extend(_end_of_msg())563 elif role == "tool":564 tool_index += 1565 tool_name = message.get("tool", message.get("name"))566 if (567 tool_name is None568 and tool_calls is not None569 and tool_index <= len(tool_calls)570 ):571 tc = tool_calls[tool_index - 1]572 fn = tc.get("function", tc)573 tool_name = fn["name"]574 if tool_name is None:575 raise ValueError(576 "Kimi K3 tool messages need a resolvable tool name: "577 "carry `tool`/`name`, or match a preceding assistant "578 "tool_call by order."579 )580 segments.extend(581 _open_tag(582 "message",583 [("role", "tool"), ("tool", tool_name), ("index", tool_index)],584 )585 )586 segments.extend(_render_content_segments(message.get("content"), image_state))587 segments.extend(_close_tag("message"))588 segments.extend(_end_of_msg())589 elif role == "assistant":590 tool_calls = message.get("tool_calls")591 tool_index = 0592 attrs = [("role", "assistant")]593 if message.get("name"):594 attrs.append(("name", message["name"]))595 segments.extend(_open_tag("message", attrs))596 segments.extend(_render_assistant_segments(message, image_state, thinking))597 segments.extend(_close_tag("message"))598 segments.extend(_end_of_msg())599 600 tool_choice = kwargs.get("tool_choice")601 if tool_choice == "required":602 segments.extend(603 _internal_system_message(604 "tool-choice",605 "The system is invoked with `tool_choice=required`.\n"606 "You MUST call tools in the next message.",607 )608 )609 elif tool_choice == "none":610 segments.extend(611 _internal_system_message(612 "tool-choice",613 "The system is invoked with `tool_choice=none`.\n"614 "You MUST NOT call any tools in the next message.",615 )616 )617 618 rf = kwargs.get("response_format")619 rf_type = _get_value(rf, "type", rf) if isinstance(rf, dict) else rf620 if rf_type == "json_object":621 segments.extend(622 _internal_system_message(623 "response-format",624 "The system is invoked with `response_format=json_object`.\n"625 "Your response must be raw JSON data without markdown code "626 "blocks (```json) or any additional formatting.",627 )628 )629 elif rf_type == "json_schema":630 schema = _json_compact(kwargs.get("response_schema"))631 segments.extend(632 _internal_system_message(633 "response-format",634 "The system is invoked with `response_format=json_schema`.\n"635 "Your response must be raw JSON data without markdown code "636 "blocks (```json) or any additional formatting.\n"637 "The JSON data must match the following schema:\n"638 f"```json\n{schema}\n```",639 )640 )641 642 if add_generation_prompt:643 segments.extend(_open_tag("message", [("role", "assistant")]))644 segments.extend(_open_tag("think" if thinking else "response"))645 646 image_state.assert_consumed()647 return segments648