53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
from collections.abc import Callable
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from mcp.server.fastmcp.exceptions import ToolError
|
|
from mcp.server.fastmcp.tools.base import Tool
|
|
from mcp.server.fastmcp.utilities.logging import get_logger
|
|
|
|
if TYPE_CHECKING:
|
|
from mcp.server.fastmcp.server import Context
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class ToolManager:
|
|
"""Manages FastMCP tools."""
|
|
|
|
def __init__(self, warn_on_duplicate_tools: bool = True):
|
|
self._tools: dict[str, Tool] = {}
|
|
self.warn_on_duplicate_tools = warn_on_duplicate_tools
|
|
|
|
def get_tool(self, name: str) -> Tool | None:
|
|
"""Get tool by name."""
|
|
return self._tools.get(name)
|
|
|
|
def list_tools(self) -> list[Tool]:
|
|
"""List all registered tools."""
|
|
return list(self._tools.values())
|
|
|
|
def add_tool(
|
|
self,
|
|
fn: Callable,
|
|
name: str | None = None,
|
|
description: str | None = None,
|
|
) -> Tool:
|
|
"""Add a tool to the server."""
|
|
tool = Tool.from_function(fn, name=name, description=description)
|
|
existing = self._tools.get(tool.name)
|
|
if existing:
|
|
if self.warn_on_duplicate_tools:
|
|
logger.warning(f"Tool already exists: {tool.name}")
|
|
return existing
|
|
self._tools[tool.name] = tool
|
|
return tool
|
|
|
|
async def call_tool(
|
|
self, name: str, arguments: dict, context: "Context | None" = None
|
|
) -> Any:
|
|
"""Call a tool by name with arguments."""
|
|
tool = self.get_tool(name)
|
|
if not tool:
|
|
raise ToolError(f"Unknown tool: {name}")
|
|
|
|
return await tool.run(arguments, context=context)
|