mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-24 06:24:18 +02:00
Support mounting FastMCPs as sub-servers
This commit is contained in:
parent
4f6aefed7a
commit
9b89715294
15 changed files with 1228 additions and 1 deletions
150
examples/modular_app.py
Normal file
150
examples/modular_app.py
Normal file
|
|
@ -0,0 +1,150 @@
|
|||
"""
|
||||
Modular FastMCP Application Example
|
||||
|
||||
This example demonstrates building a modular application with FastMCP
|
||||
by separating functionality into domain-specific modules.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from fastmcp import Context, FastMCP
|
||||
|
||||
# ----- DATA MODULE -----
|
||||
data_app = FastMCP("Data Module")
|
||||
|
||||
# Simulated database
|
||||
users_db = [
|
||||
{"id": 1, "name": "Alice", "email": "alice@example.com"},
|
||||
{"id": 2, "name": "Bob", "email": "bob@example.com"},
|
||||
{"id": 3, "name": "Charlie", "email": "charlie@example.com"},
|
||||
]
|
||||
|
||||
|
||||
@data_app.resource("users://all")
|
||||
def get_all_users() -> List[Dict[str, Any]]:
|
||||
"""Get all users in the database"""
|
||||
return users_db
|
||||
|
||||
|
||||
@data_app.resource("users://{user_id}")
|
||||
def get_user_by_id(user_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""Get a specific user by ID"""
|
||||
user_id_int = int(user_id)
|
||||
for user in users_db:
|
||||
if user["id"] == user_id_int:
|
||||
return user
|
||||
return None
|
||||
|
||||
|
||||
@data_app.tool()
|
||||
async def create_user(name: str, email: str, ctx: Context) -> Dict[str, Any]:
|
||||
"""Add a new user to the database"""
|
||||
# Simulate a slow operation
|
||||
await ctx.info(f"Creating user {name}...")
|
||||
await asyncio.sleep(1)
|
||||
|
||||
# Create user
|
||||
new_id = max(user["id"] for user in users_db) + 1
|
||||
new_user = {"id": new_id, "name": name, "email": email}
|
||||
users_db.append(new_user)
|
||||
|
||||
await ctx.info(f"User created with ID {new_id}")
|
||||
return new_user
|
||||
|
||||
|
||||
# ----- ANALYTICS MODULE -----
|
||||
analytics_app = FastMCP("Analytics Module")
|
||||
|
||||
|
||||
@analytics_app.tool()
|
||||
async def analyze_users(ctx: Context) -> Dict[str, Any]:
|
||||
"""Run analytics on user data"""
|
||||
# Get user data from the data module
|
||||
users = await ctx.read_resource("data:users://all")
|
||||
|
||||
# Perform analytics
|
||||
await ctx.info("Analyzing user data...")
|
||||
await asyncio.sleep(1)
|
||||
|
||||
# Return analytics results
|
||||
return {
|
||||
"total_users": len(users),
|
||||
"domains": {user["email"].split("@")[1] for user in users},
|
||||
}
|
||||
|
||||
|
||||
@analytics_app.resource("analytics://summary")
|
||||
def get_analytics_summary() -> Dict[str, Any]:
|
||||
"""Get a summary of analytics data"""
|
||||
return {"active_users": len(users_db), "last_updated": "2023-06-01"}
|
||||
|
||||
|
||||
# ----- FILESYSTEM MODULE -----
|
||||
files_app = FastMCP("Filesystem Module")
|
||||
|
||||
|
||||
@files_app.resource("files://desktop")
|
||||
def list_desktop_files() -> List[str]:
|
||||
"""List files on the user's desktop"""
|
||||
desktop = Path.home() / "Desktop"
|
||||
return [f.name for f in desktop.iterdir() if f.is_file()]
|
||||
|
||||
|
||||
@files_app.tool()
|
||||
async def search_files(query: str, ctx: Context) -> List[str]:
|
||||
"""Search for files matching a query"""
|
||||
await ctx.info(f"Searching for files matching '{query}'...")
|
||||
|
||||
# Simulate a file search
|
||||
desktop = Path.home() / "Desktop"
|
||||
files = [
|
||||
f.name
|
||||
for f in desktop.iterdir()
|
||||
if f.is_file() and query.lower() in f.name.lower()
|
||||
]
|
||||
|
||||
await ctx.info(f"Found {len(files)} matching files")
|
||||
return files
|
||||
|
||||
|
||||
# ----- MAIN APPLICATION -----
|
||||
# Create the main application that combines all modules
|
||||
main_app = FastMCP("Modular FastMCP Demo")
|
||||
|
||||
|
||||
@main_app.tool()
|
||||
async def get_system_info(ctx: Context) -> Dict[str, Any]:
|
||||
"""Get comprehensive system information"""
|
||||
await ctx.info("Gathering system information...")
|
||||
|
||||
# Use the mounted modules to gather info
|
||||
users = await ctx.read_resource("data:users://all")
|
||||
analytics = await ctx.read_resource("analytics:analytics://summary")
|
||||
desktop_files = await ctx.read_resource("files:files://desktop")
|
||||
|
||||
return {
|
||||
"users": {"count": len(users), "names": [user["name"] for user in users]},
|
||||
"analytics": analytics,
|
||||
"files": {"desktop_count": len(desktop_files)},
|
||||
}
|
||||
|
||||
|
||||
# Mount all modules to the main app
|
||||
main_app.mount("data", data_app)
|
||||
main_app.mount("analytics", analytics_app)
|
||||
main_app.mount("files", files_app)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Now register resources (which requires async)
|
||||
async def initialize_resources():
|
||||
await main_app.register_all_mounted_resources()
|
||||
print("Resources registered successfully!")
|
||||
|
||||
# Initialize resources
|
||||
asyncio.run(initialize_resources())
|
||||
|
||||
# Start the server
|
||||
print("Starting modular FastMCP application...")
|
||||
main_app.run()
|
||||
112
examples/mount_example.py
Normal file
112
examples/mount_example.py
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
"""Example of mounting FastMCP apps together.
|
||||
|
||||
This example demonstrates how to mount FastMCP apps together using
|
||||
the ToolManager's import_tools functionality. It shows how to:
|
||||
|
||||
1. Create sub-applications for different domains
|
||||
2. Mount those sub-applications to a main application
|
||||
3. Access tools with prefixed names and resources with prefixed URIs
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import Dict, List
|
||||
|
||||
from fastmcp import FastMCP
|
||||
|
||||
# Weather sub-application
|
||||
weather_app = FastMCP("Weather App")
|
||||
|
||||
|
||||
@weather_app.tool()
|
||||
def get_weather_forecast(location: str) -> str:
|
||||
"""Get the weather forecast for a location."""
|
||||
return f"Sunny skies for {location} today!"
|
||||
|
||||
|
||||
@weather_app.resource(uri="weather://forecast")
|
||||
async def weather_data():
|
||||
"""Return current weather data."""
|
||||
return {"temperature": 72, "conditions": "sunny", "humidity": 45, "wind_speed": 5}
|
||||
|
||||
|
||||
# News sub-application
|
||||
news_app = FastMCP("News App")
|
||||
|
||||
|
||||
@news_app.tool()
|
||||
def get_news_headlines() -> List[str]:
|
||||
"""Get the latest news headlines."""
|
||||
return [
|
||||
"Tech company launches new product",
|
||||
"Local team wins championship",
|
||||
"Scientists make breakthrough discovery",
|
||||
]
|
||||
|
||||
|
||||
@news_app.resource(uri="news://headlines")
|
||||
async def news_data():
|
||||
"""Return latest news data."""
|
||||
return {
|
||||
"top_story": "Breaking news: Important event happened",
|
||||
"categories": ["politics", "sports", "technology"],
|
||||
"sources": ["AP", "Reuters", "Local Sources"],
|
||||
}
|
||||
|
||||
|
||||
# Main application
|
||||
app = FastMCP("Main App")
|
||||
|
||||
|
||||
@app.tool()
|
||||
def check_app_status() -> Dict[str, str]:
|
||||
"""Check the status of the main application."""
|
||||
return {"status": "running", "version": "1.0.0", "uptime": "3h 24m"}
|
||||
|
||||
|
||||
# Mount sub-applications
|
||||
app.mount("weather", weather_app)
|
||||
app.mount("news", news_app)
|
||||
|
||||
|
||||
async def start_server():
|
||||
"""Print information about mounted resources."""
|
||||
# Print available tools
|
||||
tools = app._tool_manager.list_tools()
|
||||
print(f"\nAvailable tools ({len(tools)}):")
|
||||
for tool in tools:
|
||||
print(f" - {tool.name}: {tool.description}")
|
||||
|
||||
# Print available resources
|
||||
print("\nAvailable resources:")
|
||||
|
||||
# Distinguish between native and imported resources
|
||||
# Native resources would be those directly in the main app (not prefixed)
|
||||
native_resources = [
|
||||
uri
|
||||
for uri in app._resource_manager._resources
|
||||
if not (uri.startswith("weather+") or uri.startswith("news+"))
|
||||
]
|
||||
|
||||
# Imported resources - categorized by source app
|
||||
weather_resources = [
|
||||
uri for uri in app._resource_manager._resources if uri.startswith("weather+")
|
||||
]
|
||||
news_resources = [
|
||||
uri for uri in app._resource_manager._resources if uri.startswith("news+")
|
||||
]
|
||||
|
||||
print(f" - Native app resources: {native_resources}")
|
||||
print(f" - Imported from weather app: {weather_resources}")
|
||||
print(f" - Imported from news app: {news_resources}")
|
||||
|
||||
# Let's try to access resources using the prefixed URI
|
||||
weather_data = await app.read_resource("weather+weather://forecast")
|
||||
print(f"\nWeather data from prefixed URI: {weather_data}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# First run our async function to display info
|
||||
asyncio.run(start_server())
|
||||
|
||||
# Then start the server (uncomment to run the server)
|
||||
# app.run()
|
||||
3
src/fastmcp/prompts/__init__.py
Normal file
3
src/fastmcp/prompts/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
from .prompt_manager import PromptManager
|
||||
|
||||
__all__ = ["PromptManager"]
|
||||
33
src/fastmcp/prompts/prompt_manager.py
Normal file
33
src/fastmcp/prompts/prompt_manager.py
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
import logging
|
||||
|
||||
from mcp.server.fastmcp.prompts import PromptManager as BasePromptManager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PromptManager(BasePromptManager):
|
||||
"""
|
||||
Extended PromptManager that supports importing prompts from other managers.
|
||||
Adds ability to import prompts from other managers with prefixed names.
|
||||
"""
|
||||
|
||||
def import_prompts(self, manager: "PromptManager", prefix: str) -> None:
|
||||
"""
|
||||
Import all prompts from another PromptManager with prefixed names.
|
||||
|
||||
Args:
|
||||
manager: Another PromptManager instance to import prompts from
|
||||
prefix: Prefix to add to prompt names. The resulting prompt name will
|
||||
be in the format "{prefix}/{original_name}"
|
||||
For example, with prefix "weather" and prompt "forecast_prompt",
|
||||
the imported prompt would be available as "weather/forecast_prompt"
|
||||
"""
|
||||
for name, prompt in manager._prompts.items():
|
||||
# Create prefixed name - we keep the original name in the Prompt object
|
||||
prefixed_name = f"{prefix}/{name}"
|
||||
|
||||
# Log the import
|
||||
logger.debug(f"Importing prompt with name {name} as {prefixed_name}")
|
||||
|
||||
# Store the prompt with the prefixed name
|
||||
self._prompts[prefixed_name] = prompt
|
||||
3
src/fastmcp/resources/__init__.py
Normal file
3
src/fastmcp/resources/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
from .resource_manager import ResourceManager
|
||||
|
||||
__all__ = ["ResourceManager"]
|
||||
55
src/fastmcp/resources/resource_manager.py
Normal file
55
src/fastmcp/resources/resource_manager.py
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
import logging
|
||||
|
||||
from mcp.server.fastmcp.resources import (
|
||||
ResourceManager as BaseResourceManager,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ResourceManager(BaseResourceManager):
|
||||
"""ResourceManager that adds methods to import resources from other managers."""
|
||||
|
||||
def import_resources(self, manager: "ResourceManager", prefix: str) -> None:
|
||||
"""Import resources from another resource manager.
|
||||
|
||||
Resources are imported with a prefixed URI. For example, if a resource has
|
||||
URI "data://users" and you import it with prefix "app", the imported resource
|
||||
will have URI "app+data://users".
|
||||
|
||||
Args:
|
||||
manager: The ResourceManager to import from
|
||||
prefix: A prefix to apply to the resource URIs
|
||||
"""
|
||||
for uri, resource in manager._resources.items():
|
||||
# Create prefixed URI and copy the resource with the new URI
|
||||
prefixed_uri = f"{prefix}+{uri}"
|
||||
|
||||
# Log the import
|
||||
logger.debug(f"Importing resource with URI {uri} as {prefixed_uri}")
|
||||
|
||||
# Store directly in resources dictionary
|
||||
self._resources[prefixed_uri] = resource
|
||||
|
||||
def import_templates(self, manager: "ResourceManager", prefix: str) -> None:
|
||||
"""Import resource templates from another resource manager.
|
||||
|
||||
Templates are imported with a prefixed URI template. For example, if a template has
|
||||
URI template "data://users/{id}" and you import it with prefix "app", the
|
||||
imported template will have URI template "app+data://users/{id}".
|
||||
|
||||
Args:
|
||||
manager: The ResourceManager to import templates from
|
||||
prefix: A prefix to apply to the template URIs
|
||||
"""
|
||||
for uri_template, template in manager._templates.items():
|
||||
# Create prefixed URI template and copy the template with the new URI template
|
||||
prefixed_uri_template = f"{prefix}+{uri_template}"
|
||||
|
||||
# Log the import
|
||||
logger.debug(
|
||||
f"Importing resource template with URI {uri_template} as {prefixed_uri_template}"
|
||||
)
|
||||
|
||||
# Store directly in templates dictionary
|
||||
self._templates[prefixed_uri_template] = template
|
||||
|
|
@ -1,8 +1,12 @@
|
|||
from typing import Any
|
||||
from typing import Any, Dict
|
||||
|
||||
import mcp.server.fastmcp
|
||||
import mcp.types
|
||||
|
||||
from fastmcp.prompts.prompt_manager import PromptManager
|
||||
from fastmcp.resources.resource_manager import ResourceManager
|
||||
from fastmcp.server.context import Context
|
||||
from fastmcp.tools.tool_manager import ToolManager
|
||||
from fastmcp.utilities.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
|
@ -10,8 +14,23 @@ logger = get_logger(__name__)
|
|||
|
||||
class FastMCP(mcp.server.fastmcp.FastMCP):
|
||||
def __init__(self, name: str | None = None, **settings: Any):
|
||||
# First initialize with default settings
|
||||
super().__init__(name=name or "FastMCP", **settings)
|
||||
|
||||
# Replace the default managers with our extended ones
|
||||
self._tool_manager = ToolManager(
|
||||
warn_on_duplicate_tools=self.settings.warn_on_duplicate_tools
|
||||
)
|
||||
self._resource_manager = ResourceManager(
|
||||
warn_on_duplicate_resources=self.settings.warn_on_duplicate_resources
|
||||
)
|
||||
self._prompt_manager = PromptManager(
|
||||
warn_on_duplicate_prompts=self.settings.warn_on_duplicate_prompts
|
||||
)
|
||||
|
||||
# Setup for mounted apps
|
||||
self._mounted_apps: Dict[str, "FastMCP"] = {}
|
||||
|
||||
def get_context(self) -> Context:
|
||||
"""
|
||||
Returns a Context object. Note that the context will only be valid
|
||||
|
|
@ -22,3 +41,41 @@ class FastMCP(mcp.server.fastmcp.FastMCP):
|
|||
except LookupError:
|
||||
request_context = None
|
||||
return Context(request_context=request_context, fastmcp=self)
|
||||
|
||||
def mount(self, prefix: str, app: "FastMCP") -> None:
|
||||
"""Mount another FastMCP application with a given prefix.
|
||||
|
||||
When an application is mounted:
|
||||
- The tools are imported with prefixed names
|
||||
Example: If app has a tool named "get_weather", it will be available as "weather/get_weather"
|
||||
- The resources are imported with prefixed URIs
|
||||
Example: If app has a resource with URI "weather://forecast", it will be available as "weather+weather://forecast"
|
||||
- The templates are imported with prefixed URI templates
|
||||
Example: If app has a template with URI "weather://location/{id}", it will be available as "weather+weather://location/{id}"
|
||||
- The prompts are imported with prefixed names
|
||||
Example: If app has a prompt named "weather_prompt", it will be available as "weather/weather_prompt"
|
||||
|
||||
Args:
|
||||
prefix: The prefix to use for the mounted application
|
||||
app: The FastMCP application to mount
|
||||
"""
|
||||
# Mount the app in the list of mounted apps
|
||||
self._mounted_apps[prefix] = app
|
||||
|
||||
# Import tools from the mounted app
|
||||
self._tool_manager.import_tools(app._tool_manager, prefix)
|
||||
|
||||
# Import resources from the mounted app
|
||||
self._resource_manager.import_resources(app._resource_manager, prefix)
|
||||
|
||||
# Import resource templates
|
||||
self._resource_manager.import_templates(app._resource_manager, prefix)
|
||||
|
||||
# Import prompts
|
||||
self._prompt_manager.import_prompts(app._prompt_manager, prefix)
|
||||
|
||||
logger.info(f"Mounted app with prefix '{prefix}'")
|
||||
logger.debug(f"Imported tools with prefix '{prefix}/'")
|
||||
logger.debug(f"Imported resources with prefix '{prefix}+'")
|
||||
logger.debug(f"Imported templates with prefix '{prefix}+'")
|
||||
logger.debug(f"Imported prompts with prefix '{prefix}/'")
|
||||
|
|
|
|||
|
|
@ -0,0 +1,38 @@
|
|||
import mcp.server.fastmcp.tools
|
||||
from mcp.server.fastmcp.tools import Tool
|
||||
|
||||
from fastmcp.utilities.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class ToolManager(mcp.server.fastmcp.tools.ToolManager):
|
||||
"""
|
||||
Extended ToolManager that supports importing tools from other managers.
|
||||
Adds ability to import tools from other managers with prefixed names.
|
||||
"""
|
||||
|
||||
def import_tools(self, tool_manager: "ToolManager", prefix: str) -> None:
|
||||
"""
|
||||
Import all tools from another ToolManager with prefixed names.
|
||||
|
||||
Args:
|
||||
tool_manager: Another ToolManager instance to import tools from
|
||||
prefix: Prefix to add to tool names. The resulting tool name will
|
||||
be in the format "{prefix}/{original_name}"
|
||||
For example, with prefix "weather" and tool "forecast",
|
||||
the imported tool would be available as "weather/forecast"
|
||||
"""
|
||||
for name, tool in tool_manager._tools.items():
|
||||
prefixed_name = f"{prefix}/{name}"
|
||||
|
||||
# Create a shallow copy of the tool with the prefixed name
|
||||
copied_tool = Tool.from_function(
|
||||
tool.fn,
|
||||
name=prefixed_name,
|
||||
description=tool.description,
|
||||
)
|
||||
|
||||
# Store the copied tool
|
||||
self._tools[prefixed_name] = copied_tool
|
||||
logger.debug(f"Imported tool: {name} as {prefixed_name}")
|
||||
166
tests/prompts/test_prompt_manager.py
Normal file
166
tests/prompts/test_prompt_manager.py
Normal file
|
|
@ -0,0 +1,166 @@
|
|||
from mcp.server.fastmcp.prompts import Prompt
|
||||
from mcp.server.fastmcp.prompts.base import PromptArgument
|
||||
|
||||
from fastmcp.prompts.prompt_manager import PromptManager
|
||||
|
||||
|
||||
def test_import_prompts():
|
||||
"""Test importing prompts from one manager to another with a prefix."""
|
||||
# Setup source manager with prompts
|
||||
source_manager = PromptManager()
|
||||
|
||||
# Create test prompts with proper function handlers
|
||||
async def summary_fn(**kwargs):
|
||||
return [{"role": "assistant", "content": f"Summary of: {kwargs.get('text')}"}]
|
||||
|
||||
async def translate_fn(**kwargs):
|
||||
return [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": f"Translation to {kwargs.get('language')}: {kwargs.get('text')}",
|
||||
}
|
||||
]
|
||||
|
||||
summary_prompt = Prompt(
|
||||
name="summary",
|
||||
description="Generate a summary of text",
|
||||
arguments=[PromptArgument(name="text", description="Text to summarize")],
|
||||
fn=summary_fn,
|
||||
)
|
||||
source_manager._prompts["summary"] = summary_prompt
|
||||
|
||||
translate_prompt = Prompt(
|
||||
name="translate",
|
||||
description="Translate text to another language",
|
||||
arguments=[
|
||||
PromptArgument(name="text", description="Text to translate"),
|
||||
PromptArgument(name="language", description="Target language"),
|
||||
],
|
||||
fn=translate_fn,
|
||||
)
|
||||
source_manager._prompts["translate"] = translate_prompt
|
||||
|
||||
# Create target manager
|
||||
target_manager = PromptManager()
|
||||
|
||||
# Import prompts from source to target
|
||||
prefix = "nlp"
|
||||
target_manager.import_prompts(source_manager, prefix)
|
||||
|
||||
# Verify prompts were imported with prefixes
|
||||
assert "nlp/summary" in target_manager._prompts
|
||||
assert "nlp/translate" in target_manager._prompts
|
||||
|
||||
# Verify the original prompts still exist in source manager
|
||||
assert "summary" in source_manager._prompts
|
||||
assert "translate" in source_manager._prompts
|
||||
|
||||
# Verify the imported prompts have the correct properties
|
||||
assert target_manager._prompts["nlp/summary"].name == "summary"
|
||||
assert (
|
||||
target_manager._prompts["nlp/summary"].description
|
||||
== "Generate a summary of text"
|
||||
)
|
||||
|
||||
assert target_manager._prompts["nlp/translate"].name == "translate"
|
||||
assert (
|
||||
target_manager._prompts["nlp/translate"].description
|
||||
== "Translate text to another language"
|
||||
)
|
||||
|
||||
# Verify functions were properly copied
|
||||
if hasattr(target_manager._prompts["nlp/summary"], "fn"):
|
||||
assert target_manager._prompts["nlp/summary"].fn.__name__ == summary_fn.__name__
|
||||
|
||||
if hasattr(target_manager._prompts["nlp/translate"], "fn"):
|
||||
assert (
|
||||
target_manager._prompts["nlp/translate"].fn.__name__
|
||||
== translate_fn.__name__
|
||||
)
|
||||
|
||||
|
||||
def test_import_prompts_with_duplicates():
|
||||
"""Test handling of duplicate prompts during import."""
|
||||
# Setup source and target managers with same prompt names
|
||||
source_manager = PromptManager()
|
||||
target_manager = PromptManager()
|
||||
|
||||
# Add the same prompt name to both managers with functions
|
||||
async def source_fn(**kwargs):
|
||||
return [{"role": "assistant", "content": "Source content"}]
|
||||
|
||||
async def target_fn(**kwargs):
|
||||
return [{"role": "assistant", "content": "Target content"}]
|
||||
|
||||
source_prompt = Prompt(
|
||||
name="common",
|
||||
description="Source description",
|
||||
arguments=None,
|
||||
fn=source_fn,
|
||||
)
|
||||
source_manager._prompts["common"] = source_prompt
|
||||
|
||||
target_prompt = Prompt(
|
||||
name="common",
|
||||
description="Target description",
|
||||
arguments=None,
|
||||
fn=target_fn,
|
||||
)
|
||||
target_manager._prompts["common"] = target_prompt
|
||||
|
||||
# Import prompts with prefix
|
||||
prefix = "external"
|
||||
target_manager.import_prompts(source_manager, prefix)
|
||||
|
||||
# Verify both prompts exist in target manager
|
||||
assert "common" in target_manager._prompts
|
||||
assert "external/common" in target_manager._prompts
|
||||
|
||||
# Verify the functions of both prompts
|
||||
if hasattr(target_manager._prompts["common"], "fn") and hasattr(
|
||||
target_manager._prompts["external/common"], "fn"
|
||||
):
|
||||
assert target_manager._prompts["common"].fn.__name__ == target_fn.__name__
|
||||
assert (
|
||||
target_manager._prompts["external/common"].fn.__name__ == source_fn.__name__
|
||||
)
|
||||
|
||||
|
||||
def test_import_prompts_with_nested_prefixes():
|
||||
"""Test importing already prefixed prompts."""
|
||||
# Setup source manager with already prefixed prompts
|
||||
first_manager = PromptManager()
|
||||
second_manager = PromptManager()
|
||||
third_manager = PromptManager()
|
||||
|
||||
# Add prompt to first manager with a function
|
||||
async def analyze_fn(**kwargs):
|
||||
return [{"role": "assistant", "content": f"Analysis of: {kwargs.get('text')}"}]
|
||||
|
||||
original_prompt = Prompt(
|
||||
name="analyze",
|
||||
description="Analyze text",
|
||||
arguments=[PromptArgument(name="text", description="Text to analyze")],
|
||||
fn=analyze_fn,
|
||||
)
|
||||
first_manager._prompts["analyze"] = original_prompt
|
||||
|
||||
# Import to second manager with prefix
|
||||
second_manager.import_prompts(first_manager, "text")
|
||||
|
||||
# Import from second to third with another prefix
|
||||
third_manager.import_prompts(second_manager, "ai")
|
||||
|
||||
# Verify the nested prefixing
|
||||
assert "text/analyze" in second_manager._prompts
|
||||
assert "ai/text/analyze" in third_manager._prompts
|
||||
|
||||
# Verify the properties of the most nested prompt
|
||||
assert third_manager._prompts["ai/text/analyze"].name == "analyze"
|
||||
assert third_manager._prompts["ai/text/analyze"].description == "Analyze text"
|
||||
|
||||
# Verify function was properly copied through multiple imports
|
||||
if hasattr(third_manager._prompts["ai/text/analyze"], "fn"):
|
||||
assert (
|
||||
third_manager._prompts["ai/text/analyze"].fn.__name__ == analyze_fn.__name__
|
||||
)
|
||||
221
tests/resources/test_resource_manager.py
Normal file
221
tests/resources/test_resource_manager.py
Normal file
|
|
@ -0,0 +1,221 @@
|
|||
from mcp.server.fastmcp.resources import FunctionResource, ResourceTemplate
|
||||
from pydantic.networks import AnyUrl
|
||||
|
||||
from fastmcp.resources.resource_manager import ResourceManager
|
||||
|
||||
|
||||
def test_import_resources():
|
||||
"""Test importing resources from one manager to another with a prefix."""
|
||||
# Setup source manager with resources
|
||||
source_manager = ResourceManager()
|
||||
|
||||
# Create mock resource functions
|
||||
async def weather_fn():
|
||||
return "Weather data"
|
||||
|
||||
async def traffic_fn():
|
||||
return "Traffic data"
|
||||
|
||||
# Add resources to source manager
|
||||
weather_resource = FunctionResource(
|
||||
uri=AnyUrl("weather://forecast"),
|
||||
name="weather_forecast",
|
||||
description="Get weather forecast",
|
||||
mime_type="application/json",
|
||||
fn=weather_fn,
|
||||
)
|
||||
source_manager._resources["weather://forecast"] = weather_resource
|
||||
|
||||
traffic_resource = FunctionResource(
|
||||
uri=AnyUrl("traffic://status"),
|
||||
name="traffic_status",
|
||||
description="Get traffic status",
|
||||
mime_type="application/json",
|
||||
fn=traffic_fn,
|
||||
)
|
||||
source_manager._resources["traffic://status"] = traffic_resource
|
||||
|
||||
# Create target manager
|
||||
target_manager = ResourceManager()
|
||||
|
||||
# Import resources from source to target
|
||||
prefix = "data"
|
||||
target_manager.import_resources(source_manager, prefix)
|
||||
|
||||
# Verify resources were imported with prefixes
|
||||
assert "data+weather://forecast" in target_manager._resources
|
||||
assert "data+traffic://status" in target_manager._resources
|
||||
|
||||
# Verify the original resources still exist in source manager
|
||||
assert "weather://forecast" in source_manager._resources
|
||||
assert "traffic://status" in source_manager._resources
|
||||
|
||||
# Verify the imported resources have the correct properties
|
||||
assert (
|
||||
target_manager._resources["data+weather://forecast"].name == "weather_forecast"
|
||||
)
|
||||
assert (
|
||||
target_manager._resources["data+weather://forecast"].description
|
||||
== "Get weather forecast"
|
||||
)
|
||||
assert (
|
||||
target_manager._resources["data+weather://forecast"].mime_type
|
||||
== "application/json"
|
||||
)
|
||||
|
||||
assert target_manager._resources["data+traffic://status"].name == "traffic_status"
|
||||
assert (
|
||||
target_manager._resources["data+traffic://status"].description
|
||||
== "Get traffic status"
|
||||
)
|
||||
assert (
|
||||
target_manager._resources["data+traffic://status"].mime_type
|
||||
== "application/json"
|
||||
)
|
||||
|
||||
# Since we're dealing with FunctionResource type, we can safely check function attributes
|
||||
assert isinstance(
|
||||
target_manager._resources["data+weather://forecast"], FunctionResource
|
||||
)
|
||||
assert isinstance(
|
||||
target_manager._resources["data+traffic://status"], FunctionResource
|
||||
)
|
||||
|
||||
weather_resource = target_manager._resources["data+weather://forecast"]
|
||||
traffic_resource = target_manager._resources["data+traffic://status"]
|
||||
|
||||
if hasattr(weather_resource, "fn") and hasattr(traffic_resource, "fn"):
|
||||
assert weather_resource.fn.__name__ == weather_fn.__name__
|
||||
assert traffic_resource.fn.__name__ == traffic_fn.__name__
|
||||
|
||||
|
||||
def test_import_templates():
|
||||
"""Test importing resource templates from one manager to another with a prefix."""
|
||||
# Setup source manager with templates
|
||||
source_manager = ResourceManager()
|
||||
|
||||
# Create mock template functions
|
||||
async def user_fn(**params):
|
||||
return f"User data for id {params.get('id')}"
|
||||
|
||||
async def product_fn(**params):
|
||||
return f"Product data for id {params.get('id')}"
|
||||
|
||||
# Add templates to source manager
|
||||
user_template = ResourceTemplate(
|
||||
uri_template="api://users/{id}",
|
||||
name="user_template",
|
||||
description="Get user by ID",
|
||||
mime_type="application/json",
|
||||
fn=user_fn,
|
||||
parameters={"id": {"type": "string", "description": "User ID"}},
|
||||
)
|
||||
source_manager._templates["api://users/{id}"] = user_template
|
||||
|
||||
product_template = ResourceTemplate(
|
||||
uri_template="api://products/{id}",
|
||||
name="product_template",
|
||||
description="Get product by ID",
|
||||
mime_type="application/json",
|
||||
fn=product_fn,
|
||||
parameters={"id": {"type": "string", "description": "Product ID"}},
|
||||
)
|
||||
source_manager._templates["api://products/{id}"] = product_template
|
||||
|
||||
# Create target manager
|
||||
target_manager = ResourceManager()
|
||||
|
||||
# Import templates from source to target
|
||||
prefix = "shop"
|
||||
target_manager.import_templates(source_manager, prefix)
|
||||
|
||||
# Verify templates were imported with prefixes
|
||||
assert "shop+api://users/{id}" in target_manager._templates
|
||||
assert "shop+api://products/{id}" in target_manager._templates
|
||||
|
||||
# Verify the original templates still exist in source manager
|
||||
assert "api://users/{id}" in source_manager._templates
|
||||
assert "api://products/{id}" in source_manager._templates
|
||||
|
||||
# Verify the imported templates have the correct properties
|
||||
assert target_manager._templates["shop+api://users/{id}"].name == "user_template"
|
||||
assert (
|
||||
target_manager._templates["shop+api://users/{id}"].description
|
||||
== "Get user by ID"
|
||||
)
|
||||
assert (
|
||||
target_manager._templates["shop+api://users/{id}"].mime_type
|
||||
== "application/json"
|
||||
)
|
||||
assert target_manager._templates["shop+api://users/{id}"].parameters == {
|
||||
"id": {"type": "string", "description": "User ID"}
|
||||
}
|
||||
|
||||
assert (
|
||||
target_manager._templates["shop+api://products/{id}"].name == "product_template"
|
||||
)
|
||||
assert (
|
||||
target_manager._templates["shop+api://products/{id}"].description
|
||||
== "Get product by ID"
|
||||
)
|
||||
assert (
|
||||
target_manager._templates["shop+api://products/{id}"].mime_type
|
||||
== "application/json"
|
||||
)
|
||||
assert target_manager._templates["shop+api://products/{id}"].parameters == {
|
||||
"id": {"type": "string", "description": "Product ID"}
|
||||
}
|
||||
|
||||
# Verify the template functions were properly copied (only if the fn attribute exists)
|
||||
user_template = target_manager._templates["shop+api://users/{id}"]
|
||||
product_template = target_manager._templates["shop+api://products/{id}"]
|
||||
|
||||
if hasattr(user_template, "fn") and hasattr(product_template, "fn"):
|
||||
assert user_template.fn.__name__ == user_fn.__name__
|
||||
assert product_template.fn.__name__ == product_fn.__name__
|
||||
|
||||
|
||||
def test_import_multiple_resource_types():
|
||||
"""Test importing both resources and templates with the same prefix."""
|
||||
# Setup source manager with both resources and templates
|
||||
source_manager = ResourceManager()
|
||||
|
||||
# Create mock functions
|
||||
async def resource_fn():
|
||||
return "Resource data"
|
||||
|
||||
async def template_fn(**params):
|
||||
return f"Template data for id {params.get('id')}"
|
||||
|
||||
# Add a resource to source manager
|
||||
resource = FunctionResource(
|
||||
uri=AnyUrl("data://resource"),
|
||||
name="test_resource",
|
||||
description="Test resource",
|
||||
mime_type="application/json",
|
||||
fn=resource_fn,
|
||||
)
|
||||
source_manager._resources["data://resource"] = resource
|
||||
|
||||
# Add a template to source manager
|
||||
template = ResourceTemplate(
|
||||
uri_template="data://template/{id}",
|
||||
name="test_template",
|
||||
description="Test template",
|
||||
mime_type="application/json",
|
||||
fn=template_fn,
|
||||
parameters={"id": {"type": "string", "description": "ID parameter"}},
|
||||
)
|
||||
source_manager._templates["data://template/{id}"] = template
|
||||
|
||||
# Create target manager
|
||||
target_manager = ResourceManager()
|
||||
|
||||
# Import both resources and templates
|
||||
prefix = "test"
|
||||
target_manager.import_resources(source_manager, prefix)
|
||||
target_manager.import_templates(source_manager, prefix)
|
||||
|
||||
# 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
|
||||
0
tests/server.py
Normal file
0
tests/server.py
Normal file
184
tests/server/test_mount.py
Normal file
184
tests/server/test_mount.py
Normal file
|
|
@ -0,0 +1,184 @@
|
|||
from fastmcp.server.server import FastMCP
|
||||
|
||||
|
||||
async def test_mount_basic_functionality():
|
||||
"""Test that the mount method properly imports tools and other resources."""
|
||||
# Create main app and sub-app
|
||||
main_app = FastMCP("MainApp")
|
||||
sub_app = FastMCP("SubApp")
|
||||
|
||||
# Add a tool to the sub-app
|
||||
@sub_app.tool()
|
||||
def sub_tool() -> str:
|
||||
return "This is from the sub app"
|
||||
|
||||
# Mount the sub-app to the main app
|
||||
main_app.mount("sub", sub_app)
|
||||
|
||||
# Verify the tool was imported with the prefix
|
||||
assert "sub/sub_tool" in main_app._tool_manager._tools
|
||||
assert "sub_tool" in sub_app._tool_manager._tools
|
||||
|
||||
# Verify the original tool still exists in the sub-app
|
||||
tool = main_app._tool_manager._tools["sub/sub_tool"]
|
||||
assert tool.name == "sub/sub_tool"
|
||||
assert callable(tool.fn)
|
||||
|
||||
|
||||
async def test_mount_multiple_apps():
|
||||
"""Test mounting multiple apps to a main app."""
|
||||
# Create main app and multiple sub-apps
|
||||
main_app = FastMCP("MainApp")
|
||||
weather_app = FastMCP("WeatherApp")
|
||||
news_app = FastMCP("NewsApp")
|
||||
|
||||
# Add tools to each sub-app
|
||||
@weather_app.tool()
|
||||
def get_forecast() -> str:
|
||||
return "Weather forecast"
|
||||
|
||||
@news_app.tool()
|
||||
def get_headlines() -> str:
|
||||
return "News headlines"
|
||||
|
||||
# Mount both sub-apps to the main app
|
||||
main_app.mount("weather", weather_app)
|
||||
main_app.mount("news", news_app)
|
||||
|
||||
# Verify tools were imported with the correct prefixes
|
||||
assert "weather/get_forecast" in main_app._tool_manager._tools
|
||||
assert "news/get_headlines" in main_app._tool_manager._tools
|
||||
|
||||
|
||||
async def test_mount_combines_tools():
|
||||
"""Test that mounting preserves existing tools with the same prefix."""
|
||||
# Create apps
|
||||
main_app = FastMCP("MainApp")
|
||||
first_app = FastMCP("FirstApp")
|
||||
second_app = FastMCP("SecondApp")
|
||||
|
||||
# Add tools to each sub-app
|
||||
@first_app.tool()
|
||||
def first_tool() -> str:
|
||||
return "First app tool"
|
||||
|
||||
@second_app.tool()
|
||||
def second_tool() -> str:
|
||||
return "Second app tool"
|
||||
|
||||
# Mount first app
|
||||
main_app.mount("api", first_app)
|
||||
assert "api/first_tool" in main_app._tool_manager._tools
|
||||
|
||||
# Mount second app to same prefix
|
||||
main_app.mount("api", second_app)
|
||||
|
||||
# Verify second tool is there
|
||||
assert "api/second_tool" in main_app._tool_manager._tools
|
||||
|
||||
# Tools from both mounts are combined
|
||||
assert "api/first_tool" in main_app._tool_manager._tools
|
||||
|
||||
|
||||
async def test_mount_with_resources():
|
||||
"""Test mounting with resources."""
|
||||
# Create apps
|
||||
main_app = FastMCP("MainApp")
|
||||
data_app = FastMCP("DataApp")
|
||||
|
||||
# Add a resource to the data app
|
||||
@data_app.resource(uri="data://users")
|
||||
async def get_users():
|
||||
return ["user1", "user2"]
|
||||
|
||||
# Mount the data app
|
||||
main_app.mount("data", data_app)
|
||||
|
||||
# Verify the resource was imported with the prefix
|
||||
assert "data+data://users" in main_app._resource_manager._resources
|
||||
|
||||
|
||||
async def test_mount_with_resource_templates():
|
||||
"""Test mounting with resource templates."""
|
||||
# Create apps
|
||||
main_app = FastMCP("MainApp")
|
||||
user_app = FastMCP("UserApp")
|
||||
|
||||
# Add a resource template to the user app
|
||||
@user_app.resource(uri="users://{user_id}/profile")
|
||||
def get_user_profile(user_id: str) -> dict:
|
||||
return {"id": user_id, "name": f"User {user_id}"}
|
||||
|
||||
# Mount the user app
|
||||
main_app.mount("api", user_app)
|
||||
|
||||
# Verify the template was imported with the prefix
|
||||
assert "api+users://{user_id}/profile" in main_app._resource_manager._templates
|
||||
|
||||
|
||||
async def test_mount_with_prompts():
|
||||
"""Test mounting with prompts."""
|
||||
# Create apps
|
||||
main_app = FastMCP("MainApp")
|
||||
assistant_app = FastMCP("AssistantApp")
|
||||
|
||||
# Add a prompt to the assistant app
|
||||
@assistant_app.prompt()
|
||||
def greeting(name: str) -> str:
|
||||
return f"Hello, {name}!"
|
||||
|
||||
# Mount the assistant app
|
||||
main_app.mount("assistant", assistant_app)
|
||||
|
||||
# Verify the prompt was imported with the prefix
|
||||
assert "assistant/greeting" in main_app._prompt_manager._prompts
|
||||
|
||||
|
||||
async def test_mount_multiple_resource_templates():
|
||||
"""Test mounting multiple apps with resource templates."""
|
||||
# Create apps
|
||||
main_app = FastMCP("MainApp")
|
||||
weather_app = FastMCP("WeatherApp")
|
||||
news_app = FastMCP("NewsApp")
|
||||
|
||||
# Add templates to each app
|
||||
@weather_app.resource(uri="weather://{city}")
|
||||
def get_weather(city: str) -> str:
|
||||
return f"Weather for {city}"
|
||||
|
||||
@news_app.resource(uri="news://{category}")
|
||||
def get_news(category: str) -> str:
|
||||
return f"News for {category}"
|
||||
|
||||
# Mount both apps
|
||||
main_app.mount("data", weather_app)
|
||||
main_app.mount("content", news_app)
|
||||
|
||||
# Verify templates were imported with correct prefixes
|
||||
assert "data+weather://{city}" in main_app._resource_manager._templates
|
||||
assert "content+news://{category}" in main_app._resource_manager._templates
|
||||
|
||||
|
||||
async def test_mount_multiple_prompts():
|
||||
"""Test mounting multiple apps with prompts."""
|
||||
# Create apps
|
||||
main_app = FastMCP("MainApp")
|
||||
python_app = FastMCP("PythonApp")
|
||||
sql_app = FastMCP("SQLApp")
|
||||
|
||||
# Add prompts to each app
|
||||
@python_app.prompt()
|
||||
def review_python(code: str) -> str:
|
||||
return f"Reviewing Python code:\n{code}"
|
||||
|
||||
@sql_app.prompt()
|
||||
def explain_sql(query: str) -> str:
|
||||
return f"Explaining SQL query:\n{query}"
|
||||
|
||||
# Mount both apps
|
||||
main_app.mount("python", python_app)
|
||||
main_app.mount("sql", sql_app)
|
||||
|
||||
# Verify prompts were imported with correct prefixes
|
||||
assert "python/review_python" in main_app._prompt_manager._prompts
|
||||
assert "sql/explain_sql" in main_app._prompt_manager._prompts
|
||||
0
tests/tools/__init__.py
Normal file
0
tests/tools/__init__.py
Normal file
101
tests/tools/test_tool_manager.py
Normal file
101
tests/tools/test_tool_manager.py
Normal file
|
|
@ -0,0 +1,101 @@
|
|||
from fastmcp.tools.tool_manager import ToolManager
|
||||
|
||||
|
||||
def test_import_tools():
|
||||
"""Test importing tools from one manager to another with a prefix."""
|
||||
# Setup source manager with tools
|
||||
source_manager = ToolManager()
|
||||
|
||||
# Create some test tools
|
||||
def tool1_fn():
|
||||
return "Tool 1 result"
|
||||
|
||||
def tool2_fn():
|
||||
return "Tool 2 result"
|
||||
|
||||
# Add tools to source manager
|
||||
source_manager.add_tool(tool1_fn, name="get_data", description="Get some data")
|
||||
source_manager.add_tool(
|
||||
tool2_fn, name="process_data", description="Process the data"
|
||||
)
|
||||
|
||||
# Create target manager
|
||||
target_manager = ToolManager()
|
||||
|
||||
# Import tools from source to target
|
||||
prefix = "source"
|
||||
target_manager.import_tools(source_manager, prefix)
|
||||
|
||||
# Verify tools were imported with prefixes
|
||||
assert "source/get_data" in target_manager._tools
|
||||
assert "source/process_data" in target_manager._tools
|
||||
|
||||
# Verify the original tools still exist in source manager
|
||||
assert "get_data" in source_manager._tools
|
||||
assert "process_data" in source_manager._tools
|
||||
|
||||
# Verify the imported tools have the correct descriptions
|
||||
assert target_manager._tools["source/get_data"].description == "Get some data"
|
||||
assert (
|
||||
target_manager._tools["source/process_data"].description == "Process the data"
|
||||
)
|
||||
|
||||
# Verify the tool functions were properly copied
|
||||
# We can't directly compare functions, so we'll check their __name__ attribute
|
||||
assert target_manager._tools["source/get_data"].fn.__name__ == tool1_fn.__name__
|
||||
assert target_manager._tools["source/process_data"].fn.__name__ == tool2_fn.__name__
|
||||
|
||||
|
||||
def test_tool_duplicate_behavior():
|
||||
"""Test the behavior when importing tools with duplicate names."""
|
||||
# Setup source and target managers
|
||||
source_manager = ToolManager()
|
||||
target_manager = ToolManager()
|
||||
|
||||
# Add the same tool name to both managers
|
||||
def source_fn():
|
||||
return "Source result"
|
||||
|
||||
def target_fn():
|
||||
return "Target result"
|
||||
|
||||
source_manager.add_tool(source_fn, name="common_tool")
|
||||
target_manager.add_tool(
|
||||
target_fn, name="source/common_tool"
|
||||
) # Pre-create with the prefixed name
|
||||
|
||||
# Import tools from source to target
|
||||
target_manager.import_tools(source_manager, "source")
|
||||
|
||||
# The original tool in the target manager is replaced by the imported one
|
||||
assert target_manager._tools["source/common_tool"].fn.__name__ == source_fn.__name__
|
||||
|
||||
|
||||
def test_import_tools_with_multiple_prefixes():
|
||||
"""Test importing tools from multiple managers with different prefixes."""
|
||||
# Setup source managers
|
||||
weather_manager = ToolManager()
|
||||
news_manager = ToolManager()
|
||||
|
||||
# Add tools to source managers
|
||||
def forecast_fn():
|
||||
return "Weather forecast"
|
||||
|
||||
def headlines_fn():
|
||||
return "News headlines"
|
||||
|
||||
weather_manager.add_tool(forecast_fn, name="forecast")
|
||||
news_manager.add_tool(headlines_fn, name="headlines")
|
||||
|
||||
# Create target manager and import from both sources
|
||||
main_manager = ToolManager()
|
||||
main_manager.import_tools(weather_manager, "weather")
|
||||
main_manager.import_tools(news_manager, "news")
|
||||
|
||||
# Verify tools were imported with correct prefixes
|
||||
assert "weather/forecast" in main_manager._tools
|
||||
assert "news/headlines" in main_manager._tools
|
||||
|
||||
# Verify the tools are accessible and functioning
|
||||
assert main_manager._tools["weather/forecast"].fn.__name__ == forecast_fn.__name__
|
||||
assert main_manager._tools["news/headlines"].fn.__name__ == headlines_fn.__name__
|
||||
104
tests/tools/tool_manager.py
Normal file
104
tests/tools/tool_manager.py
Normal file
|
|
@ -0,0 +1,104 @@
|
|||
from fastmcp.tools.tool_manager import ToolManager
|
||||
|
||||
|
||||
def test_import_tools():
|
||||
"""Test importing tools from one manager to another with a prefix."""
|
||||
# Setup source manager with tools
|
||||
source_manager = ToolManager()
|
||||
|
||||
# Create some test tools
|
||||
def tool1_fn():
|
||||
return "Tool 1 result"
|
||||
|
||||
def tool2_fn():
|
||||
return "Tool 2 result"
|
||||
|
||||
# Add tools to source manager
|
||||
source_manager.add_tool(tool1_fn, name="get_data", description="Get some data")
|
||||
source_manager.add_tool(
|
||||
tool2_fn, name="process_data", description="Process the data"
|
||||
)
|
||||
|
||||
# Create target manager
|
||||
target_manager = ToolManager()
|
||||
|
||||
# Import tools from source to target
|
||||
prefix = "source"
|
||||
target_manager.import_tools(source_manager, prefix)
|
||||
|
||||
# Verify tools were imported with prefixes
|
||||
assert "source:get_data" in target_manager._tools
|
||||
assert "source:process_data" in target_manager._tools
|
||||
|
||||
# Verify the original tools still exist in source manager
|
||||
assert "get_data" in source_manager._tools
|
||||
assert "process_data" in source_manager._tools
|
||||
|
||||
# Verify the imported tools have the correct descriptions
|
||||
assert target_manager._tools["source:get_data"].description == "Get some data"
|
||||
assert (
|
||||
target_manager._tools["source:process_data"].description == "Process the data"
|
||||
)
|
||||
|
||||
# Verify the tool functions were properly copied
|
||||
# We can't directly compare functions, so we'll check their __name__ attribute
|
||||
assert target_manager._tools["source:get_data"].fn.__name__ == tool1_fn.__name__
|
||||
assert target_manager._tools["source:process_data"].fn.__name__ == tool2_fn.__name__
|
||||
|
||||
|
||||
def test_import_tools_duplicate_warning(caplog):
|
||||
"""Test that warning is logged when importing a tool with a name that already exists."""
|
||||
# Setup source and target managers
|
||||
source_manager = ToolManager()
|
||||
target_manager = ToolManager(warn_on_duplicate_tools=True)
|
||||
|
||||
# Add the same tool name to both managers
|
||||
def source_fn():
|
||||
return "Source result"
|
||||
|
||||
def target_fn():
|
||||
return "Target result"
|
||||
|
||||
source_manager.add_tool(source_fn, name="common_tool")
|
||||
target_manager.add_tool(
|
||||
target_fn, name="source:common_tool"
|
||||
) # Pre-create with the prefixed name
|
||||
|
||||
# Import tools from source to target
|
||||
target_manager.import_tools(source_manager, "source")
|
||||
|
||||
# Verify a warning was logged
|
||||
assert any("already exists" in record.message for record in caplog.records)
|
||||
|
||||
# The original tool in the target manager should be preserved
|
||||
assert target_manager._tools["source:common_tool"].fn.__name__ == target_fn.__name__
|
||||
|
||||
|
||||
def test_import_tools_with_multiple_prefixes():
|
||||
"""Test importing tools from multiple managers with different prefixes."""
|
||||
# Setup source managers
|
||||
weather_manager = ToolManager()
|
||||
news_manager = ToolManager()
|
||||
|
||||
# Add tools to source managers
|
||||
def forecast_fn():
|
||||
return "Weather forecast"
|
||||
|
||||
def headlines_fn():
|
||||
return "News headlines"
|
||||
|
||||
weather_manager.add_tool(forecast_fn, name="forecast")
|
||||
news_manager.add_tool(headlines_fn, name="headlines")
|
||||
|
||||
# Create target manager and import from both sources
|
||||
main_manager = ToolManager()
|
||||
main_manager.import_tools(weather_manager, "weather")
|
||||
main_manager.import_tools(news_manager, "news")
|
||||
|
||||
# Verify tools were imported with correct prefixes
|
||||
assert "weather:forecast" in main_manager._tools
|
||||
assert "news:headlines" in main_manager._tools
|
||||
|
||||
# Verify the tools are accessible and functioning
|
||||
assert main_manager._tools["weather:forecast"].fn.__name__ == forecast_fn.__name__
|
||||
assert main_manager._tools["news:headlines"].fn.__name__ == headlines_fn.__name__
|
||||
Loading…
Add table
Add a link
Reference in a new issue