refactor(block-h): move AI tool registry into core AI layer
This commit is contained in:
Agent Zero
@@ -50,7 +50,7 @@ if TYPE_CHECKING:
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.plugins.builtins.ai_assistant.tool_registry import ToolRegistry
|
||||
from app.ai.tool_registry import ToolRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -61,7 +61,7 @@ def _resolve_effective_tool_ids(
|
||||
orchestrate tools but never grant additional permissions.
|
||||
"""
|
||||
from app.ai.skill_registry import get_skill_registry
|
||||
from app.plugins.builtins.ai_assistant.contracts import get_tool_registry
|
||||
from app.ai.tool_registry import get_tool_registry
|
||||
|
||||
agent_tool_ids: list[str] = list(getattr(agent_definition, "tool_ids", None) or [])
|
||||
agent_skill_ids: list[str] = list(getattr(agent_definition, "skill_ids", None) or [])
|
||||
|
||||
@@ -221,7 +221,7 @@ async def build_agent_context(
|
||||
tool_descriptions: list[dict[str, str]] = []
|
||||
tool_ids = list(getattr(agent_definition, "tool_ids", None) or [])
|
||||
try:
|
||||
from app.plugins.builtins.ai_assistant.contracts import get_tool_registry
|
||||
from app.ai.tool_registry import get_tool_registry
|
||||
|
||||
registry = get_tool_registry()
|
||||
if tool_ids:
|
||||
|
||||
@@ -1,170 +0,0 @@
|
||||
"""Integration tools — Agent → Workflow + Agent → Knowledge.
|
||||
|
||||
Registers AI tools that allow agents to start workflows and ask knowledge questions.
|
||||
Builds on: workflow_service, knowledge services, tool_registry.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import logging
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def register_integration_tools() -> None:
|
||||
"""Register workflow + knowledge tools in the AI tool registry."""
|
||||
from app.plugins.builtins.ai_assistant.tool_registry import get_tool_registry
|
||||
registry = get_tool_registry()
|
||||
|
||||
# ── I-AW: Agent → Workflow ──
|
||||
async def _start_workflow_handler(arguments: dict[str, Any], context: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Start a workflow by ID."""
|
||||
from app.services.workflow_service import create_instance
|
||||
from app.core.db import get_worker_session_factory
|
||||
workflow_id = arguments.get("workflow_id", "")
|
||||
tenant_id = context.get("tenant_id")
|
||||
user_id = context.get("user_id")
|
||||
if not workflow_id or not tenant_id:
|
||||
return {"error": "workflow_id and tenant_id required"}
|
||||
factory = get_worker_session_factory()
|
||||
async with factory() as db:
|
||||
instance = await create_instance(
|
||||
db=db,
|
||||
tenant_id=uuid.UUID(str(tenant_id)),
|
||||
workflow_id=uuid.UUID(workflow_id),
|
||||
initiated_by=uuid.UUID(str(user_id)) if user_id else None,
|
||||
)
|
||||
await db.commit()
|
||||
return {"instance_id": str(instance.get("id", "")), "status": instance.get("status", "created")}
|
||||
|
||||
registry.register(
|
||||
name="start_workflow",
|
||||
description="Start a workflow by its ID. Returns the instance ID and status.",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"workflow_id": {"type": "string", "description": "UUID of the workflow to start"},
|
||||
},
|
||||
"required": ["workflow_id"],
|
||||
},
|
||||
handler=_start_workflow_handler,
|
||||
plugin_name="integration",
|
||||
required_permission="workflows:read",
|
||||
category="workflow",
|
||||
)
|
||||
|
||||
async def _check_workflow_status_handler(arguments: dict[str, Any], context: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Check the status of a workflow instance."""
|
||||
from sqlalchemy import select
|
||||
from app.models.workflow import WorkflowInstance
|
||||
from app.core.db import get_worker_session_factory
|
||||
instance_id = arguments.get("instance_id", "")
|
||||
tenant_id = context.get("tenant_id")
|
||||
if not instance_id or not tenant_id:
|
||||
return {"error": "instance_id and tenant_id required"}
|
||||
factory = get_worker_session_factory()
|
||||
async with factory() as db:
|
||||
result = await db.execute(
|
||||
select(WorkflowInstance).where(
|
||||
WorkflowInstance.id == uuid.UUID(instance_id),
|
||||
WorkflowInstance.tenant_id == uuid.UUID(str(tenant_id)),
|
||||
)
|
||||
)
|
||||
inst = result.scalar_one_or_none()
|
||||
if not inst:
|
||||
return {"error": "Instance not found"}
|
||||
return {
|
||||
"instance_id": str(inst.id),
|
||||
"status": inst.status,
|
||||
"current_step": inst.current_step_index,
|
||||
"completed_at": inst.completed_at.isoformat() if inst.completed_at else None,
|
||||
}
|
||||
|
||||
registry.register(
|
||||
name="check_workflow_status",
|
||||
description="Check the status of a workflow instance by its ID.",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"instance_id": {"type": "string", "description": "UUID of the workflow instance"},
|
||||
},
|
||||
"required": ["instance_id"],
|
||||
},
|
||||
handler=_check_workflow_status_handler,
|
||||
plugin_name="integration",
|
||||
required_permission="workflows:read",
|
||||
category="workflow",
|
||||
)
|
||||
|
||||
# ── I-AK: Agent → Knowledge ──
|
||||
async def _ask_knowledge_handler(arguments: dict[str, Any], context: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Ask a knowledge question."""
|
||||
from app.plugins.builtins.knowledge.services import ask_knowledge
|
||||
from app.core.db import get_worker_session_factory
|
||||
question = arguments.get("question", "")
|
||||
tenant_id = context.get("tenant_id")
|
||||
if not question or not tenant_id:
|
||||
return {"error": "question and tenant_id required"}
|
||||
factory = get_worker_session_factory()
|
||||
async with factory() as db:
|
||||
result = await ask_knowledge(db=db, tenant_id=uuid.UUID(str(tenant_id)), question=question)
|
||||
return {"answer": result.get("answer", ""), "evidence_count": len(result.get("evidence", []))}
|
||||
|
||||
registry.register(
|
||||
name="ask_knowledge",
|
||||
description="Ask a knowledge question. Searches wiki articles and graph relationships for context.",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"question": {"type": "string", "description": "The question to ask"},
|
||||
},
|
||||
"required": ["question"],
|
||||
},
|
||||
handler=_ask_knowledge_handler,
|
||||
plugin_name="integration",
|
||||
required_permission="wiki:read",
|
||||
category="knowledge",
|
||||
)
|
||||
|
||||
async def _search_knowledge_handler(arguments: dict[str, Any], context: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Search wiki articles via unified search."""
|
||||
from app.plugins.builtins.unified_search.provider_registry import get_search_registry
|
||||
from app.core.db import get_worker_session_factory
|
||||
query = arguments.get("query", "")
|
||||
tenant_id = context.get("tenant_id")
|
||||
if not query or not tenant_id:
|
||||
return {"error": "query and tenant_id required"}
|
||||
factory = get_worker_session_factory()
|
||||
async with factory() as db:
|
||||
registry = get_search_registry()
|
||||
wiki_provider = registry.get("wiki_article")
|
||||
if not wiki_provider:
|
||||
return {"error": "Wiki search provider not available"}
|
||||
results = await wiki_provider._search_fts_filtered(
|
||||
db=db, tsquery=query, tenant_id=uuid.UUID(str(tenant_id)), limit=5, visible_ids=None
|
||||
)
|
||||
return {
|
||||
"results": [
|
||||
{"title": r.get("title", ""), "summary": (r.get("summary") or r.get("content", "")[:200] or "")}
|
||||
for r in results
|
||||
],
|
||||
"count": len(results),
|
||||
}
|
||||
|
||||
registry.register(
|
||||
name="search_knowledge",
|
||||
description="Search wiki articles by keyword. Returns matching articles with title and summary.",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {"type": "string", "description": "Search query"},
|
||||
},
|
||||
"required": ["query"],
|
||||
},
|
||||
handler=_search_knowledge_handler,
|
||||
plugin_name="integration",
|
||||
required_permission="wiki:read",
|
||||
category="knowledge",
|
||||
)
|
||||
|
||||
logger.info("Registered integration tools: start_workflow, check_workflow_status, ask_knowledge, search_knowledge")
|
||||
@@ -0,0 +1,126 @@
|
||||
"""Global tool registry for AI agent tools (core platform service).
|
||||
|
||||
Plugins register tools here so AI agents can call them during chat sessions.
|
||||
Each tool declares a name, description, JSON schema for parameters,
|
||||
and an async handler. Tools can optionally require specific RBAC permissions.
|
||||
|
||||
This registry lives in the core AI layer (not inside a plugin) so that the
|
||||
agent runtime keeps working regardless of which optional plugins are active.
|
||||
Plugins contribute tools via ``register()`` / ``unregister_plugin()`` during
|
||||
their activate/deactivate lifecycle.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, 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()
|
||||
@@ -1,121 +1,19 @@
|
||||
"""Global tool registry for AI Assistant plugin tools.
|
||||
"""Compatibility shim — the tool registry moved to the core AI layer.
|
||||
|
||||
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.
|
||||
The agent runtime must not depend on this optional plugin being active,
|
||||
so ``ToolRegistry`` / ``AITool`` / ``get_tool_registry`` now live in
|
||||
``app.ai.tool_registry``. Import from here still works for existing
|
||||
plugin code; new code should import from ``app.ai.tool_registry``
|
||||
directly (or via the ai_assistant contract).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Protocol
|
||||
from app.ai.tool_registry import ( # noqa: F401
|
||||
AITool,
|
||||
ToolHandler,
|
||||
ToolRegistry,
|
||||
get_tool_registry,
|
||||
)
|
||||
|
||||
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()
|
||||
__all__ = ["AITool", "ToolHandler", "ToolRegistry", "get_tool_registry"]
|
||||
|
||||
Reference in New Issue
Block a user