mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-18 19:44:19 +02:00
Add template support
This commit is contained in:
Jeremiah Lowin
parent
659b7418cb
commit
9fbc8a3fec
3 changed files with 448 additions and 90 deletions
|
|
@ -1,13 +1,13 @@
|
|||
import pydantic.json
|
||||
|
||||
import abc
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Callable, Any, Union, Awaitable
|
||||
from typing import Dict, Optional, Callable, Any, Union
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, field_validator
|
||||
from pydantic import BaseModel, Field, TypeAdapter, validate_call, field_validator
|
||||
from pydantic.networks import _BaseUrl
|
||||
|
||||
from .utilities.logging import get_logger
|
||||
|
|
@ -15,19 +15,13 @@ from .utilities.logging import get_logger
|
|||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class Resource(BaseModel):
|
||||
"""Base class for all resources.
|
||||
class Resource(BaseModel, abc.ABC):
|
||||
"""Base class for all resources."""
|
||||
|
||||
Resources can contain either text (UTF-8 encoded) or binary data.
|
||||
Text resources are suitable for source code, logs, JSON, etc.
|
||||
Binary resources are suitable for images, PDFs, audio, etc.
|
||||
"""
|
||||
|
||||
uri: _BaseUrl
|
||||
name: str
|
||||
description: Optional[str] = None
|
||||
mime_type: Optional[str] = None
|
||||
is_binary: bool = False
|
||||
uri: _BaseUrl = Field(description="URI of the resource")
|
||||
name: str = Field(description="Name of the resource")
|
||||
description: Optional[str] = Field(description="Description of the resource")
|
||||
mime_type: Optional[str] = Field(description="MIME type of the resource content")
|
||||
|
||||
@field_validator("name", mode="before")
|
||||
@classmethod
|
||||
|
|
@ -43,112 +37,105 @@ class Resource(BaseModel):
|
|||
|
||||
@abc.abstractmethod
|
||||
async def read(self) -> Union[str, bytes]:
|
||||
"""Read the resource content.
|
||||
|
||||
Returns:
|
||||
Union[str, bytes]: Text content as str for text resources,
|
||||
binary content as bytes for binary resources
|
||||
"""
|
||||
return ""
|
||||
"""Read the resource content."""
|
||||
pass
|
||||
|
||||
|
||||
class FunctionResource(Resource):
|
||||
"""A resource that is generated by a function call.
|
||||
class TextResource(Resource):
|
||||
"""A resource containing text content."""
|
||||
|
||||
The function can be sync or async and must return a string, bytes,
|
||||
or another Resource.
|
||||
"""
|
||||
text: str = Field(description="Text content of the resource")
|
||||
mime_type: Optional[str] = Field(
|
||||
default="text/plain", description="MIME type of the resource content"
|
||||
)
|
||||
|
||||
func: Union[Callable[[], Any], Callable[[], Awaitable[Any]]]
|
||||
is_async: bool = False
|
||||
async def read(self) -> str:
|
||||
"""Read the text content."""
|
||||
return self.text
|
||||
|
||||
def __init__(self, **data):
|
||||
super().__init__(**data)
|
||||
self.is_async = asyncio.iscoroutinefunction(self.func)
|
||||
|
||||
async def read(self) -> Union[str, bytes]:
|
||||
"""Read the resource content by calling the function."""
|
||||
try:
|
||||
result = (
|
||||
await self.func()
|
||||
if self.is_async
|
||||
else await asyncio.to_thread(self.func)
|
||||
)
|
||||
class BinaryResource(Resource):
|
||||
"""A resource containing binary content."""
|
||||
|
||||
if isinstance(result, Resource):
|
||||
return await result.read()
|
||||
if isinstance(result, bytes):
|
||||
return result
|
||||
if not isinstance(result, str):
|
||||
try:
|
||||
return json.dumps(result, default=pydantic.json.pydantic_encoder)
|
||||
except json.JSONDecodeError:
|
||||
return str(result)
|
||||
return result
|
||||
except Exception as e:
|
||||
raise ValueError(f"Error calling function {self.func.__name__}: {e}")
|
||||
data: bytes = Field(description="Binary content of the resource")
|
||||
mime_type: Optional[str] = Field(
|
||||
default="application/octet-stream",
|
||||
description="MIME type of the resource content",
|
||||
)
|
||||
|
||||
async def read(self) -> bytes:
|
||||
"""Read the binary content."""
|
||||
return self.data
|
||||
|
||||
|
||||
class FileResource(Resource):
|
||||
"""A file resource."""
|
||||
"""A resource that reads from a file."""
|
||||
|
||||
path: Path
|
||||
path: Path = Field(description="Path to the file")
|
||||
mime_type: Optional[str] = Field(
|
||||
default="application/octet-stream",
|
||||
description="MIME type of the resource content",
|
||||
)
|
||||
|
||||
@field_validator("path")
|
||||
@classmethod
|
||||
def validate_absolute_path(cls, path: Path) -> Path:
|
||||
"""Ensure path is absolute."""
|
||||
if not path.is_absolute():
|
||||
raise ValueError(f"Path must be absolute: {path}")
|
||||
raise ValueError("Path must be absolute")
|
||||
return path
|
||||
|
||||
async def read(self) -> Union[str, bytes]:
|
||||
"""Read the file content."""
|
||||
try:
|
||||
if self.is_binary:
|
||||
return await asyncio.to_thread(self.path.read_bytes)
|
||||
return await asyncio.to_thread(self.path.read_text)
|
||||
except FileNotFoundError:
|
||||
raise FileNotFoundError(f"File not found: {self.path}")
|
||||
except PermissionError:
|
||||
raise PermissionError(f"Permission denied: {self.path}")
|
||||
except Exception as e:
|
||||
raise ValueError(f"Error reading file {self.path}: {e}")
|
||||
if self.mime_type and self.mime_type.startswith("text/"):
|
||||
return await self._read_text()
|
||||
return await self._read_binary()
|
||||
|
||||
async def _read_text(self) -> str:
|
||||
"""Read file as text."""
|
||||
return await asyncio.to_thread(self.path.read_text)
|
||||
|
||||
async def _read_binary(self) -> bytes:
|
||||
"""Read file as binary."""
|
||||
return await asyncio.to_thread(self.path.read_bytes)
|
||||
|
||||
|
||||
class HttpResource(Resource):
|
||||
"""An HTTP resource."""
|
||||
"""A resource that reads from an HTTP endpoint."""
|
||||
|
||||
url: str
|
||||
headers: Optional[Dict[str, str]] = None
|
||||
url: str = Field(description="URL to fetch content from")
|
||||
mime_type: Optional[str] = Field(
|
||||
default="application/json", description="MIME type of the resource content"
|
||||
)
|
||||
|
||||
async def read(self) -> Union[str, bytes]:
|
||||
"""Read the HTTP resource content."""
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(self.url, headers=self.headers)
|
||||
response.raise_for_status()
|
||||
return response.content if self.is_binary else response.text
|
||||
except httpx.HTTPStatusError as e:
|
||||
raise ValueError(f"HTTP error {e.response.status_code}: {e}")
|
||||
except httpx.RequestError as e:
|
||||
raise ValueError(f"Request failed: {e}")
|
||||
"""Read the HTTP content."""
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(self.url)
|
||||
response.raise_for_status()
|
||||
return response.text
|
||||
|
||||
|
||||
class DirectoryResource(Resource):
|
||||
"""A directory resource."""
|
||||
"""A resource that lists files in a directory."""
|
||||
|
||||
path: Path
|
||||
recursive: bool = False
|
||||
pattern: Optional[str] = None
|
||||
mime_type: Optional[str] = "application/json"
|
||||
path: Path = Field(description="Path to the directory")
|
||||
recursive: bool = Field(
|
||||
default=False, description="Whether to list files recursively"
|
||||
)
|
||||
pattern: Optional[str] = Field(
|
||||
default=None, description="Optional glob pattern to filter files"
|
||||
)
|
||||
mime_type: Optional[str] = Field(
|
||||
default="application/json", description="MIME type of the resource content"
|
||||
)
|
||||
|
||||
@field_validator("path")
|
||||
@classmethod
|
||||
def validate_absolute_path(cls, path: Path) -> Path:
|
||||
"""Ensure path is absolute."""
|
||||
if not path.is_absolute():
|
||||
raise ValueError(f"Path must be absolute: {path}")
|
||||
raise ValueError("Path must be absolute")
|
||||
return path
|
||||
|
||||
def list_files(self) -> list[Path]:
|
||||
|
|
@ -183,21 +170,132 @@ class DirectoryResource(Resource):
|
|||
raise ValueError(f"Error reading directory {self.path}: {e}")
|
||||
|
||||
|
||||
class ResourceTemplate(BaseModel):
|
||||
"""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: Optional[str] = Field(
|
||||
description="Description of what the resource does"
|
||||
)
|
||||
mime_type: Optional[str] = Field(
|
||||
default="text/plain", description="MIME type of the resource content"
|
||||
)
|
||||
func: Callable = Field(exclude=True)
|
||||
parameters: dict = Field(description="JSON schema for function parameters")
|
||||
|
||||
@classmethod
|
||||
def from_function(
|
||||
cls,
|
||||
func: Callable,
|
||||
uri_template: str,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
mime_type: Optional[str] = None,
|
||||
) -> "ResourceTemplate":
|
||||
"""Create a template from a function."""
|
||||
func_name = name or func.__name__
|
||||
if func_name == "<lambda>":
|
||||
raise ValueError("You must provide a name for lambda functions")
|
||||
|
||||
# Get schema from TypeAdapter - will fail if function isn't properly typed
|
||||
parameters = TypeAdapter(func).json_schema()
|
||||
|
||||
# ensure the arguments are properly cast
|
||||
func = validate_call(func)
|
||||
|
||||
return cls(
|
||||
uri_template=uri_template,
|
||||
name=func_name,
|
||||
description=description or func.__doc__ or "",
|
||||
mime_type=mime_type or "text/plain",
|
||||
func=func,
|
||||
parameters=parameters,
|
||||
)
|
||||
|
||||
def matches(self, uri: str) -> Optional[Dict[str, Any]]:
|
||||
"""Check if URI matches template and extract parameters."""
|
||||
# Convert template to regex pattern
|
||||
pattern = self.uri_template.replace("{", "(?P<").replace("}", ">[^/]+)")
|
||||
match = re.match(f"^{pattern}$", uri)
|
||||
if match:
|
||||
return match.groupdict()
|
||||
return None
|
||||
|
||||
async def create_resource(self, uri: str, params: Dict[str, Any]) -> Resource:
|
||||
"""Create a resource from the template with the given parameters."""
|
||||
result = await self.func(**params)
|
||||
|
||||
if isinstance(result, bytes):
|
||||
return BinaryResource(
|
||||
uri=uri,
|
||||
name=self.name,
|
||||
description=self.description,
|
||||
mime_type=self.mime_type,
|
||||
data=result,
|
||||
)
|
||||
|
||||
else:
|
||||
if not isinstance(result, str):
|
||||
try:
|
||||
result = json.dumps(result, default=pydantic.json.pydantic_encoder)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Error converting result to JSON: {e}")
|
||||
return TextResource(
|
||||
uri=uri,
|
||||
name=self.name,
|
||||
description=self.description,
|
||||
mime_type=self.mime_type,
|
||||
text=result,
|
||||
)
|
||||
|
||||
|
||||
class ResourceManager:
|
||||
"""Manages FastMCP resources."""
|
||||
|
||||
def __init__(self, warn_on_duplicate_resources: bool = True):
|
||||
self._resources: Dict[str, Resource] = {}
|
||||
self._templates: Dict[str, ResourceTemplate] = {}
|
||||
self.warn_on_duplicate_resources = warn_on_duplicate_resources
|
||||
|
||||
def get_resource(self, uri: Union[_BaseUrl, str]) -> Optional[Resource]:
|
||||
"""Get resource by URI."""
|
||||
uri = str(uri)
|
||||
logger.debug("Getting resource", extra={"uri": uri})
|
||||
def add_template(
|
||||
self,
|
||||
func: Callable,
|
||||
uri_template: str,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
mime_type: Optional[str] = None,
|
||||
) -> ResourceTemplate:
|
||||
"""Add a template from a function."""
|
||||
template = ResourceTemplate.from_function(
|
||||
func,
|
||||
uri_template=uri_template,
|
||||
name=name,
|
||||
description=description,
|
||||
mime_type=mime_type,
|
||||
)
|
||||
self._templates[template.uri_template] = template
|
||||
return template
|
||||
|
||||
if resource := self._resources.get(uri):
|
||||
async def get_resource(self, uri: Union[_BaseUrl, str]) -> Optional[Resource]:
|
||||
"""Get resource by URI, checking concrete resources first, then templates."""
|
||||
uri_str = str(uri)
|
||||
logger.debug("Getting resource", extra={"uri": uri_str})
|
||||
|
||||
# First check concrete resources
|
||||
if resource := self._resources.get(uri_str):
|
||||
return resource
|
||||
|
||||
# Then check templates
|
||||
for template in self._templates.values():
|
||||
if params := template.matches(uri_str):
|
||||
try:
|
||||
return await template.create_resource(uri_str, params)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Error creating resource from template: {e}")
|
||||
|
||||
raise ValueError(f"Unknown resource: {uri}")
|
||||
|
||||
def list_resources(self) -> list[Resource]:
|
||||
|
|
|
|||
|
|
@ -228,6 +228,50 @@ class FastMCP:
|
|||
|
||||
return decorator
|
||||
|
||||
def template(
|
||||
self,
|
||||
uri_template: str,
|
||||
*,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
mime_type: Optional[str] = None,
|
||||
) -> Callable:
|
||||
"""Decorator to register a function as a resource template.
|
||||
|
||||
Args:
|
||||
uri_template: URI template with parameters (e.g. "weather://{city}/current")
|
||||
name: Optional name for the resource
|
||||
description: Optional description of the resource
|
||||
mime_type: Optional MIME type for the resource
|
||||
|
||||
Example:
|
||||
@server.template("weather://{city}/current")
|
||||
def get_weather(city: str) -> str:
|
||||
return f"Weather for {city}"
|
||||
"""
|
||||
# Check if user passed function directly instead of calling decorator
|
||||
if callable(uri_template):
|
||||
raise TypeError(
|
||||
"The @template decorator was used incorrectly. "
|
||||
"Did you forget to call it? Use @template('uri_template') instead of @template"
|
||||
)
|
||||
|
||||
def decorator(func: Callable) -> Callable:
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
return func(*args, **kwargs)
|
||||
|
||||
self._resource_manager.add_template(
|
||||
wrapper,
|
||||
uri_template=uri_template,
|
||||
name=name,
|
||||
description=description,
|
||||
mime_type=mime_type or "text/plain",
|
||||
)
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
async def run_stdio_async(self) -> None:
|
||||
"""Run the server using stdio transport."""
|
||||
async with stdio_server() as (read_stream, write_stream):
|
||||
|
|
|
|||
|
|
@ -3,7 +3,14 @@ import pytest
|
|||
from pathlib import Path
|
||||
from tempfile import NamedTemporaryFile, TemporaryDirectory
|
||||
|
||||
from fastmcp.resources import FileResource, FunctionResource, ResourceManager
|
||||
from fastmcp.resources import (
|
||||
FileResource,
|
||||
FunctionResource,
|
||||
ResourceManager,
|
||||
TextResource,
|
||||
BinaryResource,
|
||||
ResourceTemplate,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -240,3 +247,212 @@ class TestResourceManagerList:
|
|||
assert resources[0] == added1
|
||||
assert resources[1] == added2
|
||||
assert added1 != added2
|
||||
|
||||
|
||||
class TestTextResource:
|
||||
"""Test TextResource functionality."""
|
||||
|
||||
async def test_text_resource_read(self):
|
||||
"""Test reading from a TextResource."""
|
||||
resource = TextResource(
|
||||
uri="text://test",
|
||||
name="test",
|
||||
text="Hello, world!",
|
||||
)
|
||||
content = await resource.read()
|
||||
assert content == "Hello, world!"
|
||||
assert resource.mime_type == "text/plain"
|
||||
|
||||
def test_text_resource_custom_mime(self):
|
||||
"""Test TextResource with custom mime type."""
|
||||
resource = TextResource(
|
||||
uri="text://test",
|
||||
name="test",
|
||||
text="<html></html>",
|
||||
mime_type="text/html",
|
||||
)
|
||||
assert resource.mime_type == "text/html"
|
||||
|
||||
|
||||
class TestBinaryResource:
|
||||
"""Test BinaryResource functionality."""
|
||||
|
||||
async def test_binary_resource_read(self):
|
||||
"""Test reading from a BinaryResource."""
|
||||
data = b"Hello, world!"
|
||||
resource = BinaryResource(
|
||||
uri="binary://test",
|
||||
name="test",
|
||||
data=data,
|
||||
)
|
||||
content = await resource.read()
|
||||
assert content == data
|
||||
assert resource.mime_type == "application/octet-stream"
|
||||
|
||||
def test_binary_resource_custom_mime(self):
|
||||
"""Test BinaryResource with custom mime type."""
|
||||
resource = BinaryResource(
|
||||
uri="binary://test",
|
||||
name="test",
|
||||
data=b"test",
|
||||
mime_type="image/png",
|
||||
)
|
||||
assert resource.mime_type == "image/png"
|
||||
|
||||
|
||||
class TestResourceTemplate:
|
||||
"""Test ResourceTemplate functionality."""
|
||||
|
||||
def test_template_from_function(self):
|
||||
"""Test creating a template from a function."""
|
||||
|
||||
def weather(city: str, units: str = "metric") -> str:
|
||||
return f"Weather in {city} ({units})"
|
||||
|
||||
template = ResourceTemplate.from_function(
|
||||
func=weather,
|
||||
uri_template="weather://{city}/current",
|
||||
name="weather",
|
||||
description="Get current weather",
|
||||
mime_type="text/plain",
|
||||
)
|
||||
|
||||
assert template.name == "weather"
|
||||
assert template.uri_template == "weather://{city}/current"
|
||||
assert template.mime_type == "text/plain"
|
||||
assert "city" in template.parameters["properties"]
|
||||
|
||||
def test_template_from_lambda_error(self):
|
||||
"""Test error when creating template from lambda without name."""
|
||||
with pytest.raises(
|
||||
ValueError, match="You must provide a name for lambda functions"
|
||||
):
|
||||
ResourceTemplate.from_function(
|
||||
func=lambda x: x,
|
||||
uri_template="test://{x}",
|
||||
)
|
||||
|
||||
def test_template_matches(self):
|
||||
"""Test URI matching against template."""
|
||||
|
||||
def dummy(x: str) -> str:
|
||||
return x
|
||||
|
||||
template = ResourceTemplate.from_function(
|
||||
func=dummy,
|
||||
uri_template="test://{x}/value",
|
||||
name="test",
|
||||
)
|
||||
|
||||
# Test matching URI
|
||||
params = template.matches("test://hello/value")
|
||||
assert params == {"x": "hello"}
|
||||
|
||||
# Test non-matching URI
|
||||
params = template.matches("test://hello/wrong")
|
||||
assert params is None
|
||||
|
||||
async def test_template_create_text_resource(self):
|
||||
"""Test creating a TextResource from template."""
|
||||
|
||||
def greet(name: str) -> str:
|
||||
return f"Hello, {name}!"
|
||||
|
||||
template = ResourceTemplate.from_function(
|
||||
func=greet,
|
||||
uri_template="greet://{name}",
|
||||
name="greeter",
|
||||
)
|
||||
|
||||
resource = await template.create_resource(
|
||||
"greet://world",
|
||||
{"name": "world"},
|
||||
)
|
||||
|
||||
assert isinstance(resource, TextResource)
|
||||
content = await resource.read()
|
||||
assert content == "Hello, world!"
|
||||
|
||||
async def test_template_create_binary_resource(self):
|
||||
"""Test creating a BinaryResource from template."""
|
||||
|
||||
def get_bytes(value: str) -> bytes:
|
||||
return value.encode()
|
||||
|
||||
template = ResourceTemplate.from_function(
|
||||
func=get_bytes,
|
||||
uri_template="bytes://{value}",
|
||||
name="bytes",
|
||||
mime_type="application/octet-stream",
|
||||
)
|
||||
|
||||
resource = await template.create_resource(
|
||||
"bytes://test",
|
||||
{"value": "test"},
|
||||
)
|
||||
|
||||
assert isinstance(resource, BinaryResource)
|
||||
content = await resource.read()
|
||||
assert content == b"test"
|
||||
|
||||
async def test_template_json_conversion(self):
|
||||
"""Test automatic JSON conversion of non-string/bytes results."""
|
||||
|
||||
def get_data(key: str) -> dict:
|
||||
return {"key": key, "value": 123}
|
||||
|
||||
template = ResourceTemplate.from_function(
|
||||
func=get_data,
|
||||
uri_template="data://{key}",
|
||||
name="data",
|
||||
)
|
||||
|
||||
resource = await template.create_resource(
|
||||
"data://test",
|
||||
{"key": "test"},
|
||||
)
|
||||
|
||||
assert isinstance(resource, TextResource)
|
||||
content = await resource.read()
|
||||
assert '"key": "test"' in content
|
||||
assert '"value": 123' in content
|
||||
|
||||
|
||||
class TestResourceManagerWithTemplates:
|
||||
"""Test ResourceManager template functionality."""
|
||||
|
||||
async def test_get_resource_from_template(self):
|
||||
"""Test getting a resource through a template."""
|
||||
manager = ResourceManager()
|
||||
|
||||
def greet(name: str) -> str:
|
||||
return f"Hello, {name}!"
|
||||
|
||||
template = ResourceTemplate.from_function(
|
||||
func=greet,
|
||||
uri_template="greet://{name}",
|
||||
name="greeter",
|
||||
)
|
||||
manager._templates[template.uri_template] = template
|
||||
|
||||
resource = await manager.get_resource("greet://world")
|
||||
assert isinstance(resource, TextResource)
|
||||
content = await resource.read()
|
||||
assert content == "Hello, world!"
|
||||
|
||||
async def test_template_error_handling(self):
|
||||
"""Test error handling in template resource creation."""
|
||||
manager = ResourceManager()
|
||||
|
||||
def failing_func(x: str) -> str:
|
||||
raise ValueError("Test error")
|
||||
|
||||
template = ResourceTemplate.from_function(
|
||||
func=failing_func,
|
||||
uri_template="fail://{x}",
|
||||
name="fail",
|
||||
)
|
||||
manager._templates[template.uri_template] = template
|
||||
|
||||
with pytest.raises(ValueError, match="Error creating resource from template"):
|
||||
await manager.get_resource("fail://test")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue