From df59c88e119a5853a1d04050def97d0077f85f01 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sat, 30 Nov 2024 16:36:02 -0500 Subject: [PATCH] Refactor tools + resourcemanager --- src/fastmcp/resources/__init__.py | 2 +- .../{manager.py => resource_manager.py} | 0 src/fastmcp/server.py | 3 +- src/fastmcp/tools/__init__.py | 4 ++ src/fastmcp/{tools.py => tools/base.py} | 55 ++----------------- src/fastmcp/tools/tool_manager.py | 55 +++++++++++++++++++ 6 files changed, 66 insertions(+), 53 deletions(-) rename src/fastmcp/resources/{manager.py => resource_manager.py} (100%) create mode 100644 src/fastmcp/tools/__init__.py rename src/fastmcp/{tools.py => tools/base.py} (61%) create mode 100644 src/fastmcp/tools/tool_manager.py diff --git a/src/fastmcp/resources/__init__.py b/src/fastmcp/resources/__init__.py index b89d4a4f0..92deb8735 100644 --- a/src/fastmcp/resources/__init__.py +++ b/src/fastmcp/resources/__init__.py @@ -8,7 +8,7 @@ from .types import ( DirectoryResource, ) from .templates import ResourceTemplate -from .manager import ResourceManager +from .resource_manager import ResourceManager __all__ = [ "Resource", diff --git a/src/fastmcp/resources/manager.py b/src/fastmcp/resources/resource_manager.py similarity index 100% rename from src/fastmcp/resources/manager.py rename to src/fastmcp/resources/resource_manager.py diff --git a/src/fastmcp/server.py b/src/fastmcp/server.py index 19f3c91db..699d63243 100644 --- a/src/fastmcp/server.py +++ b/src/fastmcp/server.py @@ -29,8 +29,7 @@ from pydantic_settings import BaseSettings from pydantic.networks import _BaseUrl from fastmcp.exceptions import ResourceError -from fastmcp.resources import Resource, ResourceManager -from fastmcp.resources.types import FunctionResource +from fastmcp.resources import Resource, ResourceManager, FunctionResource from fastmcp.tools import ToolManager from fastmcp.utilities.logging import configure_logging from fastmcp.utilities.types import Image diff --git a/src/fastmcp/tools/__init__.py b/src/fastmcp/tools/__init__.py new file mode 100644 index 000000000..ae9c65619 --- /dev/null +++ b/src/fastmcp/tools/__init__.py @@ -0,0 +1,4 @@ +from .base import Tool +from .tool_manager import ToolManager + +__all__ = ["Tool", "ToolManager"] diff --git a/src/fastmcp/tools.py b/src/fastmcp/tools/base.py similarity index 61% rename from src/fastmcp/tools.py rename to src/fastmcp/tools/base.py index 24a09665d..de60a0faf 100644 --- a/src/fastmcp/tools.py +++ b/src/fastmcp/tools/base.py @@ -1,19 +1,16 @@ -"""Tool management for FastMCP.""" +import fastmcp +from fastmcp.exceptions import ToolError -import inspect -from typing import Any, Callable, Dict, Optional, TYPE_CHECKING from pydantic import BaseModel, Field, TypeAdapter, validate_call -from .exceptions import ToolError -from .utilities.logging import get_logger -import fastmcp + +import inspect +from typing import TYPE_CHECKING, Any, Callable, Optional if TYPE_CHECKING: from fastmcp.server import Context -logger = get_logger(__name__) - class Tool(BaseModel): """Internal tool registration info.""" @@ -80,45 +77,3 @@ class Tool(BaseModel): return self.func(**arguments) except Exception as e: raise ToolError(f"Error executing tool {self.name}: {e}") from e - - -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) -> Optional[Tool]: - """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, - func: Callable, - name: Optional[str] = None, - description: Optional[str] = None, - ) -> Tool: - """Add a tool to the server.""" - tool = Tool.from_function(func, 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: Optional["Context"] = 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) diff --git a/src/fastmcp/tools/tool_manager.py b/src/fastmcp/tools/tool_manager.py new file mode 100644 index 000000000..c8d6d3090 --- /dev/null +++ b/src/fastmcp/tools/tool_manager.py @@ -0,0 +1,55 @@ +from fastmcp.exceptions import ToolError + +from fastmcp.tools.base import Tool + + +from typing import Any, Callable, Dict, Optional, TYPE_CHECKING + +from fastmcp.utilities.logging import get_logger + +if TYPE_CHECKING: + from 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) -> Optional[Tool]: + """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, + func: Callable, + name: Optional[str] = None, + description: Optional[str] = None, + ) -> Tool: + """Add a tool to the server.""" + tool = Tool.from_function(func, 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: Optional["Context"] = 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)