leocrm/app/plugins/builtins/ai_assistant/tool_registry.py

"""Global tool registry for AI Assistant plugin tools.

Plugins can register tools that AI agents can call during chat sessions.
Each tool declares a name, description, JSON schema for parameters,
and an async handler. Tools can optionally require specific RBAC permissions.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Any, Awaitable, Callable, Protocol

logger = logging.getLogger(__name__)


class ToolHandler(Protocol):
    async def __call__(
        self,
    arguments: dict[str, Any],
    context: dict[str, Any],
    ) -> str: ...


@dataclass
class AITool:
    """Represents a tool that an AI agent can call."""

    name: str
    description: str
    parameters: dict[str, Any]  # JSON Schema for parameters
    handler: ToolHandler
    plugin_name: str = ""
    required_permission: str | None = None  # e.g. "mail:send"
    category: str = "general"

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


class ToolRegistry:
    """Singleton registry for AI tools."""

    _instance: ToolRegistry | None = None

    def __new__(cls) -> ToolRegistry:
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._tools: dict[str, AITool] = {}
        return cls._instance

    def register(
        self,
    name: str,
    description: str,
    parameters: dict[str, Any],
    handler: ToolHandler,
    plugin_name: str = "",
    required_permission: str | None = None,
    category: str = "general",
    ) -> None:
        """Register a tool."""
        tool = AITool(
            name=name,
            description=description,
            parameters=parameters,
            handler=handler,
            plugin_name=plugin_name,
            required_permission=required_permission,
            category=category,
        )
        self._tools[name] = tool
        logger.info("AI tool registered: %s (plugin=%s)", name, plugin_name)

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

    def unregister_plugin(self, plugin_name: str) -> None:
        """Unregister all tools from a plugin."""
        to_remove = [
            name for name, tool in self._tools.items() if tool.plugin_name == plugin_name
        ]
        for name in to_remove:
            self._tools.pop(name, None)

    def get(self, name: str) -> AITool | None:
        return self._tools.get(name)

    def get_all(self) -> list[AITool]:
        return list(self._tools.values())

    def get_by_names(self, names: list[str]) -> list[AITool]:
        return [self._tools[name] for name in names if name in self._tools]

    def list_for_api(self) -> list[dict[str, Any]]:
        """Return tool list for API response."""
        return [
            {
                "name": tool.name,
                "description": tool.description,
                "parameters": tool.parameters,
                "plugin_name": tool.plugin_name,
                "required_permission": tool.required_permission,
                "category": tool.category,
            }
            for tool in self._tools.values()
        ]


def get_tool_registry() -> ToolRegistry:
    """Get the global tool registry singleton."""
    return ToolRegistry()