f71e31585b
- Single harness for all API interactions - Unified tool registry with routing - System, Git, Network tool layers - Local-first, self-healing design
266 lines
8.0 KiB
Python
266 lines
8.0 KiB
Python
"""
|
|
Uni-Wizard Tool Registry
|
|
Central registry for all tool capabilities
|
|
"""
|
|
|
|
import json
|
|
import inspect
|
|
from typing import Dict, Callable, Any, Optional
|
|
from dataclasses import dataclass, asdict
|
|
from functools import wraps
|
|
|
|
|
|
@dataclass
|
|
class ToolSchema:
|
|
"""Schema definition for a tool"""
|
|
name: str
|
|
description: str
|
|
parameters: Dict[str, Any]
|
|
returns: str
|
|
examples: list = None
|
|
|
|
def to_dict(self):
|
|
return asdict(self)
|
|
|
|
|
|
@dataclass
|
|
class ToolResult:
|
|
"""Standardized tool execution result"""
|
|
success: bool
|
|
data: Any
|
|
error: Optional[str] = None
|
|
execution_time_ms: Optional[float] = None
|
|
|
|
def to_json(self) -> str:
|
|
return json.dumps({
|
|
'success': self.success,
|
|
'data': self.data,
|
|
'error': self.error,
|
|
'execution_time_ms': self.execution_time_ms
|
|
}, indent=2)
|
|
|
|
def __str__(self) -> str:
|
|
if self.success:
|
|
return str(self.data)
|
|
return f"Error: {self.error}"
|
|
|
|
|
|
class ToolRegistry:
|
|
"""
|
|
Central registry for all uni-wizard tools.
|
|
|
|
All tools register here with their schemas.
|
|
The LLM queries available tools via get_tool_definitions().
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._tools: Dict[str, Dict] = {}
|
|
self._categories: Dict[str, list] = {}
|
|
|
|
def register(
|
|
self,
|
|
name: str,
|
|
handler: Callable,
|
|
description: str = None,
|
|
parameters: Dict = None,
|
|
category: str = "general",
|
|
examples: list = None
|
|
):
|
|
"""
|
|
Register a tool in the registry.
|
|
|
|
Args:
|
|
name: Tool name (used in tool calls)
|
|
handler: Function to execute
|
|
description: What the tool does
|
|
parameters: JSON Schema for parameters
|
|
category: Tool category (system, git, network, file)
|
|
examples: Example usages
|
|
"""
|
|
# Auto-extract description from docstring if not provided
|
|
if description is None and handler.__doc__:
|
|
description = handler.__doc__.strip().split('\n')[0]
|
|
|
|
# Auto-extract parameters from function signature
|
|
if parameters is None:
|
|
parameters = self._extract_params(handler)
|
|
|
|
self._tools[name] = {
|
|
'name': name,
|
|
'handler': handler,
|
|
'description': description or f"Execute {name}",
|
|
'parameters': parameters,
|
|
'category': category,
|
|
'examples': examples or []
|
|
}
|
|
|
|
# Add to category
|
|
if category not in self._categories:
|
|
self._categories[category] = []
|
|
self._categories[category].append(name)
|
|
|
|
return self # For chaining
|
|
|
|
def _extract_params(self, handler: Callable) -> Dict:
|
|
"""Extract parameter schema from function signature"""
|
|
sig = inspect.signature(handler)
|
|
params = {
|
|
"type": "object",
|
|
"properties": {},
|
|
"required": []
|
|
}
|
|
|
|
for name, param in sig.parameters.items():
|
|
# Skip 'self', 'cls', and params with defaults
|
|
if name in ('self', 'cls'):
|
|
continue
|
|
|
|
param_info = {"type": "string"} # Default
|
|
|
|
# Try to infer type from annotation
|
|
if param.annotation != inspect.Parameter.empty:
|
|
if param.annotation == int:
|
|
param_info["type"] = "integer"
|
|
elif param.annotation == float:
|
|
param_info["type"] = "number"
|
|
elif param.annotation == bool:
|
|
param_info["type"] = "boolean"
|
|
elif param.annotation == list:
|
|
param_info["type"] = "array"
|
|
elif param.annotation == dict:
|
|
param_info["type"] = "object"
|
|
|
|
# Add description if in docstring
|
|
if handler.__doc__:
|
|
# Simple param extraction from docstring
|
|
for line in handler.__doc__.split('\n'):
|
|
if f'{name}:' in line or f'{name} (' in line:
|
|
desc = line.split(':', 1)[-1].strip()
|
|
param_info["description"] = desc
|
|
break
|
|
|
|
params["properties"][name] = param_info
|
|
|
|
# Required if no default
|
|
if param.default == inspect.Parameter.empty:
|
|
params["required"].append(name)
|
|
|
|
return params
|
|
|
|
def execute(self, name: str, **params) -> ToolResult:
|
|
"""
|
|
Execute a tool by name with parameters.
|
|
|
|
Args:
|
|
name: Tool name
|
|
**params: Tool parameters
|
|
|
|
Returns:
|
|
ToolResult with success/failure and data
|
|
"""
|
|
import time
|
|
start = time.time()
|
|
|
|
tool = self._tools.get(name)
|
|
if not tool:
|
|
return ToolResult(
|
|
success=False,
|
|
data=None,
|
|
error=f"Tool '{name}' not found in registry",
|
|
execution_time_ms=(time.time() - start) * 1000
|
|
)
|
|
|
|
try:
|
|
handler = tool['handler']
|
|
result = handler(**params)
|
|
|
|
return ToolResult(
|
|
success=True,
|
|
data=result,
|
|
execution_time_ms=(time.time() - start) * 1000
|
|
)
|
|
|
|
except Exception as e:
|
|
return ToolResult(
|
|
success=False,
|
|
data=None,
|
|
error=f"{type(e).__name__}: {str(e)}",
|
|
execution_time_ms=(time.time() - start) * 1000
|
|
)
|
|
|
|
def get_tool(self, name: str) -> Optional[Dict]:
|
|
"""Get tool definition by name"""
|
|
tool = self._tools.get(name)
|
|
if tool:
|
|
# Return without handler (not serializable)
|
|
return {
|
|
'name': tool['name'],
|
|
'description': tool['description'],
|
|
'parameters': tool['parameters'],
|
|
'category': tool['category'],
|
|
'examples': tool['examples']
|
|
}
|
|
return None
|
|
|
|
def get_tools_by_category(self, category: str) -> list:
|
|
"""Get all tools in a category"""
|
|
tool_names = self._categories.get(category, [])
|
|
return [self.get_tool(name) for name in tool_names if self.get_tool(name)]
|
|
|
|
def list_tools(self, category: str = None) -> list:
|
|
"""List all tool names, optionally filtered by category"""
|
|
if category:
|
|
return self._categories.get(category, [])
|
|
return list(self._tools.keys())
|
|
|
|
def get_tool_definitions(self) -> str:
|
|
"""
|
|
Get all tool definitions formatted for LLM system prompt.
|
|
Returns JSON string of all tools with schemas.
|
|
"""
|
|
tools = []
|
|
for name, tool in self._tools.items():
|
|
tools.append({
|
|
"name": name,
|
|
"description": tool['description'],
|
|
"parameters": tool['parameters']
|
|
})
|
|
|
|
return json.dumps(tools, indent=2)
|
|
|
|
def get_categories(self) -> list:
|
|
"""Get all tool categories"""
|
|
return list(self._categories.keys())
|
|
|
|
|
|
# Global registry instance
|
|
registry = ToolRegistry()
|
|
|
|
|
|
def tool(name: str = None, category: str = "general", examples: list = None):
|
|
"""
|
|
Decorator to register a function as a tool.
|
|
|
|
Usage:
|
|
@tool(category="system")
|
|
def system_info():
|
|
return {...}
|
|
"""
|
|
def decorator(func: Callable):
|
|
tool_name = name or func.__name__
|
|
registry.register(
|
|
name=tool_name,
|
|
handler=func,
|
|
category=category,
|
|
examples=examples
|
|
)
|
|
return func
|
|
return decorator
|
|
|
|
|
|
# Convenience function for quick tool execution
|
|
def call_tool(name: str, **params) -> str:
|
|
"""Execute a tool and return string result"""
|
|
result = registry.execute(name, **params)
|
|
return str(result)
|