tool_registry.py•4.23 kB
"""Auto-discovery and registration system for MCP tools."""
import json
from pathlib import Path
from typing import Any, Dict, List, Optional, Type
from src.tools.base_tool import BaseTool
from src.utils.logger import get_logger
logger = get_logger(__name__)
class ToolRegistry:
"""
Auto-discovery and registration system for MCP tools.
Loads tool metadata from JSON files and registers tool instances.
"""
def __init__(self, metadata_dir: Path):
"""
Initialize tool registry.
Args:
metadata_dir: Directory containing tool metadata JSON files
"""
self.metadata_dir = metadata_dir
self.tools: Dict[str, BaseTool] = {}
self.tool_metadata: Dict[str, Dict[str, Any]] = {}
def load_metadata(self) -> Dict[str, Dict[str, Any]]:
"""
Load all tool metadata from JSON files.
Returns:
Dictionary mapping tool names to metadata
"""
metadata = {}
if not self.metadata_dir.exists():
logger.warning("metadata_directory_not_found", path=str(self.metadata_dir))
return metadata
for json_file in self.metadata_dir.glob("*.json"):
try:
with open(json_file, "r") as f:
tool_meta = json.load(f)
tool_name = tool_meta.get("name")
if tool_name:
metadata[tool_name] = tool_meta
logger.debug("metadata_loaded", tool=tool_name, file=json_file.name)
except Exception as e:
logger.error("metadata_load_error", file=json_file.name, error=str(e))
logger.info("metadata_loaded_count", count=len(metadata))
return metadata
def register_tool(self, tool_instance: BaseTool) -> None:
"""
Register a tool instance with the registry.
Args:
tool_instance: Tool instance to register
"""
tool_name = tool_instance.get_name()
# Load metadata for this tool
if tool_name not in self.tool_metadata:
logger.warning("metadata_not_found_for_tool", tool=tool_name)
return
metadata = self.tool_metadata[tool_name]
# Generate MCP schema from metadata
mcp_schema = self._generate_mcp_schema(metadata)
# Register tool
self.tools[tool_name] = tool_instance
logger.info("tool_registered", tool=tool_name)
def _generate_mcp_schema(self, metadata: Dict[str, Any]) -> Dict[str, Any]:
"""
Generate MCP-compatible tool schema from metadata.
Args:
metadata: Tool metadata
Returns:
MCP schema
"""
return {
"name": metadata["name"],
"description": metadata["description"],
"inputSchema": metadata["input_schema"],
}
def get_tool(self, tool_name: str) -> Optional[BaseTool]:
"""
Get a registered tool instance by name.
Args:
tool_name: Name of the tool
Returns:
Tool instance or None if not found
"""
return self.tools.get(tool_name)
def list_tools(self) -> List[Dict[str, Any]]:
"""
List all registered tools with their MCP schemas.
Returns:
List of tool schemas
"""
schemas = []
for tool_name, tool_instance in self.tools.items():
if tool_name in self.tool_metadata:
metadata = self.tool_metadata[tool_name]
schemas.append(self._generate_mcp_schema(metadata))
return schemas
def get_tool_metadata(self, tool_name: str) -> Optional[Dict[str, Any]]:
"""
Get metadata for a specific tool.
Args:
tool_name: Name of the tool
Returns:
Tool metadata or None if not found
"""
return self.tool_metadata.get(tool_name)
def initialize(self) -> int:
"""
Initialize registry by loading metadata.
Returns:
Number of metadata files loaded
"""
self.tool_metadata = self.load_metadata()
return len(self.tool_metadata)