"""Base classes and decorators for MCP tools."""
from typing import Any, Callable, Dict, List
from ..logger import get_logger
logger = get_logger(__name__)
# Global tool registry
TOOL_REGISTRY: Dict[str, Dict[str, Any]] = {}
def tool(name: str, description: str, input_schema: Dict[str, Any]) -> Callable:
"""
Decorator to register a tool.
Args:
name: Tool name (should be unique)
description: Tool description for LLM
input_schema: JSON Schema for tool parameters
Returns:
Decorator function
Example:
@tool(
name="calculator_add",
description="Add two numbers together",
input_schema={
"type": "object",
"properties": {
"a": {"type": "number", "description": "First number"},
"b": {"type": "number", "description": "Second number"}
},
"required": ["a", "b"]
}
)
async def add(a: float, b: float) -> Dict[str, Any]:
return {"result": a + b}
"""
def decorator(func: Callable) -> Callable:
if name in TOOL_REGISTRY:
logger.warning(f"Tool '{name}' is already registered. Overwriting.")
TOOL_REGISTRY[name] = {
"name": name,
"description": description,
"inputSchema": input_schema,
"handler": func,
}
logger.debug(f"Registered tool: {name}")
return func
return decorator
def get_all_tools() -> List[Dict[str, Any]]:
"""
Get list of all registered tools (without handlers).
Returns:
List of tool metadata dicts
"""
return [
{
"name": tool_data["name"],
"description": tool_data["description"],
"inputSchema": tool_data["inputSchema"],
}
for tool_data in TOOL_REGISTRY.values()
]
def get_tool_handler(name: str) -> Callable:
"""
Get tool handler function by name.
Args:
name: Tool name
Returns:
Tool handler function
Raises:
KeyError: If tool not found
"""
if name not in TOOL_REGISTRY:
raise KeyError(f"Tool not found: {name}")
return TOOL_REGISTRY[name]["handler"]
