tool_manager.py•2.4 kB
from typing import Any, Callable, Dict, List, Optional
class ToolDefinition:
"""Defines a tool that can be called by the MCP."""
def __init__(self, name: str, handler: Callable, description: str = "", parameters: Optional[Dict] = None):
"""Initialize a tool definition.
Args:
name: Name of the tool
handler: Async callable that handles the tool execution
description: Tool description
parameters: JSON schema for the tool parameters
"""
self.name = name
self.handler = handler
self.description = description
self.parameters = parameters or {}
class ToolManager:
"""Manages MCP tools for YDB interactions."""
def __init__(self):
"""Initialize the tool manager."""
self._tools: Dict[str, ToolDefinition] = {}
def register_tool(
self, name: str, handler: Callable, description: str = "", parameters: Optional[Dict] = None
) -> None:
"""Register a tool with the manager.
Args:
name: Name of the tool
handler: Async callable that handles the tool execution
description: Tool description
parameters: JSON schema for tool parameters
"""
self._tools[name] = ToolDefinition(name=name, handler=handler, description=description, parameters=parameters)
def get(self, name: str) -> Optional[ToolDefinition]:
"""Get a tool by name.
Args:
name: Name of the tool to retrieve
Returns:
Tool definition if found, None otherwise
"""
return self._tools.get(name)
def get_all_tools(self) -> Dict[str, ToolDefinition]:
"""Get all registered tools.
Returns:
Dictionary of tool name to tool definition
"""
return self._tools
def get_schema(self) -> List[Dict[str, Any]]:
"""Get JSON schema for all registered tools.
Returns:
List of tool schema definitions
"""
result = []
for name, tool in self._tools.items():
tool_schema: Dict[str, Any] = {
"name": name,
"description": tool.description,
}
if tool.parameters:
tool_schema["parameters"] = tool.parameters
result.append(tool_schema)
return result