From 5e6b13e986017cfba0007fb5190dc863ca33ec44 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sat, 12 Apr 2025 10:27:20 -0400 Subject: [PATCH] Add tags for all core MCP objects --- src/fastmcp/prompts/base.py | 11 ++++++++-- src/fastmcp/resources/base.py | 6 ++++++ src/fastmcp/resources/resource_manager.py | 2 ++ src/fastmcp/resources/templates.py | 11 ++++++++-- src/fastmcp/server/server.py | 26 ++++++++++++++++++----- src/fastmcp/tools/base.py | 10 +++++++-- src/fastmcp/tools/tool_manager.py | 9 ++++++-- src/fastmcp/utilities/types.py | 12 +++++++++++ 8 files changed, 74 insertions(+), 13 deletions(-) diff --git a/src/fastmcp/prompts/base.py b/src/fastmcp/prompts/base.py index 98d141786..f4fce3413 100644 --- a/src/fastmcp/prompts/base.py +++ b/src/fastmcp/prompts/base.py @@ -3,11 +3,13 @@ import inspect import json from collections.abc import Awaitable, Callable, Sequence -from typing import Any, Literal +from typing import Annotated, Any, Literal import pydantic_core from mcp.types import EmbeddedResource, ImageContent, TextContent -from pydantic import BaseModel, Field, TypeAdapter, validate_call +from pydantic import BaseModel, BeforeValidator, Field, TypeAdapter, validate_call + +from fastmcp.utilities.types import _convert_set_defaults CONTENT_TYPES = TextContent | ImageContent | EmbeddedResource @@ -71,6 +73,9 @@ class Prompt(BaseModel): description: str | None = Field( None, description="Description of what the prompt does" ) + tags: Annotated[set[str], BeforeValidator(_convert_set_defaults)] = Field( + default_factory=set, description="Tags for the prompt" + ) arguments: list[PromptArgument] | None = Field( None, description="Arguments that can be passed to the prompt" ) @@ -82,6 +87,7 @@ class Prompt(BaseModel): 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. @@ -120,6 +126,7 @@ class Prompt(BaseModel): description=description or fn.__doc__ or "", arguments=arguments, fn=fn, + tags=tags or set(), ) async def render(self, arguments: dict[str, Any] | None = None) -> list[Message]: diff --git a/src/fastmcp/resources/base.py b/src/fastmcp/resources/base.py index b2050e7f8..d489d8919 100644 --- a/src/fastmcp/resources/base.py +++ b/src/fastmcp/resources/base.py @@ -6,6 +6,7 @@ from typing import Annotated from pydantic import ( AnyUrl, BaseModel, + BeforeValidator, ConfigDict, Field, UrlConstraints, @@ -13,6 +14,8 @@ from pydantic import ( field_validator, ) +from fastmcp.utilities.types import _convert_set_defaults + class Resource(BaseModel, abc.ABC): """Base class for all resources.""" @@ -26,6 +29,9 @@ class Resource(BaseModel, abc.ABC): description: str | None = Field( description="Description of the resource", default=None ) + tags: Annotated[set[str], BeforeValidator(_convert_set_defaults)] = Field( + default_factory=set, description="Tags for the resource" + ) mime_type: str = Field( default="text/plain", description="MIME type of the resource content", diff --git a/src/fastmcp/resources/resource_manager.py b/src/fastmcp/resources/resource_manager.py index c3bc06f19..6be187645 100644 --- a/src/fastmcp/resources/resource_manager.py +++ b/src/fastmcp/resources/resource_manager.py @@ -53,6 +53,7 @@ class ResourceManager: name: str | None = None, description: str | None = None, mime_type: str | None = None, + tags: set[str] | None = None, ) -> ResourceTemplate: """Add a template from a function.""" template = ResourceTemplate.from_function( @@ -61,6 +62,7 @@ class ResourceManager: name=name, description=description, mime_type=mime_type, + tags=tags, ) self._templates[template.uri_template] = template return template diff --git a/src/fastmcp/resources/templates.py b/src/fastmcp/resources/templates.py index ac1316e73..69a085c28 100644 --- a/src/fastmcp/resources/templates.py +++ b/src/fastmcp/resources/templates.py @@ -5,11 +5,12 @@ from __future__ import annotations import inspect import re from collections.abc import Callable -from typing import Any +from typing import Annotated, Any -from pydantic import BaseModel, Field, TypeAdapter, validate_call +from pydantic import BaseModel, BeforeValidator, Field, TypeAdapter, validate_call from fastmcp.resources.types import FunctionResource, Resource +from fastmcp.utilities.types import _convert_set_defaults class ResourceTemplate(BaseModel): @@ -20,6 +21,9 @@ class ResourceTemplate(BaseModel): ) name: str = Field(description="Name of the resource") description: str | None = Field(description="Description of what the resource does") + tags: Annotated[set[str], BeforeValidator(_convert_set_defaults)] = Field( + default_factory=set, description="Tags for the resource" + ) mime_type: str = Field( default="text/plain", description="MIME type of the resource content" ) @@ -36,6 +40,7 @@ class ResourceTemplate(BaseModel): name: str | None = None, description: str | None = None, mime_type: str | None = None, + tags: set[str] | None = None, ) -> ResourceTemplate: """Create a template from a function.""" func_name = name or fn.__name__ @@ -55,6 +60,7 @@ class ResourceTemplate(BaseModel): mime_type=mime_type or "text/plain", fn=fn, parameters=parameters, + tags=tags or set(), ) def matches(self, uri: str) -> dict[str, Any] | None: @@ -80,6 +86,7 @@ class ResourceTemplate(BaseModel): description=self.description, mime_type=self.mime_type, fn=lambda: result, # Capture result in closure + tags=self.tags, ) except Exception as e: raise ValueError(f"Error creating resource from template: {e}") diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 53c95d3fe..517e041a0 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -228,6 +228,7 @@ class FastMCP(Generic[LifespanResultT]): fn: AnyFunction, name: str | None = None, description: str | None = None, + tags: set[str] | None = None, ) -> None: """Add a tool to the server. @@ -238,11 +239,15 @@ class FastMCP(Generic[LifespanResultT]): fn: The function to register as a tool name: Optional name for the tool (defaults to function name) 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) + self._tool_manager.add_tool(fn, name=name, description=description, tags=tags) def tool( - self, name: str | None = None, description: str | None = None + self, + name: str | None = None, + description: str | None = None, + tags: set[str] | None = None, ) -> Callable[[AnyFunction], AnyFunction]: """Decorator to register a tool. @@ -253,6 +258,7 @@ class FastMCP(Generic[LifespanResultT]): Args: name: Optional name for the tool (defaults to function name) description: Optional description of what the tool does + tags: Optional set of tags for categorizing the tool Example: @server.tool() @@ -277,7 +283,7 @@ class FastMCP(Generic[LifespanResultT]): ) def decorator(fn: AnyFunction) -> AnyFunction: - self.add_tool(fn, name=name, description=description) + self.add_tool(fn, name=name, description=description, tags=tags) return fn return decorator @@ -297,6 +303,7 @@ class FastMCP(Generic[LifespanResultT]): name: str | None = None, description: str | None = None, mime_type: str | None = None, + tags: set[str] | None = None, ) -> Callable[[AnyFunction], AnyFunction]: """Decorator to register a function as a resource. @@ -314,6 +321,7 @@ class FastMCP(Generic[LifespanResultT]): name: Optional name for the resource description: Optional description of the resource mime_type: Optional MIME type for the resource + tags: Optional set of tags for categorizing the resource Example: @server.resource("resource://my-resource") @@ -364,6 +372,7 @@ class FastMCP(Generic[LifespanResultT]): name=name, description=description, mime_type=mime_type or "text/plain", + tags=tags, ) else: # Register as regular resource @@ -373,6 +382,7 @@ class FastMCP(Generic[LifespanResultT]): description=description, mime_type=mime_type or "text/plain", fn=fn, + tags=tags or set(), # Default to empty set if None ) self.add_resource(resource) return fn @@ -388,13 +398,17 @@ class FastMCP(Generic[LifespanResultT]): self._prompt_manager.add_prompt(prompt) def prompt( - self, name: str | None = None, description: str | None = None + self, + name: str | None = None, + description: str | None = None, + tags: set[str] | None = None, ) -> Callable[[AnyFunction], AnyFunction]: """Decorator to register a prompt. Args: name: Optional name for the prompt (defaults to function name) description: Optional description of what the prompt does + tags: Optional set of tags for categorizing the prompt Example: @server.prompt() @@ -431,7 +445,9 @@ class FastMCP(Generic[LifespanResultT]): ) def decorator(func: AnyFunction) -> AnyFunction: - prompt = Prompt.from_function(func, name=name, description=description) + prompt = Prompt.from_function( + func, name=name, description=description, tags=tags + ) self.add_prompt(prompt) return func diff --git a/src/fastmcp/tools/base.py b/src/fastmcp/tools/base.py index 27305cb5a..cebe2c10a 100644 --- a/src/fastmcp/tools/base.py +++ b/src/fastmcp/tools/base.py @@ -2,12 +2,13 @@ from __future__ import annotations as _annotations import inspect from collections.abc import Callable -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Annotated, Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, BeforeValidator, Field from fastmcp.exceptions import ToolError from fastmcp.utilities.func_metadata import FuncMetadata, func_metadata +from fastmcp.utilities.types import _convert_set_defaults if TYPE_CHECKING: from mcp.server.session import ServerSessionT @@ -31,6 +32,9 @@ class Tool(BaseModel): context_kwarg: str | None = Field( None, description="Name of the kwarg that should receive context" ) + tags: Annotated[set[str], BeforeValidator(_convert_set_defaults)] = Field( + default_factory=set, description="Tags for the tool" + ) @classmethod def from_function( @@ -39,6 +43,7 @@ class Tool(BaseModel): name: str | None = None, description: str | None = None, context_kwarg: str | None = None, + tags: set[str] | None = None, ) -> Tool: """Create a Tool from a function.""" from fastmcp import Context @@ -72,6 +77,7 @@ class Tool(BaseModel): fn_metadata=func_arg_metadata, is_async=is_async, context_kwarg=context_kwarg, + tags=tags or set(), ) async def run( diff --git a/src/fastmcp/tools/tool_manager.py b/src/fastmcp/tools/tool_manager.py index cd8c882b8..31b02b2aa 100644 --- a/src/fastmcp/tools/tool_manager.py +++ b/src/fastmcp/tools/tool_manager.py @@ -37,9 +37,14 @@ class ToolManager: fn: Callable[..., Any], name: str | None = None, description: str | None = None, + tags: set[str] | None = None, ) -> Tool: """Add a tool to the server.""" - tool = Tool.from_function(fn, name=name, description=description) + tool = Tool.from_function(fn, name=name, description=description, tags=tags) + return self._register_tool(tool) + + def _register_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: @@ -86,5 +91,5 @@ class ToolManager: ) # Store the copied tool - self._tools[prefixed_name] = copied_tool + self._register_tool(copied_tool) logger.debug(f"Imported tool: {name} as {prefixed_name}") diff --git a/src/fastmcp/utilities/types.py b/src/fastmcp/utilities/types.py index ccaa3d69a..98197a922 100644 --- a/src/fastmcp/utilities/types.py +++ b/src/fastmcp/utilities/types.py @@ -2,9 +2,21 @@ import base64 from pathlib import Path +from typing import TypeVar from mcp.types import ImageContent +T = TypeVar("T") + + +def _convert_set_defaults(maybe_set: set[T] | list[T] | None) -> set[T]: + """Convert a set or list to a set, defaulting to an empty set if None.""" + if maybe_set is None: + return set() + if isinstance(maybe_set, set): + return maybe_set + return set(maybe_set) + class Image: """Helper class for returning images from tools."""