src/tools/base.py

from __future__ import annotations
"""Base classes for tool definitions and registry."""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any

from src.utils.exceptions import ToolExecutionError, ToolNotFoundError
from src.utils.logging import get_logger

logger = get_logger(__name__)


@dataclass
class ToolResult:
    """Result from a tool execution."""

    success: bool
    data: Any
    error: str | None = None

    @classmethod
    def ok(cls, data: Any) -> "ToolResult":
        """Create a successful result."""
        return cls(success=True, data=data, error=None)

    @classmethod
    def fail(cls, error: str) -> "ToolResult":
        """Create a failed result."""
        return cls(success=False, data=None, error=error)


@dataclass
class ToolParameter:
    """Definition of a tool parameter."""

    name: str
    type: str
    description: str
    required: bool = True
    default: Any = None
    enum: list[str] | None = None


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

    name: str
    description: str
    parameters: list[ToolParameter] = field(default_factory=list)

    @abstractmethod
    async def execute(self, **kwargs: Any) -> ToolResult:
        """Execute the tool with the given parameters.
        
        Args:
            **kwargs: Tool parameters
            
        Returns:
            ToolResult with success/failure and data
        """
        pass

    def to_openai_schema(self) -> dict[str, Any]:
        """Convert tool to OpenAI function calling schema.
        
        Returns:
            OpenAI-compatible tool definition
        """
        properties = {}
        required = []

        for param in self.parameters:
            prop: dict[str, Any] = {
                "type": param.type,
                "description": param.description,
            }
            if param.enum:
                prop["enum"] = param.enum
            properties[param.name] = prop

            if param.required:
                required.append(param.name)

        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": {
                    "type": "object",
                    "properties": properties,
                    "required": required,
                },
            },
        }

    def validate_parameters(self, **kwargs: Any) -> tuple[bool, str | None]:
        """Validate that required parameters are provided.
        
        Args:
            **kwargs: Parameters to validate
            
        Returns:
            Tuple of (is_valid, error_message)
        """
        for param in self.parameters:
            if param.required and param.name not in kwargs:
                return False, f"Missing required parameter: {param.name}"
        return True, None


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

    def __init__(self) -> None:
        """Initialize empty tool registry."""
        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 get(self, name: str) -> Tool:
        """Get a tool by name.
        
        Args:
            name: Tool name
            
        Returns:
            Tool instance
            
        Raises:
            ToolNotFoundError: If tool not found
        """
        if name not in self._tools:
            raise ToolNotFoundError(f"Tool not found: {name}")
        return self._tools[name]

    def list_tools(self) -> list[Tool]:
        """List all registered tools.
        
        Returns:
            List of all tools
        """
        return list(self._tools.values())

    def get_schemas(self) -> list[dict[str, Any]]:
        """Get OpenAI schemas for all tools.
        
        Returns:
            List of tool schemas
        """
        return [tool.to_openai_schema() for tool in self._tools.values()]

    async def execute(self, tool_name: str, **kwargs: Any) -> ToolResult:
        """Execute a tool by name.
        
        Args:
            tool_name: Name of the tool to execute
            **kwargs: Tool parameters
            
        Returns:
            ToolResult from execution
        """
        tool = self.get(tool_name)

        # Validate parameters
        is_valid, error = tool.validate_parameters(**kwargs)
        if not is_valid:
            return ToolResult.fail(error or "Invalid parameters")

        try:
            logger.info(f"Executing tool: {tool_name} with params: {kwargs}")
            result = await tool.execute(**kwargs)
            logger.info(f"Tool {tool_name} completed: success={result.success}")
            return result
        except Exception as e:
            logger.error(f"Tool {tool_name} failed: {e}")
            return ToolResult.fail(str(e))


def create_default_registry() -> ToolRegistry:
    """Create a registry with all default tools.
    
    Returns:
        ToolRegistry with default tools registered
    """
    from src.tools.web_search import WebSearchTool
    from src.tools.web_scraper import WebScraperTool
    from src.tools.calculator import CalculatorTool
    from src.tools.datetime_tool import DateTimeTool

    registry = ToolRegistry()
    registry.register(WebSearchTool())
    registry.register(WebScraperTool())
    registry.register(CalculatorTool())
    registry.register(DateTimeTool())

    return registry