Update type handling for resources

This commit is contained in:
Jeremiah Lowin 2024-11-29 19:42:46 -05:00
commit 88d34a038b
9 changed files with 133 additions and 184 deletions

View file

@ -10,22 +10,21 @@ from pathlib import Path
from fastmcp.server import FastMCP
# Create server
app = FastMCP("desktop")
# Add desktop as a directory resource
desktop = Path.home() / "Desktop"
app.add_dir_resource(
str(desktop),
recursive=True,
name="Desktop",
description="Files on the desktop",
)
mcp = FastMCP("desktop")
def main123():
# Run the server
asyncio.run(FastMCP.run_stdio(app))
@mcp.resource("desktop")
def desktop() -> list[str]:
"""List the files in the desktop directory"""
desktop = Path.home() / "Desktop"
return [str(f) for f in desktop.iterdir()]
@mcp.tool()
def add(a: int, b: int) -> int:
"""Add two numbers"""
return a + b
if __name__ == "__main__":
main123()
asyncio.run(FastMCP.run_stdio(mcp))

View file

@ -1,14 +1,14 @@
"""Resource management for FastMCP."""
import pydantic.json
import abc
import asyncio
import json
from pathlib import Path
from typing import Dict, Optional, Callable, Any
from urllib.parse import parse_qs, urlparse
from typing import Dict, Optional, Callable, Any, Union, Awaitable
import httpx
from pydantic import BaseModel, field_validator
from pydantic.networks import _BaseUrl
from .utilities.logging import get_logger
@ -18,26 +18,22 @@ logger = get_logger(__name__)
class Resource(BaseModel):
"""Base class for all resources."""
uri: str
uri: _BaseUrl
name: str
description: Optional[str] = None
mime_type: str = "text/plain"
@field_validator("uri")
@field_validator("name", mode="before")
@classmethod
def validate_uri_format(cls, uri: str) -> str:
"""Validate URI follows [protocol]://[host]/[path] format."""
parsed = urlparse(uri)
# Check protocol exists and is not empty
if not parsed.scheme:
raise ValueError("URI must have a protocol (e.g., 'http://', 'file://')")
# Check host exists and is not empty
if not parsed.netloc and not parsed.path:
raise ValueError("URI must have a host or path")
return uri
def set_default_name(cls, name: str | None, info) -> str:
"""Set default name from URI if not provided."""
if name is not None:
return name
# Extract everything after the protocol (e.g., "desktop" from "resource://desktop")
uri = info.data.get("uri")
if uri:
return str(uri).split("://", 1)[1]
raise ValueError("Either name or uri must be provided")
@abc.abstractmethod
async def read(self) -> str:
@ -48,32 +44,34 @@ class Resource(BaseModel):
class FunctionResource(Resource):
"""A resource that is generated by a function call.
The function is called with kwargs parsed from the URI query string.
For example, a URI of "fn://my_func?x=1&y=2" will call the function with
kwargs {"x": "1", "y": "2"}.
The function can be sync or async and must return a string
or another Resource.
"""
func: Callable[..., Any]
func: Union[Callable[[], Any], Callable[[], Awaitable[Any]]]
is_async: bool = False
def _parse_uri_params(self) -> Dict[str, str]:
"""Parse URI query string into kwargs."""
parsed = urlparse(self.uri)
if not parsed.query:
return {}
# parse_qs returns Dict[str, List[str]], we want Dict[str, str]
params = parse_qs(parsed.query)
return {k: v[0] for k, v in params.items()}
def __init__(self, **data):
super().__init__(**data)
self.is_async = asyncio.iscoroutinefunction(self.func)
async def read(self) -> str:
"""Read the resource content by calling the function with URI params."""
"""Read the resource content by calling the function."""
try:
kwargs = self._parse_uri_params()
result = await asyncio.to_thread(self.func, **kwargs)
result = (
await self.func()
if self.is_async
else await asyncio.to_thread(self.func)
)
if isinstance(result, Resource):
return await result.read()
if isinstance(result, bytes):
return result.decode()
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:
@ -179,33 +177,14 @@ class ResourceManager:
self._resources: Dict[str, Resource] = {}
self.warn_on_duplicate_resources = warn_on_duplicate_resources
def get_resource(self, uri: str) -> Optional[Resource]:
"""Get resource by URI.
First tries to find an exact match for the URI. If none is found,
tries to match against any FunctionResources with wildcard patterns.
"""
def get_resource(self, uri: Union[_BaseUrl, str]) -> Optional[Resource]:
"""Get resource by URI."""
uri = str(uri)
logger.debug("Getting resource", extra={"uri": uri})
# First try exact match
if resource := self._resources.get(uri):
return resource
# Then try pattern matching for FunctionResources
for resource in self._resources.values():
if isinstance(resource, FunctionResource) and hasattr(
resource, "uri_regex"
):
if resource.uri_regex.match(uri):
# Create a new instance with the actual URI
return FunctionResource(
uri=uri, # Use actual URI
name=resource.name,
description=resource.description,
mime_type=resource.mime_type,
func=resource.func,
)
raise ValueError(f"Unknown resource: {uri}")
def list_resources(self) -> list[Resource]:
@ -231,10 +210,10 @@ class ResourceManager:
"name": resource.name,
},
)
existing = self._resources.get(resource.uri)
existing = self._resources.get(str(resource.uri))
if existing:
if self.warn_on_duplicate_resources:
logger.warning(f"Resource already exists: {resource.uri}")
return existing
self._resources[resource.uri] = resource
self._resources[str(resource.uri)] = resource
return resource

View file

@ -17,6 +17,7 @@ from .exceptions import ResourceError
from .resources import Resource, FunctionResource, ResourceManager
from .tools import ToolManager
from .utilities.logging import get_logger, configure_logging
from pydantic.networks import _BaseUrl
logger = get_logger(__name__)
@ -56,15 +57,20 @@ class FastMCP:
warn_on_duplicate_resources=self.settings.warn_on_duplicate_resources
)
# Set up MCP protocol handlers
self._setup_handlers()
# Configure logging
configure_logging(self.settings.log_level)
self._setup_handlers()
@property
def name(self) -> str:
return self._mcp_server.name
async def run(self, *args, **kwargs) -> None:
"""Run the FastMCP server."""
await self._mcp_server.run(*args, **kwargs)
def _setup_handlers(self) -> None:
"""Set up core MCP protocol handlers."""
self._mcp_server.list_tools()(self.list_tools)
@ -104,7 +110,7 @@ class FastMCP:
for resource in resources
]
async def read_resource(self, uri: str) -> Union[str, bytes]:
async def read_resource(self, uri: _BaseUrl) -> Union[str, bytes]:
"""Read a resource by URI."""
resource = self._resource_manager.get_resource(uri)
if not resource:
@ -170,42 +176,40 @@ class FastMCP:
def resource(
self,
name: str,
uri: str,
*,
name: Optional[str] = None,
description: Optional[str] = None,
mime_type: Optional[str] = None,
) -> Callable:
"""Decorator to register a function as a dynamic resource.
"""Decorator to register a function as a resource.
The function will be called with kwargs parsed from the URI query string.
For example, a URI of "fn://my_func?x=1&y=2" will call the function with
kwargs {"x": "1", "y": "2"}.
The function will be called when the resource is read to generate its content.
Args:
name: Name for the resource (used in fn:// URI)
uri: URI for the resource (e.g. "resource://my-resource")
description: Optional description of the resource
mime_type: Optional MIME type for the resource
Example:
@server.resource("my_func")
def get_data(x: str, y: str) -> str:
# Called with fn://my_func?x=1&y=2
return f"x={x}, y={y}"
@server.resource("resource://my-resource")
def get_data() -> str:
return "Hello, world!"
"""
# Check if user passed function directly instead of calling decorator
if callable(name):
if callable(uri):
raise TypeError(
"The @resource decorator was used incorrectly. "
"Did you forget to call it? Use @resource('name') instead of @resource"
"Did you forget to call it? Use @resource('uri') instead of @resource"
)
def decorator(func: Callable) -> Callable:
@functools.wraps(func)
def wrapper(**kwargs) -> Any:
return func(**kwargs)
def wrapper() -> Any:
return func()
resource = FunctionResource(
uri=f"fn://{name}", # Base URI, params added when called
uri=uri,
name=name,
description=description,
mime_type=mime_type or "text/plain",
@ -216,10 +220,6 @@ class FastMCP:
return decorator
async def run(self, *args, **kwargs) -> None:
"""Run the FastMCP server."""
await self._mcp_server.run(*args, **kwargs)
@classmethod
async def run_stdio(cls, app: "FastMCP") -> None:
"""Run the server using stdio transport."""

View file

@ -3,7 +3,7 @@
import inspect
from typing import Any, Callable, Dict, Optional
from pydantic import BaseModel, Field, TypeAdapter
from pydantic import BaseModel, Field, TypeAdapter, validate_call
from .exceptions import ToolError
from .utilities.logging import get_logger
@ -37,13 +37,16 @@ class Tool(BaseModel):
is_async = inspect.iscoroutinefunction(func)
# Get schema from TypeAdapter - will fail if function isn't properly typed
schema = TypeAdapter(func).json_schema()
parameters = TypeAdapter(func).json_schema()
# ensure the arguments are properly cast
func = validate_call(func)
return cls(
func=func,
name=func_name,
description=func_doc,
parameters=schema,
parameters=parameters,
is_async=is_async,
)
@ -94,4 +97,5 @@ class ToolManager:
tool = self.get_tool(name)
if not tool:
raise ToolError(f"Unknown tool: {name}")
return await tool.run(arguments)

0
tests/__init__.py Normal file
View file

View file

View file

@ -0,0 +1,47 @@
import json
from fastmcp import FastMCP
import pytest
from pathlib import Path
@pytest.fixture(scope="session")
def test_dir(tmp_path_factory) -> Path:
"""Create a temporary directory with test files."""
tmp = tmp_path_factory.mktemp("test_files")
# Create test files
(tmp / "example.py").write_text("print('hello world')")
(tmp / "readme.md").write_text("# Test Directory\nThis is a test.")
(tmp / "config.json").write_text('{"test": true}')
return tmp
@pytest.fixture
def mcp(test_dir: Path) -> FastMCP:
mcp = FastMCP()
@mcp.resource("fs://test_dir")
def list_files() -> list[str]:
"""List the files in the test directory"""
return [str(f) for f in test_dir.iterdir()]
return mcp
async def test_list_resources(mcp: FastMCP):
resources = await mcp.list_resources()
assert len(resources) == 1
assert str(resources[0].uri) == "fs://test_dir"
assert resources[0].name == "test_dir"
async def test_read_resource(mcp: FastMCP):
files = await mcp.read_resource("fs://test_dir")
files = json.loads(files)
assert isinstance(files, list)
assert len(files) == 3
assert any("example.py" in f for f in files)
assert any("readme.md" in f for f in files)
assert any("config.json" in f for f in files)

View file

@ -54,10 +54,10 @@ class TestResourceValidation:
name="test",
func=dummy_func,
)
assert resource.uri == "http://example.com/data"
assert str(resource.uri) == "http://example.com/data"
# Missing protocol
with pytest.raises(ValueError, match="URI must have a protocol"):
with pytest.raises(ValueError, match="Input should be a valid URL"):
FunctionResource(
uri="invalid",
name="test",
@ -65,7 +65,7 @@ class TestResourceValidation:
)
# Missing host
with pytest.raises(ValueError, match="URI must have a host"):
with pytest.raises(ValueError, match="Input should be a valid URL"):
FunctionResource(
uri="http://",
name="test",
@ -85,7 +85,7 @@ class TestFileResource:
mime_type="text/plain",
path=temp_file,
)
assert resource.uri == f"file://{temp_file}"
assert str(resource.uri) == f"file://{temp_file}"
assert resource.name == "test"
assert resource.description == "test file"
assert resource.mime_type == "text/plain"
@ -162,13 +162,13 @@ class TestFunctionResource:
mime_type="text/plain",
func=my_func,
)
assert resource.uri == "fn://test"
assert str(resource.uri) == "fn://test"
assert resource.name == "test"
assert resource.description == "test function"
assert resource.mime_type == "text/plain"
assert resource.func == my_func
async def test_function_resource_read_no_params(self):
async def test_function_resource_read(self):
"""Test reading a FunctionResource with no parameters."""
def my_func() -> str:
@ -182,86 +182,6 @@ class TestFunctionResource:
content = await resource.read()
assert content == "test content"
async def test_function_resource_read_with_params(self):
"""Test reading a FunctionResource with query parameters."""
def my_func(x: str, y: str) -> str:
return f"x={x}, y={y}"
resource = FunctionResource(
uri="fn://test?x=1&y=2",
name="test",
func=my_func,
)
content = await resource.read()
assert content == "x=1, y=2"
async def test_function_resource_read_returns_resource(self, temp_file: Path):
"""Test reading a FunctionResource that returns another Resource."""
def my_func(name: str = "test") -> FileResource:
return FileResource(
uri=f"file://{temp_file}",
name=name,
path=temp_file,
)
resource = FunctionResource(
uri="fn://test?name=example",
name="test",
func=my_func,
)
content = await resource.read()
assert content == "test content"
async def test_function_resource_read_error(self):
"""Test error handling when reading a FunctionResource."""
def my_func(x: str) -> str:
raise ValueError(f"test error: {x}")
resource = FunctionResource(
uri="fn://test?x=bad",
name="test",
func=my_func,
)
with pytest.raises(
ValueError, match="Error calling function my_func: test error: bad"
):
await resource.read()
def test_parse_uri_params(self):
"""Test parsing URI parameters."""
def my_func() -> str:
return "test"
resource = FunctionResource(
uri="fn://test?x=1&y=hello&z=true",
name="test",
func=my_func,
)
params = resource._parse_uri_params()
assert params == {
"x": "1",
"y": "hello",
"z": "true",
}
def test_parse_uri_no_params(self):
"""Test parsing URI with no parameters."""
def my_func() -> str:
return "test"
resource = FunctionResource(
uri="fn://test",
name="test",
func=my_func,
)
params = resource._parse_uri_params()
assert params == {}
class TestResourceManagerAdd:
"""Test ResourceManager add functionality."""
@ -278,7 +198,7 @@ class TestResourceManagerAdd:
)
added = manager.add_resource(resource)
assert isinstance(added, FileResource)
assert added.uri == f"file://{temp_file}"
assert str(added.uri) == f"file://{temp_file}"
assert added.name == "test"
assert added.description == "test file"
assert added.mime_type == "text/plain"

View file

@ -31,7 +31,7 @@ class TestServer:
async def test_add_resource_decorator(self):
mcp = FastMCP()
@mcp.resource("data")
@mcp.resource("r://data")
def get_data(x: str) -> str:
return f"Data: {x}"