mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-22 05:24:18 +02:00
Merge pull request #170 from jlowin/resources
Add custom key for storing resources
This commit is contained in:
Jeremiah Lowin
GitHub
commit
dced861899
8 changed files with 342 additions and 85 deletions
|
|
@ -149,7 +149,15 @@ notice_resource = TextResource(
|
|||
)
|
||||
mcp.add_resource(notice_resource)
|
||||
|
||||
# 3. Exposing a directory listing
|
||||
# 3. Using a custom key different from the URI
|
||||
special_resource = TextResource(
|
||||
uri="resource://common-notice",
|
||||
name="Special Notice",
|
||||
text="This is a special notice with a custom storage key.",
|
||||
)
|
||||
mcp.add_resource(special_resource, key="resource://custom-key")
|
||||
|
||||
# 4. Exposing a directory listing
|
||||
data_dir_path = Path("./app_data").resolve()
|
||||
if data_dir_path.is_dir():
|
||||
data_listing_resource = DirectoryResource(
|
||||
|
|
@ -173,6 +181,22 @@ if data_dir_path.is_dir():
|
|||
|
||||
Use these when the content is static or sourced directly from a file/URL, bypassing the need for a dedicated Python function.
|
||||
|
||||
#### Custom Resource Keys
|
||||
|
||||
When adding resources directly with `mcp.add_resource()`, you can optionally provide a custom storage key:
|
||||
|
||||
```python
|
||||
# Creating a resource with standard URI as the key
|
||||
resource = TextResource(uri="resource://data")
|
||||
mcp.add_resource(resource) # Will be stored and accessed using "resource://data"
|
||||
|
||||
# Creating a resource with a custom key
|
||||
special_resource = TextResource(uri="resource://special-data")
|
||||
mcp.add_resource(special_resource, key="internal://data-v2") # Will be stored and accessed using "internal://data-v2"
|
||||
```
|
||||
|
||||
Note that this parameter is only available when using `add_resource()` directly and not through the `@resource` decorator, as URIs are provided explicitly when using the decorator.
|
||||
|
||||
## Defining Resource Templates
|
||||
|
||||
Resource Templates allow clients to request resources whose content depends on parameters embedded in the URI. Define a template using the **same `@mcp.resource` decorator**, but include `{parameter_name}` placeholders in the URI string and add corresponding arguments to your function signature.
|
||||
|
|
@ -289,6 +313,26 @@ In this stacked decorator pattern:
|
|||
|
||||
Templates provide a powerful way to expose parameterized data access points following REST-like principles.
|
||||
|
||||
### Custom Template Keys
|
||||
|
||||
Similar to resources, you can provide custom keys when directly adding templates:
|
||||
|
||||
```python
|
||||
from fastmcp.resources import ResourceTemplate
|
||||
|
||||
# Create a template with a function
|
||||
template = ResourceTemplate.from_function(
|
||||
my_function,
|
||||
uri_template="data://{id}/details",
|
||||
name="Data Details"
|
||||
)
|
||||
|
||||
# Register with a custom key
|
||||
mcp._resource_manager.add_template(template, key="custom://{id}/view")
|
||||
```
|
||||
|
||||
This allows accessing the same template implementation through different URI patterns.
|
||||
|
||||
## Server Behavior
|
||||
|
||||
### Duplicate Resources
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
"""Resource manager functionality."""
|
||||
|
||||
import copy
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
|
@ -9,7 +8,7 @@ from pydantic import AnyUrl
|
|||
|
||||
from fastmcp.exceptions import ResourceError
|
||||
from fastmcp.resources import FunctionResource, Resource
|
||||
from fastmcp.resources.template import ResourceTemplate
|
||||
from fastmcp.resources.template import ResourceTemplate, match_uri_template
|
||||
from fastmcp.settings import DuplicateBehavior
|
||||
from fastmcp.utilities.logging import get_logger
|
||||
|
||||
|
|
@ -109,34 +108,35 @@ class ResourceManager:
|
|||
)
|
||||
return self.add_resource(resource)
|
||||
|
||||
def add_resource(self, resource: Resource) -> Resource:
|
||||
def add_resource(self, resource: Resource, key: str | None = None) -> Resource:
|
||||
"""Add a resource to the manager.
|
||||
|
||||
Args:
|
||||
resource: A Resource instance to add
|
||||
key: Optional URI to use as the storage key (if different from resource.uri)
|
||||
"""
|
||||
uri_str = str(resource.uri)
|
||||
storage_key = key or str(resource.uri)
|
||||
logger.debug(
|
||||
"Adding resource",
|
||||
extra={
|
||||
"uri": uri_str,
|
||||
"uri": resource.uri,
|
||||
"storage_key": storage_key,
|
||||
"type": type(resource).__name__,
|
||||
"resource_name": resource.name,
|
||||
},
|
||||
)
|
||||
existing = self._resources.get(uri_str)
|
||||
existing = self._resources.get(storage_key)
|
||||
if existing:
|
||||
if self.duplicate_behavior == "warn":
|
||||
logger.warning(f"Resource already exists: {uri_str}")
|
||||
self._resources[uri_str] = resource
|
||||
logger.warning(f"Resource already exists: {storage_key}")
|
||||
self._resources[storage_key] = resource
|
||||
elif self.duplicate_behavior == "replace":
|
||||
self._resources[uri_str] = resource
|
||||
self._resources[storage_key] = resource
|
||||
elif self.duplicate_behavior == "error":
|
||||
raise ValueError(f"Resource already exists: {uri_str}")
|
||||
raise ValueError(f"Resource already exists: {storage_key}")
|
||||
elif self.duplicate_behavior == "ignore":
|
||||
return existing
|
||||
else:
|
||||
self._resources[uri_str] = resource
|
||||
self._resources[storage_key] = resource
|
||||
return resource
|
||||
|
||||
def add_template_from_fn(
|
||||
|
|
@ -160,38 +160,42 @@ class ResourceManager:
|
|||
)
|
||||
return self.add_template(template)
|
||||
|
||||
def add_template(self, template: ResourceTemplate) -> ResourceTemplate:
|
||||
def add_template(
|
||||
self, template: ResourceTemplate, key: str | None = None
|
||||
) -> ResourceTemplate:
|
||||
"""Add a template to the manager.
|
||||
|
||||
Args:
|
||||
template: A ResourceTemplate instance to add
|
||||
key: Optional URI template to use as the storage key (if different from template.uri_template)
|
||||
|
||||
Returns:
|
||||
The added template. If a template with the same URI already exists,
|
||||
returns the existing template.
|
||||
"""
|
||||
uri_template_str = str(template.uri_template)
|
||||
storage_key = key or uri_template_str
|
||||
logger.debug(
|
||||
"Adding resource",
|
||||
"Adding template",
|
||||
extra={
|
||||
"uri": uri_template_str,
|
||||
"uri_template": uri_template_str,
|
||||
"storage_key": storage_key,
|
||||
"type": type(template).__name__,
|
||||
"resource_name": template.name,
|
||||
"template_name": template.name,
|
||||
},
|
||||
)
|
||||
existing = self._templates.get(uri_template_str)
|
||||
existing = self._templates.get(storage_key)
|
||||
if existing:
|
||||
if self.duplicate_behavior == "warn":
|
||||
logger.warning(f"Resource already exists: {uri_template_str}")
|
||||
self._templates[uri_template_str] = template
|
||||
logger.warning(f"Template already exists: {storage_key}")
|
||||
self._templates[storage_key] = template
|
||||
elif self.duplicate_behavior == "replace":
|
||||
self._templates[uri_template_str] = template
|
||||
self._templates[storage_key] = template
|
||||
elif self.duplicate_behavior == "error":
|
||||
raise ValueError(f"Resource already exists: {uri_template_str}")
|
||||
raise ValueError(f"Template already exists: {storage_key}")
|
||||
elif self.duplicate_behavior == "ignore":
|
||||
return existing
|
||||
else:
|
||||
self._templates[uri_template_str] = template
|
||||
self._templates[storage_key] = template
|
||||
return template
|
||||
|
||||
async def get_resource(self, uri: AnyUrl | str) -> Resource | None:
|
||||
|
|
@ -203,9 +207,10 @@ class ResourceManager:
|
|||
if resource := self._resources.get(uri_str):
|
||||
return resource
|
||||
|
||||
# Then check templates
|
||||
for template in self._templates.values():
|
||||
if params := template.matches(uri_str):
|
||||
# Then check templates - use the utility function to match against storage keys
|
||||
for storage_key, template in self._templates.items():
|
||||
# Try to match against the storage key (which might be a custom key)
|
||||
if params := match_uri_template(uri_str, storage_key):
|
||||
try:
|
||||
return await template.create_resource(uri_str, params)
|
||||
except Exception as e:
|
||||
|
|
@ -213,11 +218,19 @@ class ResourceManager:
|
|||
|
||||
raise ResourceError(f"Unknown resource: {uri}")
|
||||
|
||||
def get_resources(self) -> dict[str, Resource]:
|
||||
"""Get all registered resources, keyed by URI."""
|
||||
return self._resources
|
||||
|
||||
def list_resources(self) -> list[Resource]:
|
||||
"""List all registered resources."""
|
||||
logger.debug("Listing resources", extra={"count": len(self._resources)})
|
||||
return list(self._resources.values())
|
||||
|
||||
def get_templates(self) -> dict[str, ResourceTemplate]:
|
||||
"""Get all registered templates, keyed by URI template."""
|
||||
return self._templates
|
||||
|
||||
def list_templates(self) -> list[ResourceTemplate]:
|
||||
"""List all registered templates."""
|
||||
logger.debug("Listing templates", extra={"count": len(self._templates)})
|
||||
|
|
@ -240,14 +253,9 @@ class ResourceManager:
|
|||
If None, the original URI is used.
|
||||
"""
|
||||
for uri, resource in manager._resources.items():
|
||||
# Create prefixed URI and copy the resource with the new URI
|
||||
# Create prefixed URI and import the resource with the new URI as the storage key
|
||||
prefixed_uri = f"{prefix}{uri}" if prefix else uri
|
||||
|
||||
new_resource = copy.copy(resource)
|
||||
new_resource.uri = AnyUrl(prefixed_uri)
|
||||
|
||||
# Store directly in resources dictionary
|
||||
self.add_resource(new_resource)
|
||||
self.add_resource(resource, key=prefixed_uri)
|
||||
logger.debug(f'Imported resource "{uri}" as "{prefixed_uri}"')
|
||||
|
||||
def import_templates(
|
||||
|
|
@ -267,16 +275,11 @@ class ResourceManager:
|
|||
If None, the original URI template is used.
|
||||
"""
|
||||
for uri_template, template in manager._templates.items():
|
||||
# Create prefixed URI template and copy the template with the new URI template
|
||||
# Create prefixed URI template and import the template with the new URI as the storage key
|
||||
prefixed_uri_template = (
|
||||
f"{prefix}{uri_template}" if prefix else uri_template
|
||||
)
|
||||
|
||||
new_template = copy.copy(template)
|
||||
new_template.uri_template = prefixed_uri_template
|
||||
|
||||
# Store directly in templates dictionary
|
||||
self.add_template(new_template)
|
||||
self.add_template(template, key=prefixed_uri_template)
|
||||
logger.debug(
|
||||
f'Imported template "{uri_template}" as "{prefixed_uri_template}"'
|
||||
)
|
||||
|
|
|
|||
|
|
@ -7,12 +7,37 @@ import re
|
|||
from collections.abc import Callable
|
||||
from typing import Annotated, Any
|
||||
|
||||
from pydantic import BaseModel, BeforeValidator, Field, TypeAdapter, validate_call
|
||||
from pydantic import (
|
||||
AnyUrl,
|
||||
BaseModel,
|
||||
BeforeValidator,
|
||||
Field,
|
||||
TypeAdapter,
|
||||
validate_call,
|
||||
)
|
||||
|
||||
from fastmcp.resources.types import FunctionResource, Resource
|
||||
from fastmcp.utilities.types import _convert_set_defaults
|
||||
|
||||
|
||||
def match_uri_template(uri: str, uri_template: str) -> dict[str, Any] | None:
|
||||
"""Match a URI against a template and extract parameters.
|
||||
|
||||
Args:
|
||||
uri: The URI to match against the template
|
||||
uri_template: The URI template to match against
|
||||
|
||||
Returns:
|
||||
A dictionary of extracted parameters if there's a match, or None if no match
|
||||
"""
|
||||
# Convert template to regex pattern
|
||||
pattern = uri_template.replace("{", "(?P<").replace("}", ">[^/]+)")
|
||||
match = re.match(f"^{pattern}$", uri)
|
||||
if match:
|
||||
return match.groupdict()
|
||||
return None
|
||||
|
||||
|
||||
class MyModel(BaseModel):
|
||||
key: str
|
||||
value: int
|
||||
|
|
@ -94,12 +119,7 @@ class ResourceTemplate(BaseModel):
|
|||
|
||||
def matches(self, uri: str) -> dict[str, Any] | None:
|
||||
"""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
|
||||
return match_uri_template(uri, self.uri_template)
|
||||
|
||||
async def create_resource(self, uri: str, params: dict[str, Any]) -> Resource:
|
||||
"""Create a resource from the template with the given parameters."""
|
||||
|
|
@ -110,7 +130,7 @@ class ResourceTemplate(BaseModel):
|
|||
result = await result
|
||||
|
||||
return FunctionResource(
|
||||
uri=uri, # type: ignore
|
||||
uri=AnyUrl(uri), # Explicitly convert to AnyUrl
|
||||
name=self.name,
|
||||
description=self.description,
|
||||
mime_type=self.mime_type,
|
||||
|
|
|
|||
|
|
@ -186,7 +186,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
self._mcp_server.list_resource_templates()(self._mcp_list_resource_templates)
|
||||
|
||||
def get_tools(self) -> dict[str, Tool]:
|
||||
"""Get all registered tools, keyed by registered name."""
|
||||
"""Get all registered tools, indexed by registered key."""
|
||||
return self._tool_manager.get_tools()
|
||||
|
||||
def list_tools(self) -> list[Tool]:
|
||||
|
|
@ -362,14 +362,14 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
|
||||
return decorator
|
||||
|
||||
def add_resource(self, resource: Resource) -> None:
|
||||
def add_resource(self, resource: Resource, key: str | None = None) -> None:
|
||||
"""Add a resource to the server.
|
||||
|
||||
Args:
|
||||
resource: A Resource instance to add
|
||||
"""
|
||||
|
||||
self._resource_manager.add_resource(resource)
|
||||
self._resource_manager.add_resource(resource, key=key)
|
||||
|
||||
def add_resource_fn(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -41,7 +41,7 @@ class ToolManager:
|
|||
return self._tools.get(name)
|
||||
|
||||
def get_tools(self) -> dict[str, Tool]:
|
||||
"""Get all registered tools, keyed by registered name."""
|
||||
"""Get all registered tools, indexed by registered key."""
|
||||
return self._tools
|
||||
|
||||
def list_tools(self) -> list[Tool]:
|
||||
|
|
@ -50,7 +50,7 @@ class ToolManager:
|
|||
|
||||
def list_mcp_tools(self) -> list[MCPTool]:
|
||||
"""List all registered tools in the format expected by the low-level MCP server."""
|
||||
return [tool.to_mcp_tool(name=name) for name, tool in self._tools.items()]
|
||||
return [tool.to_mcp_tool(name=key) for key, tool in self._tools.items()]
|
||||
|
||||
def add_tool_from_fn(
|
||||
self,
|
||||
|
|
@ -63,34 +63,34 @@ class ToolManager:
|
|||
tool = Tool.from_function(fn, name=name, description=description, tags=tags)
|
||||
return self.add_tool(tool)
|
||||
|
||||
def add_tool(self, tool: Tool, name: str | None = None) -> Tool:
|
||||
def add_tool(self, tool: Tool, key: str | None = None) -> Tool:
|
||||
"""Register a tool with the server."""
|
||||
name = name or tool.name
|
||||
existing = self._tools.get(name)
|
||||
key = key or tool.name
|
||||
existing = self._tools.get(key)
|
||||
if existing:
|
||||
if self.duplicate_behavior == "warn":
|
||||
logger.warning(f"Tool already exists: {name}")
|
||||
self._tools[name] = tool
|
||||
logger.warning(f"Tool already exists: {key}")
|
||||
self._tools[key] = tool
|
||||
elif self.duplicate_behavior == "replace":
|
||||
self._tools[name] = tool
|
||||
self._tools[key] = tool
|
||||
elif self.duplicate_behavior == "error":
|
||||
raise ValueError(f"Tool already exists: {name}")
|
||||
raise ValueError(f"Tool already exists: {key}")
|
||||
elif self.duplicate_behavior == "ignore":
|
||||
return existing
|
||||
else:
|
||||
self._tools[name] = tool
|
||||
self._tools[key] = tool
|
||||
return tool
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
name: str,
|
||||
key: str,
|
||||
arguments: dict[str, Any],
|
||||
context: Context[ServerSessionT, LifespanContextT] | None = None,
|
||||
) -> Any:
|
||||
"""Call a tool by name with arguments."""
|
||||
tool = self.get_tool(name)
|
||||
tool = self.get_tool(key)
|
||||
if not tool:
|
||||
raise ToolError(f"Unknown tool: {name}")
|
||||
raise ToolError(f"Unknown tool: {key}")
|
||||
|
||||
return await tool.run(arguments, context=context)
|
||||
|
||||
|
|
@ -110,5 +110,5 @@ class ToolManager:
|
|||
"""
|
||||
for name, tool in tool_manager._tools.items():
|
||||
prefixed_name = f"{prefix}{name}" if prefix else name
|
||||
self.add_tool(tool, name=prefixed_name)
|
||||
self.add_tool(tool, key=prefixed_name)
|
||||
logger.debug(f'Imported tool "{tool.name}" as "{prefixed_name}"')
|
||||
|
|
|
|||
|
|
@ -168,7 +168,7 @@ class TestResourceManager:
|
|||
manager.add_template(template)
|
||||
manager.add_template(template)
|
||||
|
||||
assert "Resource already exists" in caplog.text
|
||||
assert "Template already exists" in caplog.text
|
||||
# Should have the template
|
||||
assert len(manager.list_templates()) == 1
|
||||
|
||||
|
|
@ -187,7 +187,7 @@ class TestResourceManager:
|
|||
|
||||
manager.add_template(template)
|
||||
|
||||
with pytest.raises(ValueError, match="Resource already exists"):
|
||||
with pytest.raises(ValueError, match="Template already exists"):
|
||||
manager.add_template(template)
|
||||
|
||||
def test_replace_duplicate_templates(self):
|
||||
|
|
@ -688,3 +688,172 @@ class TestImports:
|
|||
# Verify both resource types were imported with prefixes
|
||||
assert "test+data://resource" in target_manager._resources
|
||||
assert "test+data://template/{id}" in target_manager._templates
|
||||
|
||||
|
||||
class TestCustomResourceKeys:
|
||||
"""Test adding resources and templates with custom keys."""
|
||||
|
||||
def test_add_resource_with_custom_key(self, temp_file: Path):
|
||||
"""Test adding a resource with a custom key different from its URI."""
|
||||
manager = ResourceManager()
|
||||
original_uri = "data://test/resource"
|
||||
custom_key = "custom://resource/key"
|
||||
|
||||
# Create a function resource instead of file resource to avoid path issues
|
||||
async def get_data():
|
||||
return "Test data"
|
||||
|
||||
resource = FunctionResource(
|
||||
uri=AnyUrl(original_uri),
|
||||
name="test_resource",
|
||||
fn=get_data,
|
||||
)
|
||||
|
||||
manager.add_resource(resource, key=custom_key)
|
||||
|
||||
# Resource should be accessible via custom key
|
||||
assert custom_key in manager._resources
|
||||
# But not via its original URI
|
||||
assert original_uri not in manager._resources
|
||||
# The resource's internal URI remains unchanged
|
||||
assert str(manager._resources[custom_key].uri) == original_uri
|
||||
|
||||
def test_add_template_with_custom_key(self):
|
||||
"""Test adding a template with a custom key different from its URI template."""
|
||||
manager = ResourceManager()
|
||||
|
||||
def template_fn(id: str) -> str:
|
||||
return f"Template {id}"
|
||||
|
||||
original_uri_template = "test://{id}"
|
||||
custom_key = "custom://{id}/template"
|
||||
|
||||
template = ResourceTemplate.from_function(
|
||||
fn=template_fn,
|
||||
uri_template=original_uri_template,
|
||||
name="test_template",
|
||||
)
|
||||
|
||||
manager.add_template(template, key=custom_key)
|
||||
|
||||
# Template should be accessible via custom key
|
||||
assert custom_key in manager._templates
|
||||
# But not via its original URI template
|
||||
assert original_uri_template not in manager._templates
|
||||
# The template's internal URI template remains unchanged
|
||||
assert str(manager._templates[custom_key].uri_template) == original_uri_template
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_get_resource_with_custom_key(self, temp_file: Path):
|
||||
"""Test that get_resource works with resources added with custom keys."""
|
||||
manager = ResourceManager()
|
||||
original_uri = "data://test/resource"
|
||||
custom_key = "custom://resource/path"
|
||||
|
||||
# Create a function resource instead of file resource to avoid path issues
|
||||
async def get_data():
|
||||
return "Test data"
|
||||
|
||||
resource = FunctionResource(
|
||||
uri=AnyUrl(original_uri),
|
||||
name="test_resource",
|
||||
fn=get_data,
|
||||
)
|
||||
|
||||
manager.add_resource(resource, key=custom_key)
|
||||
|
||||
# Should be retrievable by the custom key
|
||||
retrieved = await manager.get_resource(custom_key)
|
||||
assert retrieved is not None
|
||||
assert str(retrieved.uri) == original_uri
|
||||
|
||||
# Should NOT be retrievable by the original URI
|
||||
with pytest.raises(ResourceError, match="Unknown resource"):
|
||||
await manager.get_resource(original_uri)
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_get_resource_from_template_with_custom_key(self):
|
||||
"""Test that templates with custom keys can create resources."""
|
||||
manager = ResourceManager()
|
||||
|
||||
def greet(name: str) -> str:
|
||||
return f"Hello, {name}!"
|
||||
|
||||
original_template = "greet://{name}"
|
||||
custom_key = "custom://greet/{name}"
|
||||
|
||||
template = ResourceTemplate.from_function(
|
||||
fn=greet,
|
||||
uri_template=original_template,
|
||||
name="custom_greeter",
|
||||
)
|
||||
|
||||
manager.add_template(template, key=custom_key)
|
||||
|
||||
# Using a URI that matches the custom key pattern
|
||||
resource = await manager.get_resource("custom://greet/world")
|
||||
assert isinstance(resource, FunctionResource)
|
||||
content = await resource.read()
|
||||
assert content == "Hello, world!"
|
||||
|
||||
# Shouldn't work with the original template pattern
|
||||
with pytest.raises(ResourceError, match="Unknown resource"):
|
||||
await manager.get_resource("greet://world")
|
||||
|
||||
def test_import_resources_with_custom_keys(self):
|
||||
"""Test that import_resources properly uses custom keys."""
|
||||
source_manager = ResourceManager()
|
||||
target_manager = ResourceManager()
|
||||
|
||||
# Add a resource to source manager
|
||||
async def resource_fn():
|
||||
return "Resource data"
|
||||
|
||||
original_uri = "data://original"
|
||||
resource = FunctionResource(
|
||||
uri=AnyUrl(original_uri),
|
||||
name="original_resource",
|
||||
fn=resource_fn,
|
||||
)
|
||||
source_manager.add_resource(resource)
|
||||
|
||||
# Import with prefix which creates a new key
|
||||
prefix = "imported+"
|
||||
target_manager.import_resources(source_manager, prefix)
|
||||
|
||||
# Resource should be in target manager with prefixed URI as key
|
||||
prefixed_uri = f"{prefix}{original_uri}"
|
||||
assert prefixed_uri in target_manager._resources
|
||||
|
||||
# The resource's internal URI should remain unchanged
|
||||
stored_resource = target_manager._resources[prefixed_uri]
|
||||
assert str(stored_resource.uri) == original_uri
|
||||
|
||||
def test_import_templates_with_custom_keys(self):
|
||||
"""Test that import_templates properly uses custom keys."""
|
||||
source_manager = ResourceManager()
|
||||
target_manager = ResourceManager()
|
||||
|
||||
# Add a template to source manager
|
||||
async def template_fn(id: str):
|
||||
return f"Template {id}"
|
||||
|
||||
original_template = "template://{id}"
|
||||
template = ResourceTemplate.from_function(
|
||||
fn=template_fn,
|
||||
uri_template=original_template,
|
||||
name="original_template",
|
||||
)
|
||||
source_manager.add_template(template)
|
||||
|
||||
# Import with prefix which creates a new key
|
||||
prefix = "imported+"
|
||||
target_manager.import_templates(source_manager, prefix)
|
||||
|
||||
# Template should be in target manager with prefixed URI template as key
|
||||
prefixed_template = f"{prefix}{original_template}"
|
||||
assert prefixed_template in target_manager._templates
|
||||
|
||||
# The template's internal URI template should remain unchanged
|
||||
stored_template = target_manager._templates[prefixed_template]
|
||||
assert str(stored_template.uri_template) == original_template
|
||||
|
|
|
|||
|
|
@ -696,20 +696,26 @@ class TestMountFastMCP:
|
|||
|
||||
mcp.mount("fastapi", fastmcp_openapi_server)
|
||||
|
||||
# Check that resources are available with prefixed URIs
|
||||
resources = await mcp._mcp_list_resources()
|
||||
assert len(resources) == 1
|
||||
assert resources[0].uri == AnyUrl(
|
||||
"fastapi+resource://openapi/get_users_users_get"
|
||||
)
|
||||
# We're checking the key used by mcp to store the resource
|
||||
# The prefixed URI is used as the key, but the resource's original uri is preserved
|
||||
prefixed_uri = "fastapi+resource://openapi/get_users_users_get"
|
||||
resource = mcp._resource_manager.get_resources().get(prefixed_uri)
|
||||
assert resource is not None
|
||||
|
||||
# Check that templates are available with prefixed URIs
|
||||
templates = await mcp._mcp_list_resource_templates()
|
||||
assert len(templates) == 1
|
||||
assert templates[0].name == "get_user_users__user_id__get"
|
||||
assert (
|
||||
templates[0].uriTemplate
|
||||
== r"fastapi+resource://openapi/get_user_users__user_id__get/{user_id}"
|
||||
prefixed_template_uri = (
|
||||
r"fastapi+resource://openapi/get_user_users__user_id__get/{user_id}"
|
||||
)
|
||||
template = mcp._resource_manager.get_templates().get(prefixed_template_uri)
|
||||
assert template is not None
|
||||
|
||||
# Check that tools are available with prefixed names
|
||||
tools = await mcp._mcp_list_tools()
|
||||
assert len(tools) == 2
|
||||
assert tools[0].name == "fastapi_create_user_users_post"
|
||||
|
|
|
|||
|
|
@ -575,8 +575,8 @@ class TestCustomToolNames:
|
|||
# The tool should not be accessible via its original function name
|
||||
assert manager.get_tool("original_fn") is None
|
||||
|
||||
def test_add_tool_object_with_custom_storage_name(self):
|
||||
"""Test adding a Tool object with a custom storage name using add_tool()."""
|
||||
def test_add_tool_object_with_custom_key(self):
|
||||
"""Test adding a Tool object with a custom key using add_tool()."""
|
||||
|
||||
def fn(x: int) -> int:
|
||||
return x + 1
|
||||
|
|
@ -585,8 +585,8 @@ class TestCustomToolNames:
|
|||
tool = Tool.from_function(fn, name="my_tool")
|
||||
manager = ToolManager()
|
||||
# Store it under a different name
|
||||
manager.add_tool(tool, name="proxy_tool")
|
||||
# The tool is accessible under the storage name
|
||||
manager.add_tool(tool, key="proxy_tool")
|
||||
# The tool is accessible under the key
|
||||
stored = manager.get_tool("proxy_tool")
|
||||
assert stored is not None
|
||||
# But the tool's .name is unchanged
|
||||
|
|
@ -612,20 +612,35 @@ class TestCustomToolNames:
|
|||
with pytest.raises(ToolError):
|
||||
await manager.call_tool("multiply", {"a": 5, "b": 3})
|
||||
|
||||
def test_tool_to_mcp_tool_with_custom_name(self):
|
||||
"""Test that to_mcp_tool uses the storage name, not the internal name."""
|
||||
def test_tool_to_mcp_tool(self):
|
||||
"""Test that to_mcp_tool uses the key, not the internal name."""
|
||||
|
||||
def some_function(x: int) -> int:
|
||||
return x
|
||||
|
||||
manager = ToolManager()
|
||||
manager.add_tool_from_fn(some_function, name="api_function")
|
||||
tool = Tool.from_function(some_function, name="api_function")
|
||||
manager.add_tool(tool)
|
||||
|
||||
# When listing tools for MCP, the custom name should be used
|
||||
mcp_tools = manager.list_mcp_tools()
|
||||
assert len(mcp_tools) == 1
|
||||
assert mcp_tools[0].name == "api_function"
|
||||
|
||||
def test_tool_to_mcp_tool_with_custom_key(self):
|
||||
"""Test that to_mcp_tool uses the key, not the internal name."""
|
||||
|
||||
def some_function(x: int) -> int:
|
||||
return x
|
||||
|
||||
manager = ToolManager()
|
||||
tool = Tool.from_function(some_function, name="api_function")
|
||||
manager.add_tool(tool, key="custom-key")
|
||||
|
||||
# When listing tools for MCP, the key should be used
|
||||
mcp_tools = manager.list_mcp_tools()
|
||||
assert len(mcp_tools) == 1
|
||||
assert mcp_tools[0].name == "custom-key"
|
||||
|
||||
def test_import_tools_with_custom_names(self):
|
||||
"""Test importing tools with custom names."""
|
||||
|
||||
|
|
@ -675,20 +690,20 @@ class TestCustomToolNames:
|
|||
assert stored_tool.fn.__name__ == "replacement_fn"
|
||||
|
||||
def test_mcp_tool_name_for_add_tool(self):
|
||||
"""Test MCPTool name for add_tool (storage name != tool.name)."""
|
||||
"""Test MCPTool name for add_tool (key != tool.name)."""
|
||||
|
||||
def fn(x: int) -> int:
|
||||
return x + 1
|
||||
|
||||
tool = Tool.from_function(fn, name="my_tool")
|
||||
manager = ToolManager()
|
||||
manager.add_tool(tool, name="proxy_tool")
|
||||
manager.add_tool(tool, key="proxy_tool")
|
||||
mcp_tools = manager.list_mcp_tools()
|
||||
assert len(mcp_tools) == 1
|
||||
assert mcp_tools[0].name == "proxy_tool"
|
||||
|
||||
def test_mcp_tool_name_for_add_tool_from_fn(self):
|
||||
"""Test MCPTool name for add_tool_from_fn (storage name == tool.name)."""
|
||||
"""Test MCPTool name for add_tool_from_fn (key == tool.name)."""
|
||||
|
||||
def fn(x: int) -> int:
|
||||
return x + 1
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue