mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-21 13:04:18 +02:00
Update tool manager
This commit is contained in:
Jeremiah Lowin
parent
cafa11d3ad
commit
607125abf8
3 changed files with 182 additions and 171 deletions
|
|
@ -434,10 +434,10 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
logger.debug("Handler called: list_tools")
|
||||
|
||||
with fastmcp.server.context.Context(fastmcp=self):
|
||||
tools = await self._middleware_list_tools()
|
||||
tools = await self._list_tools()
|
||||
return [tool.to_mcp_tool(name=tool.key) for tool in tools]
|
||||
|
||||
async def _middleware_list_tools(self) -> list[Tool]:
|
||||
async def _list_tools(self) -> list[Tool]:
|
||||
"""
|
||||
List all available tools, in the format expected by the low-level MCP
|
||||
server.
|
||||
|
|
@ -447,7 +447,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
async def _handler(
|
||||
context: MiddlewareContext[mcp.types.ListToolsRequest],
|
||||
) -> list[Tool]:
|
||||
tools = await self._list_tools()
|
||||
tools = await self._tool_manager.list_tools()
|
||||
|
||||
mcp_tools: list[Tool] = []
|
||||
for tool in tools:
|
||||
|
|
@ -469,39 +469,6 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
# Apply the middleware chain.
|
||||
return await self._apply_middleware(mw_context, _handler)
|
||||
|
||||
async def _list_tools(self, apply_middleware: bool = True) -> list[Tool]:
|
||||
"""
|
||||
List all available tools.
|
||||
"""
|
||||
|
||||
if (tools := self._cache.get("tools")) is self._cache.NOT_FOUND:
|
||||
tools: dict[str, Tool] = {}
|
||||
|
||||
# iterate such that new mounts overwrite older ones
|
||||
for mounted_server in self._mounted_servers:
|
||||
try:
|
||||
if apply_middleware:
|
||||
server_tools = (
|
||||
await mounted_server.server._middleware_list_tools()
|
||||
)
|
||||
else:
|
||||
server_tools = await mounted_server.server._list_tools()
|
||||
# Apply prefix to each tool key if prefix exists and is not empty
|
||||
if mounted_server.prefix:
|
||||
for tool in server_tools:
|
||||
tool = tool.with_key(f"{mounted_server.prefix}_{tool.key}")
|
||||
tools[tool.key] = tool
|
||||
else:
|
||||
tools.update({tool.key: tool for tool in server_tools})
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to get tools from mounted server '{mounted_server.prefix}': {e}"
|
||||
)
|
||||
continue
|
||||
tools.update(self._tool_manager.get_tools())
|
||||
self._cache.set("tools", tools)
|
||||
return list(tools.values())
|
||||
|
||||
async def _mcp_list_resources(self) -> list[MCPResource]:
|
||||
logger.debug("Handler called: list_resources")
|
||||
|
||||
|
|
@ -580,7 +547,11 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
f"Failed to get resources from mounted server '{mounted_server.prefix}': {e}"
|
||||
)
|
||||
continue
|
||||
resources.update(self._resource_manager.get_resources())
|
||||
(
|
||||
local_resources,
|
||||
_,
|
||||
) = await self._resource_manager.get_resources_and_templates()
|
||||
resources.update(local_resources)
|
||||
self._cache.set("resources", resources)
|
||||
return list(resources.values())
|
||||
|
||||
|
|
@ -668,7 +639,11 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
f"'{mounted_server.prefix}': {e}"
|
||||
)
|
||||
continue
|
||||
templates.update(self._resource_manager.get_templates())
|
||||
(
|
||||
_,
|
||||
local_templates,
|
||||
) = await self._resource_manager.get_resources_and_templates()
|
||||
templates.update(local_templates)
|
||||
self._cache.set("resource_templates", templates)
|
||||
return list(templates.values())
|
||||
|
||||
|
|
@ -744,7 +719,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
f"Failed to get prompts from mounted server '{mounted_server.prefix}': {e}"
|
||||
)
|
||||
continue
|
||||
prompts.update(self._prompt_manager.get_prompts())
|
||||
prompts.update(await self._prompt_manager.get_prompts())
|
||||
self._cache.set("prompts", prompts)
|
||||
return list(prompts.values())
|
||||
|
||||
|
|
@ -767,27 +742,26 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
|
||||
with fastmcp.server.context.Context(fastmcp=self):
|
||||
try:
|
||||
return await self._middleware_call_tool(key, arguments)
|
||||
return await self._call_tool(key, arguments)
|
||||
except DisabledError:
|
||||
raise NotFoundError(f"Unknown tool: {key}")
|
||||
except NotFoundError:
|
||||
raise NotFoundError(f"Unknown tool: {key}")
|
||||
|
||||
async def _middleware_call_tool(
|
||||
self,
|
||||
key: str,
|
||||
arguments: dict[str, Any],
|
||||
) -> list[MCPContent]:
|
||||
async def _call_tool(self, key: str, arguments: dict[str, Any]) -> list[MCPContent]:
|
||||
"""
|
||||
Call a tool with middleware.
|
||||
Applies this server's middleware and delegates the filtered call to the manager.
|
||||
"""
|
||||
|
||||
async def _handler(
|
||||
context: MiddlewareContext[mcp.types.CallToolRequestParams],
|
||||
) -> list[MCPContent]:
|
||||
return await self._call_tool(
|
||||
key=context.message.name,
|
||||
arguments=context.message.arguments or {},
|
||||
tool = await self._tool_manager.get_tool(context.message.name)
|
||||
if not self._should_enable_component(tool):
|
||||
raise NotFoundError(f"Unknown tool: {context.message.name!r}")
|
||||
|
||||
return await self._tool_manager.call_tool(
|
||||
key=context.message.name, arguments=context.message.arguments or {}
|
||||
)
|
||||
|
||||
mw_context = MiddlewareContext(
|
||||
|
|
@ -799,46 +773,6 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
)
|
||||
return await self._apply_middleware(mw_context, _handler)
|
||||
|
||||
async def _call_tool(self, key: str, arguments: dict[str, Any]) -> list[MCPContent]:
|
||||
"""
|
||||
Call a tool with raw MCP arguments. FastMCP subclasses should override
|
||||
this method, not _mcp_call_tool.
|
||||
|
||||
Args:
|
||||
key: The name of the tool to call arguments: Arguments to pass to
|
||||
the tool
|
||||
|
||||
Returns:
|
||||
List of MCP Content objects containing the tool results
|
||||
"""
|
||||
|
||||
# Get tool, checking first from our tools, then from the mounted servers
|
||||
if self._tool_manager.has_tool(key):
|
||||
tool = self._tool_manager.get_tool(key)
|
||||
if not self._should_enable_component(tool):
|
||||
raise DisabledError(f"Tool {key!r} is disabled")
|
||||
return await self._tool_manager.call_tool(key, arguments)
|
||||
|
||||
# Check mounted servers to see if they have the tool
|
||||
# iterate such that new mounts take precedence over older ones
|
||||
for mounted_server in reversed(self._mounted_servers):
|
||||
tool_key = key
|
||||
try:
|
||||
# If server has a prefix, check if key matches and strip prefix
|
||||
if mounted_server.prefix:
|
||||
if tool_key.startswith(f"{mounted_server.prefix}_"):
|
||||
tool_key = tool_key.removeprefix(f"{mounted_server.prefix}_")
|
||||
else:
|
||||
continue
|
||||
return await mounted_server.server._middleware_call_tool(
|
||||
tool_key, arguments
|
||||
)
|
||||
except NotFoundError:
|
||||
# Tool not found on this server, try the next one
|
||||
continue
|
||||
|
||||
raise NotFoundError(f"Unknown tool: {key!r}")
|
||||
|
||||
async def _mcp_read_resource(self, uri: AnyUrl | str) -> list[ReadResourceContents]:
|
||||
"""
|
||||
Handle MCP 'readResource' requests.
|
||||
|
|
@ -1853,11 +1787,12 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
if as_proxy and not isinstance(server, FastMCPProxy):
|
||||
server = FastMCPProxy(Client(transport=FastMCPTransport(server)))
|
||||
|
||||
mounted_server = MountedServer(
|
||||
server=server,
|
||||
prefix=prefix,
|
||||
)
|
||||
self._mounted_servers.append(mounted_server)
|
||||
# Delegate mounting to all three managers
|
||||
mounted_server = MountedServer(prefix=prefix, server=server)
|
||||
self._tool_manager.mount(mounted_server)
|
||||
self._resource_manager.mount(mounted_server)
|
||||
self._prompt_manager.mount(mounted_server)
|
||||
|
||||
self._cache.clear()
|
||||
|
||||
async def import_server(
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
from __future__ import annotations as _annotations
|
||||
from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from typing import TYPE_CHECKING, Any, Literal
|
||||
|
||||
from mcp.types import ToolAnnotations
|
||||
|
||||
|
|
@ -14,7 +14,7 @@ from fastmcp.utilities.logging import get_logger
|
|||
from fastmcp.utilities.types import MCPContent
|
||||
|
||||
if TYPE_CHECKING:
|
||||
pass
|
||||
from fastmcp.server.server import MountedServer
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
|
@ -28,6 +28,7 @@ class ToolManager:
|
|||
mask_error_details: bool | None = None,
|
||||
):
|
||||
self._tools: dict[str, Tool] = {}
|
||||
self._mounted_sources: list[MountedServer] = []
|
||||
self.mask_error_details = mask_error_details or settings.mask_error_details
|
||||
|
||||
# Default to "warn" if None is provided
|
||||
|
|
@ -42,23 +43,78 @@ class ToolManager:
|
|||
|
||||
self.duplicate_behavior = duplicate_behavior
|
||||
|
||||
def has_tool(self, key: str) -> bool:
|
||||
def mount(self, server: MountedServer) -> None:
|
||||
"""Adds a mounted server as a source for tools."""
|
||||
self._mounted_sources.append(server)
|
||||
|
||||
async def _load_tools(
|
||||
self, *, mode: Literal["inventory", "protocol"]
|
||||
) -> dict[str, Tool]:
|
||||
"""
|
||||
The single, consolidated recursive method for fetching tools. The 'mode'
|
||||
parameter determines the communication path.
|
||||
|
||||
- mode="inventory": Manager-to-manager path for complete, unfiltered inventory
|
||||
- mode="protocol": Server-to-server path for filtered MCP requests
|
||||
"""
|
||||
all_tools: dict[str, Tool] = {}
|
||||
|
||||
for mounted in self._mounted_sources:
|
||||
try:
|
||||
if mode == "protocol":
|
||||
# PATH 2: Use the server-to-server filtered path
|
||||
child_results = await mounted.server._list_tools()
|
||||
else: # mode == "inventory"
|
||||
# PATH 1: Use the manager-to-manager unfiltered path
|
||||
child_results = await mounted.server._tool_manager.get_tools()
|
||||
|
||||
# The combination logic is the same for both paths
|
||||
child_dict = (
|
||||
{t.key: t for t in child_results}
|
||||
if isinstance(child_results, list)
|
||||
else child_results
|
||||
)
|
||||
if mounted.prefix:
|
||||
for tool in child_dict.values():
|
||||
prefixed_tool = tool.with_key(f"{mounted.prefix}_{tool.key}")
|
||||
all_tools[prefixed_tool.key] = prefixed_tool
|
||||
else:
|
||||
all_tools.update(child_dict)
|
||||
except Exception as e:
|
||||
# Skip failed mounts silently, matches existing behavior
|
||||
logger.warning(
|
||||
f"Failed to get tools from mounted server '{mounted.prefix}': {e}"
|
||||
)
|
||||
continue
|
||||
|
||||
# Finally, add local tools, which always take precedence
|
||||
all_tools.update(self._tools)
|
||||
return all_tools
|
||||
|
||||
async def has_tool(self, key: str) -> bool:
|
||||
"""Check if a tool exists."""
|
||||
return key in self._tools
|
||||
tools = await self.get_tools()
|
||||
return key in tools
|
||||
|
||||
def get_tool(self, key: str) -> Tool:
|
||||
async def get_tool(self, key: str) -> Tool:
|
||||
"""Get tool by key."""
|
||||
if key in self._tools:
|
||||
return self._tools[key]
|
||||
raise NotFoundError(f"Unknown tool: {key}")
|
||||
tools = await self.get_tools()
|
||||
if key in tools:
|
||||
return tools[key]
|
||||
raise NotFoundError(f"Tool {key!r} not found")
|
||||
|
||||
def get_tools(self) -> dict[str, Tool]:
|
||||
"""Get all registered tools, indexed by registered key."""
|
||||
return self._tools
|
||||
async def get_tools(self) -> dict[str, Tool]:
|
||||
"""
|
||||
Gets the complete, unfiltered inventory of all tools.
|
||||
"""
|
||||
return await self._load_tools(mode="inventory")
|
||||
|
||||
def list_tools(self) -> list[Tool]:
|
||||
"""List all registered tools."""
|
||||
return list(self.get_tools().values())
|
||||
async def list_tools(self) -> list[Tool]:
|
||||
"""
|
||||
Lists all tools, applying protocol filtering.
|
||||
"""
|
||||
tools_dict = await self._load_tools(mode="protocol")
|
||||
return list(tools_dict.values())
|
||||
|
||||
def add_tool_from_fn(
|
||||
self,
|
||||
|
|
@ -119,28 +175,44 @@ class ToolManager:
|
|||
if key in self._tools:
|
||||
del self._tools[key]
|
||||
else:
|
||||
raise NotFoundError(f"Unknown tool: {key}")
|
||||
raise NotFoundError(f"Tool {key!r} not found")
|
||||
|
||||
async def call_tool(self, key: str, arguments: dict[str, Any]) -> list[MCPContent]:
|
||||
"""Call a tool by name with arguments."""
|
||||
tool = self.get_tool(key)
|
||||
if not tool:
|
||||
raise NotFoundError(f"Unknown tool: {key}")
|
||||
"""
|
||||
Internal API for servers: Finds and calls a tool, respecting the
|
||||
filtered protocol path.
|
||||
"""
|
||||
# 1. Check local tools first. The server will have already applied its filter.
|
||||
if key in self._tools:
|
||||
tool = await self.get_tool(key)
|
||||
if not tool:
|
||||
raise NotFoundError(f"Tool {key!r} not found")
|
||||
|
||||
try:
|
||||
return await tool.run(arguments)
|
||||
try:
|
||||
return await tool.run(arguments)
|
||||
|
||||
# raise ToolErrors as-is
|
||||
except ToolError as e:
|
||||
logger.exception(f"Error calling tool {key!r}: {e}")
|
||||
raise e
|
||||
# raise ToolErrors as-is
|
||||
except ToolError as e:
|
||||
logger.exception(f"Error calling tool {key!r}: {e}")
|
||||
raise e
|
||||
|
||||
# Handle other exceptions
|
||||
except Exception as e:
|
||||
logger.exception(f"Error calling tool {key!r}: {e}")
|
||||
if self.mask_error_details:
|
||||
# Mask internal details
|
||||
raise ToolError(f"Error calling tool {key!r}") from e
|
||||
else:
|
||||
# Include original error details
|
||||
raise ToolError(f"Error calling tool {key!r}: {e}") from e
|
||||
# Handle other exceptions
|
||||
except Exception as e:
|
||||
logger.exception(f"Error calling tool {key!r}: {e}")
|
||||
if self.mask_error_details:
|
||||
# Mask internal details
|
||||
raise ToolError(f"Error calling tool {key!r}") from e
|
||||
else:
|
||||
# Include original error details
|
||||
raise ToolError(f"Error calling tool {key!r}: {e}") from e
|
||||
|
||||
# 2. Check mounted servers using the filtered protocol path.
|
||||
for mounted in reversed(self._mounted_sources):
|
||||
if mounted.prefix and key.startswith(f"{mounted.prefix}_"):
|
||||
key_on_child = key.removeprefix(f"{mounted.prefix}_")
|
||||
try:
|
||||
return await mounted.server._call_tool(key_on_child, arguments)
|
||||
except NotFoundError:
|
||||
continue
|
||||
|
||||
raise NotFoundError(f"Tool {key!r} not found.")
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ from fastmcp.utilities.types import Image
|
|||
|
||||
|
||||
class TestAddTools:
|
||||
def test_basic_function(self):
|
||||
async def test_basic_function(self):
|
||||
"""Test registering and running a basic function."""
|
||||
|
||||
def add(a: int, b: int) -> int:
|
||||
|
|
@ -28,7 +28,7 @@ class TestAddTools:
|
|||
tool = Tool.from_function(add)
|
||||
manager.add_tool(tool)
|
||||
|
||||
tool = manager.get_tool("add")
|
||||
tool = await manager.get_tool("add")
|
||||
assert tool is not None
|
||||
assert tool.name == "add"
|
||||
assert tool.description == "Add two numbers."
|
||||
|
|
@ -46,13 +46,13 @@ class TestAddTools:
|
|||
tool = Tool.from_function(fetch_data)
|
||||
manager.add_tool(tool)
|
||||
|
||||
tool = manager.get_tool("fetch_data")
|
||||
tool = await manager.get_tool("fetch_data")
|
||||
assert tool is not None
|
||||
assert tool.name == "fetch_data"
|
||||
assert tool.description == "Fetch data from URL."
|
||||
assert tool.parameters["properties"]["url"]["type"] == "string"
|
||||
|
||||
def test_pydantic_model_function(self):
|
||||
async def test_pydantic_model_function(self):
|
||||
"""Test registering a function that takes a Pydantic model."""
|
||||
|
||||
class UserInput(BaseModel):
|
||||
|
|
@ -67,7 +67,7 @@ class TestAddTools:
|
|||
tool = Tool.from_function(create_user)
|
||||
manager.add_tool(tool)
|
||||
|
||||
tool = manager.get_tool("create_user")
|
||||
tool = await manager.get_tool("create_user")
|
||||
assert tool is not None
|
||||
assert tool.name == "create_user"
|
||||
assert tool.description == "Create a new user."
|
||||
|
|
@ -75,7 +75,7 @@ class TestAddTools:
|
|||
assert "age" in tool.parameters["$defs"]["UserInput"]["properties"]
|
||||
assert "flag" in tool.parameters["properties"]
|
||||
|
||||
def test_callable_object(self):
|
||||
async def test_callable_object(self):
|
||||
class Adder:
|
||||
"""Adds two numbers."""
|
||||
|
||||
|
|
@ -87,7 +87,7 @@ class TestAddTools:
|
|||
tool = Tool.from_function(Adder())
|
||||
manager.add_tool(tool)
|
||||
|
||||
tool = manager.get_tool("Adder")
|
||||
tool = await manager.get_tool("Adder")
|
||||
assert tool is not None
|
||||
assert tool.name == "Adder"
|
||||
assert tool.description == "Adds two numbers."
|
||||
|
|
@ -95,7 +95,7 @@ class TestAddTools:
|
|||
assert tool.parameters["properties"]["x"]["type"] == "integer"
|
||||
assert tool.parameters["properties"]["y"]["type"] == "integer"
|
||||
|
||||
def test_async_callable_object(self):
|
||||
async def test_async_callable_object(self):
|
||||
class Adder:
|
||||
"""Adds two numbers."""
|
||||
|
||||
|
|
@ -107,7 +107,7 @@ class TestAddTools:
|
|||
tool = Tool.from_function(Adder())
|
||||
manager.add_tool(tool)
|
||||
|
||||
tool = manager.get_tool("Adder")
|
||||
tool = await manager.get_tool("Adder")
|
||||
assert tool is not None
|
||||
assert tool.name == "Adder"
|
||||
assert tool.description == "Adds two numbers."
|
||||
|
|
@ -123,7 +123,7 @@ class TestAddTools:
|
|||
tool = Tool.from_function(image_tool)
|
||||
manager.add_tool(tool)
|
||||
|
||||
tool = manager.get_tool("image_tool")
|
||||
tool = await manager.get_tool("image_tool")
|
||||
result = await tool.run({"data": "test.png"})
|
||||
assert tool.parameters["properties"]["data"]["type"] == "string"
|
||||
assert isinstance(result[0], ImageContent)
|
||||
|
|
@ -148,7 +148,7 @@ class TestAddTools:
|
|||
tool = Tool.from_function(lambda x: x)
|
||||
manager.add_tool(tool)
|
||||
|
||||
def test_remove_tool_successfully(self):
|
||||
async def test_remove_tool_successfully(self):
|
||||
"""Test removing an added tool by key."""
|
||||
manager = ToolManager()
|
||||
|
||||
|
|
@ -157,19 +157,19 @@ class TestAddTools:
|
|||
|
||||
tool = Tool.from_function(add)
|
||||
manager.add_tool(tool)
|
||||
assert manager.get_tool("add") is not None
|
||||
assert await manager.get_tool("add") is not None
|
||||
|
||||
manager.remove_tool("add")
|
||||
with pytest.raises(NotFoundError):
|
||||
manager.get_tool("add")
|
||||
await manager.get_tool("add")
|
||||
|
||||
def test_remove_tool_missing_key(self):
|
||||
"""Test removing a tool that does not exist raises NotFoundError."""
|
||||
manager = ToolManager()
|
||||
with pytest.raises(NotFoundError, match=f"Unknown tool: {'missing'}"):
|
||||
with pytest.raises(NotFoundError, match="Tool 'missing' not found"):
|
||||
manager.remove_tool("missing")
|
||||
|
||||
def test_warn_on_duplicate_tools(self, caplog):
|
||||
async def test_warn_on_duplicate_tools(self, caplog):
|
||||
"""Test warning on duplicate tools."""
|
||||
manager = ToolManager(duplicate_behavior="warn")
|
||||
|
||||
|
|
@ -183,7 +183,7 @@ class TestAddTools:
|
|||
|
||||
assert "Tool already exists: test_tool" in caplog.text
|
||||
# Should have the tool
|
||||
assert manager.get_tool("test_tool") is not None
|
||||
assert await manager.get_tool("test_tool") is not None
|
||||
|
||||
def test_disable_warn_on_duplicate_tools(self, caplog):
|
||||
"""Test disabling warning on duplicate tools."""
|
||||
|
|
@ -213,7 +213,7 @@ class TestAddTools:
|
|||
tool2 = Tool.from_function(test_fn, name="test_tool")
|
||||
manager.add_tool(tool2)
|
||||
|
||||
def test_replace_duplicate_tools(self):
|
||||
async def test_replace_duplicate_tools(self):
|
||||
"""Test replacing duplicate tools."""
|
||||
manager = ToolManager(duplicate_behavior="replace")
|
||||
|
||||
|
|
@ -229,12 +229,12 @@ class TestAddTools:
|
|||
manager.add_tool(result)
|
||||
|
||||
# Should have replaced with the new tool
|
||||
tool = manager.get_tool("test_tool")
|
||||
tool = await manager.get_tool("test_tool")
|
||||
assert tool is not None
|
||||
assert isinstance(tool, FunctionTool)
|
||||
assert tool.fn.__name__ == "replacement_fn"
|
||||
|
||||
def test_ignore_duplicate_tools(self):
|
||||
async def test_ignore_duplicate_tools(self):
|
||||
"""Test ignoring duplicate tools."""
|
||||
manager = ToolManager(duplicate_behavior="ignore")
|
||||
|
||||
|
|
@ -250,7 +250,7 @@ class TestAddTools:
|
|||
manager.add_tool(result)
|
||||
|
||||
# Should keep the original
|
||||
tool = manager.get_tool("test_tool")
|
||||
tool = await manager.get_tool("test_tool")
|
||||
assert tool is not None
|
||||
assert isinstance(tool, FunctionTool)
|
||||
assert tool.fn.__name__ == "original_fn"
|
||||
|
|
@ -262,7 +262,7 @@ class TestAddTools:
|
|||
class TestToolTags:
|
||||
"""Test functionality related to tool tags."""
|
||||
|
||||
def test_add_tool_with_tags(self):
|
||||
async def test_add_tool_with_tags(self):
|
||||
"""Test adding tags to a tool."""
|
||||
|
||||
def example_tool(x: int) -> int:
|
||||
|
|
@ -274,11 +274,11 @@ class TestToolTags:
|
|||
manager.add_tool(tool)
|
||||
|
||||
assert tool.tags == {"math", "utility"}
|
||||
tool = manager.get_tool("example_tool")
|
||||
tool = await manager.get_tool("example_tool")
|
||||
assert tool is not None
|
||||
assert tool.tags == {"math", "utility"}
|
||||
|
||||
def test_add_tool_with_empty_tags(self):
|
||||
async def test_add_tool_with_empty_tags(self):
|
||||
"""Test adding a tool with empty tags set."""
|
||||
|
||||
def example_tool(x: int) -> int:
|
||||
|
|
@ -291,7 +291,7 @@ class TestToolTags:
|
|||
|
||||
assert tool.tags == set()
|
||||
|
||||
def test_add_tool_with_none_tags(self):
|
||||
async def test_add_tool_with_none_tags(self):
|
||||
"""Test adding a tool with None tags."""
|
||||
|
||||
def example_tool(x: int) -> int:
|
||||
|
|
@ -304,7 +304,7 @@ class TestToolTags:
|
|||
|
||||
assert tool.tags == set()
|
||||
|
||||
def test_list_tools_with_tags(self):
|
||||
async def test_list_tools_with_tags(self):
|
||||
"""Test listing tools with specific tags."""
|
||||
|
||||
def math_tool(x: int) -> int:
|
||||
|
|
@ -328,12 +328,16 @@ class TestToolTags:
|
|||
manager.add_tool(tool3)
|
||||
|
||||
# Check if we can filter by tags when listing tools
|
||||
math_tools = [tool for tool in manager.list_tools() if "math" in tool.tags]
|
||||
math_tools = [
|
||||
tool for tool in (await manager.get_tools()).values() if "math" in tool.tags
|
||||
]
|
||||
assert len(math_tools) == 2
|
||||
assert {tool.name for tool in math_tools} == {"math_tool", "mixed_tool"}
|
||||
|
||||
utility_tools = [
|
||||
tool for tool in manager.list_tools() if "utility" in tool.tags
|
||||
tool
|
||||
for tool in (await manager.get_tools()).values()
|
||||
if "utility" in tool.tags
|
||||
]
|
||||
assert len(utility_tools) == 2
|
||||
assert {tool.name for tool in utility_tools} == {"string_tool", "mixed_tool"}
|
||||
|
|
@ -416,7 +420,7 @@ class TestCallTools:
|
|||
|
||||
async def test_call_unknown_tool(self):
|
||||
manager = ToolManager()
|
||||
with pytest.raises(NotFoundError, match="Unknown tool: unknown"):
|
||||
with pytest.raises(NotFoundError, match="Tool 'unknown' not found"):
|
||||
await manager.call_tool("unknown", {"a": 1})
|
||||
|
||||
async def test_call_tool_with_list_int_input(self):
|
||||
|
|
@ -728,7 +732,7 @@ class TestContextHandling:
|
|||
class TestCustomToolNames:
|
||||
"""Test adding tools with custom names that differ from their function names."""
|
||||
|
||||
def test_add_tool_with_custom_name(self):
|
||||
async def test_add_tool_with_custom_name(self):
|
||||
"""Test adding a tool with a custom name parameter using add_tool_from_fn."""
|
||||
|
||||
def original_fn(x: int) -> int:
|
||||
|
|
@ -739,15 +743,15 @@ class TestCustomToolNames:
|
|||
manager.add_tool(tool)
|
||||
|
||||
# The tool is stored under the custom name and its .name is also set to custom_name
|
||||
assert manager.get_tool("custom_name") is not None
|
||||
assert await manager.get_tool("custom_name") is not None
|
||||
assert tool.name == "custom_name"
|
||||
assert isinstance(tool, FunctionTool)
|
||||
assert tool.fn.__name__ == "original_fn"
|
||||
# The tool should not be accessible via its original function name
|
||||
with pytest.raises(NotFoundError, match="Unknown tool: original_fn"):
|
||||
manager.get_tool("original_fn")
|
||||
with pytest.raises(NotFoundError, match="Tool 'original_fn' not found"):
|
||||
await manager.get_tool("original_fn")
|
||||
|
||||
def test_add_tool_object_with_custom_key(self):
|
||||
async 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:
|
||||
|
|
@ -759,13 +763,13 @@ class TestCustomToolNames:
|
|||
# Store it under a different name
|
||||
manager.add_tool(tool, key="proxy_tool")
|
||||
# The tool is accessible under the key
|
||||
stored = manager.get_tool("proxy_tool")
|
||||
stored = await manager.get_tool("proxy_tool")
|
||||
assert stored is not None
|
||||
# But the tool's .name is unchanged
|
||||
assert stored.name == "my_tool"
|
||||
# The tool is not accessible under its original name
|
||||
with pytest.raises(NotFoundError, match="Unknown tool: my_tool"):
|
||||
manager.get_tool("my_tool")
|
||||
with pytest.raises(NotFoundError, match="Tool 'my_tool' not found"):
|
||||
await manager.get_tool("my_tool")
|
||||
|
||||
async def test_call_tool_with_custom_name(self):
|
||||
"""Test calling a tool added with a custom name."""
|
||||
|
|
@ -783,10 +787,10 @@ class TestCustomToolNames:
|
|||
assert result[0].text == "15" # type: ignore[attr-defined]
|
||||
|
||||
# Original name should not be registered
|
||||
with pytest.raises(NotFoundError, match="Unknown tool: multiply"):
|
||||
with pytest.raises(NotFoundError, match="Tool 'multiply' not found"):
|
||||
await manager.call_tool("multiply", {"a": 5, "b": 3})
|
||||
|
||||
def test_replace_tool_keeps_original_name(self):
|
||||
async def test_replace_tool_keeps_original_name(self):
|
||||
"""Test that replacing a tool with "replace" keeps the original name."""
|
||||
|
||||
def original_fn(x: int) -> int:
|
||||
|
|
@ -808,7 +812,7 @@ class TestCustomToolNames:
|
|||
manager.add_tool(replacement_tool)
|
||||
|
||||
# The tool object should have been replaced
|
||||
stored_tool = manager.get_tool("test_tool")
|
||||
stored_tool = await manager.get_tool("test_tool")
|
||||
assert stored_tool is not None
|
||||
assert stored_tool == replacement_tool
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue