Refactor tools + resourcemanager

This commit is contained in:
Jeremiah Lowin 2024-11-30 16:36:02 -05:00
commit df59c88e11
6 changed files with 66 additions and 53 deletions

View file

@ -8,7 +8,7 @@ from .types import (
DirectoryResource,
)
from .templates import ResourceTemplate
from .manager import ResourceManager
from .resource_manager import ResourceManager
__all__ = [
"Resource",

View file

@ -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

View file

@ -0,0 +1,4 @@
from .base import Tool
from .tool_manager import ToolManager
__all__ = ["Tool", "ToolManager"]

View file

@ -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)

View file

@ -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)