From 430610ca3dfdaaf4c5eb370fac63a920604f9ab0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 27 May 2026 13:15:24 +0000 Subject: [PATCH 1/6] Studio: re-introduce multi-format tool calling with parser bug fixes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resubmits the work from #5615 (reverted in #5619) with the parser / route bug fixes that were subsequently developed on #5620 folded in. The healing-parity package in #5620 -- the GGUF canonical heal key in `llama_cpp.py` and the safetensors_agentic re-prompt loop -- is deliberately left to #5620 so this redo stays scoped to the same 4 files that #5615 originally touched. Adds multi-format tool-call parsing for the safetensors / MLX agentic loop so Llama-3, Llama-3.2 bare JSON, Mistral pre-v11 / v11+ / Ministral, and Gemma 4 tool emissions are normalised to OpenAI shape instead of leaking as prose, plus a route-layer strip that removes the same shapes from streamed and non-streamed completions. Formats: Qwen / Hermes {json} Qwen3.5 / Hermes v Llama-3 built-in <|python_tag|>NAME.call(k="v", ...) Llama-3 custom <|python_tag|>{"name":..., "parameters":...} Llama-3.2 bare {"name":..., "parameters":...} (no marker) Mistral pre-v11 [TOOL_CALLS] [{"name":..., "arguments":...}, ...] Mistral v11+ [TOOL_CALLS]name{json} (may chain) Ministral / Large 3 [TOOL_CALLS]name[ARGS]{json} Gemma 4 <|tool_call>call:NAME{k:<|"|>v<|"|>} The four parser bugs that motivated the revert are fixed here: 1. Mistral nested-JSON truncation. The closed-pair Mistral regex `\[TOOL_CALLS\]...\{.*?\}` was non-greedy on `}`, so `[TOOL_CALLS]search{"filters":{"date":"2024"},"q":"foo"}` was stripped only up to the inner `}`, leaking `,"q":"foo"}` to the user. Replaced with `_strip_mistral_closed_calls` + the balanced-brace / balanced-bracket helpers that ignore braces inside JSON strings. 2. `<|python_tag|>` stop-on-`<`. The route-layer strip clause `<\|python_tag\|>[^\n<]*` stopped at any literal `<`, so `<|python_tag|>python.call(code="if x < 10: pass")` was sliced to `< 10: pass")`. Replaced with `<\|python_tag\|>(?:[^<]|<(?!\|))*` so the strip consumes any character that is not a Llama-3 `<|sentinel|>` start -- literal `<`, newlines, and embedded JSON all stay inside. 3. Llama-3 sentinel single-pass loop. The fixed-order `for sentinel in (...)` loop in the bare-JSON parser silently dropped calls when the stream contained `<|eot_id|><|begin_of_text|>{json}` because `begin_of_text` was tested before `eot_id` consumed its prefix. Replaced with a `while True / matched` loop so the order of sentinels in the stream no longer matters. 4. UTF-8 corruption in Llama-3 KV decoder. `bytes(s, "utf-8").decode("unicode_escape")` mangles non-ASCII bytes (`"café日本"` -> `'caféæ\x97¥æ\x9c¬'`). Replaced with `json.loads('"' + value + '"')` which handles `\n` / `\t` / `\uXXXX` escapes correctly while preserving literal UTF-8 bytes (emoji, CJK, etc.). `_TOOL_XML_RE` keeps the orphan-handling clauses that #5735 added for the speculative buffer leak shapes (closed pair OR orphan-open to EOF, bare orphan close, tail-only ``) so the route layer continues to strip in-flight tool markup as well as the multi-format closed pairs. The new `_strip_tool_xml(text)` helper composes `_TOOL_XML_RE` with `_strip_mistral_closed_calls` so the Mistral nested-JSON shape gets balanced-brace handling at every call site (8 sites updated). Capability gating in `_detect_safetensors_features` now allows templates whose tool-call format is any of the seven supported markers; the gate still suppresses `supports_tools` for templates that advertise tools but use a shape the parser cannot honour, so the UI never enables a pill the loop will not return. Tests in scope: - tests/test_safetensors_tool_loop.py: full multi-format parser coverage (Qwen/Hermes, Llama-3 python_tag and bare JSON, Mistral all variants, Gemma 4), plus `TestRoutesPythonTagStrip` (8 tests) pinning the multi-line / less-than-in-code / sentinel-stop behaviour of bug 2's regex. - tests/test_safetensors_capability_advertise.py: capability gate keeps tools enabled for Llama-3 / Mistral / Gemma 4 / Llama-3.2 bare-JSON templates while still suppressing tools for unknown emission formats. Tests deliberately out of scope (they belong to #5620 because they exercise `llama_cpp.py` / `safetensors_agentic.py`): - TestLoopRePrompt (6) -- safetensors_agentic re-prompt loop. - TestLoopCanonicalHealKey (3) -- canonical heal key under loop. - TestGGUFSafetensorsHealingParity (5) -- GGUF / safetensors parity assertions on shared constants and `_MAX_REPROMPTS`. `pytest studio/backend/tests/test_safetensors_tool_loop.py studio/backend/tests/test_safetensors_capability_advertise.py -q` -> 104 passed. --- .../core/inference/tool_call_parser.py | 866 +++++++++++++++--- studio/backend/routes/inference.py | 84 +- .../test_safetensors_capability_advertise.py | 83 +- .../tests/test_safetensors_tool_loop.py | 440 +++++++++ 4 files changed, 1313 insertions(+), 160 deletions(-) diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index a0ab8a2a53..9f603b7839 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -2,32 +2,57 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """ -Backend-neutral tool-call XML parser shared by GGUF and safetensors. -Tolerates missing closing tags in either ``{json}`` -or ``v...`` shape. +Backend-neutral tool-call parser shared by GGUF, safetensors, and MLX. + +Covers the emission formats so the safetensors + MLX agentic loop sees +the same call shape llama-server normalises for GGUF: + + - ``{json}`` (Qwen / Hermes) + - ``v`` (Qwen3.5 xml) + - ``<|python_tag|>NAME.call(k="v", ...)`` (Llama-3 built-in tools) + - ``<|python_tag|>{"name":..., "parameters":...}`` (Llama-3 custom) + - ``{"name":..., "parameters":...}`` (Llama-3.2 bare JSON) + - ``[TOOL_CALLS] [{...}, ...]`` (Mistral v0.3 / Nemo / Small) + - ``[TOOL_CALLS]name{json}`` (Mistral v11+ / Magistral) + - ``[TOOL_CALLS]name[ARGS]{json}`` (Ministral / Mistral Large 3) + - ``<|tool_call>call:NAME{k:<|"|>v<|"|>}`` (Gemma 4) + +Closing tags / brackets are tolerated when missing because models +frequently truncate them mid-stream. """ import json import re +from typing import Any -# _TOOL_CLOSED_PATS: closed pairs only. _TOOL_ALL_PATS: also trailing -# unclosed runs so truncated tails don't leak markup. +# Markers that flip the streaming buffer from STREAMING to DRAINING so +# partial markup never leaks before the parser sees it. +TOOL_XML_SIGNALS = ( + "", + "", + "[TOOL_CALLS]", + "<|tool_call>", +) + + +# Closed pairs only (mid-stream); _TOOL_ALL_PATS also eats unclosed +# tails for end-of-turn cleanup. _TOOL_CLOSED_PATS = [ re.compile(r".*?", re.DOTALL), re.compile(r".*?", re.DOTALL), + re.compile(r"<\|tool_call>.*?", re.DOTALL), ] _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ re.compile(r".*$", re.DOTALL), re.compile(r".*$", re.DOTALL), + re.compile(r"<\|tool_call>.*$", re.DOTALL), + re.compile(r"\[TOOL_CALLS\].*$", re.DOTALL), + re.compile(r"<\|python_tag\|>.*$", re.DOTALL), ] -# Prefixes the streaming buffer watches for to gate in-progress text. -TOOL_XML_SIGNALS = ("", "{json}``. _TC_JSON_START_RE = re.compile(r"\s*\{") -_TC_FUNC_START_RE = re.compile(r"\s*") +# Qwen3.5 / Hermes XML ``v``. +_TC_FUNC_START_RE = re.compile(r"\s*") _TC_END_TAG_RE = re.compile(r"") _TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") -_TC_PARAM_START_RE = re.compile(r"\s*") +_TC_PARAM_START_RE = re.compile(r"\s*") _TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") +# Llama-3 ``<|python_tag|>NAME.call(...)``. +_LLAMA3_PYTHON_TAG = "<|python_tag|>" +_LLAMA3_PY_CALL_RE = re.compile( + r"<\|python_tag\|>\s*([\w\.\-]+)\s*\.\s*call\s*\(", +) +_LLAMA3_KV_RE = re.compile( + r"""(\w+)\s*=\s*(?:"((?:\\.|[^"\\])*)"|(-?\d+(?:\.\d+)?)|(true|false|null))""", + re.VERBOSE, +) + +# Mistral ``[TOOL_CALLS]`` trigger. v11+ chains them, each followed by +# a bare name plus ``{json}`` (Magistral) or ``[ARGS]{json}`` (Ministral +# / Large 3). +_MISTRAL_TRIGGER = "[TOOL_CALLS]" +_MISTRAL_ARGS_MARKER = "[ARGS]" +_MISTRAL_V11_NAME_RE = re.compile(r"\s*([\w\.\-]+)\s*") + +# Gemma 4: ``<|tool_call>call:NAME{...}``, ``<|"|>`` wraps strings. +_GEMMA_TC_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w\.\-]+)\s*\{") +_GEMMA_STR_BEGIN = '<|"|>' +_GEMMA_STR_END = '<|"|>' +_GEMMA_TC_END = "" + + +def _balanced_bracket_end(text: str, start: int) -> int | None: + """Index of `]` matching `[` at ``text[start]``; ignores brackets + in JSON strings. None if unmatched.""" + if start >= len(text) or text[start] != "[": + return None + depth = 0 + in_string = False + esc = False + i = start + while i < len(text): + ch = text[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "[": + depth += 1 + elif ch == "]": + depth -= 1 + if depth == 0: + return i + i += 1 + return None + + +def _strip_mistral_closed_calls(text: str) -> str: + """Strip cleanly-closed ``[TOOL_CALLS]`` blocks (array, ``name{json}``, + or ``name[ARGS]{json}``) via balanced brace/bracket scanning. + + A non-greedy ``\\{.*?\\}`` would truncate at the first ``}`` and lose + nested JSON. Unclosed runs are left for ``final=True`` cleanup. + """ + n = len(text) + out = [] + cursor = 0 + while cursor < n: + idx = text.find(_MISTRAL_TRIGGER, cursor) + if idx == -1: + out.append(text[cursor:]) + break + out.append(text[cursor:idx]) + body_start = idx + len(_MISTRAL_TRIGGER) + i = body_start + while i < n and text[i] in " \t\n\r": + i += 1 + # Array shape: ``[TOOL_CALLS] [...]``. + if i < n and text[i] == "[": + end = _balanced_bracket_end(text, i) + if end is None: + # Truncated; let caller buffer / final-strip. + out.append(text[idx:]) + break + cursor = end + 1 + if text.startswith("", cursor): + cursor += len("") + continue + # Named shape: ``[TOOL_CALLS] name [ARGS]? { json }``. + name_match = _MISTRAL_V11_NAME_RE.match(text, i) + if not name_match: + out.append(text[idx:body_start]) + cursor = body_start + continue + i = name_match.end() + while i < n and text[i] in " \t\n\r": + i += 1 + if text.startswith(_MISTRAL_ARGS_MARKER, i): + i += len(_MISTRAL_ARGS_MARKER) + while i < n and text[i] in " \t\n\r": + i += 1 + if i >= n or text[i] != "{": + out.append(text[idx:i]) + cursor = i + continue + end = _balanced_brace_end(text, i) + if end is None: + out.append(text[idx:]) + break + cursor = end + 1 + return "".join(out) + def strip_tool_markup(text: str, *, final: bool = False) -> str: - """Strip tool-call XML from streamed text. - - ``final=False`` only removes closed pairs (used during streaming so - in-progress XML stays buffered). ``final=True`` also removes a - trailing unclosed run and trims the result. - """ + """Strip tool-call markup. ``final=False`` keeps in-progress + markup buffered; ``final=True`` also drops trailing unclosed runs + and trims.""" + text = _strip_mistral_closed_calls(text) pats = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS for pat in pats: text = pat.sub("", text) return text.strip() if final else text +def has_tool_signal(text: str) -> bool: + return any(s in text for s in TOOL_XML_SIGNALS) + + def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict]: - """Parse OpenAI-format ``tool_calls`` from model text. + """Return OpenAI-format tool calls. Tries each format and returns + as soon as one matches so we never double-count.""" + for parser in ( + _parse_tool_call_json, # Qwen / Hermes + _parse_function_xml, # Qwen3.5 / Hermes XML + _parse_llama3_python_tag, # Llama-3 + _parse_mistral_tool_calls, # Mistral + _parse_gemma_tool_calls, # Gemma 4 + ): + calls = parser(content, id_offset = id_offset) + if calls: + return calls - Returns a list of ``{"id", "type", "function": {"name", "arguments"}}`` - dicts. ``arguments`` is always a JSON string so callers can hand it - straight back into an OpenAI-style response. + # Llama-3.2 bare ``{"name":..., "parameters":...}``. Strict: only + # fires on content that starts with ``{`` and parses as the right + # shape, so plain prose stays untouched. + return _parse_llama3_bare_json(content, id_offset = id_offset) - Handles two shapes: - - JSON inside ```` tags: - ``{"name":"web_search","arguments":{"query":"..."}}`` - - XML-style function blocks: - ``v`` - - Closing tags (````, ````, ````) - are all optional since models frequently omit them. - """ - tool_calls: list[dict] = [] - - # Pattern 1: {json}. Balanced-brace scan that skips - # braces inside JSON strings. +def _parse_tool_call_json(content: str, *, id_offset: int) -> list[dict]: + out: list[dict] = [] for m in _TC_JSON_START_RE.finditer(content): - brace_start = m.end() - 1 # position of the opening { - depth, i = 0, brace_start + brace_start = m.end() - 1 + end = _balanced_brace_end(content, brace_start) + if end is None: + continue + try: + obj = json.loads(content[brace_start : end + 1]) + except (json.JSONDecodeError, ValueError): + continue + name = obj.get("name", "") + args = obj.get("arguments", {}) + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + args_str = json.dumps({"value": args}) + if not name: + continue + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + return out + + +def _parse_function_xml(content: str, *, id_offset: int) -> list[dict]: + out: list[dict] = [] + func_starts = list(_TC_FUNC_START_RE.finditer(content)) + for idx, fm in enumerate(func_starts): + func_name = fm.group(1) + body_start = fm.end() + next_func = ( + func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) + ) + end_tag = _TC_END_TAG_RE.search(content[body_start:]) + if end_tag: + body_end = body_start + end_tag.start() + else: + body_end = len(content) + body_end = min(body_end, next_func) + body = _TC_FUNC_CLOSE_RE.sub("", content[body_start:body_end]) + + args: dict = {} + param_starts = list(_TC_PARAM_START_RE.finditer(body)) + if len(param_starts) == 1: + pm = param_starts[0] + val = _TC_PARAM_CLOSE_RE.sub("", body[pm.end() :]) + args[pm.group(1)] = val.strip() + else: + for pidx, pm in enumerate(param_starts): + val_start = pm.end() + next_param = ( + param_starts[pidx + 1].start() + if pidx + 1 < len(param_starts) + else len(body) + ) + val = _TC_PARAM_CLOSE_RE.sub("", body[val_start:next_param]) + args[pm.group(1)] = val.strip() + + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": func_name, "arguments": json.dumps(args)}, + } + ) + return out + + +def _parse_llama3_python_tag(content: str, *, id_offset: int) -> list[dict]: + """Parse the four Llama-3 emissions: ``<|python_tag|>NAME.call(...)`` + (built-in), ``<|python_tag|>{"name":..., "parameters":...}`` (custom), + multi-call via ``; `` separators, ``parameters`` or ``arguments`` key. + """ + out: list[dict] = [] + if _LLAMA3_PYTHON_TAG not in content: + return out + + # 1. ``NAME.call(...)`` built-in form. + for m in _LLAMA3_PY_CALL_RE.finditer(content): + name = m.group(1) + i = m.end() + depth = 1 in_string = False - while i < len(content): + esc = False + while i < len(content) and depth > 0: ch = content[i] if in_string: - if ch == "\\" and i + 1 < len(content): - i += 2 - continue - if ch == '"': + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': in_string = False + else: + if ch == '"': + in_string = True + elif ch == "(": + depth += 1 + elif ch == ")": + depth -= 1 + if depth == 0: + break + i += 1 + body = content[m.end() : i] + args: dict[str, Any] = {} + for kv in _LLAMA3_KV_RE.finditer(body): + k = kv.group(1) + if kv.group(2) is not None: + # ``json.loads`` on a quoted string handles \n/\t/\uXXXX + # escapes correctly AND keeps literal UTF-8 bytes (emoji + # / CJK) intact -- the older ``bytes.decode('unicode_escape')`` + # path mangled non-ASCII. + try: + args[k] = json.loads('"' + kv.group(2) + '"') + except (json.JSONDecodeError, ValueError): + args[k] = kv.group(2) + elif kv.group(3) is not None: + v = kv.group(3) + args[k] = float(v) if "." in v else int(v) + elif kv.group(4) is not None: + args[k] = {"true": True, "false": False, "null": None}[kv.group(4)] + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": json.dumps(args)}, + } + ) + + # 2. ``<|python_tag|>{"name":..., "parameters":...}``. ``raw_decode`` + # peels multiple ``; ``-separated objects from one emission. + if not out: + decoder = json.JSONDecoder() + idx = content.find(_LLAMA3_PYTHON_TAG) + while idx >= 0: + search_from = idx + len(_LLAMA3_PYTHON_TAG) + cursor = search_from + while cursor < len(content): + brace = content.find("{", cursor) + if brace < 0: + break + # Stop at the next ``<|python_tag|>``. + next_tag = content.find(_LLAMA3_PYTHON_TAG, search_from, brace) + if next_tag >= 0: + break + try: + obj, end_offset = decoder.raw_decode(content[brace:]) + except (json.JSONDecodeError, ValueError): + cursor = brace + 1 + continue + if not isinstance(obj, dict): + cursor = brace + end_offset + continue + name = obj.get("name") or obj.get("function") or "" + args = ( + obj.get("parameters") + if "parameters" in obj + else obj.get("arguments", {}) + ) + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + args_str = json.dumps({"value": args}) + if name: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + cursor = brace + end_offset + idx = content.find(_LLAMA3_PYTHON_TAG, cursor) + return out + + +def _parse_llama3_bare_json(content: str, *, id_offset: int) -> list[dict]: + """Llama-3.2 ``custom_tools``: bare ``{"name":..., "parameters":{...}}`` + without ``<|python_tag|>``. Strict (must start with ``{`` after sentinel + strip; ``name`` non-empty; ``parameters`` or ``arguments`` is a dict) so + plain prose and tool-message echoes don't trigger.""" + out: list[dict] = [] + stripped = content.lstrip() + # Sentinels can chain in any order, so loop until none match. + _sentinels = ( + "<|begin_of_text|>", + "<|eot_id|>", + "<|start_header_id|>", + "<|end_header_id|>", + "<|eom_id|>", + ) + while True: + stripped = stripped.lstrip() + matched = False + for sentinel in _sentinels: + if stripped.startswith(sentinel): + stripped = stripped[len(sentinel) :] + matched = True + break + if not matched: + break + if not stripped.startswith("{"): + return out + + decoder = json.JSONDecoder() + cursor = 0 + n = len(stripped) + while cursor < n: + # Skip whitespace and the Llama-3 ``;`` inter-call separator. + while cursor < n and stripped[cursor] in " \t\n\r;": + cursor += 1 + if cursor >= n or stripped[cursor] != "{": + break + try: + obj, end_offset = decoder.raw_decode(stripped[cursor:]) + except (json.JSONDecodeError, ValueError): + break + if not isinstance(obj, dict): + break + name = obj.get("name") or obj.get("function") or "" + if not isinstance(name, str) or not name: + break + if "parameters" in obj: + args = obj.get("parameters") + elif "arguments" in obj: + args = obj.get("arguments") + else: + break + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + break + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + cursor += end_offset + return out + + +def _parse_mistral_tool_calls(content: str, *, id_offset: int) -> list[dict]: + """Parse all Mistral emissions: pre-v11 ``[TOOL_CALLS][...]`` / + ``[TOOL_CALLS]{...}`` and v11+ ``[TOOL_CALLS]name{json}`` / + ``[TOOL_CALLS]name[ARGS]{json}`` (parallel-friendly).""" + out: list[dict] = [] + idx = content.find(_MISTRAL_TRIGGER) + if idx < 0: + return out + + # Disambiguate the first occurrence: array (pre-v11), single object + # (pre-v11), or bare-name (v11+). + j = idx + len(_MISTRAL_TRIGGER) + k = j + while k < len(content) and content[k] in " \t\n\r": + k += 1 + if k >= len(content): + return out + + if content[k] == "[": + return _parse_mistral_array(content, k, id_offset) + + if content[k] == "{": + # Pre-v11 single ``{"name":...}``; fall through if it doesn't + # carry a ``name`` so v11+ handling still gets a chance. + end = _balanced_brace_end(content, k) + if end is not None: + try: + obj = json.loads(content[k : end + 1]) + if isinstance(obj, dict) and obj.get("name"): + _consume_mistral_call(content[k : end + 1], out, id_offset) + return out + except (json.JSONDecodeError, ValueError): + pass + + # v11+: walk every ``[TOOL_CALLS]``, parsing ``name{json}`` or + # ``name[ARGS]{json}`` after each trigger. + pos = idx + while pos >= 0: + cur = pos + len(_MISTRAL_TRIGGER) + nm = _MISTRAL_V11_NAME_RE.match(content, cur) + if not nm: + pos = content.find(_MISTRAL_TRIGGER, cur) + continue + name = nm.group(1) + after_name = nm.end() + if content.startswith(_MISTRAL_ARGS_MARKER, after_name): + after_name += len(_MISTRAL_ARGS_MARKER) + while after_name < len(content) and content[after_name] in " \t\n\r": + after_name += 1 + if after_name >= len(content) or content[after_name] != "{": + pos = content.find(_MISTRAL_TRIGGER, cur) + continue + end = _balanced_brace_end(content, after_name) + if end is None: + break + try: + args = json.loads(content[after_name : end + 1]) + except (json.JSONDecodeError, ValueError): + pos = content.find(_MISTRAL_TRIGGER, end + 1) + continue + if not isinstance(args, dict): + pos = content.find(_MISTRAL_TRIGGER, end + 1) + continue + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + pos = content.find(_MISTRAL_TRIGGER, end + 1) + return out + + +def _parse_mistral_array(content: str, start: int, id_offset: int) -> list[dict]: + """Pre-v11 ``[TOOL_CALLS] [{...}, ...]`` array form.""" + out: list[dict] = [] + j = start + depth = 0 + in_string = False + esc = False + while j < len(content): + ch = content[j] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "[": + depth += 1 + elif ch == "]": + depth -= 1 + if depth == 0: + break + j += 1 + body = content[start : j + 1] if depth == 0 else content[start:] + + try: + arr = json.loads(body) + if isinstance(arr, list): + for obj in arr: + if isinstance(obj, dict): + _consume_mistral_call(json.dumps(obj), out, id_offset) + return out + except (json.JSONDecodeError, ValueError): + pass + + # Healing path for unclosed arrays: walk objects by hand. + for m in re.finditer(r"\{", body): + end = _balanced_brace_end(body, m.start()) + if end is None: + continue + _consume_mistral_call(body[m.start() : end + 1], out, id_offset) + return out + + +def _consume_mistral_call(obj_text: str, out: list[dict], id_offset: int) -> None: + try: + obj = json.loads(obj_text) + except (json.JSONDecodeError, ValueError): + return + if not isinstance(obj, dict): + return + name = obj.get("name") or "" + args = obj.get("arguments") or {} + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + args_str = json.dumps({"value": args}) + if name: + out.append( + { + "id": obj.get("id") or f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + + +def _parse_gemma_tool_calls(content: str, *, id_offset: int) -> list[dict]: + """Gemma 4: ``<|tool_call>call:NAME{k:<|"|>v<|"|>, ...}``.""" + out: list[dict] = [] + for m in _GEMMA_TC_RE.finditer(content): + name = m.group(1) + body_start = m.end() - 1 + end_marker = content.find(_GEMMA_TC_END, body_start) + scan_end = end_marker if end_marker >= 0 else len(content) + end = _gemma_balanced_brace_end(content, body_start, scan_end) + if end is None: + continue + body = content[body_start + 1 : end] + try: + args = _gemma_parse_mapping_body(body) + except Exception: + args = {} + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": json.dumps(args)}, + } + ) + return out + + +def _balanced_brace_end(text: str, brace_pos: int) -> int | None: + """Index of `}` matching `{` at ``brace_pos``; ignores braces inside + JSON strings. None if unmatched.""" + if brace_pos >= len(text) or text[brace_pos] != "{": + return None + depth = 0 + in_string = False + esc = False + i = brace_pos + while i < len(text): + ch = text[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': in_string = True elif ch == "{": depth += 1 elif ch == "}": depth -= 1 if depth == 0: + return i + i += 1 + return None + + +def _gemma_balanced_brace_end(text: str, brace_pos: int, hard_stop: int) -> int | None: + """Like ``_balanced_brace_end`` but skips ``<|"|>`` strings and + matches `{`/`[` symmetrically.""" + if brace_pos >= len(text) or text[brace_pos] != "{": + return None + depth = 0 + i = brace_pos + while i < hard_stop: + if text.startswith(_GEMMA_STR_BEGIN, i): + close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + if close < 0: + return None + i = close + len(_GEMMA_STR_END) + continue + ch = text[i] + if ch == "{" or ch == "[": + depth += 1 + elif ch == "}" or ch == "]": + depth -= 1 + if depth == 0: + return i + i += 1 + return None + + +def _gemma_parse_value(text: str, i: int): + """Parse one Gemma arg value at ``i``; returns ``(value, next_index)``.""" + if text.startswith(_GEMMA_STR_BEGIN, i): + close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + if close < 0: + return text[i + len(_GEMMA_STR_BEGIN) :], len(text) + return text[i + len(_GEMMA_STR_BEGIN) : close], close + len(_GEMMA_STR_END) + if text[i] == "{": + end = _gemma_balanced_brace_end(text, i, len(text)) + if end is None: + return {}, len(text) + return _gemma_parse_mapping_body(text[i + 1 : end]), end + 1 + if text[i] == "[": + j, depth = i, 0 + while j < len(text): + if text.startswith(_GEMMA_STR_BEGIN, j): + k = text.find(_GEMMA_STR_END, j + len(_GEMMA_STR_BEGIN)) + if k < 0: + j = len(text) break + j = k + len(_GEMMA_STR_END) + continue + ch = text[j] + if ch == "[": + depth += 1 + elif ch == "]": + depth -= 1 + if depth == 0: + break + j += 1 + body = text[i + 1 : j] + items: list[Any] = [] + k = 0 + while k < len(body): + if body[k] in " \t\n\r,": + k += 1 + continue + v, k = _gemma_parse_value(body, k) + items.append(v) + return items, j + 1 + # Primitive: number / true/false/null / bare identifier. + end = i + while ( + end < len(text) + and text[end] not in ",}]" + and not text.startswith(_GEMMA_STR_BEGIN, end) + ): + end += 1 + raw = text[i:end].strip() + if raw == "true": + return True, end + if raw == "false": + return False, end + if raw == "null": + return None, end + try: + return int(raw), end + except ValueError: + pass + try: + return float(raw), end + except ValueError: + pass + return raw, end + + +def _gemma_parse_mapping_body(body: str) -> dict[str, Any]: + """Parse a Gemma argument mapping (content between `{` and `}`).""" + out: dict[str, Any] = {} + i = 0 + n = len(body) + while i < n: + while i < n and body[i] in " \t\n\r,": i += 1 - if depth == 0: - json_str = content[brace_start : i + 1] - try: - obj = json.loads(json_str) - tc = { - "id": f"call_{id_offset + len(tool_calls)}", - "type": "function", - "function": { - "name": obj.get("name", ""), - "arguments": obj.get("arguments", {}), - }, - } - if isinstance(tc["function"]["arguments"], dict): - tc["function"]["arguments"] = json.dumps( - tc["function"]["arguments"] - ) - tool_calls.append(tc) - except (json.JSONDecodeError, ValueError): - pass - - # Pattern 2: v... -- closing tags - # optional; don't use as body boundary because code - # values can contain that literal. - if not tool_calls: - func_starts = list(_TC_FUNC_START_RE.finditer(content)) - for idx, fm in enumerate(func_starts): - func_name = fm.group(1) - body_start = fm.end() - next_func = ( - func_starts[idx + 1].start() - if idx + 1 < len(func_starts) - else len(content) - ) - end_tag = _TC_END_TAG_RE.search(content[body_start:]) - if end_tag: - body_end = body_start + end_tag.start() - else: - body_end = len(content) - body_end = min(body_end, next_func) - body = content[body_start:body_end] - body = _TC_FUNC_CLOSE_RE.sub("", body) - - arguments: dict = {} - param_starts = list(_TC_PARAM_START_RE.finditer(body)) - if len(param_starts) == 1: - # Single param: take everything to body end so - # embedded in code strings is preserved. - pm = param_starts[0] - val = body[pm.end() :] - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[pm.group(1)] = val.strip() - else: - for pidx, pm in enumerate(param_starts): - param_name = pm.group(1) - val_start = pm.end() - next_param = ( - param_starts[pidx + 1].start() - if pidx + 1 < len(param_starts) - else len(body) - ) - val = body[val_start:next_param] - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[param_name] = val.strip() - - tc = { - "id": f"call_{id_offset + len(tool_calls)}", - "type": "function", - "function": { - "name": func_name, - "arguments": json.dumps(arguments), - }, - } - tool_calls.append(tc) - - return tool_calls - - -def has_tool_signal(text: str) -> bool: - """Return True if ``text`` contains any tool-call XML signal.""" - return any(s in text for s in TOOL_XML_SIGNALS) + if i >= n: + break + if body.startswith(_GEMMA_STR_BEGIN, i): + close = body.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + if close < 0: + break + key = body[i + len(_GEMMA_STR_BEGIN) : close] + i = close + len(_GEMMA_STR_END) + else: + kstart = i + while i < n and body[i] != ":": + i += 1 + key = body[kstart:i].strip() + while i < n and body[i] in " \t\n\r": + i += 1 + if i < n and body[i] == ":": + i += 1 + while i < n and body[i] in " \t\n\r": + i += 1 + if i >= n: + out[key] = None + break + v, i = _gemma_parse_value(body, i) + out[key] = v + return out diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 92498b8a9a..abbc9d7a50 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -260,16 +260,24 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: "supports_tools": False, } ) - # Our safetensors loop only parses {json} - # and .... Llama uses <|python_tag|>, - # Mistral uses [TOOL_CALLS]; advertising tools for those would - # enable a pill the parser cannot honour. GGUF is unaffected -- - # llama-server normalises every format into structured deltas. + # Markers the safetensors / MLX parser recognises. If the template + # advertises tools but uses none of them, drop the pill (parser + # can't honour the emission). The two ``{"name":`` variants cover + # Llama-3.2 ``custom_tools`` whose template prompts the bare-JSON + # form without a ``<|python_tag|>`` prefix. + _PARSER_MARKERS = ( + "", + "", + "[TOOL_CALLS]", + "<|tool_call>", + '{"name":', + '{\\"name\\":', + ) if ( flags.get("supports_tools") and chat_template - and "" not in chat_template - and "...` / `...` -# 2. orphan opening to EOF (close was DRAINED) -# 3. bare orphan close (open was DRAINED) -# 4. tail-only `` (outer close truncated by EOS); anchored to -# `\Z` so mid-text `` in user code samples survives. +# Strip leaked tool-call markup. Covers every shared-parser format AND +# the four leak shapes the speculative buffer in ``llama_cpp.py`` splits +# across the visible/DRAIN boundary (closed pair, orphan open to EOF, +# bare orphan close, tail-only ````). Mistral ``[TOOL_CALLS]`` +# is delegated to the parser's balanced-brace helper -- a non-greedy +# ``\{.*?\}`` here would truncate nested JSON at the first ``}``. _TOOL_XML_RE = _re.compile( - r"<(?:tool_call|function=\w+)>.*?(?:|\Z)" - r"|" - r"|\s*\Z", + "|".join( + [ + # Tool-call / function XML: closed pair OR orphan open to EOF. + r"<(?:tool_call|function=\w+)>.*?(?:|\Z)", + # Bare orphan close (open was DRAINED upstream). + r"", + # Gemma 4. + r"<\|tool_call>.*?", + # Llama-3 ``<|python_tag|>...`` to the next ``<|`` sentinel + # or EOF. ``(?:[^<]|<(?!\|))*`` (not ``[^\n<]*`` or + # ``[^\n]*``) keeps literal ``<``, newlines, and embedded + # JSON inside the strip. + r"<\|python_tag\|>(?:[^<]|<(?!\|))*", + # Tail-only ```` (anchored so mid-text survives). + r"\s*\Z", + ] + ), _re.DOTALL, ) + + +def _strip_tool_xml(text: str) -> str: + """Combine the Mistral balanced-brace helper with ``_TOOL_XML_RE``.""" + from studio.backend.core.inference.tool_call_parser import ( + _strip_mistral_closed_calls, + ) + + return _TOOL_XML_RE.sub("", _strip_mistral_closed_calls(text)) + + logger = get_logger(__name__) @@ -2736,7 +2768,7 @@ async def openai_chat_completions( if _msg.get("role") == "assistant" and isinstance( _msg.get("content"), str ): - _msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip() + _msg["content"] = _strip_tool_xml(_msg["content"]).strip() def gguf_generate_with_tools(): return llama_backend.generate_chat_completion_with_tools( @@ -2837,7 +2869,7 @@ async def openai_chat_completions( # the last sanitized snapshot so cross-chunk XML # tags are handled correctly. raw_cumulative = event.get("text", "") - clean_cumulative = _TOOL_XML_RE.sub("", raw_cumulative) + clean_cumulative = _strip_tool_xml(raw_cumulative) new_text = clean_cumulative[len(prev_text) :] prev_text = clean_cumulative if not new_text: @@ -3222,7 +3254,7 @@ async def openai_chat_completions( _sf_chat_messages.append( { **_msg, - "content": _TOOL_XML_RE.sub("", _msg["content"]).strip(), + "content": _strip_tool_xml(_msg["content"]).strip(), } ) else: @@ -3309,7 +3341,7 @@ async def openai_chat_completions( # Diff cumulative cleaned text against last snapshot. raw_cumulative = event.get("text", "") - clean_cumulative = _TOOL_XML_RE.sub("", raw_cumulative) + clean_cumulative = _strip_tool_xml(raw_cumulative) new_text = clean_cumulative[len(prev_text) :] prev_text = clean_cumulative if not new_text: @@ -3381,7 +3413,7 @@ async def openai_chat_completions( if cancel_event.is_set(): break if event.get("type") == "content": - full_text = _TOOL_XML_RE.sub("", event.get("text", "")) + full_text = _strip_tool_xml(event.get("text", "")) return full_text content_text = await asyncio.to_thread(_drain_to_text) @@ -4913,7 +4945,7 @@ async def anthropic_messages( # Strip stale tool-call XML from conversation for _msg in openai_messages: if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): - _msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip() + _msg["content"] = _strip_tool_xml(_msg["content"]).strip() def _run_tool_gen(): return llama_backend.generate_chat_completion_with_tools( @@ -5005,7 +5037,7 @@ async def _anthropic_tool_stream( # Strip leaked tool-call XML from content events if event.get("type") == "content": event = dict(event) - event["text"] = _TOOL_XML_RE.sub("", event["text"]) + event["text"] = _strip_tool_xml(event["text"]) for line in emitter.feed(event): yield line except Exception as e: @@ -5096,7 +5128,7 @@ async def _anthropic_tool_non_streaming(run_gen, message_id, model_name): etype = event.get("type", "") if etype == "content": # Strip leaked tool-call XML - clean = _TOOL_XML_RE.sub("", event["text"]) + clean = _strip_tool_xml(event["text"]) new = clean[len(prev_text) :] prev_text = clean if new: @@ -5433,7 +5465,7 @@ async def _anthropic_passthrough_non_streaming( content_blocks = [] text = message.get("content") or "" if text: - text = _TOOL_XML_RE.sub("", text).strip() + text = _strip_tool_xml(text).strip() if text: content_blocks.append(AnthropicResponseTextBlock(text = text)) diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index c3ee5b9ff1..b63e835d2a 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -129,11 +129,11 @@ def test_detect_safetensors_features_gptoss_disables_tools(): assert flags["supports_tools"] is False -# Llama-3 / Mistral templates advertise tool handling but the model emits -# tool calls in <|python_tag|> / [TOOL_CALLS] format -- not the -# / , [TOOL_CALLS], and +# <|tool_call>). The route helper must surface supports_tools=True for +# all of them so the UI enables the pill. Only templates whose tool +# format is NONE of the five known markers should be suppressed. LLAMA3_TEMPLATE = """ {%- if tools %} @@ -165,27 +165,88 @@ MISTRAL_TEMPLATE = """ {%- endfor %} """ +GEMMA4_TEMPLATE = """ +{%- if tools %} + {{- 'Tools available. Emit calls as ' }} + {{- '<|tool_call>call:NAME{key:<|"|>val<|"|>}' }} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +""" -def test_detect_safetensors_features_llama3_template_suppresses_tools(): - """Llama-3 emits <|python_tag|>; safetensors loop cannot parse it.""" + +def test_detect_safetensors_features_llama3_template_keeps_tools_on(): + """Llama-3 emits <|python_tag|>; parser now supports it.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") flags = _detect_safetensors_features(backend, LLAMA3_TEMPLATE) - assert flags["supports_tools"] is False + assert flags["supports_tools"] is True -def test_detect_safetensors_features_mistral_template_suppresses_tools(): - """Mistral emits [TOOL_CALLS]; safetensors loop cannot parse it.""" +def test_detect_safetensors_features_mistral_template_keeps_tools_on(): + """Mistral emits [TOOL_CALLS]; parser now supports it.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/mistral-7b-instruct-v0.3") flags = _detect_safetensors_features(backend, MISTRAL_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_gemma4_template_keeps_tools_on(): + """Gemma 4 emits <|tool_call>; parser now supports it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/gemma-4-E2B-it-UD-MLX-4bit") + flags = _detect_safetensors_features(backend, GEMMA4_TEMPLATE) + assert flags["supports_tools"] is True + + +LLAMA3_2_BARE_JSON_TEMPLATE = """ +{%- if tools %} + {{- 'Given the following functions, respond with JSON for a function call.' }} + {{- 'Respond in the format {"name": function name, "parameters": dictionary}.' }} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +{%- for message in messages %} + {%- if 'tool_calls' in message %} + {{- '{"name": "' + message.tool_calls[0].function.name + '", '}} + {{- '"parameters": ' + (message.tool_calls[0].function.arguments | tojson) + '}' }} + {%- endif %} +{%- endfor %} +""" + + +def test_detect_safetensors_features_llama3_2_bare_json_keeps_tools_on(): + """Llama-3.2 emits bare JSON ``{"name":..., "parameters":...}`` -- the + parser now handles that path, so the pill must stay enabled.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, LLAMA3_2_BARE_JSON_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_unknown_format_suppresses_tools(): + """A template that advertises tools but uses no known marker must + be suppressed so the UI does not enable an unsupported pill.""" + from routes.inference import _detect_safetensors_features + + tpl = ( + "{%- if tools %}<|im_start|>system\n" + "Emit tool calls as JSON-RPC notifications inside the response." + "<|im_end|>{%- endif %}" + ) + backend = SimpleNamespace(active_model_name = "custom/unknown-tool-format") + flags = _detect_safetensors_features(backend, tpl) assert flags["supports_tools"] is False def test_detect_safetensors_features_qwen_tool_call_keeps_tools_on(): - """Sanity check: gate only suppresses non-Qwen formats.""" + """Sanity check: Qwen marker still flips supports_tools.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/Qwen3-0.6B") diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 923af87c4f..a5e8413a1c 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -130,6 +130,292 @@ class TestParser: assert "partial" in strip_tool_markup(text) +class TestParserMultiFormat: + """Parser coverage for Llama-3 / Mistral / Gemma 4 emission formats. + + Each model family upstream of GGUF emits a different tool-call + shape. The shared parser must turn all of them into the same + OpenAI ``{name, arguments}`` shape so the safetensors / MLX + agentic loop is family-agnostic. + """ + + # ── Llama-3 ──────────────────────────────────────────────────── + + def test_llama3_python_tag_dot_call(self): + # Llama-3 built-in tools: <|python_tag|>NAME.call(k="v", ...). + import json + + text = '<|python_tag|>brave_search.call(query="weather in Tokyo")' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "brave_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "weather in Tokyo"} + + def test_llama3_python_tag_dot_call_multi_arg(self): + import json + + text = ( + "<|python_tag|>get_weather.call(" + 'location="Tokyo", units="celsius", days=5)' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"location": "Tokyo", "units": "celsius", "days": 5} + + def test_llama3_python_tag_json_form(self): + import json + + text = ( + '<|python_tag|>{"name":"web_search",' '"parameters":{"query":"hi","n":5}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "hi", "n": 5} + + def test_llama3_python_tag_json_form_with_eom(self): + # Llama-3 emits ``<|eom_id|>`` after the JSON; must not break parsing. + import json + + text = ( + '<|python_tag|>{"name":"python",' + '"parameters":{"code":"print(2+2)"}}<|eom_id|>' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"code": "print(2+2)"} + + def test_llama3_strip_markup_final(self): + text = '<|python_tag|>brave_search.call(query="x")' + assert strip_tool_markup(text, final = True) == "" + + # ── Llama-3.2 bare JSON ``custom_tools`` ───────────────────── + + def test_llama3_2_bare_json_parameters(self): + # Llama-3.2-Instruct emits bare JSON directly as content; no + # <|python_tag|> prefix per its training template. + import json + + text = '{"name":"web_search","parameters":{"query":"Tokyo weather"}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "Tokyo weather"} + + def test_llama3_2_bare_json_arguments_key(self): + import json + + text = '{"name":"add","arguments":{"a":1,"b":2}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"a": 1, "b": 2} + + def test_llama3_2_bare_json_multi_call(self): + # Llama-3 may chain calls with ``; `` per training template. + text = '{"name":"a","parameters":{}}; ' '{"name":"b","parameters":{}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_llama3_2_bare_json_with_eom_sentinel(self): + text = '{"name":"x","parameters":{"y":1}}<|eom_id|>' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "x" + + def test_llama3_2_bare_json_leading_sentinel_skipped(self): + # Sometimes prior <|eot_id|> leaks into the next turn. + text = '<|eot_id|>{"name":"x","parameters":{}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "x" + + def test_llama3_2_bare_json_plain_prose_does_not_fire(self): + # Defensive: must NOT fire on plain assistant prose. + text = "Hello world, how are you today?" + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_embedded_in_prose_does_not_fire(self): + # Defensive: JSON embedded in prose must NOT fire (parser is + # strict about content STARTING with `{`). + text = 'The tool result was: {"name":"foo"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_missing_name_does_not_fire(self): + text = '{"result":"ok","data":[1,2,3]}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_missing_args_does_not_fire(self): + text = '{"name":"x"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_args_not_dict_does_not_fire(self): + text = '{"name":"x","parameters":42}' + assert parse_tool_calls_from_text(text) == [] + + # ── Mistral pre-v11 ─────────────────────────────────────────── + + def test_mistral_pre_v11_array(self): + import json + + text = ( + '[TOOL_CALLS] [{"name":"web_search",' + '"arguments":{"query":"hello"},"id":"abc"}]' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + # Mistral provides its own id; preserve it. + assert result[0]["id"] == "abc" + assert json.loads(result[0]["function"]["arguments"]) == {"query": "hello"} + + def test_mistral_pre_v11_array_multi(self): + text = ( + '[TOOL_CALLS] [{"name":"a","arguments":{"x":1},"id":"id1"},' + '{"name":"b","arguments":{"y":2},"id":"id2"}]' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_mistral_pre_v11_unclosed_array(self): + # Closing ``]`` truncated -- parser must heal off individual objects. + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"q":"x"},"id":"id"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + # ── Mistral v11+ ─────────────────────────────────────────────── + + def test_mistral_v11_single(self): + # Magistral / Mistral Small 3.1: bare ``name{json}`` after trigger. + import json + + text = '[TOOL_CALLS]add{"a":3.5,"b":4}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "add" + assert json.loads(result[0]["function"]["arguments"]) == {"a": 3.5, "b": 4} + + def test_mistral_v11_parallel(self): + # v11+ parallel: ``[TOOL_CALLS]a{...}[TOOL_CALLS]b{...}``. + text = '[TOOL_CALLS]add{"a":1}[TOOL_CALLS]sub{"b":2}' + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "add" + assert result[1]["function"]["name"] == "sub" + + def test_mistral_v11_with_args_marker(self): + # Ministral / Mistral Large 3: ``[TOOL_CALLS]name[ARGS]{json}``. + import json + + text = '[TOOL_CALLS]add[ARGS]{"a":1,"b":2}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "add" + assert json.loads(result[0]["function"]["arguments"]) == {"a": 1, "b": 2} + + def test_mistral_strip_markup_v11(self): + text = '[TOOL_CALLS]add{"a":1}' + assert strip_tool_markup(text, final = True) == "" + + # ── Gemma 4 ─────────────────────────────────────────────────── + + def test_gemma4_simple_call(self): + import json + + text = ( + "<|tool_call>call:get_weather{" + 'location:<|"|>Tokyo<|"|>,units:<|"|>celsius<|"|>}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"location": "Tokyo", "units": "celsius"} + + def test_gemma4_with_primitives(self): + import json + + text = ( + "<|tool_call>call:set_pref{" + "enabled:true,attempts:5,threshold:1.5,nickname:null}" + ) + result = parse_tool_calls_from_text(text) + args = json.loads(result[0]["function"]["arguments"]) + assert args == { + "enabled": True, + "attempts": 5, + "threshold": 1.5, + "nickname": None, + } + + def test_gemma4_nested_args(self): + # Gemma 4 nests dicts / lists with bare keys and ``<|"|>`` strings. + import json + + text = ( + "<|tool_call>call:search{" + 'query:<|"|>foo<|"|>,filters:{site:<|"|>example.com<|"|>,recent:true},' + 'tags:[<|"|>a<|"|>,<|"|>b<|"|>]}' + ) + result = parse_tool_calls_from_text(text) + args = json.loads(result[0]["function"]["arguments"]) + assert args["query"] == "foo" + assert args["filters"] == {"site": "example.com", "recent": True} + assert args["tags"] == ["a", "b"] + + def test_gemma4_multi_call(self): + text = ( + "<|tool_call>call:a{x:1}" "<|tool_call>call:b{y:2}" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_gemma4_unclosed_does_not_raise(self): + # Truncated mid-stream; must not raise. + text = '<|tool_call>call:foo{x:<|"|>bar<|"|>' + result = parse_tool_calls_from_text(text) + assert isinstance(result, list) + + def test_gemma4_strip_markup_final(self): + text = "<|tool_call>call:foo{x:1}" + assert strip_tool_markup(text, final = True) == "" + + # ── Cross-format sentinels ──────────────────────────────────── + + def test_all_markers_in_tool_xml_signals(self): + # Streaming buffer wakes up on every emission marker. + from core.inference.tool_call_parser import TOOL_XML_SIGNALS + + for marker in ( + "", + "", + "[TOOL_CALLS]", + "<|tool_call>", + ): + assert ( + marker in TOOL_XML_SIGNALS + ), f"streaming loop would not wake on {marker!r}" + + def test_has_tool_signal_for_all_formats(self): + assert has_tool_signal('<|python_tag|>brave_search.call(q="x")') + assert has_tool_signal('[TOOL_CALLS] [{"name":"x"}]') + assert has_tool_signal('[TOOL_CALLS]add{"a":1}') + assert has_tool_signal("<|tool_call>call:foo{}") + + # ──────────────────────────────────────────────────────────────────── # run_safetensors_tool_loop # ──────────────────────────────────────────────────────────────────── @@ -280,6 +566,71 @@ class TestLoopBasic: contents = [e for e in events if e["type"] == "content"] assert "Result: 1" in contents[-1]["text"] + def test_llama3_python_tag_form(self): + # The agentic loop must recognise Llama-3's <|python_tag|> + # marker, drain the rest of the turn, and execute the call. + loop, exec_fn = _make_loop( + turns = [ + [ + "<|python_tag|>web_search.call(", + 'query="weather in Tokyo"', + ")", + ], + ["The weather is sunny."], + ], + exec_results = ["Sunny, 22C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather in Tokyo"})] + contents = [e for e in events if e["type"] == "content"] + assert "sunny" in contents[-1]["text"].lower() + + def test_mistral_pre_v11_form(self): + # Pre-v11 Mistral emission: ``[TOOL_CALLS] [{...}]``. + loop, exec_fn = _make_loop( + turns = [ + [ + '[TOOL_CALLS] [{"name":"web_search",', + '"arguments":{"query":"hi"},"id":"abc"}]', + ], + ["done"], + ], + exec_results = ["ok"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "hi"})] + # Mistral-provided ids must propagate to tool_start events. + tool_start = next(e for e in events if e["type"] == "tool_start") + assert tool_start["tool_call_id"] == "abc" + + def test_mistral_v11_form(self): + # v11+ Mistral emission: bare ``name{json}`` after the trigger. + loop, exec_fn = _make_loop( + turns = [ + ['[TOOL_CALLS]web_search{"query":"hi"}'], + ["done"], + ], + exec_results = ["ok"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "hi"})] + + def test_gemma4_form(self): + # Gemma 4 emission: ``<|tool_call>call:NAME{...}``. + loop, exec_fn = _make_loop( + turns = [ + [ + "<|tool_call>call:web_search{", + 'query:<|"|>weather<|"|>', + "}", + ], + ["sunny"], + ], + exec_results = ["Sunny, 22C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather"})] + def test_truncated_unclosed_tool_call(self): loop, exec_fn = _make_loop( turns = [ @@ -784,5 +1135,94 @@ class TestGptOssNameDetection: assert is_gpt_oss_model_name(None) is False +# ──────────────────────────────────────────────────────────────────── +# Routes-level python_tag strip (multi-line; stop on next sentinel) +# ──────────────────────────────────────────────────────────────────── + + +class TestRoutesPythonTagStrip: + """Earlier revisions of ``_TOOL_XML_RE`` in + ``studio.backend.routes.inference`` used either ``[^\\n<]*`` (5615 -- + leaked the tail of any tool call whose argument contained a literal + ``<`` like ``code="if x < 10"``) or ``[^\\n]*`` (5620 round one -- + single-line only, so the second line of + ``python.call(code="line1\\nline2")`` leaked). The current pattern + ``(?:[^<]|<(?!\\|))*`` consumes any character that is not a Llama-3 + ``<|`` sentinel start, so multi-line code, embedded JSON, and bare + ``<`` characters in code all stay inside the strip. + + The fully resolved strip is also exposed via + ``strip_tool_markup(text, final=True)`` in the parser; the + streaming path's routes-level strip is the regression-prone one + because it runs on every cumulative emission while content is + still arriving. + """ + + def _strip(self, text: str) -> str: + # Import inside the test so a routes-module import error does + # not blow up the entire test file at collection time. + from routes.inference import _strip_tool_xml + + return _strip_tool_xml(text) + + def test_single_line_python_tag_stripped(self): + # Floor: the original 5620 single-line behaviour still works. + text = '<|python_tag|>brave_search.call(query="weather")' + assert self._strip(text) == "" + + def test_python_tag_with_less_than_in_code(self): + # 5615 regression: literal ``<`` inside code must NOT terminate + # the strip early. + text = '<|python_tag|>python.call(code="if x < 10: pass")' + assert self._strip(text) == "" + + def test_python_tag_multiline_code_stripped(self): + # 5620 round-1 regression: multi-line code's second line leaked. + text = '<|python_tag|>python.call(code="line1\nline2\nline3")' + assert self._strip(text) == "" + + def test_python_tag_multiline_with_less_than(self): + # Combined: multi-line code AND literal ``<`` in code. + text = ( + '<|python_tag|>python.call(code="for i in range(10):\n' + " if i < 5:\n" + ' print(i)")' + ) + assert self._strip(text) == "" + + def test_python_tag_stops_at_eom_sentinel(self): + # Strip stops at the next Llama-3 ``<|`` sentinel so any + # trailing assistant content survives. + text = ( + '<|python_tag|>python.call(code="multi\nline")' + "<|eom_id|>final answer text" + ) + assert self._strip(text) == "<|eom_id|>final answer text" + + def test_python_tag_stops_at_eot_sentinel(self): + text = '<|python_tag|>brave_search.call(query="x")' "<|eot_id|>after" + assert self._strip(text) == "<|eot_id|>after" + + def test_python_tag_json_form_multiline_stripped(self): + # The JSON form of python_tag with newlines inside string args. + text = ( + '<|python_tag|>{"name":"python",' + '"parameters":{"code":"a = 1\nb = 2\nprint(a+b)"}}' + ) + assert self._strip(text) == "" + + def test_python_tag_with_eom_then_trailing_python_tag(self): + # Two python_tag emissions back-to-back across a sentinel: both + # should strip independently. + text = ( + '<|python_tag|>brave_search.call(query="a")' + "<|eom_id|>" + '<|python_tag|>python.call(code="x=1")' + ) + # ``<|eom_id|>`` between the two strips remains; both + # python_tag blocks are fully consumed. + assert self._strip(text) == "<|eom_id|>" + + if __name__ == "__main__": pytest.main([__file__, "-v"]) From 4d69577d9761159058fd55ea2df7252c8056a946 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 27 May 2026 13:16:18 +0000 Subject: [PATCH 2/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_cpu_threads.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/studio/backend/tests/test_cpu_threads.py b/studio/backend/tests/test_cpu_threads.py index 0dbcbdb74b..1224941622 100644 --- a/studio/backend/tests/test_cpu_threads.py +++ b/studio/backend/tests/test_cpu_threads.py @@ -63,16 +63,18 @@ def test_cpu_thread_cap_is_opt_in(raw): # Anything that is not a positive integer raises a clear ValueError. -@pytest.mark.parametrize("raw", ["zero", "0", "-3", "1.5", "abc", "8a", "0x4", "1e3", "4 0"]) +@pytest.mark.parametrize( + "raw", ["zero", "0", "-3", "1.5", "abc", "8a", "0x4", "1e3", "4 0"] +) def test_cpu_thread_cap_requires_positive_integer(raw): - with pytest.raises(ValueError, match="must be a positive integer"): + with pytest.raises(ValueError, match = "must be a positive integer"): configure_cpu_threads({"UNSLOTH_CPU_THREADS": raw}) # env=None path uses real os.environ (production call from run.py / main.py). def test_cpu_thread_cap_uses_os_environ_when_env_is_none(monkeypatch): for variable in (*_THREAD_POOL_ENV_VARS, "UNSLOTH_CPU_THREADS"): - monkeypatch.delenv(variable, raising=False) + monkeypatch.delenv(variable, raising = False) monkeypatch.setenv("UNSLOTH_CPU_THREADS", "3") configure_cpu_threads() @@ -84,7 +86,7 @@ def test_cpu_thread_cap_uses_os_environ_when_env_is_none(monkeypatch): # Calling twice must not flip any seeded value. def test_cpu_thread_cap_idempotent(monkeypatch): for variable in (*_THREAD_POOL_ENV_VARS, "UNSLOTH_CPU_THREADS"): - monkeypatch.delenv(variable, raising=False) + monkeypatch.delenv(variable, raising = False) monkeypatch.setenv("UNSLOTH_CPU_THREADS", "5") configure_cpu_threads() @@ -138,9 +140,9 @@ def test_invalid_cpu_thread_cap_exits_without_traceback(entry_point): result = subprocess.run( [sys.executable, str(entry_point)], - env=env, - capture_output=True, - text=True, + env = env, + capture_output = True, + text = True, ) assert result.returncode == 1 From 518d0a5570d12e3afde8bc02084f6814dd8360f0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 27 May 2026 14:15:04 +0000 Subject: [PATCH 3/6] Studio: correct tool-call parser docstrings for PR #5811 Three label corrections surfaced while cross-checking the parser against official HF chat templates. Comment-only, no behaviour change; 104/104 tests still pass. 1. ```` XML was labelled "Qwen3.5 xml". The canonical emitter is Qwen3-Coder and it always wraps the block inside an outer ..., never bare. Updated labels in the module docstring and the three inline comments. 2. Mistral grouping was wrong on two counts: Ministral-8B-2410 uses Tekken V3 and emits the ``[TOOL_CALLS] [...]`` array form, not ``[ARGS]``. Mistral-Large-2411 ships tokenizer.model.v7 (no ``[ARGS]`` token) and also emits the array form. ``[ARGS]`` only enters with Tekken V13 (Devstral, Magistral-Small-2509). Regrouped the doc lines accordingly. 3. Gemma 4 ``<|tool_call>`` is forward-looking. Neither gemma-3-12b-it nor gemma-3n-E4B-it emit this shape in their chat templates today. Noted that the capability gate correctly suppresses the tools pill on real Gemma 3 templates, so the parser stays as ready-when-Google-ships scaffolding. --- .../core/inference/tool_call_parser.py | 31 +++++++++++-------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 9f603b7839..13156ab520 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -8,14 +8,14 @@ Covers the emission formats so the safetensors + MLX agentic loop sees the same call shape llama-server normalises for GGUF: - ``{json}`` (Qwen / Hermes) - - ``v`` (Qwen3.5 xml) + - ``v`` (Qwen3-Coder XML; nested in ) - ``<|python_tag|>NAME.call(k="v", ...)`` (Llama-3 built-in tools) - ``<|python_tag|>{"name":..., "parameters":...}`` (Llama-3 custom) - ``{"name":..., "parameters":...}`` (Llama-3.2 bare JSON) - - ``[TOOL_CALLS] [{...}, ...]`` (Mistral v0.3 / Nemo / Small) - - ``[TOOL_CALLS]name{json}`` (Mistral v11+ / Magistral) - - ``[TOOL_CALLS]name[ARGS]{json}`` (Ministral / Mistral Large 3) - - ``<|tool_call>call:NAME{k:<|"|>v<|"|>}`` (Gemma 4) + - ``[TOOL_CALLS] [{...}, ...]`` (Mistral V3 tokenizer: v0.3, Nemo, Small, Ministral-8B-2410, Large-2411) + - ``[TOOL_CALLS]name{json}`` (Mistral V11+ Tekken: Magistral) + - ``[TOOL_CALLS]name[ARGS]{json}`` (Mistral V13 Tekken: Devstral, Magistral-Small-2509) + - ``<|tool_call>call:NAME{k:<|"|>v<|"|>}`` (Gemma 4; forward-looking, no shipping Gemma model emits this yet) Closing tags / brackets are tolerated when missing because models frequently truncate them mid-stream. @@ -85,7 +85,7 @@ BUDGET_EXHAUSTED_NUDGE = ( # Qwen / Hermes ``{json}``. _TC_JSON_START_RE = re.compile(r"\s*\{") -# Qwen3.5 / Hermes XML ``v``. +# Qwen3-Coder XML ``v``, nested inside . _TC_FUNC_START_RE = re.compile(r"\s*") _TC_END_TAG_RE = re.compile(r"") _TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") @@ -102,14 +102,17 @@ _LLAMA3_KV_RE = re.compile( re.VERBOSE, ) -# Mistral ``[TOOL_CALLS]`` trigger. v11+ chains them, each followed by -# a bare name plus ``{json}`` (Magistral) or ``[ARGS]{json}`` (Ministral -# / Large 3). +# Mistral ``[TOOL_CALLS]`` trigger. V11+ Tekken chains them, each +# followed by a bare name plus ``{json}`` (Magistral) or +# ``[ARGS]{json}`` (V13 Tekken: Devstral, Magistral-Small-2509). +# V3 tokenizer models (v0.3, Nemo, Small, Ministral-8B-2410, +# Large-2411) emit the ``[TOOL_CALLS] [...]`` array form instead. _MISTRAL_TRIGGER = "[TOOL_CALLS]" _MISTRAL_ARGS_MARKER = "[ARGS]" _MISTRAL_V11_NAME_RE = re.compile(r"\s*([\w\.\-]+)\s*") -# Gemma 4: ``<|tool_call>call:NAME{...}``, ``<|"|>`` wraps strings. +# Gemma 4 (forward-looking; Gemma 3 does not emit this shape yet): +# ``<|tool_call>call:NAME{...}``, ``<|"|>`` wraps strings. _GEMMA_TC_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w\.\-]+)\s*\{") _GEMMA_STR_BEGIN = '<|"|>' _GEMMA_STR_END = '<|"|>' @@ -223,10 +226,10 @@ def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict as soon as one matches so we never double-count.""" for parser in ( _parse_tool_call_json, # Qwen / Hermes - _parse_function_xml, # Qwen3.5 / Hermes XML + _parse_function_xml, # Qwen3-Coder XML _parse_llama3_python_tag, # Llama-3 _parse_mistral_tool_calls, # Mistral - _parse_gemma_tool_calls, # Gemma 4 + _parse_gemma_tool_calls, # Gemma 4 (forward-looking) ): calls = parser(content, id_offset = id_offset) if calls: @@ -642,7 +645,9 @@ def _consume_mistral_call(obj_text: str, out: list[dict], id_offset: int) -> Non def _parse_gemma_tool_calls(content: str, *, id_offset: int) -> list[dict]: - """Gemma 4: ``<|tool_call>call:NAME{k:<|"|>v<|"|>, ...}``.""" + """Gemma 4 (forward-looking; no shipping Gemma 3 model emits this): + ``<|tool_call>call:NAME{k:<|"|>v<|"|>, ...}``. + Capability gate suppresses the tools pill on real Gemma 3 templates.""" out: list[dict] = [] for m in _GEMMA_TC_RE.finditer(content): name = m.group(1) From 615b86080372847f7fd62f3617de4f8f115284e4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 27 May 2026 14:28:08 +0000 Subject: [PATCH 4/6] Studio: tighten Llama-3.2 bare-JSON guard for PR #5811 A fuzz pass turned up that ``_parse_llama3_bare_json`` accepted ``parameters`` as a string, contradicting the docstring's "parameters or arguments is a dict" guard. Prose like ``{"name":"foo","parameters":"a sentence"}`` would wrongly fire the parser, which the agentic loop would then heal into a real ``foo(query="a sentence")`` call. Tightened guard: - ``parameters`` must be a dict (Llama-3 spec). - ``arguments`` may be a dict, or a JSON-encoded string that decodes to a dict (OpenAI shape, e.g. ``"arguments":"{\"q\":\"x\"}"``). Plain non-JSON strings, JSON-strings of lists / scalars / null no longer pass. Added 4 regression tests under TestParserMultiFormat: - test_llama3_2_bare_json_string_parameters_does_not_fire - test_llama3_2_bare_json_string_arguments_not_json_does_not_fire - test_llama3_2_bare_json_string_arguments_json_dict_fires - test_llama3_2_bare_json_string_arguments_json_non_dict_does_not_fire Existing tests stay green (104 -> 108 passing) and a 50-case cross-version fuzz suite passes on Python 3.10 / 3.11 / 3.12 / 3.13. --- .../core/inference/tool_call_parser.py | 24 ++++++++++---- .../tests/test_safetensors_tool_loop.py | 31 +++++++++++++++++++ 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 13156ab520..e089421fd6 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -472,16 +472,28 @@ def _parse_llama3_bare_json(content: str, *, id_offset: int) -> list[dict]: name = obj.get("name") or obj.get("function") or "" if not isinstance(name, str) or not name: break + # ``parameters`` must be a dict (Llama-3 spec). + # ``arguments`` may be a dict or a JSON-string of a dict (OpenAI shape). + # Anything looser would fire on prose like ``{"name":"x","parameters":"sentence"}``. if "parameters" in obj: args = obj.get("parameters") + if not isinstance(args, dict): + break + args_str = json.dumps(args) elif "arguments" in obj: args = obj.get("arguments") - else: - break - if isinstance(args, dict): - args_str = json.dumps(args) - elif isinstance(args, str): - args_str = args + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + try: + parsed = json.loads(args) + except (json.JSONDecodeError, ValueError): + break + if not isinstance(parsed, dict): + break + args_str = args + else: + break else: break out.append( diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index a5e8413a1c..84d3817d77 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -260,6 +260,37 @@ class TestParserMultiFormat: text = '{"name":"x","parameters":42}' assert parse_tool_calls_from_text(text) == [] + def test_llama3_2_bare_json_string_parameters_does_not_fire(self): + # Llama-3 spec: parameters must be a dict. Prose like + # ``{"name":"foo","parameters":"a sentence"}`` must NOT trigger. + text = '{"name":"foo","parameters":"this is a sentence"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_string_arguments_not_json_does_not_fire(self): + # OpenAI ``arguments`` may be a JSON-string of a dict, but a + # plain non-JSON string must not pass the guard. + text = '{"name":"foo","arguments":"not json"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_string_arguments_json_dict_fires(self): + # OpenAI shape: arguments is a JSON-encoded string of a dict. + text = '{"name":"foo","arguments":"{\\"q\\":\\"x\\"}"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "foo" + # arguments stays as the original JSON-string. + assert result[0]["function"]["arguments"] == '{"q":"x"}' + + def test_llama3_2_bare_json_string_arguments_json_non_dict_does_not_fire(self): + # JSON-string that parses to a list / scalar / null must NOT fire. + for bad in ( + '{"name":"foo","arguments":"[1,2,3]"}', + '{"name":"foo","arguments":"\\"plain\\""}', + '{"name":"foo","arguments":"null"}', + '{"name":"foo","arguments":"42"}', + ): + assert parse_tool_calls_from_text(bad) == [], bad + # ── Mistral pre-v11 ─────────────────────────────────────────── def test_mistral_pre_v11_array(self): From ade6613e99fde2d936ee5ff368879b00ad74f9f7 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 2 Jul 2026 07:25:40 +0000 Subject: [PATCH 5/6] Studio: trim redundant comments (comment-only, AST-verified) --- .../core/inference/tool_call_parser.py | 13 ++++------ studio/backend/routes/inference.py | 24 +++++++------------ .../tests/test_safetensors_tool_loop.py | 14 +++++------ 3 files changed, 19 insertions(+), 32 deletions(-) diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index e089421fd6..283a3f648a 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -26,8 +26,7 @@ import re from typing import Any -# Markers that flip the streaming buffer from STREAMING to DRAINING so -# partial markup never leaks before the parser sees it. +# Flip the stream buffer to DRAINING so partial markup never leaks. TOOL_XML_SIGNALS = ( "", ".*?", re.DOTALL), re.compile(r".*?", re.DOTALL), @@ -102,11 +100,8 @@ _LLAMA3_KV_RE = re.compile( re.VERBOSE, ) -# Mistral ``[TOOL_CALLS]`` trigger. V11+ Tekken chains them, each -# followed by a bare name plus ``{json}`` (Magistral) or -# ``[ARGS]{json}`` (V13 Tekken: Devstral, Magistral-Small-2509). -# V3 tokenizer models (v0.3, Nemo, Small, Ministral-8B-2410, -# Large-2411) emit the ``[TOOL_CALLS] [...]`` array form instead. +# Mistral ``[TOOL_CALLS]``: V11+ Tekken chains ``name{json}`` / ``name[ARGS]{json}`` +# (Magistral, Devstral); V3 tokenizer models emit the ``[...]`` array form. _MISTRAL_TRIGGER = "[TOOL_CALLS]" _MISTRAL_ARGS_MARKER = "[ARGS]" _MISTRAL_V11_NAME_RE = re.compile(r"\s*([\w\.\-]+)\s*") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index abbc9d7a50..0896fcf5ce 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -260,11 +260,9 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: "supports_tools": False, } ) - # Markers the safetensors / MLX parser recognises. If the template - # advertises tools but uses none of them, drop the pill (parser - # can't honour the emission). The two ``{"name":`` variants cover - # Llama-3.2 ``custom_tools`` whose template prompts the bare-JSON - # form without a ``<|python_tag|>`` prefix. + # Markers the safetensors / MLX parser recognises; drop the pill if the + # template advertises tools but emits none. The ``{"name":`` variants cover + # Llama-3.2 ``custom_tools`` bare-JSON (no ``<|python_tag|>`` prefix). _PARSER_MARKERS = ( "", "``). Mistral ``[TOOL_CALLS]`` -# is delegated to the parser's balanced-brace helper -- a non-greedy -# ``\{.*?\}`` here would truncate nested JSON at the first ``}``. +# Strip leaked tool-call markup: every shared-parser format plus the four +# leak shapes the llama_cpp.py buffer splits across the visible/DRAIN boundary. +# Mistral ``[TOOL_CALLS]`` goes via the parser's balanced-brace helper (a +# non-greedy ``\{.*?\}`` would truncate nested JSON at the first ``}``). _TOOL_XML_RE = _re.compile( "|".join( [ @@ -450,10 +446,8 @@ _TOOL_XML_RE = _re.compile( r"", # Gemma 4. r"<\|tool_call>.*?", - # Llama-3 ``<|python_tag|>...`` to the next ``<|`` sentinel - # or EOF. ``(?:[^<]|<(?!\|))*`` (not ``[^\n<]*`` or - # ``[^\n]*``) keeps literal ``<``, newlines, and embedded - # JSON inside the strip. + # Llama-3 ``<|python_tag|>...`` to the next ``<|`` sentinel or EOF; + # ``(?:[^<]|<(?!\|))*`` keeps literal ``<``, newlines, embedded JSON. r"<\|python_tag\|>(?:[^<]|<(?!\|))*", # Tail-only ```` (anchored so mid-text survives). r"\s*\Z", diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 84d3817d77..e8268ab0de 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -139,7 +139,7 @@ class TestParserMultiFormat: agentic loop is family-agnostic. """ - # ── Llama-3 ──────────────────────────────────────────────────── + # Llama-3 def test_llama3_python_tag_dot_call(self): # Llama-3 built-in tools: <|python_tag|>NAME.call(k="v", ...). @@ -193,7 +193,7 @@ class TestParserMultiFormat: text = '<|python_tag|>brave_search.call(query="x")' assert strip_tool_markup(text, final = True) == "" - # ── Llama-3.2 bare JSON ``custom_tools`` ───────────────────── + # Llama-3.2 bare JSON ``custom_tools`` def test_llama3_2_bare_json_parameters(self): # Llama-3.2-Instruct emits bare JSON directly as content; no @@ -291,7 +291,7 @@ class TestParserMultiFormat: ): assert parse_tool_calls_from_text(bad) == [], bad - # ── Mistral pre-v11 ─────────────────────────────────────────── + # Mistral pre-v11 def test_mistral_pre_v11_array(self): import json @@ -324,7 +324,7 @@ class TestParserMultiFormat: assert len(result) == 1 assert result[0]["function"]["name"] == "web_search" - # ── Mistral v11+ ─────────────────────────────────────────────── + # Mistral v11+ def test_mistral_v11_single(self): # Magistral / Mistral Small 3.1: bare ``name{json}`` after trigger. @@ -358,7 +358,7 @@ class TestParserMultiFormat: text = '[TOOL_CALLS]add{"a":1}' assert strip_tool_markup(text, final = True) == "" - # ── Gemma 4 ─────────────────────────────────────────────────── + # Gemma 4 def test_gemma4_simple_call(self): import json @@ -423,7 +423,7 @@ class TestParserMultiFormat: text = "<|tool_call>call:foo{x:1}" assert strip_tool_markup(text, final = True) == "" - # ── Cross-format sentinels ──────────────────────────────────── + # Cross-format sentinels def test_all_markers_in_tool_xml_signals(self): # Streaming buffer wakes up on every emission marker. @@ -1166,9 +1166,7 @@ class TestGptOssNameDetection: assert is_gpt_oss_model_name(None) is False -# ──────────────────────────────────────────────────────────────────── # Routes-level python_tag strip (multi-line; stop on next sentinel) -# ──────────────────────────────────────────────────────────────────── class TestRoutesPythonTagStrip: From 698e76e3231ebde5761ec9e3459b810fb44c052c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 06:59:31 +0000 Subject: [PATCH 6/6] Studio: stop the Gemma tool-call array parser hanging on malformed input --- .../core/inference/tool_call_parser.py | 8 +++++++- .../tests/test_safetensors_tool_loop.py | 19 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 283a3f648a..a38c5b1108 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -772,7 +772,13 @@ def _gemma_parse_value(text: str, i: int): if body[k] in " \t\n\r,": k += 1 continue - v, k = _gemma_parse_value(body, k) + v, next_k = _gemma_parse_value(body, k) + if next_k <= k: + # Malformed array (e.g. a stray ``}`` with no opening + # ``{``): the primitive branch consumed nothing, so stop + # rather than spin forever. + break + k = next_k items.append(v) return items, j + 1 # Primitive: number / true/false/null / bare identifier. diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index e8268ab0de..bbb75090ac 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -419,6 +419,25 @@ class TestParserMultiFormat: result = parse_tool_calls_from_text(text) assert isinstance(result, list) + def test_gemma4_malformed_array_terminates(self): + # Regression: a stray ``}`` inside a Gemma array (unbalanced + # braces) once spun the list loop forever because + # ``_gemma_parse_value`` returned without advancing the index. + # Parsing must terminate on this malformed input. + import threading + + text = "<|tool_call>call:foo{items:[1}}" + box: dict = {} + + def _run(): + box["result"] = parse_tool_calls_from_text(text) + + worker = threading.Thread(target = _run, daemon = True) + worker.start() + worker.join(timeout = 5) + assert not worker.is_alive(), "Gemma parser hung on malformed array input" + assert isinstance(box["result"], list) + def test_gemma4_strip_markup_final(self): text = "<|tool_call>call:foo{x:1}" assert strip_tool_markup(text, final = True) == ""