Create common base class for components

This commit is contained in:
Jeremiah Lowin 2025-06-10 09:22:57 -04:00
commit b74bce6ff7
6 changed files with 47 additions and 53 deletions

View file

@ -5,13 +5,13 @@ from __future__ import annotations as _annotations
import inspect
from abc import ABC, abstractmethod
from collections.abc import Awaitable, Callable, Sequence
from typing import TYPE_CHECKING, Annotated, Any
from typing import TYPE_CHECKING, Any
import pydantic_core
from mcp.types import EmbeddedResource, ImageContent, PromptMessage, Role, TextContent
from mcp.types import Prompt as MCPPrompt
from mcp.types import PromptArgument as MCPPromptArgument
from pydantic import BeforeValidator, Field, TypeAdapter, validate_call
from pydantic import Field, TypeAdapter, validate_call
from fastmcp.exceptions import PromptError
from fastmcp.server.dependencies import get_context
@ -19,7 +19,7 @@ from fastmcp.utilities.json_schema import compress_schema
from fastmcp.utilities.logging import get_logger
from fastmcp.utilities.types import (
FastMCPBaseModel,
_convert_set_defaults,
FastMCPComponent,
find_kwarg_by_type,
get_cached_typeadapter,
)
@ -66,26 +66,13 @@ class PromptArgument(FastMCPBaseModel):
)
class Prompt(FastMCPBaseModel, ABC):
class Prompt(FastMCPComponent, ABC):
"""A prompt template that can be rendered with parameters."""
name: str = Field(description="Name of the prompt")
description: str | None = Field(
default=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(
default=None, description="Arguments that can be passed to the prompt"
)
def __eq__(self, other: object) -> bool:
if type(self) is not type(other):
return False
assert isinstance(other, type(self))
return self.model_dump() == other.model_dump()
def to_mcp_prompt(self, **overrides: Any) -> MCPPrompt:
"""Convert the prompt to an MCP prompt."""
arguments = [

View file

@ -22,7 +22,7 @@ from pydantic import (
from fastmcp.server.dependencies import get_context
from fastmcp.utilities.types import (
FastMCPBaseModel,
_convert_set_defaults,
_convert_set_default_none,
find_kwarg_by_type,
)
@ -42,7 +42,7 @@ class Resource(FastMCPBaseModel, abc.ABC):
description: str | None = Field(
default=None, description="Description of the resource"
)
tags: Annotated[set[str], BeforeValidator(_convert_set_defaults)] = Field(
tags: Annotated[set[str], BeforeValidator(_convert_set_default_none)] = Field(
default_factory=set, description="Tags for the resource"
)
mime_type: str = Field(

View file

@ -5,12 +5,11 @@ from __future__ import annotations
import inspect
import re
from collections.abc import Callable
from typing import Annotated, Any
from typing import Any
from urllib.parse import unquote
from mcp.types import ResourceTemplate as MCPResourceTemplate
from pydantic import (
BeforeValidator,
Field,
field_validator,
validate_call,
@ -20,8 +19,7 @@ from fastmcp.resources.types import Resource
from fastmcp.server.dependencies import get_context
from fastmcp.utilities.json_schema import compress_schema
from fastmcp.utilities.types import (
FastMCPBaseModel,
_convert_set_defaults,
FastMCPComponent,
find_kwarg_by_type,
get_cached_typeadapter,
)
@ -51,17 +49,12 @@ def match_uri_template(uri: str, uri_template: str) -> dict[str, str] | None:
return None
class ResourceTemplate(FastMCPBaseModel):
class ResourceTemplate(FastMCPComponent):
"""A template for dynamically creating resources."""
uri_template: str = Field(
description="URI template with parameters (e.g. weather://{city}/current)"
)
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"
)

View file

@ -5,21 +5,20 @@ import json
from abc import ABC, abstractmethod
from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Annotated, Any
from typing import TYPE_CHECKING, Any
import pydantic_core
from mcp.types import EmbeddedResource, ImageContent, TextContent, ToolAnnotations
from mcp.types import Tool as MCPTool
from pydantic import BeforeValidator, Field
from pydantic import Field
import fastmcp
from fastmcp.server.dependencies import get_context
from fastmcp.utilities.json_schema import compress_schema
from fastmcp.utilities.logging import get_logger
from fastmcp.utilities.types import (
FastMCPBaseModel,
FastMCPComponent,
Image,
_convert_set_defaults,
find_kwarg_by_type,
get_cached_typeadapter,
)
@ -34,17 +33,10 @@ def default_serializer(data: Any) -> str:
return pydantic_core.to_json(data, fallback=str, indent=2).decode()
class Tool(FastMCPBaseModel, ABC):
class Tool(FastMCPComponent, ABC):
"""Internal tool registration info."""
name: str = Field(description="Name of the tool")
description: str | None = Field(
default=None, description="Description of what the tool does"
)
parameters: dict[str, Any] = Field(description="JSON schema for tool parameters")
tags: Annotated[set[str], BeforeValidator(_convert_set_defaults)] = Field(
default_factory=set, description="Tags for the tool"
)
annotations: ToolAnnotations | None = Field(
default=None, description="Additional annotations about the tool"
)

View file

@ -2,24 +2,55 @@
import base64
import inspect
from collections.abc import Callable
from collections.abc import Callable, Sequence
from functools import lru_cache
from pathlib import Path
from types import UnionType
from typing import Annotated, TypeVar, Union, get_args, get_origin
from mcp.types import ImageContent
from pydantic import BaseModel, ConfigDict, TypeAdapter
from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, TypeAdapter
T = TypeVar("T")
def _convert_set_default_none(maybe_set: set[T] | Sequence[T] | None) -> set[T]:
"""Convert a sequence 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 FastMCPBaseModel(BaseModel):
"""Base model for FastMCP models."""
model_config = ConfigDict(extra="forbid")
class FastMCPComponent(FastMCPBaseModel):
"""Base class for FastMCP tools, prompts, resources, and resource templates."""
name: str = Field(
description="The name of the component.",
)
description: str | None = Field(
default=None,
description="The description of the component.",
)
tags: Annotated[set[str], BeforeValidator(_convert_set_default_none)] = Field(
default_factory=set,
description="Tags for the component.",
)
def __eq__(self, other: object) -> bool:
if type(self) is not type(other):
return False
assert isinstance(other, type(self))
return self.model_dump() == other.model_dump()
@lru_cache(maxsize=5000)
def get_cached_typeadapter(cls: T) -> TypeAdapter[T]:
"""
@ -80,15 +111,6 @@ def find_kwarg_by_type(fn: Callable, kwarg_type: type) -> str | None:
return None
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."""

View file

@ -140,7 +140,7 @@ class TestTools:
assert proxy_result[0].text == "3" # type: ignore[attr-defined]
async def test_error_tool_raises_error(self, proxy_server):
with pytest.raises(ToolError, match=""):
with pytest.raises(ToolError, match="This is a test error"):
async with Client(proxy_server) as client:
await client.call_tool("error_tool", {})