ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: coerce tool_call arguments to dict before chat templating (fixes MLX tool follow-up error) (#6807)

Browse source 
* Studio: coerce tool_call arguments to dict before chat templating

Strict tool chat templates (e.g. mlx-community Qwen3.5 checkpoints) iterate
arguments.items() and raise "TypeError: Can only get item pairs from a mapping"
when a prior assistant tool call is re-rendered on the next turn. The agentic
loop stores arguments in the OpenAI JSON-string form (as_assistant_tool_call),
which is correct on the wire and for llama-server, but the transformers / MLX
paths apply_chat_template directly and hit the strict Jinja templates.

Normalize each assistant tool_call's function.arguments from a JSON string to a
dict inside apply_chat_template_for_generation (shared by both the MLX and
safetensors paths). A dict renders on strict and lenient templates alike;
non-JSON / non-dict values are left untouched, and the OpenAI-format
as_assistant_tool_call (used by the GGUF path + API responses) is unchanged.

Verified against the real mlx-community/Qwen3.5-2B-8bit template: string args
raised the tester's error, the fix renders cleanly, and the lenient
unsloth/Qwen3.5-0.8B template still works.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: make tool-arg coercion a string-first fallback (non-regressive)

Render the original OpenAI string-arg form first and only coerce arguments to a
dict when the template raises the mapping TypeError, instead of always coercing.
Any template that already renders is now byte-identical (a template that emits
arguments verbatim keeps the JSON string, not a Python dict repr).

Verified across Llama-3, Qwen2.5, Qwen3, Qwen3.5, Phi-3.5 (byte-identical) and
mlx-community/Qwen3.5-2B-8bit (strict -> fixed). Gemma-3 / Mistral tool-template
errors are unrelated (role alternation / tool-id length) and identical with or
without the change.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Make core.inference package init lazy so dependency-light helpers import standalone

Importing any core.inference submodule ran the package __init__, which
eagerly imported orchestrator and llama_cpp; both pull loggers ->
structlog (and httpx), so a dependency-light helper like
chat_template_helpers dragged in the full heavy stack and its unit test
failed to collect in a backend env without structlog. Defer those
imports to attribute access via PEP 562 __getattr__, mirroring the lazy
pattern already in core/__init__.py. The re-exports resolve unchanged on
first access.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Retry dict-coercion for strict templates that raise non-TypeError

apply_chat_template_for_generation only retried the OpenAI JSON-string arguments
coercion when the first render raised TypeError (the arguments.items() form). The
bundled gemma-4.jinja instead rejects string arguments with raise_exception, which
surfaces as a Jinja error, so a second tool turn with string function.arguments
propagated and failed rather than retrying with the parsed dict.

Broaden the outer catch to Exception, still gated on there being a string arg to
normalize (normalized is messages -> re-raise), so unrelated template errors and
templates that already render are unaffected.

* Tighten comments in tool-call argument coercion helper and tests

* Tighten tool-call argument coercion comments

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-07-06 10:12:22 -07:00 • committed by GitHub
commit e9ea45b6a5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 229 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

90
studio/backend/core/inference/chat_template_helpers.py
View file

@ -6,9 +6,51 @@ Dependency-light wrapper around tokenizer.apply_chat_template with a kwarg
fallback for templates that reject reasoning/tools args.
"""
import json
from typing import Optional
def _normalize_tool_call_arguments(messages: list) -> list:
"""Coerce each assistant ``tool_calls[].function.arguments`` from a JSON
string to a dict.
The OpenAI wire format carries ``arguments`` as a JSON string, but some chat
templates (e.g. the stricter Qwen tool templates shipped with mlx-community
checkpoints) iterate ``arguments.items()`` and raise
``TypeError: Can only get item pairs from a mapping.`` on the string form
when a prior tool call is re-rendered on the next turn. A dict works on both
strict and lenient templates, so parse the string; leave non-JSON or non-dict
values untouched. Returns the original list unchanged when nothing needed
coercing (no copy)."""
mutated = False
out: list = []
for msg in messages:
tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None
if not tool_calls:
out.append(msg)
continue
new_calls = []
msg_changed = False
for call in tool_calls:
fn = call.get("function") if isinstance(call, dict) else None
args = fn.get("arguments") if isinstance(fn, dict) else None
if isinstance(args, str):
try:
parsed = json.loads(args)
except (ValueError, TypeError):
parsed = None
if isinstance(parsed, dict):
call = {**call, "function": {**fn, "arguments": parsed}}
msg_changed = True
new_calls.append(call)
if msg_changed:
out.append({**msg, "tool_calls": new_calls})
mutated = True
else:
out.append(msg)
return out if mutated else messages
def apply_chat_template_for_generation(
tokenizer,
messages: list,
@ -38,21 +80,33 @@ def apply_chat_template_for_generation(
attempts.append(dict(reasoning_kwargs))
attempts.append({})
last_exc: Optional[Exception] = None
for kwargs in attempts:
try:
return tokenizer.apply_chat_template(
messages,
tokenize = False,
add_generation_prompt = True,
**kwargs,
)
except TypeError as e:
last_exc = e
continue
except Exception as e:
last_exc = e
break
if last_exc is not None:
raise last_exc
raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result")
def _render(msgs: list) -> str:
last_exc: Optional[Exception] = None
for kwargs in attempts:
try:
return tokenizer.apply_chat_template(
msgs,
tokenize = False,
add_generation_prompt = True,
**kwargs,
)
except TypeError as e:
last_exc = e
continue
except Exception as e:
last_exc = e
break
if last_exc is not None:
raise last_exc
raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result")
try:
return _render(messages)
except Exception:
# Strict tool templates reject the JSON-string ``arguments`` form via
# TypeError or a broad Jinja raise_exception, so retry with dicts coerced.
# Original messages render first, so working templates stay byte-identical.
normalized = _normalize_tool_call_arguments(messages)
if normalized is messages:
raise
return _render(normalized)

157
studio/backend/tests/test_chat_template_tool_arguments.py Normal file
View file

@ -0,0 +1,157 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""apply_chat_template_for_generation must coerce assistant tool_call arguments
from the OpenAI JSON-string form to a dict before rendering. Strict tool
templates (e.g. mlx-community Qwen3.5 checkpoints) iterate arguments.items() and
raise "Can only get item pairs from a mapping." on the string form when a prior
tool call is re-rendered on the next turn (MLX + transformers paths).
"""
from __future__ import annotations
import sys
from pathlib import Path
import pytest
_BACKEND = Path(__file__).resolve().parent.parent
if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
from core.inference.chat_template_helpers import ( # noqa: E402
_normalize_tool_call_arguments,
apply_chat_template_for_generation,
)
def _conv(arguments):
return [
{"role": "user", "content": "weather?"},
{
"role": "assistant",
"content": "",
"tool_calls": [
{
"type": "function",
"id": "c1",
"function": {"name": "web_search", "arguments": arguments},
}
],
},
{"role": "tool", "name": "web_search", "content": "21C sunny"},
]
class _StrictTemplateTokenizer:
"""Mimics a strict Qwen tool template: rejects string tool_call arguments."""
def apply_chat_template(
self,
messages,
*,
tokenize = False,
add_generation_prompt = True,
**kw,
):
for msg in messages:
for call in msg.get("tool_calls", []) or []:
args = call.get("function", {}).get("arguments")
if isinstance(args, str):
raise TypeError("Can only get item pairs from a mapping.")
return "RENDERED"
def test_string_arguments_are_parsed_to_dict():
out = _normalize_tool_call_arguments(_conv('{"query": "sweden"}'))
args = out[1]["tool_calls"][0]["function"]["arguments"]
assert args == {"query": "sweden"}
def test_dict_arguments_untouched_and_no_copy():
conv = _conv({"query": "sweden"})
assert _normalize_tool_call_arguments(conv) is conv
def test_non_json_string_left_as_is():
out = _normalize_tool_call_arguments(_conv("not json"))
assert out[1]["tool_calls"][0]["function"]["arguments"] == "not json"
def test_render_succeeds_on_strict_template_with_string_arguments():
# Regression: strict template + string args used to raise.
result = apply_chat_template_for_generation(_StrictTemplateTokenizer(), _conv('{"query": "x"}'))
assert result == "RENDERED"
class _RecordingTokenizer:
"""Lenient template: renders whatever arguments it is given (string or dict)."""
def __init__(self):
self.seen_arguments = None
def apply_chat_template(
self,
messages,
*,
tokenize = False,
add_generation_prompt = True,
**kw,
):
for msg in messages:
for call in msg.get("tool_calls", []) or []:
self.seen_arguments = call.get("function", {}).get("arguments")
return "RENDERED"
def test_lenient_template_receives_original_string_untouched():
# Lenient template must see the exact original string, not a coerced dict.
tok = _RecordingTokenizer()
apply_chat_template_for_generation(tok, _conv('{"query": "x"}'))
assert tok.seen_arguments == '{"query": "x"}'
def test_messages_without_tool_calls_pass_through_unchanged():
conv = [{"role": "user", "content": "hi"}]
assert _normalize_tool_call_arguments(conv) is conv
class _RaiseExceptionTemplateTokenizer:
"""Mimics the bundled gemma-4.jinja: rejects string tool_call arguments via
``raise_exception(...)``, which surfaces as a Jinja error, NOT a TypeError."""
def apply_chat_template(
self,
messages,
*,
tokenize = False,
add_generation_prompt = True,
**kw,
):
for msg in messages:
for call in msg.get("tool_calls", []) or []:
args = call.get("function", {}).get("arguments")
if isinstance(args, str):
raise ValueError(
"chat_template: tool_calls[].function.arguments must be a "
"JSON object (mapping), not a string."
)
return "RENDERED"
def test_render_succeeds_on_raise_exception_template_with_string_arguments():
# Regression: gemma-4.jinja rejects string args via a non-TypeError; retry must still coerce.
result = apply_chat_template_for_generation(
_RaiseExceptionTemplateTokenizer(), _conv('{"query": "x"}')
)
assert result == "RENDERED"
def test_unrelated_template_error_still_propagates_with_dict_args():
# Failure unrelated to string args (dict args, nothing to coerce) must propagate.
class _AlwaysRaises:
def apply_chat_template(self, messages, **kw):
raise ValueError("template is broken")
with pytest.raises(ValueError, match = "broken"):
apply_chat_template_for_generation(_AlwaysRaises(), _conv({"query": "x"}))