Initial Implementation

This commit is contained in:
strawgate 2025-04-20 17:37:29 -05:00
commit 9e5c5d0210
5 changed files with 407 additions and 0 deletions

View file

@ -0,0 +1,35 @@
# Bulk Tool Caller
This module provides the `BulkToolCaller` class, which extends the `MCPMixin` to offer tools for performing multiple tool calls in a single request to a FastMCP server. This can be useful for optimizing interactions with the server by reducing the overhead of individual tool calls.
## Usage
To use the `BulkToolCaller`, see the example [example.py](./example.py) file. The `BulkToolCaller` can be instantiated and then registered with a FastMCP server URL. It provides methods to call multiple tools in bulk, either different tools or the same tool with different arguments.
## Provided Tools
The `BulkToolCaller` provides the following tools:
### `call_tools_bulk`
Calls multiple different tools registered on the MCP server in a single request.
- **Arguments:**
- `tool_calls` (list of `CallToolRequest`): A list of objects, where each object specifies the `tool` name and `arguments` for an individual tool call.
- `continue_on_error` (bool, optional): If `True`, continue executing subsequent tool calls even if a previous one resulted in an error. Defaults to `True`.
- **Returns:**
A list of `CallToolRequestResult` objects, each containing the result (`isError`, `content`) and the original `tool` name and `arguments` for each call.
### `call_tool_bulk`
Calls a single tool registered on the MCP server multiple times with different arguments in a single request.
- **Arguments:**
- `tool` (str): The name of the tool to call.
- `tool_arguments` (list of dict): A list of dictionaries, where each dictionary contains the arguments for an individual run of the tool.
- `continue_on_error` (bool, optional): If `True`, continue executing subsequent tool calls even if a previous one resulted in an error. Defaults to `True`.
- **Returns:**
A list of `CallToolRequestResult` objects, each containing the result (`isError`, `content`) and the original `tool` name and `arguments` for each call.

View file

@ -0,0 +1,3 @@
from .bulk_tool_caller import BulkToolCaller
__all__ = ["BulkToolCaller"]

View file

@ -0,0 +1,131 @@
from typing import Any
from mcp.types import CallToolResult
from pydantic import BaseModel, Field
from contrib.mcp_mixin.mcp_mixin import _DEFAULT_SEPARATOR_TOOL, MCPMixin, mcp_tool
from fastmcp import FastMCP
from fastmcp.client import Client
from fastmcp.client.transports import FastMCPTransport
class CallToolRequest(BaseModel):
"""A class to represent a request to call a tool with specific arguments."""
tool: str = Field(description="The name of the tool to call.")
arguments: dict[str, Any] = Field(
description="A dictionary containing the arguments for the tool call."
)
class CallToolRequestResult(CallToolResult):
"""
A class to represent the result of a bulk tool call.
It extends CallToolResult to include information about the requested tool call.
"""
tool: str = Field(description="The name of the tool that was called.")
arguments: dict[str, Any] = Field(
description="The arguments used for the tool call."
)
@classmethod
def from_call_tool_result(
cls, result: CallToolResult, tool: str, arguments: dict[str, Any]
) -> "CallToolRequestResult":
"""
Create a CallToolRequestResult from a CallToolResult.
"""
return cls(
tool=tool,
arguments=arguments,
isError=result.isError,
content=result.content,
)
class BulkToolCaller(MCPMixin):
"""
A class to provide a "bulk tool call" tool for a FastMCP server
"""
def register_tools(
self,
mcp_server: "FastMCP",
prefix: str | None = None,
separator: str = _DEFAULT_SEPARATOR_TOOL,
) -> None:
"""
Register the tools provided by this class with the given MCP server.
"""
self.connection = FastMCPTransport(mcp_server)
super().register_tools(mcp_server=mcp_server)
@mcp_tool()
async def call_tools_bulk(
self, tool_calls: list[CallToolRequest], continue_on_error: bool = True
) -> list[CallToolRequestResult]:
"""
Call multiple tools registered on this MCP server in a single request. Each call can
be for a different tool and can include different arguments. Useful for speeding up
what would otherwise take several individual tool calls.
"""
results = []
for tool_call in tool_calls:
result = await self._call_tool(tool_call.tool, tool_call.arguments)
results.append(result)
if result.isError and not continue_on_error:
return results
return results
@mcp_tool()
async def call_tool_bulk(
self,
tool: str,
tool_arguments: list[dict[str, str | int | float | bool | None]],
continue_on_error: bool = True,
) -> list[CallToolRequestResult]:
"""
Call a single tool registered on this MCP server multiple times with a single request.
Each call can include different arguments. Useful for speeding up what would otherwise
take several individual tool calls.
Args:
tool: The name of the tool to call.
tool_arguments: A list of dictionaries, where each dictionary contains the arguments for an individual run of the tool.
"""
results = []
for tool_call_arguments in tool_arguments:
result = await self._call_tool(tool, tool_call_arguments)
results.append(result)
if result.isError and not continue_on_error:
return results
return results
async def _call_tool(
self, tool: str, arguments: dict[str, Any]
) -> CallToolRequestResult:
"""
Helper method to call a tool with the provided arguments.
"""
async with Client(self.connection) as client:
result = await client.call_tool(
name=tool, arguments=arguments, _return_raw_result=True
)
return CallToolRequestResult(
tool=tool,
arguments=arguments,
isError=result.isError,
content=result.content,
)

View file

@ -0,0 +1,17 @@
"""Sample code for FastMCP using MCPMixin."""
from contrib.bulk_tool_caller import BulkToolCaller
from fastmcp import FastMCP
mcp = FastMCP()
@mcp.tool()
def echo_tool(text: str) -> str:
"""Echo the input text"""
return text
bulk_tool_caller = BulkToolCaller()
bulk_tool_caller.register_tools(mcp)

View file

@ -0,0 +1,221 @@
from typing import Any
import pytest
from mcp.types import EmbeddedResource, ImageContent, TextContent
from contrib.bulk_tool_caller.bulk_tool_caller import (
BulkToolCaller,
CallToolRequest,
CallToolRequestResult,
)
from fastmcp import FastMCP
ContentType = TextContent | ImageContent | EmbeddedResource
class ToolException(Exception):
"""Custom exception for tool errors."""
pass
async def error_tool(arg1: str) -> dict[str, Any]:
"""A tool that raises an error for testing purposes."""
raise ToolException(f"Error in tool with arg1: {arg1}")
def error_tool_result_factory(arg1: str) -> CallToolRequestResult:
"""Generates the expected error result for error_tool."""
# Mimic the error message format generated by BulkToolCaller when catching ToolException
exception_message = f"Error in tool with arg1: {arg1}"
formatted_error_text = f"Error executing tool error_tool: {exception_message}"
return CallToolRequestResult(
isError=True,
content=[TextContent(text=formatted_error_text, type="text")],
tool="error_tool",
arguments={"arg1": arg1},
)
async def echo_tool(arg1: str) -> str:
"""A simple tool that echoes arguments or raises an error."""
return arg1
def echo_tool_result_factory(arg1: str) -> CallToolRequestResult:
"""A tool that returns a result based on the input arguments."""
return CallToolRequestResult(
isError=False,
content=[TextContent(text=f"{arg1}", type="text")],
tool="echo_tool",
arguments={"arg1": arg1},
)
async def no_return_tool(arg1: str) -> None:
"""A simple tool that echoes arguments or raises an error."""
def no_return_tool_result_factory(arg1: str) -> CallToolRequestResult:
"""A tool that returns a result based on the input arguments."""
return CallToolRequestResult(
isError=False, content=[], tool="no_return_tool", arguments={"arg1": arg1}
)
@pytest.fixture(scope="module")
def live_server_with_tool() -> FastMCP:
"""Fixture to create a FastMCP server instance with the echo_tool registered."""
server = FastMCP()
server.add_tool(echo_tool)
server.add_tool(error_tool)
server.add_tool(no_return_tool)
return server
@pytest.fixture
def bulk_caller_live(live_server_with_tool: FastMCP) -> BulkToolCaller:
"""Fixture to create a BulkToolCaller instance connected to the live server."""
bulk_tool_caller = BulkToolCaller()
bulk_tool_caller.register_tools(live_server_with_tool)
return bulk_tool_caller
ECHO_TOOL_NAME = "echo_tool"
ERROR_TOOL_NAME = "error_tool"
NO_RETURN_TOOL_NAME = "no_return_tool"
@pytest.mark.asyncio
async def test_call_tool_bulk_single_success(bulk_caller_live: BulkToolCaller):
"""Test single successful call via call_tool_bulk using echo_tool."""
tool_arguments = [{"arg1": "value1"}]
expected_result = echo_tool_result_factory(**tool_arguments[0])
results = await bulk_caller_live.call_tool_bulk(ECHO_TOOL_NAME, tool_arguments)
assert len(results) == 1
result = results[0]
assert result == expected_result
@pytest.mark.asyncio
async def test_call_tool_bulk_multiple_success(bulk_caller_live: BulkToolCaller):
"""Test multiple successful calls via call_tool_bulk using echo_tool."""
tool_arguments = [{"arg1": "value1"}, {"arg1": "value2"}]
expected_results = [echo_tool_result_factory(**args) for args in tool_arguments]
results = await bulk_caller_live.call_tool_bulk(ECHO_TOOL_NAME, tool_arguments)
assert len(results) == 2
assert results == expected_results
@pytest.mark.asyncio
async def test_call_tool_bulk_error_stops(bulk_caller_live: BulkToolCaller):
"""Test call_tool_bulk stops on first error using error_tool."""
tool_arguments = [{"arg1": "error_value"}, {"arg1": "value2"}]
expected_result = error_tool_result_factory(**tool_arguments[0])
results = await bulk_caller_live.call_tool_bulk(
ERROR_TOOL_NAME, tool_arguments, continue_on_error=False
)
assert len(results) == 1
result = results[0]
assert result == expected_result
@pytest.mark.asyncio
async def test_call_tool_bulk_error_continues(bulk_caller_live: BulkToolCaller):
"""Test call_tool_bulk continues on error using error_tool and echo_tool."""
tool_arguments = [{"arg1": "error_value"}, {"arg1": "success_value"}]
expected_error_result = error_tool_result_factory(**tool_arguments[0])
expected_success_result = echo_tool_result_factory(**tool_arguments[1])
tool_calls = [
CallToolRequest(tool=ERROR_TOOL_NAME, arguments=tool_arguments[0]),
CallToolRequest(tool=ECHO_TOOL_NAME, arguments=tool_arguments[1]),
]
results = await bulk_caller_live.call_tools_bulk(tool_calls, continue_on_error=True)
assert len(results) == 2
error_result = results[0]
assert error_result == expected_error_result
success_result = results[1]
assert success_result == expected_success_result
@pytest.mark.asyncio
async def test_call_tools_bulk_single_success(bulk_caller_live: BulkToolCaller):
"""Test single successful call via call_tools_bulk using echo_tool."""
tool_calls = [CallToolRequest(tool=ECHO_TOOL_NAME, arguments={"arg1": "value1"})]
expected_result = echo_tool_result_factory(**tool_calls[0].arguments)
results = await bulk_caller_live.call_tools_bulk(tool_calls)
assert len(results) == 1
result = results[0]
assert result == expected_result
@pytest.mark.asyncio
async def test_call_tools_bulk_multiple_success(bulk_caller_live: BulkToolCaller):
"""Test multiple successful calls via call_tools_bulk with different tools."""
tool_calls = [
CallToolRequest(tool=ECHO_TOOL_NAME, arguments={"arg1": "echo_value"}),
CallToolRequest(
tool=NO_RETURN_TOOL_NAME, arguments={"arg1": "no_return_value"}
),
]
expected_results = [
echo_tool_result_factory(**tool_calls[0].arguments),
no_return_tool_result_factory(**tool_calls[1].arguments),
]
results = await bulk_caller_live.call_tools_bulk(tool_calls)
assert len(results) == 2
assert results == expected_results
@pytest.mark.asyncio
async def test_call_tools_bulk_error_stops(bulk_caller_live: BulkToolCaller):
"""Test call_tools_bulk stops on first error using error_tool."""
tool_calls = [
CallToolRequest(tool=ERROR_TOOL_NAME, arguments={"arg1": "error_value"}),
CallToolRequest(tool=ECHO_TOOL_NAME, arguments={"arg1": "skipped_value"}),
]
expected_result = error_tool_result_factory(**tool_calls[0].arguments)
results = await bulk_caller_live.call_tools_bulk(
tool_calls, continue_on_error=False
)
assert len(results) == 1
result = results[0]
assert result == expected_result
@pytest.mark.asyncio
async def test_call_tools_bulk_error_continues(bulk_caller_live: BulkToolCaller):
"""Test call_tools_bulk continues on error using error_tool and echo_tool."""
tool_calls = [
CallToolRequest(tool=ERROR_TOOL_NAME, arguments={"arg1": "error_value"}),
CallToolRequest(tool=ECHO_TOOL_NAME, arguments={"arg1": "success_value"}),
]
expected_error_result = error_tool_result_factory(**tool_calls[0].arguments)
expected_success_result = echo_tool_result_factory(**tool_calls[1].arguments)
results = await bulk_caller_live.call_tools_bulk(tool_calls, continue_on_error=True)
assert len(results) == 2
error_result = results[0]
assert error_result == expected_error_result
success_result = results[1]
assert success_result == expected_success_result