From f0d07e5f04118d0699755742baaff0616cf50e89 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sat, 12 Apr 2025 11:20:25 -0400 Subject: [PATCH] Add tags, improve duplicate import behavior --- .cursor/rules/core-mcp-objects.mdc | 13 + src/fastmcp/prompts/__init__.py | 2 +- src/fastmcp/prompts/{base.py => prompt.py} | 15 +- src/fastmcp/prompts/prompt_manager.py | 36 ++- src/fastmcp/resources/__init__.py | 4 +- .../resources/{base.py => resource.py} | 15 +- src/fastmcp/resources/resource_manager.py | 64 ++++- .../resources/{templates.py => template.py} | 15 +- src/fastmcp/resources/types.py | 4 +- src/fastmcp/server/openapi.py | 2 +- src/fastmcp/server/proxy.py | 2 +- src/fastmcp/server/server.py | 12 +- src/fastmcp/settings.py | 14 +- src/fastmcp/tools/__init__.py | 2 +- src/fastmcp/tools/{base.py => tool.py} | 15 +- src/fastmcp/tools/tool_manager.py | 33 +-- tests/prompts/test_base.py | 2 +- tests/prompts/test_prompt_manager.py | 222 +++++++++++------- tests/resources/test_resource_manager.py | 198 +++++++++++++++- tests/server/test_server.py | 4 +- tests/tools/test_tool_manager.py | 200 ++++++++++++---- 21 files changed, 691 insertions(+), 183 deletions(-) create mode 100644 .cursor/rules/core-mcp-objects.mdc rename src/fastmcp/prompts/{base.py => prompt.py} (92%) rename src/fastmcp/resources/{base.py => resource.py} (76%) rename src/fastmcp/resources/{templates.py => template.py} (86%) rename src/fastmcp/tools/{base.py => tool.py} (87%) diff --git a/.cursor/rules/core-mcp-objects.mdc b/.cursor/rules/core-mcp-objects.mdc new file mode 100644 index 000000000..c8cc92818 --- /dev/null +++ b/.cursor/rules/core-mcp-objects.mdc @@ -0,0 +1,13 @@ +--- +description: +globs: +alwaysApply: true +--- +There are four major MCP object types: + +- Tools (src/tools/) +- Resources (src/resources/) +- Resource Templates (src/resources/) +- Prompts (src/prompts) + +While these have slightly different semantics and implementations, in general changes that affect interactions with any one (like adding tags, importing, etc.) will need to be adopted, applied, and tested on all others. Be sure to look at not only the object definition but also the related `Manager` (e.g. `ToolManager`, `ResourceManager`, and `PromptManager`). Also note that while resources and resource templates are different objects, they both are handled by the `ResourceManager`. \ No newline at end of file diff --git a/src/fastmcp/prompts/__init__.py b/src/fastmcp/prompts/__init__.py index bacb4c37d..cdc6702bd 100644 --- a/src/fastmcp/prompts/__init__.py +++ b/src/fastmcp/prompts/__init__.py @@ -1,4 +1,4 @@ -from .base import Prompt +from .prompt import Prompt from .prompt_manager import PromptManager __all__ = ["Prompt", "PromptManager"] diff --git a/src/fastmcp/prompts/base.py b/src/fastmcp/prompts/prompt.py similarity index 92% rename from src/fastmcp/prompts/base.py rename to src/fastmcp/prompts/prompt.py index f4fce3413..153d23a0a 100644 --- a/src/fastmcp/prompts/base.py +++ b/src/fastmcp/prompts/prompt.py @@ -8,6 +8,7 @@ from typing import Annotated, Any, Literal import pydantic_core from mcp.types import EmbeddedResource, ImageContent, TextContent from pydantic import BaseModel, BeforeValidator, Field, TypeAdapter, validate_call +from typing_extensions import Self from fastmcp.utilities.types import _convert_set_defaults @@ -79,7 +80,7 @@ class Prompt(BaseModel): arguments: list[PromptArgument] | None = Field( None, description="Arguments that can be passed to the prompt" ) - fn: Callable[..., PromptResult | Awaitable[PromptResult]] = Field(exclude=True) + fn: Callable[..., PromptResult | Awaitable[PromptResult]] @classmethod def from_function( @@ -171,3 +172,15 @@ class Prompt(BaseModel): return messages except Exception as e: raise ValueError(f"Error rendering prompt {self.name}: {e}") + + def copy(self, updates: dict[str, Any] | None = None) -> Self: + """Copy the prompt with optional updates.""" + data = self.model_dump() + if updates: + data.update(updates) + return type(self)(**data) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Prompt): + return False + return self.model_dump() == other.model_dump() diff --git a/src/fastmcp/prompts/prompt_manager.py b/src/fastmcp/prompts/prompt_manager.py index e15da430e..c1653a50d 100644 --- a/src/fastmcp/prompts/prompt_manager.py +++ b/src/fastmcp/prompts/prompt_manager.py @@ -1,8 +1,10 @@ """Prompt management functionality.""" +from collections.abc import Awaitable, Callable from typing import Any -from fastmcp.prompts.base import Message, Prompt +from fastmcp.prompts.prompt import Message, Prompt, PromptResult +from fastmcp.settings import DuplicateBehavior from fastmcp.utilities.logging import get_logger logger = get_logger(__name__) @@ -11,9 +13,9 @@ logger = get_logger(__name__) class PromptManager: """Manages FastMCP prompts.""" - def __init__(self, warn_on_duplicate_prompts: bool = True): + def __init__(self, duplicate_behavior: DuplicateBehavior = DuplicateBehavior.WARN): self._prompts: dict[str, Prompt] = {} - self.warn_on_duplicate_prompts = warn_on_duplicate_prompts + self.duplicate_behavior = duplicate_behavior def get_prompt(self, name: str) -> Prompt | None: """Get prompt by name.""" @@ -23,18 +25,32 @@ class PromptManager: """List all registered prompts.""" return list(self._prompts.values()) - def add_prompt( + def add_prompt_from_fn( self, - prompt: Prompt, + fn: Callable[..., PromptResult | Awaitable[PromptResult]], + name: str | None = None, + description: str | None = None, + tags: set[str] | None = None, ) -> Prompt: + """Create a prompt from a function.""" + prompt = Prompt.from_function(fn, name=name, description=description, tags=tags) + return self.add_prompt(prompt) + + def add_prompt(self, prompt: Prompt) -> Prompt: """Add a prompt to the manager.""" # Check for duplicates existing = self._prompts.get(prompt.name) if existing: - if self.warn_on_duplicate_prompts: + if self.duplicate_behavior == DuplicateBehavior.WARN: logger.warning(f"Prompt already exists: {prompt.name}") - return existing + self._prompts[prompt.name] = prompt + elif self.duplicate_behavior == DuplicateBehavior.REPLACE: + self._prompts[prompt.name] = prompt + elif self.duplicate_behavior == DuplicateBehavior.ERROR: + raise ValueError(f"Prompt already exists: {prompt.name}") + elif self.duplicate_behavior == DuplicateBehavior.IGNORE: + pass self._prompts[prompt.name] = prompt return prompt @@ -64,11 +80,13 @@ class PromptManager: the imported prompt would be available as "weather/forecast_prompt" """ for name, prompt in manager._prompts.items(): - # Create prefixed name - we keep the original name in the Prompt object + # Create prefixed name prefixed_name = f"{prefix}{name}" if prefix else name + new_prompt = prompt.copy(updates=dict(name=prefixed_name)) + # Log the import logger.debug(f"Importing prompt with name {name} as {prefixed_name}") # Store the prompt with the prefixed name - self._prompts[prefixed_name] = prompt + self.add_prompt(new_prompt) diff --git a/src/fastmcp/resources/__init__.py b/src/fastmcp/resources/__init__.py index b5805fb34..b7f7b5985 100644 --- a/src/fastmcp/resources/__init__.py +++ b/src/fastmcp/resources/__init__.py @@ -1,6 +1,6 @@ -from .base import Resource +from .resource import Resource from .resource_manager import ResourceManager -from .templates import ResourceTemplate +from .template import ResourceTemplate from .types import ( BinaryResource, DirectoryResource, diff --git a/src/fastmcp/resources/base.py b/src/fastmcp/resources/resource.py similarity index 76% rename from src/fastmcp/resources/base.py rename to src/fastmcp/resources/resource.py index d489d8919..01e0995ac 100644 --- a/src/fastmcp/resources/base.py +++ b/src/fastmcp/resources/resource.py @@ -1,7 +1,7 @@ """Base classes and interfaces for FastMCP resources.""" import abc -from typing import Annotated +from typing import Annotated, Any from pydantic import ( AnyUrl, @@ -13,6 +13,7 @@ from pydantic import ( ValidationInfo, field_validator, ) +from typing_extensions import Self from fastmcp.utilities.types import _convert_set_defaults @@ -52,3 +53,15 @@ class Resource(BaseModel, abc.ABC): async def read(self) -> str | bytes: """Read the resource content.""" pass + + def copy(self, updates: dict[str, Any] | None = None) -> Self: + """Copy the resource with optional updates.""" + data = self.model_dump() + if updates: + data.update(updates) + return type(self)(**data) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Resource): + return False + return self.model_dump() == other.model_dump() diff --git a/src/fastmcp/resources/resource_manager.py b/src/fastmcp/resources/resource_manager.py index 6be187645..afb26aecf 100644 --- a/src/fastmcp/resources/resource_manager.py +++ b/src/fastmcp/resources/resource_manager.py @@ -5,8 +5,9 @@ from typing import Any from pydantic import AnyUrl -from fastmcp.resources.base import Resource -from fastmcp.resources.templates import ResourceTemplate +from fastmcp.resources.resource import Resource +from fastmcp.resources.template import ResourceTemplate +from fastmcp.settings import DuplicateBehavior from fastmcp.utilities.logging import get_logger logger = get_logger(__name__) @@ -15,10 +16,10 @@ logger = get_logger(__name__) class ResourceManager: """Manages FastMCP resources.""" - def __init__(self, warn_on_duplicate_resources: bool = True): + def __init__(self, duplicate_behavior: DuplicateBehavior = DuplicateBehavior.WARN): self._resources: dict[str, Resource] = {} self._templates: dict[str, ResourceTemplate] = {} - self.warn_on_duplicate_resources = warn_on_duplicate_resources + self.duplicate_behavior = duplicate_behavior def add_resource(self, resource: Resource) -> Resource: """Add a resource to the manager. @@ -40,13 +41,19 @@ class ResourceManager: ) existing = self._resources.get(str(resource.uri)) if existing: - if self.warn_on_duplicate_resources: + if self.duplicate_behavior == DuplicateBehavior.WARN: logger.warning(f"Resource already exists: {resource.uri}") - return existing + self._resources[str(resource.uri)] = resource + elif self.duplicate_behavior == DuplicateBehavior.REPLACE: + self._resources[str(resource.uri)] = resource + elif self.duplicate_behavior == DuplicateBehavior.ERROR: + raise ValueError(f"Resource already exists: {resource.uri}") + elif self.duplicate_behavior == DuplicateBehavior.IGNORE: + pass self._resources[str(resource.uri)] = resource return resource - def add_template( + def add_template_from_fn( self, fn: Callable[..., Any], uri_template: str, @@ -55,7 +62,7 @@ class ResourceManager: mime_type: str | None = None, tags: set[str] | None = None, ) -> ResourceTemplate: - """Add a template from a function.""" + """Create a template from a function.""" template = ResourceTemplate.from_function( fn, uri_template=uri_template, @@ -64,6 +71,37 @@ class ResourceManager: mime_type=mime_type, tags=tags, ) + return self.add_template(template) + + def add_template(self, template: ResourceTemplate) -> ResourceTemplate: + """Add a template to the manager. + + Args: + template: A ResourceTemplate instance to add + + Returns: + The added template. If a template with the same URI already exists, + returns the existing template. + """ + logger.debug( + "Adding resource", + extra={ + "uri": template.uri_template, + "type": type(template).__name__, + "resource_name": template.name, + }, + ) + existing = self._templates.get(str(template.uri_template)) + if existing: + if self.duplicate_behavior == DuplicateBehavior.WARN: + logger.warning(f"Resource already exists: {template.uri_template}") + self._templates[str(template.uri_template)] = template + elif self.duplicate_behavior == DuplicateBehavior.REPLACE: + self._templates[str(template.uri_template)] = template + elif self.duplicate_behavior == DuplicateBehavior.ERROR: + raise ValueError(f"Resource already exists: {template.uri_template}") + elif self.duplicate_behavior == DuplicateBehavior.IGNORE: + pass self._templates[template.uri_template] = template return template @@ -116,11 +154,13 @@ class ResourceManager: # Create prefixed URI and copy the resource with the new URI prefixed_uri = f"{prefix}{uri}" if prefix else uri + new_resource = resource.copy(updates=dict(uri=prefixed_uri)) + # Log the import logger.debug(f"Importing resource with URI {uri} as {prefixed_uri}") # Store directly in resources dictionary - self._resources[prefixed_uri] = resource + self.add_resource(new_resource) def import_templates( self, manager: "ResourceManager", prefix: str | None = None @@ -144,10 +184,14 @@ class ResourceManager: f"{prefix}{uri_template}" if prefix else uri_template ) + new_template = template.copy( + updates=dict(uri_template=prefixed_uri_template) + ) + # Log the import logger.debug( f"Importing resource template with URI {uri_template} as {prefixed_uri_template}" ) # Store directly in templates dictionary - self._templates[prefixed_uri_template] = template + self.add_template(new_template) diff --git a/src/fastmcp/resources/templates.py b/src/fastmcp/resources/template.py similarity index 86% rename from src/fastmcp/resources/templates.py rename to src/fastmcp/resources/template.py index 69a085c28..289b23050 100644 --- a/src/fastmcp/resources/templates.py +++ b/src/fastmcp/resources/template.py @@ -8,6 +8,7 @@ from collections.abc import Callable from typing import Annotated, Any from pydantic import BaseModel, BeforeValidator, Field, TypeAdapter, validate_call +from typing_extensions import Self from fastmcp.resources.types import FunctionResource, Resource from fastmcp.utilities.types import _convert_set_defaults @@ -27,7 +28,7 @@ class ResourceTemplate(BaseModel): mime_type: str = Field( default="text/plain", description="MIME type of the resource content" ) - fn: Callable[..., Any] = Field(exclude=True) + fn: Callable[..., Any] parameters: dict[str, Any] = Field( description="JSON schema for function parameters" ) @@ -90,3 +91,15 @@ class ResourceTemplate(BaseModel): ) except Exception as e: raise ValueError(f"Error creating resource from template: {e}") + + def copy(self, updates: dict[str, Any] | None = None) -> Self: + """Copy the resource template with optional updates.""" + data = self.model_dump() + if updates: + data.update(updates) + return type(self)(**data) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, ResourceTemplate): + return False + return self.model_dump() == other.model_dump() diff --git a/src/fastmcp/resources/types.py b/src/fastmcp/resources/types.py index 89a142395..30d168bee 100644 --- a/src/fastmcp/resources/types.py +++ b/src/fastmcp/resources/types.py @@ -13,7 +13,7 @@ import pydantic.json import pydantic_core from pydantic import Field, ValidationInfo -from fastmcp.resources.base import Resource +from fastmcp.resources.resource import Resource class TextResource(Resource): @@ -49,7 +49,7 @@ class FunctionResource(Resource): - other types will be converted to JSON """ - fn: Callable[[], Any] = Field(exclude=True) + fn: Callable[[], Any] async def read(self) -> str | bytes: """Read the resource by calling the wrapped function.""" diff --git a/src/fastmcp/server/openapi.py b/src/fastmcp/server/openapi.py index a37c6dba4..19cac311d 100644 --- a/src/fastmcp/server/openapi.py +++ b/src/fastmcp/server/openapi.py @@ -12,7 +12,7 @@ from pydantic.networks import AnyUrl from fastmcp.resources import Resource, ResourceTemplate from fastmcp.server.server import FastMCP -from fastmcp.tools.base import Tool +from fastmcp.tools.tool import Tool from fastmcp.utilities import openapi from fastmcp.utilities.func_metadata import func_metadata from fastmcp.utilities.logging import get_logger diff --git a/src/fastmcp/server/proxy.py b/src/fastmcp/server/proxy.py index 0b3311251..048d1d1e0 100644 --- a/src/fastmcp/server/proxy.py +++ b/src/fastmcp/server/proxy.py @@ -9,7 +9,7 @@ from fastmcp.prompts import Prompt from fastmcp.resources import Resource, ResourceTemplate from fastmcp.server.context import Context from fastmcp.server.server import FastMCP -from fastmcp.tools.base import Tool +from fastmcp.tools.tool import Tool from fastmcp.utilities.func_metadata import func_metadata from fastmcp.utilities.logging import get_logger diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 517e041a0..fd50f2450 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -88,13 +88,13 @@ class FastMCP(Generic[LifespanResultT]): lifespan=lifespan_wrapper(self, lifespan) if lifespan else default_lifespan, # type: ignore ) self._tool_manager = ToolManager( - warn_on_duplicate_tools=self.settings.warn_on_duplicate_tools + duplicate_behavior=self.settings.on_duplicate_tools ) self._resource_manager = ResourceManager( - warn_on_duplicate_resources=self.settings.warn_on_duplicate_resources + duplicate_behavior=self.settings.on_duplicate_resources ) self._prompt_manager = PromptManager( - warn_on_duplicate_prompts=self.settings.warn_on_duplicate_prompts + duplicate_behavior=self.settings.on_duplicate_prompts ) self.dependencies = self.settings.dependencies @@ -241,7 +241,9 @@ class FastMCP(Generic[LifespanResultT]): description: Optional description of what the tool does tags: Optional set of tags for categorizing the tool """ - self._tool_manager.add_tool(fn, name=name, description=description, tags=tags) + self._tool_manager.add_tool_from_fn( + fn, name=name, description=description, tags=tags + ) def tool( self, @@ -366,7 +368,7 @@ class FastMCP(Generic[LifespanResultT]): ) # Register as template - self._resource_manager.add_template( + self._resource_manager.add_template_from_fn( fn=fn, uri_template=uri, name=name, diff --git a/src/fastmcp/settings.py b/src/fastmcp/settings.py index 1561a39a2..09f7e4fdc 100644 --- a/src/fastmcp/settings.py +++ b/src/fastmcp/settings.py @@ -1,5 +1,6 @@ from __future__ import annotations as _annotations +from enum import Enum from typing import TYPE_CHECKING, Literal from pydantic import Field @@ -11,6 +12,13 @@ if TYPE_CHECKING: LOG_LEVEL = Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] +class DuplicateBehavior(Enum): + WARN = "warn" + ERROR = "error" + REPLACE = "replace" + IGNORE = "ignore" + + class Settings(BaseSettings): """FastMCP settings.""" @@ -47,13 +55,13 @@ class ServerSettings(BaseSettings): debug: bool = False # resource settings - warn_on_duplicate_resources: bool = True + on_duplicate_resources: DuplicateBehavior = DuplicateBehavior.WARN # tool settings - warn_on_duplicate_tools: bool = True + on_duplicate_tools: DuplicateBehavior = DuplicateBehavior.WARN # prompt settings - warn_on_duplicate_prompts: bool = True + on_duplicate_prompts: DuplicateBehavior = DuplicateBehavior.WARN dependencies: list[str] = Field( default_factory=list, diff --git a/src/fastmcp/tools/__init__.py b/src/fastmcp/tools/__init__.py index ae9c65619..22b69a0c7 100644 --- a/src/fastmcp/tools/__init__.py +++ b/src/fastmcp/tools/__init__.py @@ -1,4 +1,4 @@ -from .base import Tool +from .tool import Tool from .tool_manager import ToolManager __all__ = ["Tool", "ToolManager"] diff --git a/src/fastmcp/tools/base.py b/src/fastmcp/tools/tool.py similarity index 87% rename from src/fastmcp/tools/base.py rename to src/fastmcp/tools/tool.py index cebe2c10a..b40d088f3 100644 --- a/src/fastmcp/tools/base.py +++ b/src/fastmcp/tools/tool.py @@ -5,6 +5,7 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Annotated, Any from pydantic import BaseModel, BeforeValidator, Field +from typing_extensions import Self from fastmcp.exceptions import ToolError from fastmcp.utilities.func_metadata import FuncMetadata, func_metadata @@ -20,7 +21,7 @@ if TYPE_CHECKING: class Tool(BaseModel): """Internal tool registration info.""" - fn: Callable[..., Any] = Field(exclude=True) + fn: Callable[..., Any] name: str = Field(description="Name of the tool") description: str = Field(description="Description of what the tool does") parameters: dict[str, Any] = Field(description="JSON schema for tool parameters") @@ -97,3 +98,15 @@ class Tool(BaseModel): ) except Exception as e: raise ToolError(f"Error executing tool {self.name}: {e}") from e + + def copy(self, updates: dict[str, Any] | None = None) -> Self: + """Copy the tool with optional updates.""" + data = self.model_dump() + if updates: + data.update(updates) + return type(self)(**data) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Tool): + return False + return self.model_dump() == other.model_dump() diff --git a/src/fastmcp/tools/tool_manager.py b/src/fastmcp/tools/tool_manager.py index 31b02b2aa..00279ed89 100644 --- a/src/fastmcp/tools/tool_manager.py +++ b/src/fastmcp/tools/tool_manager.py @@ -6,7 +6,8 @@ from typing import TYPE_CHECKING, Any from mcp.shared.context import LifespanContextT from fastmcp.exceptions import ToolError -from fastmcp.tools.base import Tool +from fastmcp.settings import DuplicateBehavior +from fastmcp.tools.tool import Tool from fastmcp.utilities.logging import get_logger if TYPE_CHECKING: @@ -20,9 +21,9 @@ logger = get_logger(__name__) class ToolManager: """Manages FastMCP tools.""" - def __init__(self, warn_on_duplicate_tools: bool = True): + def __init__(self, duplicate_behavior: DuplicateBehavior = DuplicateBehavior.WARN): self._tools: dict[str, Tool] = {} - self.warn_on_duplicate_tools = warn_on_duplicate_tools + self.duplicate_behavior = duplicate_behavior def get_tool(self, name: str) -> Tool | None: """Get tool by name.""" @@ -32,7 +33,7 @@ class ToolManager: """List all registered tools.""" return list(self._tools.values()) - def add_tool( + def add_tool_from_fn( self, fn: Callable[..., Any], name: str | None = None, @@ -41,15 +42,21 @@ class ToolManager: ) -> Tool: """Add a tool to the server.""" tool = Tool.from_function(fn, name=name, description=description, tags=tags) - return self._register_tool(tool) + return self.add_tool(tool) - def _register_tool(self, tool: Tool) -> Tool: + def add_tool(self, tool: Tool) -> Tool: """Register a tool with the server.""" existing = self._tools.get(tool.name) if existing: - if self.warn_on_duplicate_tools: + if self.duplicate_behavior == DuplicateBehavior.WARN: logger.warning(f"Tool already exists: {tool.name}") - return existing + self._tools[tool.name] = tool + elif self.duplicate_behavior == DuplicateBehavior.REPLACE: + self._tools[tool.name] = tool + elif self.duplicate_behavior == DuplicateBehavior.ERROR: + raise ValueError(f"Tool already exists: {tool.name}") + elif self.duplicate_behavior == DuplicateBehavior.IGNORE: + pass self._tools[tool.name] = tool return tool @@ -83,13 +90,7 @@ class ToolManager: for name, tool in tool_manager._tools.items(): prefixed_name = f"{prefix}{name}" if prefix else name - # Create a shallow copy of the tool with the prefixed name - copied_tool = Tool.from_function( - tool.fn, - name=prefixed_name, - description=tool.description, - ) - + new_tool = tool.copy(updates=dict(name=prefixed_name)) # Store the copied tool - self._register_tool(copied_tool) + self.add_tool(new_tool) logger.debug(f"Imported tool: {name} as {prefixed_name}") diff --git a/tests/prompts/test_base.py b/tests/prompts/test_base.py index a3b591858..fb02c93d1 100644 --- a/tests/prompts/test_base.py +++ b/tests/prompts/test_base.py @@ -2,7 +2,7 @@ import pytest from mcp.types import EmbeddedResource, TextResourceContents from pydantic import FileUrl -from fastmcp.prompts.base import ( +from fastmcp.prompts.prompt import ( AssistantMessage, Message, Prompt, diff --git a/tests/prompts/test_prompt_manager.py b/tests/prompts/test_prompt_manager.py index 1964db74d..cac976202 100644 --- a/tests/prompts/test_prompt_manager.py +++ b/tests/prompts/test_prompt_manager.py @@ -1,8 +1,9 @@ import pytest from fastmcp.prompts import Prompt -from fastmcp.prompts.base import PromptArgument, TextContent, UserMessage +from fastmcp.prompts.prompt import PromptArgument, TextContent, UserMessage from fastmcp.prompts.prompt_manager import PromptManager +from fastmcp.settings import DuplicateBehavior class TestPromptManager: @@ -24,7 +25,7 @@ class TestPromptManager: def fn() -> str: return "Hello, world!" - manager = PromptManager() + manager = PromptManager(duplicate_behavior=DuplicateBehavior.WARN) prompt = Prompt.from_function(fn) first = manager.add_prompt(prompt) second = manager.add_prompt(prompt) @@ -37,7 +38,7 @@ class TestPromptManager: def fn() -> str: return "Hello, world!" - manager = PromptManager(warn_on_duplicate_prompts=False) + manager = PromptManager(duplicate_behavior=DuplicateBehavior.IGNORE) prompt = Prompt.from_function(fn) first = manager.add_prompt(prompt) second = manager.add_prompt(prompt) @@ -112,6 +113,131 @@ class TestPromptManager: with pytest.raises(ValueError, match="Missing required arguments"): await manager.render_prompt("fn") + def test_error_on_duplicate_prompts(self): + """Test error on duplicate prompts.""" + + def fn() -> str: + return "Hello, world!" + + manager = PromptManager(duplicate_behavior=DuplicateBehavior.ERROR) + prompt = Prompt.from_function(fn) + manager.add_prompt(prompt) + + with pytest.raises(ValueError, match="Prompt already exists"): + manager.add_prompt(prompt) + + def test_replace_duplicate_prompts(self): + """Test replacing duplicate prompts.""" + + def fn1() -> str: + return "Original" + + def fn2() -> str: + return "Replacement" + + manager = PromptManager(duplicate_behavior=DuplicateBehavior.REPLACE) + prompt1 = Prompt.from_function(fn1, name="test_prompt") + prompt2 = Prompt.from_function(fn2, name="test_prompt") + + manager.add_prompt(prompt1) + manager.add_prompt(prompt2) + + # Should have replaced the first prompt with the second + stored_prompt = manager.get_prompt("test_prompt") + assert stored_prompt == prompt2 + + +class TestPromptTags: + """Test functionality related to prompt tags.""" + + def test_add_prompt_with_tags(self): + """Test adding a prompt with tags.""" + + def greeting() -> str: + return "Hello, world!" + + manager = PromptManager() + prompt = Prompt.from_function(greeting, tags={"greeting", "simple"}) + manager.add_prompt(prompt) + + prompt = manager.get_prompt("greeting") + assert prompt is not None + assert prompt.tags == {"greeting", "simple"} + + def test_add_prompt_with_empty_tags(self): + """Test adding a prompt with empty tags.""" + + def greeting() -> str: + return "Hello, world!" + + manager = PromptManager() + prompt = Prompt.from_function(greeting, tags=set()) + manager.add_prompt(prompt) + + prompt = manager.get_prompt("greeting") + assert prompt is not None + assert prompt.tags == set() + + def test_add_prompt_with_none_tags(self): + """Test adding a prompt with None tags.""" + + def greeting() -> str: + return "Hello, world!" + + manager = PromptManager() + prompt = Prompt.from_function(greeting, tags=None) + manager.add_prompt(prompt) + + prompt = manager.get_prompt("greeting") + assert prompt is not None + assert prompt.tags == set() + + def test_list_prompts_with_tags(self): + """Test listing prompts with specific tags.""" + + def greeting() -> str: + return "Hello, world!" + + def weather(location: str) -> str: + return f"Weather for {location}" + + def summary(text: str) -> str: + return f"Summary of: {text}" + + manager = PromptManager() + manager.add_prompt(Prompt.from_function(greeting, tags={"greeting", "simple"})) + manager.add_prompt(Prompt.from_function(weather, tags={"weather", "location"})) + manager.add_prompt( + Prompt.from_function(summary, tags={"summary", "nlp", "simple"}) + ) + + # Filter prompts by tags + simple_prompts = [p for p in manager.list_prompts() if "simple" in p.tags] + assert len(simple_prompts) == 2 + assert {p.name for p in simple_prompts} == {"greeting", "summary"} + + nlp_prompts = [p for p in manager.list_prompts() if "nlp" in p.tags] + assert len(nlp_prompts) == 1 + assert nlp_prompts[0].name == "summary" + + def test_import_prompts_preserves_tags(self): + """Test that importing prompts preserves their tags.""" + source_manager = PromptManager() + + def sample_prompt() -> str: + return "Sample prompt" + + source_manager.add_prompt( + Prompt.from_function(sample_prompt, tags={"example", "test"}) + ) + + target_manager = PromptManager() + target_manager.import_prompts(source_manager, "imported/") + + imported_prompt = target_manager.get_prompt("imported/sample_prompt") + assert imported_prompt is not None + assert imported_prompt.tags == {"example", "test"} + class TestImports: def test_import_prompts(self): @@ -119,27 +245,13 @@ class TestImports: # Setup source manager with prompts source_manager = PromptManager() - # Create test prompts with proper function handlers - async def summary_fn(**kwargs): - return [ - {"role": "assistant", "content": f"Summary of: {kwargs.get('text')}"} - ] - - async def translate_fn(**kwargs): - return [ - { - "role": "assistant", - "content": f"Translation to {kwargs.get('language')}: {kwargs.get('text')}", - } - ] - summary_prompt = Prompt( name="summary", description="Generate a summary of text", arguments=[PromptArgument(name="text", description="Text to summarize")], - fn=summary_fn, + fn=lambda: None, # type: ignore ) - source_manager._prompts["summary"] = summary_prompt + source_manager.add_prompt(summary_prompt) translate_prompt = Prompt( name="translate", @@ -148,9 +260,9 @@ class TestImports: PromptArgument(name="text", description="Text to translate"), PromptArgument(name="language", description="Target language"), ], - fn=translate_fn, + fn=lambda: None, # type: ignore ) - source_manager._prompts["translate"] = translate_prompt + source_manager.add_prompt(translate_prompt) # Create target manager target_manager = PromptManager() @@ -167,31 +279,8 @@ class TestImports: assert "summary" in source_manager._prompts assert "translate" in source_manager._prompts - # Verify the imported prompts have the correct properties - assert target_manager._prompts["nlp/summary"].name == "summary" - assert ( - target_manager._prompts["nlp/summary"].description - == "Generate a summary of text" - ) - - assert target_manager._prompts["nlp/translate"].name == "translate" - assert ( - target_manager._prompts["nlp/translate"].description - == "Translate text to another language" - ) - - # Verify functions were properly copied - if hasattr(target_manager._prompts["nlp/summary"], "fn"): - assert ( - target_manager._prompts["nlp/summary"].fn.__name__ - == summary_fn.__name__ - ) - - if hasattr(target_manager._prompts["nlp/translate"], "fn"): - assert ( - target_manager._prompts["nlp/translate"].fn.__name__ - == translate_fn.__name__ - ) + assert target_manager._prompts["nlp/summary"].fn == summary_prompt.fn + assert target_manager._prompts["nlp/translate"].fn == translate_prompt.fn def test_import_prompts_with_duplicates(self): """Test handling of duplicate prompts during import.""" @@ -199,18 +288,11 @@ class TestImports: source_manager = PromptManager() target_manager = PromptManager() - # Add the same prompt name to both managers with functions - async def source_fn(**kwargs): - return [{"role": "assistant", "content": "Source content"}] - - async def target_fn(**kwargs): - return [{"role": "assistant", "content": "Target content"}] - source_prompt = Prompt( name="common", description="Source description", arguments=None, - fn=source_fn, + fn=lambda: None, # type: ignore ) source_manager._prompts["common"] = source_prompt @@ -218,7 +300,7 @@ class TestImports: name="common", description="Target description", arguments=None, - fn=target_fn, + fn=lambda: None, # type: ignore ) target_manager._prompts["common"] = target_prompt @@ -230,15 +312,8 @@ class TestImports: assert "common" in target_manager._prompts assert "external/common" in target_manager._prompts - # Verify the functions of both prompts - if hasattr(target_manager._prompts["common"], "fn") and hasattr( - target_manager._prompts["external/common"], "fn" - ): - assert target_manager._prompts["common"].fn.__name__ == target_fn.__name__ - assert ( - target_manager._prompts["external/common"].fn.__name__ - == source_fn.__name__ - ) + assert target_manager._prompts["common"].fn == target_prompt.fn + assert target_manager._prompts["external/common"].fn == source_prompt.fn def test_import_prompts_with_nested_prefixes(self): """Test importing already prefixed prompts.""" @@ -247,17 +322,11 @@ class TestImports: second_manager = PromptManager() third_manager = PromptManager() - # Add prompt to first manager with a function - async def analyze_fn(**kwargs): - return [ - {"role": "assistant", "content": f"Analysis of: {kwargs.get('text')}"} - ] - original_prompt = Prompt( name="analyze", description="Analyze text", arguments=[PromptArgument(name="text", description="Text to analyze")], - fn=analyze_fn, + fn=lambda: None, # type: ignore ) first_manager._prompts["analyze"] = original_prompt @@ -271,13 +340,4 @@ class TestImports: assert "text/analyze" in second_manager._prompts assert "ai/text/analyze" in third_manager._prompts - # Verify the properties of the most nested prompt - assert third_manager._prompts["ai/text/analyze"].name == "analyze" - assert third_manager._prompts["ai/text/analyze"].description == "Analyze text" - - # Verify function was properly copied through multiple imports - if hasattr(third_manager._prompts["ai/text/analyze"], "fn"): - assert ( - third_manager._prompts["ai/text/analyze"].fn.__name__ - == analyze_fn.__name__ - ) + assert third_manager._prompts["ai/text/analyze"].fn == original_prompt.fn diff --git a/tests/resources/test_resource_manager.py b/tests/resources/test_resource_manager.py index 12f7cb59d..87d95718c 100644 --- a/tests/resources/test_resource_manager.py +++ b/tests/resources/test_resource_manager.py @@ -10,6 +10,7 @@ from fastmcp.resources import ( ResourceManager, ResourceTemplate, ) +from fastmcp.settings import DuplicateBehavior @pytest.fixture @@ -59,7 +60,7 @@ class TestResourceManager: def test_warn_on_duplicate_resources(self, temp_file: Path, caplog): """Test warning on duplicate resources.""" - manager = ResourceManager() + manager = ResourceManager(duplicate_behavior=DuplicateBehavior.WARN) resource = FileResource( uri=FileUrl(f"file://{temp_file}"), name="test", @@ -71,7 +72,7 @@ class TestResourceManager: def test_disable_warn_on_duplicate_resources(self, temp_file: Path, caplog): """Test disabling warning on duplicate resources.""" - manager = ResourceManager(warn_on_duplicate_resources=False) + manager = ResourceManager(duplicate_behavior=DuplicateBehavior.IGNORE) resource = FileResource( uri=FileUrl(f"file://{temp_file}"), name="test", @@ -81,6 +82,43 @@ class TestResourceManager: manager.add_resource(resource) assert "Resource already exists" not in caplog.text + def test_error_on_duplicate_resources(self, temp_file: Path): + """Test error on duplicate resources.""" + manager = ResourceManager(duplicate_behavior=DuplicateBehavior.ERROR) + resource = FileResource( + uri=FileUrl(f"file://{temp_file}"), + name="test", + path=temp_file, + ) + manager.add_resource(resource) + + with pytest.raises(ValueError, match="Resource already exists"): + manager.add_resource(resource) + + def test_replace_duplicate_resources(self, temp_file: Path): + """Test replacing duplicate resources.""" + manager = ResourceManager(duplicate_behavior=DuplicateBehavior.REPLACE) + + resource1 = FileResource( + uri=FileUrl(f"file://{temp_file}"), + name="test1", + path=temp_file, + ) + + resource2 = FileResource( + uri=FileUrl(f"file://{temp_file}"), + name="test2", # Different name + path=temp_file, + ) + + manager.add_resource(resource1) + manager.add_resource(resource2) + + # Should have replaced the first resource with the second + resources = manager.list_resources() + assert len(resources) == 1 + assert resources[0].name == "test2" + @pytest.mark.anyio async def test_get_resource(self, temp_file: Path): """Test getting a resource by URI.""" @@ -141,6 +179,162 @@ class TestResourceManager: assert resources == [resource1, resource2] +class TestResourceTags: + """Test functionality related to resource tags.""" + + def test_add_resource_with_tags(self, temp_file: Path): + """Test adding a resource with tags.""" + manager = ResourceManager() + resource = FileResource( + uri=FileUrl(f"file://{temp_file}"), + name="weather_data", + path=temp_file, + tags={"weather", "data"}, + ) + manager.add_resource(resource) + + # Check that tags are preserved + resources = manager.list_resources() + assert len(resources) == 1 + assert resources[0].tags == {"weather", "data"} + + def test_add_function_resource_with_tags(self): + """Test adding a function resource with tags.""" + manager = ResourceManager() + + async def get_data(): + return "Sample data" + + resource = FunctionResource( + uri=AnyUrl("data://sample"), + name="sample_data", + description="Sample data resource", + mime_type="text/plain", + fn=get_data, + tags={"sample", "test", "data"}, + ) + + manager.add_resource(resource) + resources = manager.list_resources() + assert len(resources) == 1 + assert resources[0].tags == {"sample", "test", "data"} + + def test_add_template_with_tags(self): + """Test adding a resource template with tags.""" + manager = ResourceManager() + + def user_data(user_id: str) -> str: + return f"Data for user {user_id}" + + template = ResourceTemplate.from_function( + fn=user_data, + uri_template="users://{user_id}", + name="user_template", + description="Get user data by ID", + tags={"users", "template", "data"}, + ) + + manager.add_template(template) + templates = manager.list_templates() + assert len(templates) == 1 + assert templates[0].tags == {"users", "template", "data"} + + def test_filter_resources_by_tags(self, temp_file: Path): + """Test filtering resources by tags.""" + manager = ResourceManager() + + # Create multiple resources with different tags + resource1 = FileResource( + uri=FileUrl(f"file://{temp_file}1"), + name="weather_data", + path=temp_file, + tags={"weather", "external"}, + ) + + async def get_user_data(): + return "User data" + + resource2 = FunctionResource( + uri=AnyUrl("data://users"), + name="user_data", + fn=get_user_data, + tags={"users", "internal"}, + ) + + async def get_system_data(): + return "System data" + + resource3 = FunctionResource( + uri=AnyUrl("data://system"), + name="system_data", + fn=get_system_data, + tags={"system", "internal"}, + ) + + manager.add_resource(resource1) + manager.add_resource(resource2) + manager.add_resource(resource3) + + # Filter resources by tags + internal_resources = [ + r for r in manager.list_resources() if "internal" in r.tags + ] + assert len(internal_resources) == 2 + assert {r.name for r in internal_resources} == {"user_data", "system_data"} + + external_resources = [ + r for r in manager.list_resources() if "external" in r.tags + ] + assert len(external_resources) == 1 + assert external_resources[0].name == "weather_data" + + def test_import_resources_preserves_tags(self): + """Test that importing resources preserves their tags.""" + source_manager = ResourceManager() + + async def get_data(): + return "Tagged data" + + resource = FunctionResource( + uri=AnyUrl("data://tagged"), + name="tagged_data", + fn=get_data, + tags={"test", "example", "data"}, + ) + + source_manager.add_resource(resource) + + target_manager = ResourceManager() + target_manager.import_resources(source_manager, "imported+") + + imported_resources = target_manager.list_resources() + assert len(imported_resources) == 1 + assert imported_resources[0].tags == {"test", "example", "data"} + + def test_import_templates_preserves_tags(self): + """Test that importing templates preserves their tags.""" + source_manager = ResourceManager() + + def user_template(user_id: str) -> str: + return f"User {user_id}" + + template = ResourceTemplate.from_function( + fn=user_template, + uri_template="users://{user_id}", + name="user_template", + tags={"users", "template", "test"}, + ) + + source_manager.add_template(template) + + target_manager = ResourceManager() + target_manager.import_templates(source_manager, "imported+") + + imported_templates = target_manager.list_templates() + assert len(imported_templates) == 1 + assert imported_templates[0].tags == {"users", "template", "test"} + + class TestImports: def test_import_resources(self): """Test importing resources from one manager to another with a prefix.""" diff --git a/tests/server/test_server.py b/tests/server/test_server.py index 8d2df726c..0ad2af419 100644 --- a/tests/server/test_server.py +++ b/tests/server/test_server.py @@ -17,7 +17,7 @@ from mcp.types import ( from pydantic import AnyUrl, Field from fastmcp import Context, FastMCP -from fastmcp.prompts.base import EmbeddedResource, Message, UserMessage +from fastmcp.prompts.prompt import EmbeddedResource, Message, UserMessage from fastmcp.resources import FileResource, FunctionResource from fastmcp.utilities.types import Image @@ -482,7 +482,7 @@ class TestContextInjection: def tool_with_context(x: int, ctx: Context) -> str: return f"Request {ctx.request_id}: {x}" - tool = mcp._tool_manager.add_tool(tool_with_context) + tool = mcp._tool_manager.add_tool_from_fn(tool_with_context) assert tool.context_kwarg == "ctx" async def test_context_injection(self): diff --git a/tests/tools/test_tool_manager.py b/tests/tools/test_tool_manager.py index 4115b00c9..c30078375 100644 --- a/tests/tools/test_tool_manager.py +++ b/tests/tools/test_tool_manager.py @@ -5,6 +5,7 @@ import pytest from pydantic import BaseModel from fastmcp.exceptions import ToolError +from fastmcp.settings import DuplicateBehavior from fastmcp.tools import ToolManager @@ -17,7 +18,7 @@ class TestAddTools: return a + b manager = ToolManager() - manager.add_tool(add) + manager.add_tool_from_fn(add) tool = manager.get_tool("add") assert tool is not None @@ -36,7 +37,7 @@ class TestAddTools: return f"Data from {url}" manager = ToolManager() - manager.add_tool(fetch_data) + manager.add_tool_from_fn(fetch_data) tool = manager.get_tool("fetch_data") assert tool is not None @@ -57,7 +58,7 @@ class TestAddTools: return {"id": 1, **user.model_dump()} manager = ToolManager() - manager.add_tool(create_user) + manager.add_tool_from_fn(create_user) tool = manager.get_tool("create_user") assert tool is not None @@ -71,11 +72,11 @@ class TestAddTools: def test_add_invalid_tool(self): manager = ToolManager() with pytest.raises(AttributeError): - manager.add_tool(1) # type: ignore + manager.add_tool_from_fn(1) # type: ignore def test_add_lambda(self): manager = ToolManager() - tool = manager.add_tool(lambda x: x, name="my_tool") + tool = manager.add_tool_from_fn(lambda x: x, name="my_tool") assert tool.name == "my_tool" def test_add_lambda_with_no_name(self): @@ -83,7 +84,7 @@ class TestAddTools: with pytest.raises( ValueError, match="You must provide a name for lambda functions" ): - manager.add_tool(lambda x: x) + manager.add_tool_from_fn(lambda x: x) def test_warn_on_duplicate_tools(self, caplog): """Test warning on duplicate tools.""" @@ -91,10 +92,10 @@ class TestAddTools: def f(x: int) -> int: return x - manager = ToolManager() - manager.add_tool(f) + manager = ToolManager(duplicate_behavior=DuplicateBehavior.WARN) + manager.add_tool_from_fn(f) with caplog.at_level(logging.WARNING): - manager.add_tool(f) + manager.add_tool_from_fn(f) assert "Tool already exists: f" in caplog.text def test_disable_warn_on_duplicate_tools(self, caplog): @@ -103,13 +104,132 @@ class TestAddTools: def f(x: int) -> int: return x - manager = ToolManager() - manager.add_tool(f) - manager.warn_on_duplicate_tools = False + manager = ToolManager(duplicate_behavior=DuplicateBehavior.IGNORE) + manager.add_tool_from_fn(f) with caplog.at_level(logging.WARNING): - manager.add_tool(f) + manager.add_tool_from_fn(f) assert "Tool already exists: f" not in caplog.text + def test_error_on_duplicate_tools(self): + """Test error on duplicate tools.""" + + def f(x: int) -> int: + return x + + manager = ToolManager(duplicate_behavior=DuplicateBehavior.ERROR) + manager.add_tool_from_fn(f) + + with pytest.raises(ValueError, match="Tool already exists"): + manager.add_tool_from_fn(f) + + def test_replace_duplicate_tools(self): + """Test replacing duplicate tools.""" + + def original_fn(x: int) -> int: + return x + + def replacement_fn(x: int) -> int: + return x * 2 + + manager = ToolManager(duplicate_behavior=DuplicateBehavior.REPLACE) + manager.add_tool_from_fn(original_fn, name="test_tool") + replacement_tool = manager.add_tool_from_fn(replacement_fn, name="test_tool") + + # Should have replaced the first tool with the second + stored_tool = manager.get_tool("test_tool") + assert stored_tool == replacement_tool + + +class TestToolTags: + """Test functionality related to tool tags.""" + + def test_add_tool_with_tags(self): + """Test adding tags to a tool.""" + + def example_tool(x: int) -> int: + """An example tool with tags.""" + return x * 2 + + manager = ToolManager() + tool = manager.add_tool_from_fn(example_tool, tags={"math", "utility"}) + + assert tool.tags == {"math", "utility"} + tool = manager.get_tool("example_tool") + assert tool is not None + assert tool.tags == {"math", "utility"} + + def test_add_tool_with_empty_tags(self): + """Test adding a tool with empty tags set.""" + + def example_tool(x: int) -> int: + """An example tool with empty tags.""" + return x * 2 + + manager = ToolManager() + tool = manager.add_tool_from_fn(example_tool, tags=set()) + + assert tool.tags == set() + + def test_add_tool_with_none_tags(self): + """Test adding a tool with None tags.""" + + def example_tool(x: int) -> int: + """An example tool with None tags.""" + return x * 2 + + manager = ToolManager() + tool = manager.add_tool_from_fn(example_tool, tags=None) + + assert tool.tags == set() + + def test_list_tools_with_tags(self): + """Test listing tools with specific tags.""" + + def math_tool(x: int) -> int: + """A math tool.""" + return x * 2 + + def string_tool(x: str) -> str: + """A string tool.""" + return x.upper() + + def mixed_tool(x: int) -> str: + """A tool with multiple tags.""" + return str(x) + + manager = ToolManager() + manager.add_tool_from_fn(math_tool, tags={"math"}) + manager.add_tool_from_fn(string_tool, tags={"string", "utility"}) + manager.add_tool_from_fn(mixed_tool, tags={"math", "utility", "string"}) + + # Check if we can filter by tags when listing tools + math_tools = [tool for tool in manager.list_tools() if "math" in tool.tags] + assert len(math_tools) == 2 + assert {tool.name for tool in math_tools} == {"math_tool", "mixed_tool"} + + utility_tools = [ + tool for tool in manager.list_tools() if "utility" in tool.tags + ] + assert len(utility_tools) == 2 + assert {tool.name for tool in utility_tools} == {"string_tool", "mixed_tool"} + + def test_import_tools_preserves_tags(self): + """Test that importing tools preserves their tags.""" + + def tagged_tool(x: int) -> int: + """A tool with tags.""" + return x + + source_manager = ToolManager() + source_manager.add_tool_from_fn(tagged_tool, tags={"test", "example"}) + + target_manager = ToolManager() + target_manager.import_tools(source_manager, "source/") + + imported_tool = target_manager.get_tool("source/tagged_tool") + assert imported_tool is not None + assert imported_tool.tags == {"test", "example"} + class TestCallTools: @pytest.mark.anyio @@ -119,7 +239,7 @@ class TestCallTools: return a + b manager = ToolManager() - manager.add_tool(add) + manager.add_tool_from_fn(add) result = await manager.call_tool("add", {"a": 1, "b": 2}) assert result == 3 @@ -130,7 +250,7 @@ class TestCallTools: return n * 2 manager = ToolManager() - manager.add_tool(double) + manager.add_tool_from_fn(double) result = await manager.call_tool("double", {"n": 5}) assert result == 10 @@ -141,7 +261,7 @@ class TestCallTools: return a + b manager = ToolManager() - manager.add_tool(add) + manager.add_tool_from_fn(add) result = await manager.call_tool("add", {"a": 1}) assert result == 2 @@ -152,7 +272,7 @@ class TestCallTools: return a + b manager = ToolManager() - manager.add_tool(add) + manager.add_tool_from_fn(add) with pytest.raises(ToolError): await manager.call_tool("add", {"a": 1}) @@ -168,7 +288,7 @@ class TestCallTools: return sum(vals) manager = ToolManager() - manager.add_tool(sum_vals) + manager.add_tool_from_fn(sum_vals) # Try both with plain list and with JSON list result = await manager.call_tool("sum_vals", {"vals": "[1, 2, 3]"}) assert result == 6 @@ -181,7 +301,7 @@ class TestCallTools: return vals if isinstance(vals, str) else "".join(vals) manager = ToolManager() - manager.add_tool(concat_strs) + manager.add_tool_from_fn(concat_strs) # Try both with plain python object and with JSON list result = await manager.call_tool("concat_strs", {"vals": ["a", "b", "c"]}) assert result == "abc" @@ -207,7 +327,7 @@ class TestCallTools: return [x.name for x in tank.shrimp] manager = ToolManager() - manager.add_tool(name_shrimp) + manager.add_tool_from_fn(name_shrimp) result = await manager.call_tool( "name_shrimp", {"tank": {"x": None, "shrimp": [{"name": "rex"}, {"name": "gertrude"}]}}, @@ -229,7 +349,7 @@ class TestToolSchema: return a manager = ToolManager() - tool = manager.add_tool(something) + tool = manager.add_tool_from_fn(something) assert "ctx" not in json.dumps(tool.parameters) assert "Context" not in json.dumps(tool.parameters) assert "ctx" not in tool.fn_metadata.arg_model.model_fields @@ -247,13 +367,13 @@ class TestContextHandling: return str(x) manager = ToolManager() - tool = manager.add_tool(tool_with_context) + tool = manager.add_tool_from_fn(tool_with_context) assert tool.context_kwarg == "ctx" def tool_without_context(x: int) -> str: return str(x) - tool = manager.add_tool(tool_without_context) + tool = manager.add_tool_from_fn(tool_without_context) assert tool.context_kwarg is None @pytest.mark.anyio @@ -266,7 +386,7 @@ class TestContextHandling: return str(x) manager = ToolManager() - manager.add_tool(tool_with_context) + manager.add_tool_from_fn(tool_with_context) mcp = FastMCP() ctx = mcp.get_context() @@ -283,7 +403,7 @@ class TestContextHandling: return str(x) manager = ToolManager() - manager.add_tool(async_tool) + manager.add_tool_from_fn(async_tool) mcp = FastMCP() ctx = mcp.get_context() @@ -299,7 +419,7 @@ class TestContextHandling: return str(x) manager = ToolManager() - manager.add_tool(tool_with_context) + manager.add_tool_from_fn(tool_with_context) # Should not raise an error when context is not provided result = await manager.call_tool("tool_with_context", {"x": 42}) assert result == "42" @@ -313,7 +433,7 @@ class TestContextHandling: raise ValueError("Test error") manager = ToolManager() - manager.add_tool(tool_with_context) + manager.add_tool_from_fn(tool_with_context) mcp = FastMCP() ctx = mcp.get_context() @@ -335,8 +455,10 @@ class TestImportTools: return "Tool 2 result" # Add tools to source manager - source_manager.add_tool(tool1_fn, name="get_data", description="Get some data") - source_manager.add_tool( + source_manager.add_tool_from_fn( + tool1_fn, name="get_data", description="Get some data" + ) + source_manager.add_tool_from_fn( tool2_fn, name="process_data", description="Process the data" ) @@ -364,11 +486,8 @@ class TestImportTools: # Verify the tool functions were properly copied # We can't directly compare functions, so we'll check their __name__ attribute - assert target_manager._tools["source/get_data"].fn.__name__ == tool1_fn.__name__ - assert ( - target_manager._tools["source/process_data"].fn.__name__ - == tool2_fn.__name__ - ) + assert target_manager._tools["source/get_data"].fn == tool1_fn + assert target_manager._tools["source/process_data"].fn == tool2_fn def test_tool_duplicate_behavior(self): """Test the behavior when importing tools with duplicate names.""" @@ -383,8 +502,8 @@ class TestImportTools: def target_fn(): return "Target result" - source_manager.add_tool(source_fn, name="common_tool") - target_manager.add_tool( + source_manager.add_tool_from_fn(source_fn, name="common_tool") + target_manager.add_tool_from_fn( target_fn, name="source/common_tool" ) # Pre-create with the prefixed name @@ -392,10 +511,7 @@ class TestImportTools: target_manager.import_tools(source_manager, "source/") # The original tool in the target manager is replaced by the imported one - assert ( - target_manager._tools["source/common_tool"].fn.__name__ - == source_fn.__name__ - ) + assert target_manager._tools["source/common_tool"].fn == source_fn def test_import_tools_with_multiple_prefixes(self): """Test importing tools from multiple managers with different prefixes.""" @@ -410,8 +526,8 @@ class TestImportTools: def headlines_fn(): return "News headlines" - weather_manager.add_tool(forecast_fn, name="forecast") - news_manager.add_tool(headlines_fn, name="headlines") + weather_manager.add_tool_from_fn(forecast_fn, name="forecast") + news_manager.add_tool_from_fn(headlines_fn, name="headlines") # Create target manager and import from both sources main_manager = ToolManager()