SupplyAgent/supply_agent/tools/registry.py

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