OwusuBlessing
81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
"""
|
|
Tool Registry for easy tool management and combination
|
|
"""
|
|
from typing import List, Dict, Any, Optional
|
|
from ..core.types import MCPTool
|
|
from .math_tools import MathTools
|
|
from .text_tools import TextTools
|
|
from .system_tools import SystemTools
|
|
from .web_tools import WebTools
|
|
|
|
|
|
class ToolRegistry:
|
|
"""Registry for managing and combining MCP tools from different categories"""
|
|
|
|
def __init__(self):
|
|
self._tool_categories = {
|
|
'math': MathTools,
|
|
'text': TextTools,
|
|
'system': SystemTools,
|
|
'web': WebTools,
|
|
}
|
|
self._custom_tools: List[MCPTool] = []
|
|
|
|
def get_tools_by_category(self, category: str) -> List[MCPTool]:
|
|
"""Get all tools from a specific category"""
|
|
if category not in self._tool_categories:
|
|
raise ValueError(f"Unknown tool category: {category}")
|
|
|
|
return self._tool_categories[category].get_tools()
|
|
|
|
def get_tools_by_categories(self, categories: List[str]) -> List[MCPTool]:
|
|
"""Get tools from multiple categories"""
|
|
tools = []
|
|
for category in categories:
|
|
tools.extend(self.get_tools_by_category(category))
|
|
return tools
|
|
|
|
def get_all_tools(self) -> List[MCPTool]:
|
|
"""Get all tools from all categories"""
|
|
tools = []
|
|
for category in self._tool_categories.values():
|
|
tools.extend(category.get_tools())
|
|
tools.extend(self._custom_tools)
|
|
return tools
|
|
|
|
def add_custom_tool(self, tool: MCPTool) -> None:
|
|
"""Add a custom tool to the registry"""
|
|
self._custom_tools.append(tool)
|
|
|
|
def add_custom_tools(self, tools: List[MCPTool]) -> None:
|
|
"""Add multiple custom tools to the registry"""
|
|
self._custom_tools.extend(tools)
|
|
|
|
def get_available_categories(self) -> List[str]:
|
|
"""Get list of available tool categories"""
|
|
return list(self._tool_categories.keys())
|
|
|
|
def get_category_info(self) -> Dict[str, Dict[str, Any]]:
|
|
"""Get information about each category"""
|
|
info = {}
|
|
for category_name, category_class in self._tool_categories.items():
|
|
tools = category_class.get_tools()
|
|
info[category_name] = {
|
|
'tool_count': len(tools),
|
|
'tool_names': [tool.name for tool in tools]
|
|
}
|
|
return info
|
|
|
|
def create_server_config(self, categories: Optional[List[str]] = None) -> Dict[str, Any]:
|
|
"""Create a server configuration with tools from specified categories"""
|
|
if categories is None:
|
|
tools = self.get_all_tools()
|
|
else:
|
|
tools = self.get_tools_by_categories(categories)
|
|
|
|
return {
|
|
'tools': tools,
|
|
'tool_count': len(tools),
|
|
'categories': categories or self.get_available_categories()
|
|
}
|