AI Assistant plugin: backend with LiteLLM, agents, tools, streaming chat

app/plugins/builtins/ai_assistant/tool_registry.py

@@ -0,0 +1,121 @@
+"""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()