AI_engine/core/tool_registry.py

"""
Tool and Skill Registry for LogiFlow AI / Doormile Agent System.

Provides a unified Tool/Skill abstraction, JSONSchema parameter definition,
argument validation, and central registry for agent tool-use and human-in-the-loop interactions.
"""
from dataclasses import dataclass, field
from typing import Callable, Dict, Any, List, Optional, Awaitable, Union
import inspect
import json
from core.logger import logger


@dataclass
class ToolParameter:
    name: str
    type: str  # "string", "number", "integer", "boolean", "array", "object"
    description: str
    required: bool = True
    enum: Optional[List[Any]] = None
    default: Optional[Any] = None
    items: Optional[Dict[str, Any]] = None


@dataclass
class Tool:
    name: str
    description: str
    handler: Callable[..., Awaitable[Any]]
    parameters: List[ToolParameter] = field(default_factory=list)
    requires_confirmation: bool = False
    category: str = "general"
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_schema(self) -> Dict[str, Any]:
        """Generate an Anthropic Messages API tool definition (name / description / input_schema)."""
        properties = {}
        required = []

        for p in self.parameters:
            prop = {
                "type": p.type,
                "description": p.description,
            }
            if p.enum:
                prop["enum"] = p.enum
            if p.default is not None:
                prop["default"] = p.default
            if p.items:
                prop["items"] = p.items

            properties[p.name] = prop
            if p.required:
                required.append(p.name)

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

    async def execute(self, **kwargs) -> Any:
        """Validate required arguments and execute the handler."""
        for p in self.parameters:
            if p.required and p.name not in kwargs and p.default is None:
                raise ValueError(f"Missing required parameter '{p.name}' for tool '{self.name}'")

        # Inject default values if missing
        for p in self.parameters:
            if p.name not in kwargs and p.default is not None:
                kwargs[p.name] = p.default

        if inspect.iscoroutinefunction(self.handler):
            return await self.handler(**kwargs)
        return self.handler(**kwargs)


class ToolRegistry:
    """Central registry where tools and domain skills are registered and discovered."""

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

    def register(self, tool: Tool) -> Tool:
        """Register a Tool instance."""
        if tool.name in self._tools:
            logger.warning(f"Overwriting existing tool registration: {tool.name}")
        self._tools[tool.name] = tool
        self._categories.setdefault(tool.category, []).append(tool.name)
        logger.info(f"Registered tool: {tool.name} (category: {tool.category})")
        return tool

    def get(self, name: str) -> Optional[Tool]:
        """Retrieve a tool by name."""
        return self._tools.get(name)

    def list_tools(self, category: Optional[str] = None) -> List[Tool]:
        """List registered tools, optionally filtered by category."""
        if category:
            return [self._tools[name] for name in self._categories.get(category, [])]
        return list(self._tools.values())

    def get_schemas(self, category: Optional[str] = None) -> List[Dict[str, Any]]:
        """Return JSON Schemas for registered tools."""
        return [tool.to_schema() for tool in self.list_tools(category)]

    async def execute_tool(self, name: str, **kwargs) -> Any:
        """Execute a tool by name with provided arguments."""
        tool = self.get(name)
        if not tool:
            raise KeyError(f"Tool '{name}' is not registered in ToolRegistry")
        return await tool.execute(**kwargs)


# Global singleton registry
tool_registry = ToolRegistry()


def register_tool(
    name: str,
    description: str,
    parameters: Optional[List[ToolParameter]] = None,
    requires_confirmation: bool = False,
    category: str = "general",
    metadata: Optional[Dict[str, Any]] = None,
):
    """Decorator to easily register functions as tools."""
    def decorator(fn: Callable):
        tool = Tool(
            name=name,
            description=description,
            handler=fn,
            parameters=parameters or [],
            requires_confirmation=requires_confirmation,
            category=category,
            metadata=metadata or {},
        )
        tool_registry.register(tool)
        return fn
    return decorator