MCP Server Hero

"""Tool registry for managing MCP tools.""" import inspect from typing import Any, Callable import mcp.types as types from .base import BaseTool class ToolRegistry: """Registry for managing MCP tools.""" def __init__(self) -> None: """Initialize the tool registry.""" self._tools: dict[str, BaseTool | Callable[..., Any]] = {} self._schemas: dict[str, dict[str, Any]] = {} self._descriptions: dict[str, str] = {} def register(self, name: str, tool_func: Any, description: str = "", schema: dict[str, Any] | None = None) -> None: """Register a tool. Args: name: Tool name tool_func: Tool function or BaseTool instance description: Tool description schema: JSON schema for tool arguments """ if isinstance(tool_func, BaseTool): self._tools[name] = tool_func else: self._tools[name] = tool_func self._descriptions[name] = description or self._extract_description(tool_func) self._schemas[name] = schema or self._generate_schema(tool_func) def _extract_description(self, func: Callable[..., Any]) -> str: """Extract description from function docstring. Args: func: Function to extract description from Returns: Function description """ return (func.__doc__ or "").strip() def _generate_schema(self, func: Callable[..., Any]) -> dict[str, Any]: """Generate JSON schema from function signature. Args: func: Function to generate schema for Returns: JSON schema """ sig = inspect.signature(func) properties = {} required = [] for param_name, param in sig.parameters.items(): param_type = "string" # Default type if param.annotation != inspect.Parameter.empty: if param.annotation == int: param_type = "integer" elif param.annotation == float: param_type = "number" elif param.annotation == bool: param_type = "boolean" elif param.annotation == list: param_type = "array" elif param.annotation == dict: param_type = "object" properties[param_name] = {"type": param_type} if param.default == inspect.Parameter.empty: required.append(param_name) return { "type": "object", "properties": properties, "required": required, } async def list_tools(self) -> list[types.Tool]: """List all registered tools. Returns: List of MCP tools """ tools = [] for name, tool in self._tools.items(): if isinstance(tool, BaseTool): tools.append(tool.to_mcp_tool()) else: tools.append( types.Tool( name=name, description=self._descriptions.get(name, ""), inputSchema=self._schemas.get(name, {}), ) ) return tools async def call_tool(self, name: str, arguments: dict[str, Any]) -> list[types.ContentBlock]: """Call a tool by name. Args: name: Tool name arguments: Tool arguments Returns: Tool execution result Raises: ValueError: If tool is not found """ if name not in self._tools: raise ValueError(f"Unknown tool: {name}") tool = self._tools[name] if isinstance(tool, BaseTool): return await tool.execute(arguments) else: # Call regular function try: if inspect.iscoroutinefunction(tool): result = await tool(**arguments) else: result = tool(**arguments) if isinstance(result, str): return [types.TextContent(type="text", text=result)] elif isinstance(result, list) and all(isinstance(item, types.ContentBlock) for item in result): return result else: return [types.TextContent(type="text", text=str(result))] except Exception as e: return [types.TextContent(type="text", text=f"Error executing tool {name}: {str(e)}")]