Consolidate resource methods

This commit is contained in:
Jeremiah Lowin 2024-11-29 17:13:26 -05:00
commit 6ede7a1870
5 changed files with 196 additions and 263 deletions

View file

@ -25,7 +25,7 @@ class Resource(BaseModel):
@abc.abstractmethod
async def read(self) -> str:
"""Read the resource content."""
...
return ""
class FileResource(Resource):
@ -139,111 +139,26 @@ class ResourceManager:
logger.debug("Listing resources", extra={"count": len(self._resources)})
return list(self._resources.values())
def add_file_resource(
self,
path: str,
*,
name: Optional[str] = None,
description: Optional[str] = None,
mime_type: Optional[str] = None,
) -> FileResource:
"""Add a file as a resource.
def add_resource(self, resource: Resource) -> Resource:
"""Add a resource to the manager.
Args:
path: Absolute path to the file
name: Optional name for the resource
description: Optional description of the resource
mime_type: Optional MIME type for the resource
resource: A Resource instance to add
Returns:
The created resource
Raises:
ValueError: If the path is not absolute or the file does not exist
The added resource. If a resource with the same URI already exists,
returns the existing resource.
"""
logger.debug(
"Adding file resource",
"Adding resource",
extra={
"path": path,
"name": name,
"mime_type": mime_type,
"uri": resource.uri,
"type": type(resource).__name__,
"name": resource.name,
},
)
file = Path(path)
if not file.is_absolute():
raise ValueError(f"Path must be absolute: {path}")
if not file.is_file():
raise FileNotFoundError(f"File does not exist: {path}")
resource = FileResource(
uri=f"file://{str(file)}",
name=name or file.name,
description=description,
mime_type=mime_type or "text/plain",
path=file,
)
self._resources[resource.uri] = resource
return resource
def add_http_resource(
self,
url: str,
*,
name: Optional[str] = None,
description: Optional[str] = None,
mime_type: Optional[str] = None,
headers: Optional[Dict[str, str]] = None,
) -> HttpResource:
"""Add an HTTP endpoint as a resource."""
logger.debug(
"Adding HTTP resource",
extra={
"url": url,
"name": name,
"mime_type": mime_type,
},
)
resource = HttpResource(
uri=f"http://{url}",
name=name or url.split("/")[-1],
description=description,
mime_type=mime_type or "text/plain",
url=url,
headers=headers,
)
self._resources[resource.uri] = resource
return resource
def add_dir_resource(
self,
path: str,
*,
recursive: bool = False,
pattern: Optional[str] = None,
name: Optional[str] = None,
description: Optional[str] = None,
) -> DirectoryResource:
"""Add a directory as a resource."""
logger.debug(
"Adding directory resource",
extra={
"path": path,
"recursive": recursive,
"pattern": pattern,
"name": name,
},
)
dir_path = Path(path).expanduser().resolve()
if not dir_path.is_dir():
raise ValueError(f"Directory does not exist: {path}")
resource = DirectoryResource(
uri=f"dir://{str(dir_path)}",
name=name or dir_path.name,
description=description,
path=dir_path,
recursive=recursive,
pattern=pattern,
)
existing = self._resources.get(resource.uri)
if existing:
return existing
self._resources[resource.uri] = resource
return resource

View file

@ -12,7 +12,7 @@ from mcp.types import Tool, TextContent, ImageContent, EmbeddedResource
from pydantic import BaseModel
from .exceptions import ResourceError
from .resources import ResourceManager
from .resources import Resource, ResourceManager
from .tools import ToolManager
logger = logging.getLogger("mcp")
@ -110,6 +110,14 @@ class FastMCPServer:
return decorator
def add_resource(self, resource: Resource) -> None:
"""Add a resource to the server.
Args:
resource: A Resource instance to add
"""
self._resource_manager.add_resource(resource)
def add_file_resource(
self,
path: str,
@ -118,13 +126,28 @@ class FastMCPServer:
description: Optional[str] = None,
mime_type: Optional[str] = None,
) -> None:
"""Add a file as a resource."""
self._resource_manager.add_file_resource(
path,
name=name,
"""Add a file as a resource.
This is a convenience method that constructs and adds a FileResource.
For more control, use add_resource() directly.
"""
from pathlib import Path
from .resources import FileResource
file = Path(path)
if not file.is_absolute():
raise ValueError(f"Path must be absolute: {path}")
if not file.is_file():
raise FileNotFoundError(f"File does not exist: {path}")
resource = FileResource(
uri=f"file://{str(file)}",
name=name or file.name,
description=description,
mime_type=mime_type,
mime_type=mime_type or "text/plain",
path=file,
)
self.add_resource(resource)
def add_http_resource(
self,
@ -135,14 +158,22 @@ class FastMCPServer:
mime_type: Optional[str] = None,
headers: Optional[Dict[str, str]] = None,
) -> None:
"""Add an HTTP endpoint as a resource."""
self._resource_manager.add_http_resource(
url,
name=name,
"""Add an HTTP endpoint as a resource.
This is a convenience method that constructs and adds an HttpResource.
For more control, use add_resource() directly.
"""
from .resources import HttpResource
resource = HttpResource(
uri=f"http://{url}",
name=name or url.split("/")[-1],
description=description,
mime_type=mime_type,
mime_type=mime_type or "text/plain",
url=url,
headers=headers,
)
self.add_resource(resource)
def add_dir_resource(
self,
@ -153,14 +184,27 @@ class FastMCPServer:
name: Optional[str] = None,
description: Optional[str] = None,
) -> None:
"""Add a directory as a resource."""
self._resource_manager.add_dir_resource(
path,
"""Add a directory as a resource.
This is a convenience method that constructs and adds a DirectoryResource.
For more control, use add_resource() directly.
"""
from pathlib import Path
from .resources import DirectoryResource
dir_path = Path(path).expanduser().resolve()
if not dir_path.is_dir():
raise ValueError(f"Directory does not exist: {path}")
resource = DirectoryResource(
uri=f"dir://{str(dir_path)}",
name=name or dir_path.name,
description=description,
path=dir_path,
recursive=recursive,
pattern=pattern,
name=name,
description=description,
)
self.add_resource(resource)
async def run(self, *args, **kwargs) -> None:
"""Run the FastMCP server."""

View file

@ -123,40 +123,43 @@ class TestFileResource:
temp_file.chmod(0o644) # Restore permissions
class TestResourceManager:
"""Test ResourceManager functionality."""
class TestResourceManagerAdd:
"""Test ResourceManager add functionality."""
def test_add_file_resource(
self, resource_manager: ResourceManager, temp_file: Path
):
"""Test adding a file resource."""
resource = resource_manager.add_file_resource(
str(temp_file),
resource = FileResource(
uri=f"file://{temp_file}",
name="test",
description="test file",
mime_type="text/plain",
path=temp_file,
)
assert isinstance(resource, FileResource)
assert resource.uri == f"file://{temp_file}"
assert resource.name == "test"
assert resource.description == "test file"
assert resource.mime_type == "text/plain"
assert resource.path == temp_file
added = resource_manager.add_resource(resource)
assert isinstance(added, FileResource)
assert added.uri == f"file://{temp_file}"
assert added.name == "test"
assert added.description == "test file"
assert added.mime_type == "text/plain"
assert added.path == temp_file
def test_add_file_resource_relative_path_error(
self, resource_manager: ResourceManager
):
"""Test ResourceManager rejects relative paths."""
with pytest.raises(ValueError, match="Path must be absolute"):
resource_manager.add_file_resource("test.txt")
resource = FileResource(
uri="file://test.txt",
name="test",
path=Path("test.txt"),
)
resource_manager.add_resource(resource)
def test_add_file_resource_missing_file_error(
self, resource_manager: ResourceManager, temp_dir: Path
):
"""Test ResourceManager rejects non-existent files."""
missing_file = temp_dir / "missing.txt"
with pytest.raises(FileNotFoundError):
resource_manager.add_file_resource(str(missing_file))
class TestResourceManagerRead:
"""Test ResourceManager read functionality."""
def test_get_resource_unknown_uri(self, resource_manager: ResourceManager):
"""Test getting a non-existent resource."""
@ -165,23 +168,26 @@ class TestResourceManager:
def test_get_resource(self, resource_manager: ResourceManager, temp_file: Path):
"""Test getting a resource by URI."""
added = resource_manager.add_file_resource(str(temp_file))
resource = FileResource(
uri=f"file://{temp_file}",
name="test",
path=temp_file,
)
added = resource_manager.add_resource(resource)
retrieved = resource_manager.get_resource(added.uri)
assert retrieved == added
def test_list_resources(self, resource_manager: ResourceManager, temp_file: Path):
"""Test listing all resources."""
resource = resource_manager.add_file_resource(str(temp_file))
resources = resource_manager.list_resources()
assert len(resources) == 1
assert resources[0] == resource
async def test_resource_read_through_manager(
self, resource_manager: ResourceManager, temp_file: Path
):
"""Test reading a resource through the manager."""
resource = resource_manager.add_file_resource(str(temp_file))
retrieved = resource_manager.get_resource(resource.uri)
resource = FileResource(
uri=f"file://{temp_file}",
name="test",
path=temp_file,
)
added = resource_manager.add_resource(resource)
retrieved = resource_manager.get_resource(added.uri)
assert retrieved is not None
content = await retrieved.read()
assert content == "test content"
@ -191,11 +197,75 @@ class TestResourceManager:
):
"""Test error handling when reading through manager."""
# Create resource while file exists
resource = resource_manager.add_file_resource(str(temp_file_no_cleanup))
retrieved = resource_manager.get_resource(resource.uri)
resource = FileResource(
uri=f"file://{temp_file_no_cleanup}",
name="test",
path=temp_file_no_cleanup,
)
added = resource_manager.add_resource(resource)
retrieved = resource_manager.get_resource(added.uri)
assert retrieved is not None
# Delete file and verify read fails
temp_file_no_cleanup.unlink()
with pytest.raises(FileNotFoundError):
await retrieved.read()
class TestResourceManagerList:
"""Test ResourceManager list functionality."""
def test_list_resources(self, resource_manager: ResourceManager, temp_file: Path):
"""Test listing all resources."""
resource = FileResource(
uri=f"file://{temp_file}",
name="test",
path=temp_file,
)
added = resource_manager.add_resource(resource)
resources = resource_manager.list_resources()
assert len(resources) == 1
assert resources[0] == added
def test_list_resources_duplicate(
self, resource_manager: ResourceManager, temp_file: Path
):
"""Test that adding the same resource twice only stores it once."""
resource = FileResource(
uri=f"file://{temp_file}",
name="test",
path=temp_file,
)
resource1 = resource_manager.add_resource(resource)
resource2 = resource_manager.add_resource(resource)
resources = resource_manager.list_resources()
assert len(resources) == 1
assert resources[0] == resource1
assert resource1 == resource2
def test_list_multiple_resources(
self,
resource_manager: ResourceManager,
temp_file: Path,
temp_file_no_cleanup: Path,
):
"""Test listing multiple different resources."""
resource1 = FileResource(
uri=f"file://{temp_file}",
name="test1",
path=temp_file,
)
resource2 = FileResource(
uri=f"file://{temp_file_no_cleanup}",
name="test2",
path=temp_file_no_cleanup,
)
added1 = resource_manager.add_resource(resource1)
added2 = resource_manager.add_resource(resource2)
resources = resource_manager.list_resources()
assert len(resources) == 2
assert resources[0] == added1
assert resources[1] == added2
assert added1 != added2

23
tests/test_server.py Normal file
View file

@ -0,0 +1,23 @@
from mcp.shared.memory import create_connected_server_and_client_session
from fastmcp.server import FastMCPServer
async def test_list_tools():
server = FastMCPServer("test_server")
server.add_tool(lambda x: x)
async with create_connected_server_and_client_session(
server._mcp_server
) as client_session:
tools = await client_session.list_tools()
assert len(tools.tools) == 1
async def test_call_tool():
server = FastMCPServer("test_server")
server.add_tool(lambda x: x)
async with create_connected_server_and_client_session(
server._mcp_server
) as client_session:
result = await client_session.call_tool("my_tool", {"arg1": "value"})
assert "error" not in result
assert len(result.content) > 0

View file

@ -1,119 +0,0 @@
"""Test tool registration and execution."""
import pytest
from pydantic import BaseModel
from fastmcp.exceptions import ToolError
from fastmcp.tools import ToolManager
class TestAddTools:
def test_basic_function(self):
"""Test registering and running a basic function."""
def add(a: int, b: int) -> int:
"""Add two numbers."""
return a + b
manager = ToolManager()
manager.add_tool(add)
tool = manager.get_tool("add")
assert tool is not None
assert tool.name == "add"
assert tool.description == "Add two numbers."
assert tool.is_async is False
assert tool.parameters["properties"]["a"]["type"] == "integer"
assert tool.parameters["properties"]["b"]["type"] == "integer"
async def test_async_function(self):
"""Test registering and running an async function."""
async def fetch_data(url: str) -> str:
"""Fetch data from URL."""
return f"Data from {url}"
manager = ToolManager()
manager.add_tool(fetch_data)
tool = manager.get_tool("fetch_data")
assert tool is not None
assert tool.name == "fetch_data"
assert tool.description == "Fetch data from URL."
assert tool.is_async is True
assert tool.parameters["properties"]["url"]["type"] == "string"
def test_pydantic_model_function(self):
"""Test registering a function that takes a Pydantic model."""
class UserInput(BaseModel):
name: str
age: int
def create_user(user: UserInput, flag: bool) -> dict:
"""Create a new user."""
return {"id": 1, **user.model_dump()}
manager = ToolManager()
manager.add_tool(create_user)
tool = manager.get_tool("create_user")
assert tool is not None
assert tool.name == "create_user"
assert tool.description == "Create a new user."
assert tool.is_async is False
assert "name" in tool.parameters["$defs"]["UserInput"]["properties"]
assert "age" in tool.parameters["$defs"]["UserInput"]["properties"]
assert "flag" in tool.parameters["properties"]
def test_add_invalid_tool(self):
manager = ToolManager()
with pytest.raises(AttributeError):
manager.add_tool(1)
class TestCallTools:
async def test_call_tool(self):
def add(a: int, b: int) -> int:
"""Add two numbers."""
return a + b
manager = ToolManager()
manager.add_tool(add)
result = await manager.call_tool("add", {"a": 1, "b": 2})
assert result == 3
async def test_call_async_tool(self):
async def double(n: int) -> int:
"""Double a number."""
return n * 2
manager = ToolManager()
manager.add_tool(double)
result = await manager.call_tool("double", {"n": 5})
assert result == 10
async def test_call_tool_with_default_args(self):
def add(a: int, b: int = 1) -> int:
"""Add two numbers."""
return a + b
manager = ToolManager()
manager.add_tool(add)
result = await manager.call_tool("add", {"a": 1})
assert result == 2
async def test_call_tool_with_missing_args(self):
def add(a: int, b: int) -> int:
"""Add two numbers."""
return a + b
manager = ToolManager()
manager.add_tool(add)
with pytest.raises(ToolError):
await manager.call_tool("add", {"a": 1})
async def test_call_unknown_tool(self):
manager = ToolManager()
with pytest.raises(ToolError):
await manager.call_tool("unknown", {"a": 1})