mirror of https://github.com/kortix-ai/suna.git
69 lines
2.4 KiB
Python
69 lines
2.4 KiB
Python
|
from .base_tool import AgentBuilderBaseTool
|
||
|
|
from .agent_config_tool import AgentConfigTool
|
||
|
|
from .mcp_search_tool import MCPSearchTool
|
||
|
|
from .credential_profile_tool import CredentialProfileTool
|
||
|
|
from .workflow_tool import WorkflowTool
|
||
|
|
from typing import List, Type, Dict, Any
|
||
|
|
from agentpress.thread_manager import ThreadManager
|
||
|
|
from utils.logger import logger
|
||
|
|
|
||
|
|
|
||
|
|
class AgentBuilderToolRegistry:
|
||
|
|
"""Registry for managing and registering agent builder tools."""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self.tools: Dict[str, Type[AgentBuilderBaseTool]] = {
|
||
|
|
'agent_config': AgentConfigTool,
|
||
|
|
'mcp_search': MCPSearchTool,
|
||
|
|
'credential_profile': CredentialProfileTool,
|
||
|
|
'workflow': WorkflowTool,
|
||
|
|
}
|
||
|
|
|
||
|
|
def register_tool(self, name: str, tool_class: Type[AgentBuilderBaseTool]):
|
||
|
|
"""Register a new agent builder tool."""
|
||
|
|
self.tools[name] = tool_class
|
||
|
|
logger.info(f"Registered agent builder tool: {name}")
|
||
|
|
|
||
|
|
def get_tool(self, name: str) -> Type[AgentBuilderBaseTool]:
|
||
|
|
"""Get a tool class by name."""
|
||
|
|
return self.tools.get(name)
|
||
|
|
|
||
|
|
def get_all_tools(self) -> Dict[str, Type[AgentBuilderBaseTool]]:
|
||
|
|
"""Get all registered tools."""
|
||
|
|
return self.tools.copy()
|
||
|
|
|
||
|
|
def register_all_tools(self, thread_manager: ThreadManager, db_connection, agent_id: str):
|
||
|
|
"""Register all agent builder tools with the thread manager."""
|
||
|
|
logger.info(f"Registering {len(self.tools)} agent builder tools")
|
||
|
|
|
||
|
|
for tool_name, tool_class in self.tools.items():
|
||
|
|
try:
|
||
|
|
thread_manager.add_tool(
|
||
|
|
tool_class,
|
||
|
|
thread_manager=thread_manager,
|
||
|
|
db_connection=db_connection,
|
||
|
|
agent_id=agent_id
|
||
|
|
)
|
||
|
|
logger.info(f"Successfully registered agent builder tool: {tool_name}")
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Failed to register agent builder tool {tool_name}: {e}")
|
||
|
|
|
||
|
|
def list_available_tools(self) -> List[str]:
|
||
|
|
"""List all available tool names."""
|
||
|
|
return list(self.tools.keys())
|
||
|
|
|
||
|
|
|
||
|
|
# Create a global registry instance
|
||
|
|
agent_builder_registry = AgentBuilderToolRegistry()
|
||
|
|
|
||
|
|
# Export commonly used items
|
||
|
|
__all__ = [
|
||
|
|
'AgentBuilderBaseTool',
|
||
|
|
'AgentConfigTool',
|
||
|
|
'MCPSearchTool',
|
||
|
|
'CredentialProfileTool',
|
||
|
|
'WorkflowTool',
|
||
|
|
'AgentBuilderToolRegistry',
|
||
|
|
'agent_builder_registry'
|
||
|
|
]
|