-
Notifications
You must be signed in to change notification settings - Fork 1.6k
[Draft] Convert BulkToolCaller to middleware #2267
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 1 commit
e606d36
009772a
916c8e6
a11e3ab
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,194 @@ | ||
| """Middleware for bulk tool calling functionality.""" | ||
|
|
||
| from typing import Annotated | ||
|
|
||
| from mcp.types import CallToolResult, TextContent | ||
|
|
||
| from fastmcp.server.context import Context | ||
| from fastmcp.server.middleware.bulk_tool_caller_types import ( | ||
| CallToolRequest, | ||
| CallToolRequestResult, | ||
| ) | ||
| from fastmcp.server.middleware.tool_injection import ToolInjectionMiddleware | ||
| from fastmcp.tools.tool import Tool | ||
|
|
||
|
|
||
| async def call_tools_bulk( | ||
| context: Context, | ||
| tool_calls: Annotated[ | ||
| list[CallToolRequest], | ||
| "List of tool calls to execute. Each call can be for a different tool with different arguments.", | ||
| ], | ||
| continue_on_error: Annotated[ | ||
| bool, | ||
| "If True, continue executing remaining tools even if one fails. If False, stop on first error.", | ||
| ] = 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. | ||
|
|
||
| Args: | ||
| context: The request context providing access to the server | ||
| tool_calls: List of tool calls to execute | ||
| continue_on_error: Whether to continue on errors (default: True) | ||
|
|
||
| Returns: | ||
| List of results, one per tool call | ||
| """ | ||
| results = [] | ||
|
|
||
| for tool_call in tool_calls: | ||
| try: | ||
| # Call the tool directly through the tool manager | ||
| tool_result = await context.fastmcp._tool_manager.call_tool( | ||
| key=tool_call.tool, arguments=tool_call.arguments | ||
| ) | ||
|
|
||
| # Convert ToolResult to CallToolResult | ||
| # Don't set isError - it defaults to None for successful calls | ||
| call_tool_result = CallToolResult( | ||
| content=tool_result.content, | ||
| ) | ||
|
|
||
| # Convert to CallToolRequestResult | ||
| results.append( | ||
| CallToolRequestResult.from_call_tool_result( | ||
| call_tool_result, tool_call.tool, tool_call.arguments | ||
| ) | ||
| ) | ||
| except Exception as e: | ||
| # Create error result | ||
| error_message = f"Error calling tool '{tool_call.tool}': {e}" | ||
| results.append( | ||
| CallToolRequestResult( | ||
| tool=tool_call.tool, | ||
| arguments=tool_call.arguments, | ||
| isError=True, | ||
| content=[TextContent(text=error_message, type="text")], | ||
| ) | ||
| ) | ||
|
|
||
| if not continue_on_error: | ||
| break | ||
|
|
||
| return results | ||
|
|
||
|
|
||
| async def call_tool_bulk( | ||
| context: Context, | ||
| tool: Annotated[str, "The name of the tool to call multiple times."], | ||
| tool_arguments: Annotated[ | ||
| list[dict[str, str | int | float | bool | None]], | ||
| "List of argument dictionaries. Each dictionary contains the arguments for one tool invocation.", | ||
| ], | ||
| continue_on_error: Annotated[ | ||
| bool, | ||
| "If True, continue executing remaining calls even if one fails. If False, stop on first error.", | ||
| ] = 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: | ||
| context: The request context providing access to the server | ||
| tool: The name of the tool to call | ||
| tool_arguments: List of argument dictionaries for each invocation | ||
| continue_on_error: Whether to continue on errors (default: True) | ||
|
|
||
| Returns: | ||
| List of results, one per invocation | ||
| """ | ||
| results = [] | ||
|
|
||
| for args in tool_arguments: | ||
| try: | ||
| # Call the tool directly through the tool manager | ||
| tool_result = await context.fastmcp._tool_manager.call_tool( | ||
| key=tool, arguments=args | ||
| ) | ||
|
|
||
| # Convert ToolResult to CallToolResult | ||
| # Don't set isError - it defaults to None for successful calls | ||
| call_tool_result = CallToolResult( | ||
| content=tool_result.content, | ||
| ) | ||
|
||
|
|
||
| # Convert to CallToolRequestResult | ||
| results.append( | ||
| CallToolRequestResult.from_call_tool_result( | ||
| call_tool_result, tool, args | ||
| ) | ||
| ) | ||
| except Exception as e: | ||
| # Create error result | ||
| error_message = f"Error calling tool '{tool}': {e}" | ||
| results.append( | ||
| CallToolRequestResult( | ||
| tool=tool, | ||
| arguments=args, | ||
| isError=True, | ||
| content=[TextContent(text=error_message, type="text")], | ||
| ) | ||
| ) | ||
|
|
||
| if not continue_on_error: | ||
| break | ||
|
|
||
| return results | ||
|
|
||
|
|
||
| class BulkToolCallerMiddleware(ToolInjectionMiddleware): | ||
| """Middleware for injecting bulk tool calling capabilities into the server. | ||
|
|
||
| This middleware adds two tools to the server: | ||
| - call_tools_bulk: Call multiple different tools in a single request | ||
| - call_tool_bulk: Call a single tool multiple times with different arguments | ||
|
|
||
| Example: | ||
| ```python | ||
| from fastmcp import FastMCP | ||
| from fastmcp.server.middleware import BulkToolCallerMiddleware | ||
|
|
||
| mcp = FastMCP("MyServer", middleware=[BulkToolCallerMiddleware()]) | ||
|
|
||
| @mcp.tool | ||
| def greet(name: str) -> str: | ||
| return f"Hello, {name}!" | ||
|
|
||
| @mcp.tool | ||
| def add(a: int, b: int) -> int: | ||
| return a + b | ||
| ``` | ||
|
|
||
| Now clients can use bulk calling: | ||
| ```python | ||
| # Call multiple different tools | ||
| result = await client.call_tool("call_tools_bulk", { | ||
| "tool_calls": [ | ||
| {"tool": "greet", "arguments": {"name": "Alice"}}, | ||
| {"tool": "add", "arguments": {"a": 1, "b": 2}} | ||
| ] | ||
| }) | ||
|
|
||
| # Call same tool multiple times | ||
| result = await client.call_tool("call_tool_bulk", { | ||
| "tool": "greet", | ||
| "tool_arguments": [ | ||
| {"name": "Alice"}, | ||
| {"name": "Bob"} | ||
| ] | ||
| }) | ||
| ``` | ||
| """ | ||
|
|
||
| def __init__(self) -> None: | ||
| """Initialize the bulk tool caller middleware.""" | ||
| tools: list[Tool] = [ | ||
| Tool.from_function(call_tools_bulk), | ||
| Tool.from_function(call_tool_bulk), | ||
| ] | ||
| super().__init__(tools=tools) | ||
| Original file line number | Diff line number | Diff line change | ||||||||
|---|---|---|---|---|---|---|---|---|---|---|
| @@ -0,0 +1,39 @@ | ||||||||||
| """Types for bulk tool caller.""" | ||||||||||
|
|
||||||||||
| from typing import Any | ||||||||||
|
|
||||||||||
| from mcp.types import CallToolResult | ||||||||||
| from pydantic import BaseModel, Field | ||||||||||
|
|
||||||||||
|
|
||||||||||
| 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, | ||||||||||
|
||||||||||
| content=result.content, | |
| content=result.content, | |
| _meta=getattr(result, "_meta", None), | |
| structuredContent=getattr(result, "structuredContent", None), |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The intermediate
CallToolResultcreation is unnecessary. You can passtool_result.contentdirectly toCallToolRequestResult.from_call_tool_resultor create theCallToolRequestResultdirectly fromtool_result, eliminating this temporary object.