From 9e5c5d02109926e3d70b0cfe49d6fe41906be91b Mon Sep 17 00:00:00 2001 From: strawgate Date: Sun, 20 Apr 2025 17:37:29 -0500 Subject: [PATCH] Initial Implementation --- src/contrib/bulk_tool_caller/README.md | 35 +++ src/contrib/bulk_tool_caller/__init__.py | 3 + .../bulk_tool_caller/bulk_tool_caller.py | 131 +++++++++++ src/contrib/bulk_tool_caller/example.py | 17 ++ tests/contrib/test_bulk_tool_caller.py | 221 ++++++++++++++++++ 5 files changed, 407 insertions(+) create mode 100644 src/contrib/bulk_tool_caller/README.md create mode 100644 src/contrib/bulk_tool_caller/__init__.py create mode 100644 src/contrib/bulk_tool_caller/bulk_tool_caller.py create mode 100644 src/contrib/bulk_tool_caller/example.py create mode 100644 tests/contrib/test_bulk_tool_caller.py diff --git a/src/contrib/bulk_tool_caller/README.md b/src/contrib/bulk_tool_caller/README.md new file mode 100644 index 000000000..41daa7780 --- /dev/null +++ b/src/contrib/bulk_tool_caller/README.md @@ -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. \ No newline at end of file diff --git a/src/contrib/bulk_tool_caller/__init__.py b/src/contrib/bulk_tool_caller/__init__.py new file mode 100644 index 000000000..b8f66ed8b --- /dev/null +++ b/src/contrib/bulk_tool_caller/__init__.py @@ -0,0 +1,3 @@ +from .bulk_tool_caller import BulkToolCaller + +__all__ = ["BulkToolCaller"] diff --git a/src/contrib/bulk_tool_caller/bulk_tool_caller.py b/src/contrib/bulk_tool_caller/bulk_tool_caller.py new file mode 100644 index 000000000..aa02c5e61 --- /dev/null +++ b/src/contrib/bulk_tool_caller/bulk_tool_caller.py @@ -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, + ) diff --git a/src/contrib/bulk_tool_caller/example.py b/src/contrib/bulk_tool_caller/example.py new file mode 100644 index 000000000..76170c197 --- /dev/null +++ b/src/contrib/bulk_tool_caller/example.py @@ -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) diff --git a/tests/contrib/test_bulk_tool_caller.py b/tests/contrib/test_bulk_tool_caller.py new file mode 100644 index 000000000..5e5519852 --- /dev/null +++ b/tests/contrib/test_bulk_tool_caller.py @@ -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