Spaces:
Sleeping
Sleeping
File size: 4,999 Bytes
4e316d6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | """
Robust tool-call parser for small language models.
Extracts structured tool calls from free-form model output using
regex-based parsing with multiple fallback strategies. Designed for
models that cannot reliably produce valid JSON (e.g. Qwen3-0.6B).
The expected model output format::
Thought: I need to search the documents for information about X.
<tool>retrieve_documents</tool>
<param name="query">information about X</param>
Or for a final answer::
Thought: I have enough information to answer.
<answer>The answer is Y.</answer>
If parsing fails entirely, the raw output is treated as a direct answer.
"""
from __future__ import annotations
import re
from src.agents.schemas import ToolCall
# ---------------------------------------------------------------------------
# Regex patterns
# ---------------------------------------------------------------------------
_TOOL_RE = re.compile(r"<tool>\s*(\w+)\s*</tool[!.]?>", re.IGNORECASE)
_PARAM_RE = re.compile(
r'<param\s+name=["\'](\w+)["\']>(.*?)</param[!.]?>',
re.IGNORECASE | re.DOTALL,
)
_ANSWER_RE = re.compile(r"<answer>(.*?)</answer[!.]?>", re.IGNORECASE | re.DOTALL)
_THOUGHT_RE = re.compile(
r"(?:Thought|Reasoning|Think):\s*(.+?)(?=<tool>|<answer>|$)",
re.IGNORECASE | re.DOTALL,
)
# Fallback: model uses <tool_name>key: value\n...</tool_name> instead of
# the canonical <tool>/<param> format. Common with small models.
_ALT_TOOL_RE = re.compile(
r"<(\w+)>\s*(.*?)\s*</\1>",
re.IGNORECASE | re.DOTALL,
)
_ALT_PARAM_RE = re.compile(r"(\w+)\s*:\s*(.+)", re.MULTILINE)
# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------
def parse_tool_call(text: str) -> tuple[str, ToolCall | None, str | None]:
"""Parse model output into *(thought, tool_call_or_none, final_answer_or_none)*.
Exactly one of *tool_call* or *final_answer* will be non-``None``.
If neither is parseable the raw text is returned as the final answer.
"""
thought = _extract_thought(text)
# Try final answer first — if the model says <answer>, we're done.
final_answer = _extract_final_answer(text)
if final_answer is not None:
return thought, None, final_answer
# Try tool call.
tool_call = _extract_tool_call(text)
if tool_call is not None:
return thought, tool_call, None
# Fallback: treat entire output as a direct answer.
return thought, None, text.strip()
# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------
def _extract_thought(text: str) -> str:
"""Extract the ``Thought:`` line, or everything before the first tag."""
m = _THOUGHT_RE.search(text)
if m:
return m.group(1).strip()
# Fallback: everything before the first < tag.
idx = text.find("<")
if idx > 0:
return text[:idx].strip()
return ""
def _extract_tool_call(text: str) -> ToolCall | None:
"""Try to extract tool calls, with fallback for alternative formats.
Primary format: ``<tool>name</tool>`` + ``<param name="k">v</param>``
Fallback format: ``<tool_name>key: value</tool_name>`` (common with small models)
"""
# Primary format: <tool>name</tool> + <param> tags
m = _TOOL_RE.search(text)
if m is not None:
tool_name = m.group(1)
arguments: dict[str, str] = {}
for pm in _PARAM_RE.finditer(text):
arguments[pm.group(1)] = pm.group(2).strip()
return ToolCall(tool_name=tool_name, arguments=arguments, raw_text=text)
# Fallback: <tool_name>key: value\n...</tool_name>
# e.g. <web_search>query: latest PyTorch\nnum_results: 3</web_search>
for am in _ALT_TOOL_RE.finditer(text):
tag_name = am.group(1)
# Skip known non-tool tags
if tag_name.lower() in ("answer", "think", "param"):
continue
body = am.group(2)
arguments = {}
for pm in _ALT_PARAM_RE.finditer(body):
key = pm.group(1).strip().strip('"').strip("'")
val = pm.group(2).strip().strip('"').strip("'")
arguments[key] = val
return ToolCall(tool_name=tag_name, arguments=arguments, raw_text=text)
return None
def _extract_final_answer(text: str) -> str | None:
"""Try to extract ``<answer>...</answer>`` content."""
m = _ANSWER_RE.search(text)
if m is None:
return None
return m.group(1).strip()
def coerce_param(value: str, expected_type: str):
"""Best-effort type coercion for tool parameters."""
if expected_type == "int":
try:
return int(value)
except (ValueError, TypeError):
return value
if expected_type == "float":
try:
return float(value)
except (ValueError, TypeError):
return value
return value
|