from __future__ import annotations from collections.abc import Callable from typing import Any from supply_agent.tools.base import Tool from supply_agent.types import ToolDefinition, ToolResult class ToolRegistry: """Registry for managing and executing agent tools.""" def __init__(self) -> None: self._tools: dict[str, Tool] = {} def register( self, func: Callable[..., Any], name: str | None = None, description: str | None = None, ) -> Tool: """Register a function as a tool.""" tool = Tool(func, name=name, description=description) self._tools[tool.name] = tool return tool def register_tool(self, tool: Tool) -> None: """Register an existing Tool instance.""" self._tools[tool.name] = tool def unregister(self, name: str) -> None: self._tools.pop(name, None) def get(self, name: str) -> Tool | None: return self._tools.get(name) def list_tools(self) -> list[str]: return list(self._tools.keys()) @property def definitions(self) -> list[ToolDefinition]: return [t.definition for t in self._tools.values()] def execute(self, tool_call_id: str, name: str, arguments: str) -> ToolResult: """Execute a tool by name and return the result.""" tool = self._tools.get(name) if not tool: return ToolResult( tool_call_id=tool_call_id, name=name, content=f'{{"error": "Tool \'{name}\' not found"}}', is_error=True, ) content = tool.execute(arguments) is_error = content.startswith('{"error"') return ToolResult( tool_call_id=tool_call_id, name=name, content=content, is_error=is_error, ) async def aexecute(self, tool_call_id: str, name: str, arguments: str) -> ToolResult: """Async execute a tool.""" tool = self._tools.get(name) if not tool: return ToolResult( tool_call_id=tool_call_id, name=name, content=f'{{"error": "Tool \'{name}\' not found"}}', is_error=True, ) content = await tool.aexecute(arguments) is_error = content.startswith('{"error"') return ToolResult( tool_call_id=tool_call_id, name=name, content=content, is_error=is_error, ) def from_decorated(self, *funcs: Callable[..., Any]) -> ToolRegistry: """Register multiple @tool-decorated functions.""" for func in funcs: if getattr(func, "_is_tool", False): self.register(func) return self def __len__(self) -> int: return len(self._tools) def __contains__(self, name: str) -> bool: return name in self._tools