feat: initial release - Hunzi agent framework v0.1.0
This commit is contained in:
@@ -0,0 +1,361 @@
|
||||
"""Tool system — registration, discovery, and execution."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import subprocess
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tool
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class Tool:
|
||||
"""A callable tool with a JSON-schema description.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
Unique identifier for the tool.
|
||||
description : str
|
||||
Human-readable description (used in the system prompt).
|
||||
parameters : dict
|
||||
JSON Schema ``properties`` for the tool's arguments.
|
||||
func : callable
|
||||
The actual implementation — accepts **kwargs matching *parameters*.
|
||||
"""
|
||||
|
||||
name: str
|
||||
description: str
|
||||
func: Callable[..., Any] = field(repr=False, default=None)
|
||||
parameters: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# -- JSON Schema helpers --------------------------------------------------------
|
||||
|
||||
@property
|
||||
def schema(self) -> Dict[str, Any]:
|
||||
"""Full JSON Schema object (``type: object`` wrapper)."""
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": self.parameters,
|
||||
"required": [
|
||||
k for k, v in self.parameters.items() if v.get("required", False)
|
||||
],
|
||||
}
|
||||
|
||||
def to_openai_format(self) -> Dict[str, Any]:
|
||||
"""Return the tool in OpenAI function-calling format."""
|
||||
return {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": self.name,
|
||||
"description": self.description,
|
||||
"parameters": self.schema,
|
||||
},
|
||||
}
|
||||
|
||||
# -- execution ------------------------------------------------------------------
|
||||
|
||||
def execute(self, **kwargs: Any) -> Any:
|
||||
"""Run the tool function with the given keyword arguments."""
|
||||
logger.debug("Executing tool %s with %s", self.name, kwargs)
|
||||
try:
|
||||
result = self.func(**kwargs)
|
||||
logger.debug("Tool %s returned: %s", self.name, str(result)[:200])
|
||||
return result
|
||||
except Exception as exc:
|
||||
logger.error("Tool %s failed: %s", self.name, exc)
|
||||
raise ToolExecutionError(self.name, str(exc)) from exc
|
||||
|
||||
|
||||
class ToolExecutionError(Exception):
|
||||
"""Raised when a tool execution fails."""
|
||||
|
||||
def __init__(self, tool_name: str, message: str):
|
||||
self.tool_name = tool_name
|
||||
super().__init__(f"Tool '{tool_name}' execution failed: {message}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Decorator
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_default_registry: Optional[ToolRegistry] = None
|
||||
|
||||
|
||||
def tool(
|
||||
name: str,
|
||||
description: str,
|
||||
*,
|
||||
parameters: Optional[Dict[str, Any]] = None,
|
||||
registry: Optional[ToolRegistry] = None,
|
||||
) -> Callable[[Callable[..., Any]], Tool]:
|
||||
"""Decorator to register a function as a tool.
|
||||
|
||||
Usage
|
||||
-----
|
||||
>>> @tool("add", description="Add two numbers",
|
||||
... parameters={"a": {"type": "number"}, "b": {"type": "number"}})
|
||||
... def add(a: float, b: float) -> float:
|
||||
... return a + b
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
Unique tool name.
|
||||
description : str
|
||||
One-line description of what the tool does.
|
||||
parameters : dict, optional
|
||||
JSON Schema properties dict. If omitted the decorated function's
|
||||
signature is introspected to generate a basic schema.
|
||||
registry : ToolRegistry, optional
|
||||
Registry to add the tool to. Uses the global default if omitted.
|
||||
"""
|
||||
|
||||
def wrapper(func: Callable[..., Any]) -> Tool:
|
||||
params = parameters or _infer_schema(func)
|
||||
t = Tool(name=name, description=description, parameters=params, func=func)
|
||||
|
||||
target = registry if registry is not None else global_registry()
|
||||
target.register(t)
|
||||
return t
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def _infer_schema(func: Callable[..., Any]) -> Dict[str, Any]:
|
||||
"""Build a minimal JSON Schema from a function's signature."""
|
||||
import inspect
|
||||
|
||||
sig = inspect.signature(func)
|
||||
schema: Dict[str, Any] = {}
|
||||
for param_name, param in sig.parameters.items():
|
||||
ptype = "string"
|
||||
if param.annotation != inspect.Parameter.empty:
|
||||
ann = param.annotation.__name__ if hasattr(param.annotation, "__name__") else str(param.annotation)
|
||||
mapping = {
|
||||
"str": "string",
|
||||
"int": "integer",
|
||||
"float": "number",
|
||||
"bool": "boolean",
|
||||
"list": "array",
|
||||
"dict": "object",
|
||||
}
|
||||
ptype = mapping.get(ann.lower(), "string")
|
||||
schema[param_name] = {"type": ptype}
|
||||
return schema
|
||||
|
||||
|
||||
def global_registry() -> ToolRegistry:
|
||||
"""Return the singleton global ToolRegistry."""
|
||||
global _default_registry
|
||||
if _default_registry is None:
|
||||
_default_registry = ToolRegistry()
|
||||
return _default_registry
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ToolRegistry:
|
||||
"""Central registry for all available tools.
|
||||
|
||||
Tools are stored by name and can be looked up or enumerated.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._tools: Dict[str, Tool] = {}
|
||||
|
||||
# -- CRUD -------------------------------------------------------------------
|
||||
|
||||
def register(self, tool: Tool) -> None:
|
||||
"""Add or overwrite a tool in the registry."""
|
||||
self._tools[tool.name] = tool
|
||||
logger.debug("Registered tool: %s", tool.name)
|
||||
|
||||
def get(self, name: str) -> Tool:
|
||||
"""Retrieve a tool by name. Raises ``KeyError`` if not found."""
|
||||
return self._tools[name]
|
||||
|
||||
def list(self) -> List[str]:
|
||||
"""Return all registered tool names."""
|
||||
return list(self._tools.keys())
|
||||
|
||||
def remove(self, name: str) -> None:
|
||||
"""Remove a tool from the registry."""
|
||||
self._tools.pop(name, None)
|
||||
|
||||
# -- schema helpers -----------------------------------------------------------
|
||||
|
||||
def get_schema(self) -> List[Dict[str, Any]]:
|
||||
"""Return all tool definitions in OpenAI function-calling format."""
|
||||
return [t.to_openai_format() for t in self._tools.values()]
|
||||
|
||||
# -- execution ----------------------------------------------------------------
|
||||
|
||||
def execute(self, tool_name: str, **kwargs: Any) -> Any:
|
||||
"""Look up and execute a tool by name."""
|
||||
tool = self.get(tool_name)
|
||||
return tool.execute(**kwargs)
|
||||
|
||||
# -- convenience --------------------------------------------------------------
|
||||
|
||||
def register_from_func(
|
||||
self,
|
||||
func: Callable[..., Any],
|
||||
name: str,
|
||||
description: str,
|
||||
parameters: Optional[Dict[str, Any]] = None,
|
||||
) -> Tool:
|
||||
"""Register a plain function as a tool."""
|
||||
params = parameters or _infer_schema(func)
|
||||
t = Tool(name=name, description=description, parameters=params, func=func)
|
||||
self.register(t)
|
||||
return t
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Built-in tools
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _python_repl(code: str) -> str:
|
||||
"""Execute Python code and return stdout / stderr."""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["python", "-c", code],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
)
|
||||
output = result.stdout or ""
|
||||
if result.stderr:
|
||||
output += f"\nSTDERR:\n{result.stderr}"
|
||||
if result.returncode != 0:
|
||||
output = f"EXIT CODE {result.returncode}\n" + output
|
||||
return output.strip()
|
||||
except subprocess.TimeoutExpired:
|
||||
return "ERROR: execution timed out (30s)"
|
||||
except Exception as exc:
|
||||
return f"ERROR: {exc}"
|
||||
|
||||
|
||||
def _web_search(query: str) -> str:
|
||||
"""Perform a web search and return results.
|
||||
|
||||
This is a stub implementation — replace with DuckDuckGo, Tavily,
|
||||
Google Custom Search, or any preferred backend.
|
||||
"""
|
||||
# Stub: return a placeholder; swap for a real search backend.
|
||||
logger.warning("web_search is a stub — query: %s", query)
|
||||
return json.dumps({
|
||||
"query": query,
|
||||
"results": [
|
||||
{
|
||||
"title": "Stub result",
|
||||
"url": "https://example.com",
|
||||
"snippet": f"Search for '{query}' — replace with a real backend.",
|
||||
}
|
||||
],
|
||||
}, indent=2)
|
||||
|
||||
|
||||
def _file_read(file_path: str) -> str:
|
||||
"""Read the contents of a file."""
|
||||
path = Path(file_path)
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
if not path.is_file():
|
||||
raise IsADirectoryError(f"Not a file: {file_path}")
|
||||
return path.read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def _file_write(file_path: str, content: str) -> str:
|
||||
"""Write content to a file (creates directories if needed)."""
|
||||
path = Path(file_path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(content, encoding="utf-8")
|
||||
return f"Successfully wrote {len(content)} bytes to {file_path}"
|
||||
|
||||
|
||||
# -- Register built-in tools ----------------------------------------------------
|
||||
|
||||
_BUILTIN_PARAMS: Dict[str, Dict[str, Any]] = {
|
||||
"python_repl": {
|
||||
"code": {
|
||||
"type": "string",
|
||||
"description": "Python code to execute.",
|
||||
"required": True,
|
||||
},
|
||||
},
|
||||
"web_search": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "Search query string.",
|
||||
"required": True,
|
||||
},
|
||||
},
|
||||
"file_read": {
|
||||
"file_path": {
|
||||
"type": "string",
|
||||
"description": "Absolute or relative path to the file.",
|
||||
"required": True,
|
||||
},
|
||||
},
|
||||
"file_write": {
|
||||
"file_path": {
|
||||
"type": "string",
|
||||
"description": "Absolute or relative path to write to.",
|
||||
"required": True,
|
||||
},
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": "File content to write.",
|
||||
"required": True,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _register_builtins(registry: Optional[ToolRegistry] = None) -> ToolRegistry:
|
||||
"""Populate a registry with the standard built-in tools."""
|
||||
target = registry or global_registry()
|
||||
|
||||
target.register_from_func(
|
||||
_python_repl,
|
||||
name="python_repl",
|
||||
description="Execute Python code in a sandboxed subprocess and return the output.",
|
||||
parameters=_BUILTIN_PARAMS["python_repl"],
|
||||
)
|
||||
target.register_from_func(
|
||||
_web_search,
|
||||
name="web_search",
|
||||
description="Search the web for information. Returns a JSON list of results.",
|
||||
parameters=_BUILTIN_PARAMS["web_search"],
|
||||
)
|
||||
target.register_from_func(
|
||||
_file_read,
|
||||
name="file_read",
|
||||
description="Read the full contents of a file.",
|
||||
parameters=_BUILTIN_PARAMS["file_read"],
|
||||
)
|
||||
target.register_from_func(
|
||||
_file_write,
|
||||
name="file_write",
|
||||
description="Write content to a file, creating parent directories as needed.",
|
||||
parameters=_BUILTIN_PARAMS["file_write"],
|
||||
)
|
||||
return target
|
||||
|
||||
|
||||
# Ensure builtins are available on the global registry.
|
||||
_register_builtins()
|
||||
Reference in New Issue
Block a user