cab3ed41e9
The web_search tool previously returned only titles, URLs, and short snippets from DuckDuckGo. Small models would loop repeatedly trying different search queries because the snippets never contained enough data to answer the question. Now fetches the actual page content for the top 1-2 search results (best-effort, capped at 6000 chars and 10s timeout per fetch) and includes it in the tool response. This gives the model real data to work with on the first search instead of endlessly retrying. Uses only stdlib (urllib + re) so no new dependencies are needed.
535 lines
18 KiB
Python
535 lines
18 KiB
Python
# 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 definitions and executors for LLM tool calling.
|
|
|
|
Supports web search (DuckDuckGo), Python code execution, and terminal commands.
|
|
"""
|
|
|
|
import ast
|
|
import os
|
|
|
|
os.environ["UNSLOTH_IS_PRESENT"] = "1"
|
|
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
|
|
from loggers import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
_EXEC_TIMEOUT = 300 # 5 minutes
|
|
_MAX_OUTPUT_CHARS = 8000 # truncate long output
|
|
_BASH_BLOCKED_WORDS = {"rm", "sudo", "dd", "chmod", "mkfs", "shutdown", "reboot"}
|
|
|
|
# Per-session working directories so each chat thread gets its own sandbox.
|
|
# Falls back to a shared ~/studio_sandbox/ for API callers without a session_id.
|
|
_workdirs: dict[str, str] = {}
|
|
|
|
|
|
def _get_workdir(session_id: str | None = None) -> str:
|
|
"""Return (and lazily create) a persistent working directory for tool execution."""
|
|
global _workdirs
|
|
key = session_id or "_default"
|
|
if key not in _workdirs or not os.path.isdir(_workdirs[key]):
|
|
home = os.path.expanduser("~")
|
|
sandbox_root = os.path.join(home, "studio_sandbox")
|
|
if session_id:
|
|
# Sanitize: strip path separators and parent-dir references
|
|
safe_id = os.path.basename(session_id.replace("..", ""))
|
|
if not safe_id:
|
|
safe_id = "_invalid"
|
|
workdir = os.path.join(sandbox_root, safe_id)
|
|
# Verify resolved path stays under sandbox root
|
|
if not os.path.realpath(workdir).startswith(os.path.realpath(sandbox_root)):
|
|
workdir = os.path.join(sandbox_root, "_invalid")
|
|
else:
|
|
workdir = sandbox_root
|
|
os.makedirs(workdir, exist_ok = True)
|
|
_workdirs[key] = workdir
|
|
return _workdirs[key]
|
|
|
|
|
|
WEB_SEARCH_TOOL = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "web_search",
|
|
"description": "Search the web for current information, recent events, or facts you are uncertain about.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {
|
|
"type": "string",
|
|
"description": "The search query",
|
|
}
|
|
},
|
|
"required": ["query"],
|
|
},
|
|
},
|
|
}
|
|
|
|
PYTHON_TOOL = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "python",
|
|
"description": "Execute Python code in a sandbox and return stdout/stderr.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"code": {
|
|
"type": "string",
|
|
"description": "The Python code to run",
|
|
}
|
|
},
|
|
"required": ["code"],
|
|
},
|
|
},
|
|
}
|
|
|
|
TERMINAL_TOOL = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "terminal",
|
|
"description": "Execute a terminal command and return stdout/stderr.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"command": {
|
|
"type": "string",
|
|
"description": "The command to run",
|
|
}
|
|
},
|
|
"required": ["command"],
|
|
},
|
|
},
|
|
}
|
|
|
|
ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL]
|
|
|
|
|
|
_TIMEOUT_UNSET = object()
|
|
|
|
|
|
def execute_tool(
|
|
name: str,
|
|
arguments: dict,
|
|
cancel_event = None,
|
|
timeout: int | None = _TIMEOUT_UNSET,
|
|
session_id: str | None = None,
|
|
) -> str:
|
|
"""Execute a tool by name with the given arguments. Returns result as a string.
|
|
|
|
``timeout``: int sets per-call limit in seconds, ``None`` means no limit,
|
|
unset (default) uses ``_EXEC_TIMEOUT`` (300 s).
|
|
``session_id``: optional thread/session ID for per-conversation sandbox isolation.
|
|
"""
|
|
logger.info(
|
|
f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}"
|
|
)
|
|
effective_timeout = _EXEC_TIMEOUT if timeout is _TIMEOUT_UNSET else timeout
|
|
if name == "web_search":
|
|
return _web_search(arguments.get("query", ""), timeout = effective_timeout)
|
|
if name == "python":
|
|
return _python_exec(
|
|
arguments.get("code", ""), cancel_event, effective_timeout, session_id
|
|
)
|
|
if name == "terminal":
|
|
return _bash_exec(
|
|
arguments.get("command", ""), cancel_event, effective_timeout, session_id
|
|
)
|
|
return f"Unknown tool: {name}"
|
|
|
|
|
|
def _fetch_page_text(url: str, max_chars: int = 4000, timeout: int = 10) -> str:
|
|
"""Fetch a URL and extract plain text content (best-effort)."""
|
|
import urllib.request
|
|
import urllib.error
|
|
try:
|
|
req = urllib.request.Request(url, headers = {"User-Agent": "Mozilla/5.0"})
|
|
with urllib.request.urlopen(req, timeout = timeout) as resp:
|
|
raw = resp.read(200_000).decode("utf-8", errors = "replace")
|
|
# Strip HTML tags (lightweight, no extra deps)
|
|
import re
|
|
text = re.sub(r"<script[^>]*>.*?</script>", "", raw, flags = re.DOTALL | re.IGNORECASE)
|
|
text = re.sub(r"<style[^>]*>.*?</style>", "", text, flags = re.DOTALL | re.IGNORECASE)
|
|
text = re.sub(r"<[^>]+>", " ", text)
|
|
text = re.sub(r"\s+", " ", text).strip()
|
|
# Decode HTML entities
|
|
import html
|
|
text = html.unescape(text)
|
|
if len(text) > max_chars:
|
|
text = text[:max_chars] + "..."
|
|
return text
|
|
except Exception:
|
|
return ""
|
|
|
|
|
|
def _web_search(query: str, max_results: int = 5, timeout: int = _EXEC_TIMEOUT) -> str:
|
|
"""Search the web using DuckDuckGo and return formatted results.
|
|
|
|
For the top result, also fetches the actual page content so the
|
|
model has real data to work with instead of just snippets.
|
|
"""
|
|
if not query.strip():
|
|
return "No query provided."
|
|
try:
|
|
from ddgs import DDGS
|
|
|
|
results = DDGS(timeout = timeout).text(query, max_results = max_results)
|
|
if not results:
|
|
return "No results found."
|
|
|
|
# Fetch full page content for the top result (best-effort,
|
|
# capped at 10s so it does not block the agentic loop).
|
|
top_page = ""
|
|
for r in results[:2]:
|
|
href = r.get("href", "")
|
|
if href:
|
|
top_page = _fetch_page_text(href, max_chars = 6000, timeout = 10)
|
|
if len(top_page) > 200:
|
|
break
|
|
parts = []
|
|
for i, r in enumerate(results):
|
|
entry = (
|
|
f"Title: {r.get('title', '')}\n"
|
|
f"URL: {r.get('href', '')}\n"
|
|
f"Snippet: {r.get('body', '')}"
|
|
)
|
|
if i == 0 and top_page:
|
|
entry += f"\n\nPage content:\n{top_page}"
|
|
parts.append(entry)
|
|
return "\n\n---\n\n".join(parts)
|
|
except Exception as e:
|
|
return f"Search failed: {e}"
|
|
|
|
|
|
def _check_signal_escape_patterns(code: str):
|
|
"""
|
|
Check if code contains patterns that could escape signal-based timeouts.
|
|
|
|
Vendored from unsloth_zoo.rl_environments to avoid importing unsloth_zoo
|
|
(which requires GPU drivers and fails on Mac/Apple Silicon).
|
|
|
|
Returns (safe: bool, details: dict)
|
|
"""
|
|
try:
|
|
tree = ast.parse(code)
|
|
except SyntaxError as e:
|
|
return False, {
|
|
"error": f"SyntaxError: {e}",
|
|
"signal_tampering": [],
|
|
"exception_catching": [],
|
|
"warnings": [],
|
|
}
|
|
|
|
signal_tampering = []
|
|
exception_catching = []
|
|
warnings = []
|
|
|
|
def _ast_name_matches(node, names):
|
|
if isinstance(node, ast.Name):
|
|
return node.id in names
|
|
elif isinstance(node, ast.Attribute):
|
|
full_name = []
|
|
current = node
|
|
while isinstance(current, ast.Attribute):
|
|
full_name.append(current.attr)
|
|
current = current.value
|
|
if isinstance(current, ast.Name):
|
|
full_name.append(current.id)
|
|
full_name = ".".join(reversed(full_name))
|
|
return full_name in names
|
|
return False
|
|
|
|
class SignalEscapeVisitor(ast.NodeVisitor):
|
|
def __init__(self):
|
|
self.imports_signal = False
|
|
self.signal_aliases = {"signal"}
|
|
self.loop_depth = 0
|
|
|
|
def visit_Import(self, node):
|
|
for alias in node.names:
|
|
if alias.name == "signal":
|
|
self.imports_signal = True
|
|
if alias.asname:
|
|
self.signal_aliases.add(alias.asname)
|
|
self.generic_visit(node)
|
|
|
|
def visit_ImportFrom(self, node):
|
|
if node.module == "signal":
|
|
self.imports_signal = True
|
|
for alias in node.names:
|
|
if alias.name in (
|
|
"signal",
|
|
"SIGALRM",
|
|
"SIG_IGN",
|
|
"setitimer",
|
|
"ITIMER_REAL",
|
|
"pthread_sigmask",
|
|
"SIG_BLOCK",
|
|
"alarm",
|
|
):
|
|
self.signal_aliases.add(alias.asname or alias.name)
|
|
self.generic_visit(node)
|
|
|
|
def visit_While(self, node):
|
|
self.loop_depth += 1
|
|
self.generic_visit(node)
|
|
self.loop_depth -= 1
|
|
|
|
def visit_For(self, node):
|
|
self.loop_depth += 1
|
|
self.generic_visit(node)
|
|
self.loop_depth -= 1
|
|
|
|
def visit_Call(self, node):
|
|
func = node.func
|
|
func_name = None
|
|
if isinstance(func, ast.Attribute):
|
|
if isinstance(func.value, ast.Name):
|
|
if func.value.id in self.signal_aliases:
|
|
func_name = f"signal.{func.attr}"
|
|
elif isinstance(func, ast.Name):
|
|
if func.id in ("signal", "setitimer", "alarm", "pthread_sigmask"):
|
|
func_name = func.id
|
|
|
|
if func_name:
|
|
if func_name in ("signal.signal", "signal"):
|
|
if len(node.args) >= 1:
|
|
if _ast_name_matches(
|
|
node.args[0], ("SIGALRM", "signal.SIGALRM")
|
|
):
|
|
signal_tampering.append(
|
|
{
|
|
"type": "signal_handler_override",
|
|
"line": node.lineno,
|
|
"description": "Overrides SIGALRM handler",
|
|
}
|
|
)
|
|
elif func_name in ("signal.setitimer", "setitimer"):
|
|
if len(node.args) >= 1:
|
|
if _ast_name_matches(
|
|
node.args[0], ("ITIMER_REAL", "signal.ITIMER_REAL")
|
|
):
|
|
signal_tampering.append(
|
|
{
|
|
"type": "timer_manipulation",
|
|
"line": node.lineno,
|
|
"description": "Manipulates ITIMER_REAL timer",
|
|
}
|
|
)
|
|
elif func_name in ("signal.alarm", "alarm"):
|
|
signal_tampering.append(
|
|
{
|
|
"type": "alarm_manipulation",
|
|
"line": node.lineno,
|
|
"description": "Manipulates alarm timer",
|
|
}
|
|
)
|
|
elif func_name in ("signal.pthread_sigmask", "pthread_sigmask"):
|
|
signal_tampering.append(
|
|
{
|
|
"type": "signal_mask",
|
|
"line": node.lineno,
|
|
"description": "Modifies signal mask (may block SIGALRM)",
|
|
}
|
|
)
|
|
self.generic_visit(node)
|
|
|
|
def visit_ExceptHandler(self, node):
|
|
if self.loop_depth == 0:
|
|
self.generic_visit(node)
|
|
return
|
|
if node.type is None:
|
|
exception_catching.append(
|
|
{
|
|
"type": "bare_except_in_loop",
|
|
"line": node.lineno,
|
|
"description": "Bare except in loop catches TimeoutError and continues looping",
|
|
}
|
|
)
|
|
elif isinstance(node.type, ast.Name):
|
|
if node.type.id in ("TimeoutError", "BaseException", "Exception"):
|
|
exception_catching.append(
|
|
{
|
|
"type": f"catches_{node.type.id}_in_loop",
|
|
"line": node.lineno,
|
|
"description": f"Catches {node.type.id} in loop - may suppress timeout and continue",
|
|
}
|
|
)
|
|
elif isinstance(node.type, ast.Tuple):
|
|
for elt in node.type.elts:
|
|
if isinstance(elt, ast.Name):
|
|
if elt.id in ("TimeoutError", "BaseException", "Exception"):
|
|
exception_catching.append(
|
|
{
|
|
"type": f"catches_{elt.id}_in_loop",
|
|
"line": node.lineno,
|
|
"description": f"Catches {elt.id} in loop - may suppress timeout and continue",
|
|
}
|
|
)
|
|
self.generic_visit(node)
|
|
|
|
visitor = SignalEscapeVisitor()
|
|
visitor.visit(tree)
|
|
|
|
if visitor.imports_signal and not signal_tampering:
|
|
warnings.append("Code imports 'signal' module - review manually for safety")
|
|
|
|
is_safe = len(signal_tampering) == 0 and len(exception_catching) == 0
|
|
return is_safe, {
|
|
"signal_tampering": signal_tampering,
|
|
"exception_catching": exception_catching,
|
|
"warnings": warnings,
|
|
}
|
|
|
|
|
|
def _check_code_safety(code: str) -> str | None:
|
|
"""Validate code safety via static analysis.
|
|
|
|
Returns an error message string if the code is unsafe, or None if OK.
|
|
"""
|
|
safe, info = _check_signal_escape_patterns(code)
|
|
if not safe:
|
|
reasons = [
|
|
item.get("description", "") for item in info.get("signal_tampering", [])
|
|
]
|
|
return (
|
|
f"Error: unsafe code detected ({'; '.join(reasons)}). "
|
|
f"Please remove signal manipulation from your code."
|
|
)
|
|
|
|
return None
|
|
|
|
|
|
def _cancel_watcher(proc, cancel_event, poll_interval = 0.2):
|
|
"""Daemon thread that kills a process when cancel_event is set."""
|
|
while proc.poll() is None:
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
proc.kill()
|
|
return
|
|
cancel_event.wait(poll_interval) if cancel_event else None
|
|
|
|
|
|
def _truncate(text: str, limit: int = _MAX_OUTPUT_CHARS) -> str:
|
|
if len(text) > limit:
|
|
return text[:limit] + f"\n\n... (truncated, {len(text)} chars total)"
|
|
return text
|
|
|
|
|
|
def _python_exec(
|
|
code: str,
|
|
cancel_event = None,
|
|
timeout: int = _EXEC_TIMEOUT,
|
|
session_id: str | None = None,
|
|
) -> str:
|
|
"""Execute Python code in a subprocess sandbox."""
|
|
if not code or not code.strip():
|
|
return "No code provided."
|
|
|
|
# Validate imports and code safety
|
|
error = _check_code_safety(code)
|
|
if error:
|
|
return error
|
|
|
|
tmp_path = None
|
|
workdir = _get_workdir(session_id)
|
|
try:
|
|
fd, tmp_path = tempfile.mkstemp(
|
|
suffix = ".py", prefix = "studio_exec_", dir = workdir
|
|
)
|
|
with os.fdopen(fd, "w") as f:
|
|
f.write(code)
|
|
|
|
proc = subprocess.Popen(
|
|
[sys.executable, tmp_path],
|
|
stdout = subprocess.PIPE,
|
|
stderr = subprocess.STDOUT,
|
|
text = True,
|
|
cwd = workdir,
|
|
)
|
|
|
|
# Spawn cancel watcher if we have a cancel event
|
|
if cancel_event is not None:
|
|
watcher = threading.Thread(
|
|
target = _cancel_watcher, args = (proc, cancel_event), daemon = True
|
|
)
|
|
watcher.start()
|
|
|
|
try:
|
|
output, _ = proc.communicate(timeout = timeout)
|
|
except subprocess.TimeoutExpired:
|
|
proc.kill()
|
|
proc.communicate()
|
|
return _truncate(f"Execution timed out after {timeout} seconds.")
|
|
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
return "Execution cancelled."
|
|
|
|
result = output or ""
|
|
if proc.returncode != 0:
|
|
result = f"Exit code {proc.returncode}:\n{result}"
|
|
return _truncate(result) if result.strip() else "(no output)"
|
|
|
|
except Exception as e:
|
|
return f"Execution error: {e}"
|
|
finally:
|
|
if tmp_path and os.path.exists(tmp_path):
|
|
try:
|
|
os.unlink(tmp_path)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def _bash_exec(
|
|
command: str,
|
|
cancel_event = None,
|
|
timeout: int = _EXEC_TIMEOUT,
|
|
session_id: str | None = None,
|
|
) -> str:
|
|
"""Execute a bash command in a subprocess sandbox."""
|
|
if not command or not command.strip():
|
|
return "No command provided."
|
|
|
|
# Block dangerous commands
|
|
tokens = set(command.lower().split())
|
|
blocked = tokens & _BASH_BLOCKED_WORDS
|
|
if blocked:
|
|
return f"Blocked command(s) for safety: {', '.join(sorted(blocked))}"
|
|
|
|
try:
|
|
workdir = _get_workdir(session_id)
|
|
proc = subprocess.Popen(
|
|
["bash", "-c", command],
|
|
stdout = subprocess.PIPE,
|
|
stderr = subprocess.STDOUT,
|
|
text = True,
|
|
cwd = workdir,
|
|
)
|
|
|
|
if cancel_event is not None:
|
|
watcher = threading.Thread(
|
|
target = _cancel_watcher, args = (proc, cancel_event), daemon = True
|
|
)
|
|
watcher.start()
|
|
|
|
try:
|
|
output, _ = proc.communicate(timeout = timeout)
|
|
except subprocess.TimeoutExpired:
|
|
proc.kill()
|
|
proc.communicate()
|
|
return _truncate(f"Execution timed out after {timeout} seconds.")
|
|
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
return "Execution cancelled."
|
|
|
|
result = output or ""
|
|
if proc.returncode != 0:
|
|
result = f"Exit code {proc.returncode}:\n{result}"
|
|
return _truncate(result) if result.strip() else "(no output)"
|
|
|
|
except Exception as e:
|
|
return f"Execution error: {e}"
|