Update encoding_k3.py

encoding_k3.py

@@ -187,21 +187,117 @@ def deep_sort_dict(obj: Any) -> Any:
     return obj
 
 
-def normalize_tool_arguments(arguments: Any) -> tuple[dict[str, Any], Optional[str]]:
+# One normalized tool-call argument: (key, XTML type, rendered text). The
+# text keeps the original JSON literal for non-string values and is the
+# decoded string for string values.
+XtmlArgument = tuple[str, str, str]
+
+
+def _parse_arguments_object(s: str) -> list[XtmlArgument]:
+    """Parse a JSON object string one level deep.
+
+    Each top-level key-value pair yields one ``(key, type, text)`` triple. The
+    text keeps the original JSON literal for non-string values (``1e2`` stays
+    ``1e2``, ``[1,2]`` keeps its exact bytes) and is the decoded (unescaped)
+    string for string values. Nested values are never re-serialized. Anything
+    after the closing ``}`` is ignored. Raises ``ValueError`` on malformed
+    input (including a valid non-object JSON document).
+    """
+    idx = 0
+
+    def _skip_whitespaces() -> None:
+        nonlocal idx
+        while idx < len(s) and s[idx] in (" ", "\t", "\n", "\r"):
+            idx += 1
+
+    def _next_char() -> Optional[str]:
+        nonlocal idx
+        if idx < len(s):
+            c = s[idx]
+            idx += 1
+            return c
+        return None
+
+    _skip_whitespaces()
+    if _next_char() != "{":
+        raise ValueError("JSON arguments must be an object")
+    _skip_whitespaces()
+
+    parsed: list[XtmlArgument] = []
+
+    if idx >= len(s):
+        raise ValueError("Unexpected end of JSON object")
+    if s[idx] == "}":
+        return parsed
+
+    json_decoder = json.JSONDecoder(strict=False)
+
+    def _raw_decode(start_idx: int) -> tuple[Any, int]:
+        try:
+            return json_decoder.raw_decode(s, idx=start_idx)
+        except json.JSONDecodeError as e:
+            raise ValueError(str(e)) from e
+
+    while True:
+        decoded_key, idx = _raw_decode(idx)
+        if not isinstance(decoded_key, str):
+            raise ValueError(f"JSON object key must be a string, got {decoded_key!r}")
+        _skip_whitespaces()
+        if _next_char() != ":":
+            raise ValueError(f"Expects ':' after {decoded_key}")
+        _skip_whitespaces()
+
+        value_start_idx = idx
+        decoded_value, idx = _raw_decode(idx)
+        value_end_idx = idx
+
+        text = (
+            decoded_value
+            if isinstance(decoded_value, str)
+            else s[value_start_idx:value_end_idx]
+        )
+        parsed.append((decoded_key, _xtml_type(decoded_value), text))
+
+        _skip_whitespaces()
+        c = _next_char()
+        _skip_whitespaces()
+
+        if c == "}":
+            break
+        elif c == ",":
+            continue
+        else:
+            raise ValueError(f"Expect '}}' or ',', got {c!r}")
+
+    return parsed
+
+
+def normalize_tool_arguments(
+    arguments: Any,
+) -> tuple[list[XtmlArgument], Optional[str]]:
+    """Normalize tool call arguments for XTML rendering.
+
+    Returns ``(argument_triples, raw_json_block)``. String arguments are parsed
+    one level deep so non-string values keep their original JSON literal text,
+    and anything after the closing ``}`` is discarded; dict arguments (already
+    Python objects) are serialized on the spot. Any string that is not a
+    well-formed JSON object -- unparseable, whitespace only, or valid
+    non-object JSON -- falls back to the raw ``<json>`` block.
+    """
     if arguments is None:
-        return {}, None
+        return [], None
     if isinstance(arguments, dict):
-        return arguments, None
+        return [
+            (str(key), _xtml_type(value), _xtml_value(value))
+            for key, value in arguments.items()
+        ], None
     if isinstance(arguments, str):
-        if not arguments.strip():
-            return {}, None
+        if arguments == "":
+            return [], None
         try:
-            parsed = json.loads(arguments)
-        except json.JSONDecodeError:
-            return {}, arguments
-        if not isinstance(parsed, dict):
-            raise ValueError("Kimi K3 tool call arguments must be a JSON object.")
-        return parsed, None
+            return _parse_arguments_object(arguments), None
+        except ValueError:
+            return [], arguments
     raise TypeError(
         "Kimi K3 tool call arguments must be a dict or a JSON object string."
     )
@@ -421,7 +517,9 @@ def _render_assistant_segments(
             "reasoning"
         )
         segments.extend(_open_tag("think"))
-        if reasoning_content is not None and str(reasoning_content).strip():
+        # Only an empty string counts as no reasoning, so whitespace-only
+        # reasoning is still rendered into the think channel.
+        if reasoning_content is not None and str(reasoning_content) != "":
             _append_text(segments, reasoning_content, image_state)
         segments.extend(_close_tag("think"))
 
@@ -437,21 +535,18 @@ def _render_assistant_segments(
             segments.extend(
                 _open_tag("call", [("tool", fn["name"]), ("index", index)])
             )
-            args = fn.get("arguments", {})
+            args = fn.get("arguments", [])
             json_block = fn.get("_xtml_json_block")
             if json_block is not None:
                 segments.extend(_open_tag("json", [("type", "object")]))
                 _append_text(segments, json_block, image_state)
                 segments.extend(_close_tag("json"))
-            elif _is_mapping(args):
-                for key, value in args.items():
+            else:
+                for key, arg_type, arg_text in args:
                     segments.extend(
-                        _open_tag(
-                            "argument",
-                            [("key", key), ("type", _xtml_type(value))],
-                        )
+                        _open_tag("argument", [("key", key), ("type", arg_type)])
                     )
-                    _append_text(segments, _xtml_value(value), image_state)
+                    _append_text(segments, arg_text, image_state)
                     segments.extend(_close_tag("argument"))
             segments.extend(_close_tag("call"))
         segments.extend(_close_tag("tools"))
@@ -538,10 +633,14 @@ def build_chat_segments(
         )
 
     for message_index, message in enumerate(messages):
+        # Malformed messages are rejected instead of being silently skipped.
         if not isinstance(message, dict):
-            continue
+            raise ValueError(
+                f"Kimi K3 messages must be dicts, got {type(message).__name__} "
+                f"at index {message_index}."
+            )
 
-        role = message["role"]
+        role = message.get("role")
         if role == "user":
             attrs = [("role", "user")]
             if message.get("name"):
@@ -596,6 +695,10 @@ def build_chat_segments(
             segments.extend(_render_assistant_segments(message, image_state, thinking))
             segments.extend(_close_tag("message"))
             segments.extend(_end_of_msg())
+        else:
+            raise ValueError(
+                f"Unknown message role {role!r} at index {message_index}."
+            )
 
     tool_choice = kwargs.get("tool_choice")
     if tool_choice == "required":