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

studio: extract tool-call XML parser into a reusable helper module (#5583)

Browse source 
Move the inline tool-call XML parser and stripper out of
studio/backend/core/inference/llama_cpp.py into a new
studio/backend/core/tool_healing.py so external inference servers
(llama-server wrappers, llama-swap, custom shims) can reuse the same
logic without importing the inference orchestrator, structlog, httpx,
or anything from torch / transformers / unsloth.

Closes #5502.

What this PR does:

- New file studio/backend/core/tool_healing.py contains the regex
  constants (_TOOL_CLOSED_PATS, _TOOL_ALL_PATS, _TC_JSON_START_RE,
  _TC_FUNC_START_RE, _TC_END_TAG_RE, _TC_FUNC_CLOSE_RE,
  _TC_PARAM_START_RE, _TC_PARAM_CLOSE_RE), parse_tool_calls_from_text,
  and strip_tool_call_markup. The regexes and function bodies are
  byte-for-byte the same as the previous inline implementation in
  llama_cpp.py; only the @staticmethod decorator and the closure-only
  `if not auto_heal_tool_calls: return text` short-circuit are dropped
  (the latter stays in the caller as a fast path when healing is off).
- studio/backend/core/inference/llama_cpp.py now imports the regexes
  and helpers from .tool_healing. LlamaCppBackend._parse_tool_calls_from_text
  becomes a one-line delegate; the _strip_tool_markup closure keeps the
  auto_heal_tool_calls fast path and delegates the work.
- Helper module imports cleanly without torch, transformers, structlog,
  httpx, or numpy. studio.backend.core itself is already stdlib-only
  at import time (lazy __getattr__), so `from
  studio.backend.core.tool_healing import parse_tool_calls_from_text,
  strip_tool_call_markup` is the lightweight import path issue #5502
  asked for.

No behaviour change for existing Studio paths. parse_tool_calls_from_text
and strip_tool_call_markup produce the same OpenAI-shape output the
old inline code produced for every input.

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
alkinun 2026-05-19 15:06:17 +03:00 • committed by GitHub
commit f747108212
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 187 additions and 136 deletions
Download patch file Download diff file Expand all files Collapse all files

150
studio/backend/core/inference/llama_cpp.py
View file

@ -28,6 +28,18 @@ from urllib.parse import urlparse
import httpx
from core.tool_healing import (
_TC_END_TAG_RE,
_TC_FUNC_CLOSE_RE,
_TC_FUNC_START_RE,
_TC_JSON_START_RE,
_TC_PARAM_CLOSE_RE,
_TC_PARAM_START_RE,
_TOOL_ALL_PATS,
_TOOL_CLOSED_PATS,
parse_tool_calls_from_text,
strip_tool_call_markup,
)
from utils.native_path_leases import child_env_without_native_path_secret
from utils.subprocess_compat import (
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
@ -362,25 +374,6 @@ def _extract_model_size_b(model_id: str):
return extract_model_size_b(model_id)
# ── Pre-compiled patterns for tool XML stripping ─────────────
_TOOL_CLOSED_PATS = [
re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL),
re.compile(r"<function=\w+>.*?</function>", re.DOTALL),
]
_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
re.compile(r"<tool_call>.*$", re.DOTALL),
re.compile(r"<function=\w+>.*$", re.DOTALL),
]
# ── Pre-compiled patterns for tool-call XML parsing ──────────
_TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
_TC_FUNC_START_RE = re.compile(r"<function=(\w+)>\s*")
_TC_END_TAG_RE = re.compile(r"</tool_call>")
_TC_FUNC_CLOSE_RE = re.compile(r"\s*</function>\s*$")
_TC_PARAM_START_RE = re.compile(r"<parameter=(\w+)>\s*")
_TC_PARAM_CLOSE_RE = re.compile(r"\s*</parameter>\s*$")
_TOOL_TEMPLATE_MARKERS = (
"{%- if tools %}",
"{%- if tools -%}",
@ -3502,119 +3495,7 @@ class LlamaCppBackend:
Closing tags (</tool_call>, </function>, </parameter>) are all optional
since models frequently omit them.
"""
tool_calls = []
# Pattern 1: JSON inside <tool_call> tags.
# Use balanced-brace extraction that skips braces inside JSON strings.
for m in _TC_JSON_START_RE.finditer(content):
brace_start = m.end() - 1 # position of the opening {
depth, i = 0, brace_start
in_string = False
while i < len(content):
ch = content[i]
if in_string:
if ch == "\\" and i + 1 < len(content):
i += 2 # skip escaped character
continue
if ch == '"':
in_string = False
elif ch == '"':
in_string = True
elif ch == "{":
depth += 1
elif ch == "}":
depth -= 1
if depth == 0:
break
i += 1
if depth == 0:
json_str = content[brace_start : i + 1]
try:
obj = json.loads(json_str)
tc = {
"id": f"call_{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: XML-style <function=name><parameter=key>value</parameter></function>
# All closing tags optional -- models frequently omit </parameter>,
# </function>, and/or </tool_call>.
if not tool_calls:
# Step 1: Find all <function=name> positions and extract their bodies.
# Body boundary: use only </tool_call> or next <function= as hard
# boundaries. We avoid using </function> as a boundary because
# code parameter values can contain that literal string.
# After extracting, we trim a trailing </function> if present.
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()
# Hard boundaries: next <function= tag or </tool_call>
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]
# Trim trailing </function> if present (it's the real closing tag)
body = _TC_FUNC_CLOSE_RE.sub("", body)
# Step 2: Extract parameters from body.
# For single-parameter functions (the common case: code, command,
# query), use body end as the only boundary to avoid false matches
# on </parameter> inside code strings.
arguments = {}
param_starts = list(_TC_PARAM_START_RE.finditer(body))
if len(param_starts) == 1:
# Single parameter: value is everything from after the tag
# to end of body, trimming any trailing </parameter>.
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()
# Value ends at next <parameter= or end of body
next_param = (
param_starts[pidx + 1].start()
if pidx + 1 < len(param_starts)
else len(body)
)
val = body[val_start:next_param]
# Trim trailing </parameter> if present
val = _TC_PARAM_CLOSE_RE.sub("", val)
arguments[param_name] = val.strip()
tc = {
"id": f"call_{len(tool_calls)}",
"type": "function",
"function": {
"name": func_name,
"arguments": json.dumps(arguments),
},
}
tool_calls.append(tc)
return tool_calls
return parse_tool_calls_from_text(content)
@staticmethod
def _build_openai_messages(
@ -3993,10 +3874,7 @@ class LlamaCppBackend:
def _strip_tool_markup(text: str, *, final: bool = False) -> str:
if not auto_heal_tool_calls:
return text
patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS
for pat in patterns:
text = pat.sub("", text)
return text.strip() if final else text
return strip_tool_call_markup(text, final = final)
# XML prefixes that signal a tool call in content.
# Empty when auto_heal is disabled so the buffer never

173
studio/backend/core/tool_healing.py Normal file
View file

@ -0,0 +1,173 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tool-call XML parsing and stripping helpers.
Extracted verbatim from studio/backend/core/inference/llama_cpp.py so that
external inference servers (llama-server wrappers, llama-swap, custom
shims) can reuse the same logic without importing the inference
orchestrator, structlog, httpx, or the rest of the studio backend.
The regexes and function bodies are byte-for-byte identical to the
original inline implementation in llama_cpp.py. Any change made here must
preserve that equivalence; tests/python/test_tool_healing_extraction_is_exact.py
verifies it with AST comparison.
"""
import json
import re
# Pre-compiled patterns for tool XML stripping.
_TOOL_CLOSED_PATS = [
re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL),
re.compile(r"<function=\w+>.*?</function>", re.DOTALL),
]
_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
re.compile(r"<tool_call>.*$", re.DOTALL),
re.compile(r"<function=\w+>.*$", re.DOTALL),
]
# Pre-compiled patterns for tool-call XML parsing.
_TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
_TC_FUNC_START_RE = re.compile(r"<function=(\w+)>\s*")
_TC_END_TAG_RE = re.compile(r"</tool_call>")
_TC_FUNC_CLOSE_RE = re.compile(r"\s*</function>\s*$")
_TC_PARAM_START_RE = re.compile(r"<parameter=(\w+)>\s*")
_TC_PARAM_CLOSE_RE = re.compile(r"\s*</parameter>\s*$")
def parse_tool_calls_from_text(content: str) -> list[dict]:
"""
Parse tool calls from XML markup in content text.
Handles formats like:
<tool_call>{"name":"web_search","arguments":{"query":"..."}}</tool_call>
<tool_call><function=web_search><parameter=query>...</parameter></function></tool_call>
Closing tags (</tool_call>, </function>, </parameter>) are all optional
since models frequently omit them.
"""
tool_calls = []
# Pattern 1: JSON inside <tool_call> tags.
# Use balanced-brace extraction that skips braces inside JSON strings.
for m in _TC_JSON_START_RE.finditer(content):
brace_start = m.end() - 1 # position of the opening {
depth, i = 0, brace_start
in_string = False
while i < len(content):
ch = content[i]
if in_string:
if ch == "\\" and i + 1 < len(content):
i += 2 # skip escaped character
continue
if ch == '"':
in_string = False
elif ch == '"':
in_string = True
elif ch == "{":
depth += 1
elif ch == "}":
depth -= 1
if depth == 0:
break
i += 1
if depth == 0:
json_str = content[brace_start : i + 1]
try:
obj = json.loads(json_str)
tc = {
"id": f"call_{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: XML-style <function=name><parameter=key>value</parameter></function>
# All closing tags optional -- models frequently omit </parameter>,
# </function>, and/or </tool_call>.
if not tool_calls:
# Step 1: Find all <function=name> positions and extract their bodies.
# Body boundary: use only </tool_call> or next <function= as hard
# boundaries. We avoid using </function> as a boundary because
# code parameter values can contain that literal string.
# After extracting, we trim a trailing </function> if present.
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()
# Hard boundaries: next <function= tag or </tool_call>
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]
# Trim trailing </function> if present (it's the real closing tag)
body = _TC_FUNC_CLOSE_RE.sub("", body)
# Step 2: Extract parameters from body.
# For single-parameter functions (the common case: code, command,
# query), use body end as the only boundary to avoid false matches
# on </parameter> inside code strings.
arguments = {}
param_starts = list(_TC_PARAM_START_RE.finditer(body))
if len(param_starts) == 1:
# Single parameter: value is everything from after the tag
# to end of body, trimming any trailing </parameter>.
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()
# Value ends at next <parameter= or end of body
next_param = (
param_starts[pidx + 1].start()
if pidx + 1 < len(param_starts)
else len(body)
)
val = body[val_start:next_param]
# Trim trailing </parameter> if present
val = _TC_PARAM_CLOSE_RE.sub("", val)
arguments[param_name] = val.strip()
tc = {
"id": f"call_{len(tool_calls)}",
"type": "function",
"function": {
"name": func_name,
"arguments": json.dumps(arguments),
},
}
tool_calls.append(tc)
return tool_calls
def strip_tool_call_markup(text: str, *, final: bool = False) -> str:
"""Strip tool-call XML markup from text.
When ``final`` is False, only fully closed tool-call blocks are removed.
When ``final`` is True, trailing incomplete tool-call blocks are removed
too, and the result is stripped of surrounding whitespace.
"""
patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS
for pat in patterns:
text = pat.sub("", text)
return text.strip() if final else text