""" 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)