X-Agents/core/agents/tools.py

"""Tool system for agent capabilities."""

import asyncio
import logging
from abc import ABC, abstractmethod
from typing import Any

logger = logging.getLogger(__name__)


class Tool(ABC):
    """Abstract base class for agent tools."""

    @property
    @abstractmethod
    def name(self) -> str:
        """Tool name used in function calls."""
        pass

    @property
    @abstractmethod
    def description(self) -> str:
        """Description of what the tool does."""
        pass

    @property
    @abstractmethod
    def parameters(self) -> dict[str, Any]:
        """JSON Schema for tool parameters."""
        pass

    @abstractmethod
    async def execute(self, **kwargs: Any) -> str:
        """Execute the tool with given parameters.

        Returns:
            String result of the tool execution.
        """
        pass

    def to_schema(self) -> dict[str, Any]:
        """Convert tool to function schema format."""
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters,
            },
        }


class ToolRegistry:
    """Registry for managing agent tools."""

    def __init__(self):
        self._tools: dict[str, Tool] = {}

    def register(self, tool: Tool) -> None:
        """Register a tool.

        Args:
            tool: Tool instance to register
        """
        self._tools[tool.name] = tool
        logger.info(f"Registered tool: {tool.name}")

    def unregister(self, name: str) -> None:
        """Unregister a tool.

        Args:
            name: Tool name to unregister
        """
        if name in self._tools:
            del self._tools[name]
            logger.info(f"Unregistered tool: {name}")

    def get(self, name: str) -> Tool | None:
        """Get a tool by name.

        Args:
            name: Tool name

        Returns:
            Tool instance or None
        """
        return self._tools.get(name)

    def get_definitions(self) -> list[dict[str, Any]]:
        """Get all tool definitions for LLM.

        Returns:
            List of tool schemas
        """
        return [tool.to_schema() for tool in self._tools.values()]

    async def execute(self, name: str, arguments: dict[str, Any]) -> str:
        """Execute a tool.

        Args:
            name: Tool name
            arguments: Tool arguments

        Returns:
            Tool execution result
        """
        tool = self.get(name)
        if not tool:
            return f'{{"error": "Unknown tool: {name}"}}'

        try:
            # Validate parameters
            validated = tool.cast_params(arguments)
            errors = tool.validate_params(validated)
            if errors:
                return f'{{"error": "Parameter validation failed: {errors}"}}'

            # Execute with timeout
            result = await asyncio.wait_for(
                tool.execute(**validated),
                timeout=60.0,
            )
            return result
        except asyncio.TimeoutError:
            return f'{{"error": "Tool execution timed out: {name}"}}'
        except Exception as exc:
            logger.exception(f"Tool execution error: {name}")
            return f'{{"error": "Tool execution failed: {exc}"}}'

    def list_tools(self) -> list[str]:
        """List all registered tool names.

        Returns:
            List of tool names
        """
        return list(self._tools.keys())


# Built-in placeholder tools
class EchoTool(Tool):
    """Echo tool for testing."""

    @property
    def name(self) -> str:
        return "echo"

    @property
    def description(self) -> str:
        return "Echo back the input text. Useful for testing."

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "text": {
                    "type": "string",
                    "description": "Text to echo back",
                }
            },
            "required": ["text"],
        }

    async def execute(self, **kwargs: Any) -> str:
        text = kwargs.get("text", "")
        return f'{{"echo": "{text}"}}'


class TimeTool(Tool):
    """Get current time tool."""

    @property
    def name(self) -> str:
        return "get_time"

    @property
    def description(self) -> str:
        return "Get the current date and time."

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {},
        }

    async def execute(self, **kwargs: Any) -> str:
        from datetime import datetime
        now = datetime.now()
        return f'{{"time": "{now.isoformat()}"}}'


def create_default_registry() -> ToolRegistry:
    """Create a tool registry with default tools.

    Returns:
        Tool registry with built-in tools
    """
    registry = ToolRegistry()
    registry.register(EchoTool())
    registry.register(TimeTool())
    return registry