a64378956a
Some checks failed
Pre-commit / run (ubuntu-latest) (push) Has been cancelled
Deploy Sphinx documentation to Pages / build_en (ubuntu-latest, 3.10) (push) Has been cancelled
Deploy Sphinx documentation to Pages / build_zh (ubuntu-latest, 3.10) (push) Has been cancelled
Python Unittest Coverage / test (macos-15, 3.10) (push) Has been cancelled
Python Unittest Coverage / test (macos-15, 3.11) (push) Has been cancelled
Python Unittest Coverage / test (macos-15, 3.12) (push) Has been cancelled
Python Unittest Coverage / test (ubuntu-latest, 3.10) (push) Has been cancelled
Python Unittest Coverage / test (ubuntu-latest, 3.11) (push) Has been cancelled
Python Unittest Coverage / test (ubuntu-latest, 3.12) (push) Has been cancelled
Python Unittest Coverage / test (windows-latest, 3.10) (push) Has been cancelled
Python Unittest Coverage / test (windows-latest, 3.11) (push) Has been cancelled
Python Unittest Coverage / test (windows-latest, 3.12) (push) Has been cancelled
404 lines
12 KiB
Python
404 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
.. _middleware:
|
|
|
|
Middleware
|
|
===========================
|
|
|
|
AgentScope provides a flexible middleware system that allows developers to intercept and modify the execution of various operations.
|
|
Currently, middleware support is available for **tool execution** in the ``Toolkit`` class.
|
|
|
|
The middleware system follows an **onion model**, where each middleware wraps around the previous one, forming layers.
|
|
This allows developers to:
|
|
|
|
- Perform **pre-processing** before the operation
|
|
- **Intercept and modify** responses during execution
|
|
- Perform **post-processing** after the operation completes
|
|
- **Skip** the operation execution entirely based on conditions
|
|
|
|
.. tip:: Future versions of AgentScope will expand middleware support to other components such as agents and models.
|
|
|
|
"""
|
|
import asyncio
|
|
from typing import AsyncGenerator, Callable
|
|
|
|
from agentscope.message import TextBlock, ToolUseBlock
|
|
from agentscope.tool import ToolResponse, Toolkit
|
|
|
|
|
|
# %%
|
|
# Tool Execution Middleware
|
|
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
|
#
|
|
# The ``Toolkit`` class supports middleware for tool execution via the ``register_middleware`` method.
|
|
# Each middleware can intercept the tool call and modify the input or output.
|
|
#
|
|
# Middleware Signature
|
|
# ------------------------------
|
|
#
|
|
# A middleware function should have the following signature:
|
|
#
|
|
# .. code-block:: python
|
|
#
|
|
# async def middleware(
|
|
# kwargs: dict,
|
|
# next_handler: Callable,
|
|
# ) -> AsyncGenerator[ToolResponse, None]:
|
|
# # Access parameters from kwargs
|
|
# tool_call = kwargs["tool_call"]
|
|
#
|
|
# # Pre-processing
|
|
# # ...
|
|
#
|
|
# # Call the next middleware or tool function
|
|
# async for response in await next_handler(**kwargs):
|
|
# # Post-processing
|
|
# yield response
|
|
#
|
|
# .. list-table:: Middleware Parameters
|
|
# :header-rows: 1
|
|
#
|
|
# * - Parameter
|
|
# - Type
|
|
# - Description
|
|
# * - ``kwargs``
|
|
# - ``dict``
|
|
# - Context parameters. Currently, includes ``tool_call`` (ToolUseBlock). May include additional parameters in future versions.
|
|
# * - ``next_handler``
|
|
# - ``Callable``
|
|
# - A callable that accepts kwargs dict and returns a coroutine yielding AsyncGenerator of ToolResponse objects
|
|
# * - **Returns**
|
|
# - ``AsyncGenerator[ToolResponse, None]``
|
|
# - An async generator that yields ToolResponse objects
|
|
#
|
|
# Basic Example
|
|
# ------------------------------
|
|
#
|
|
# Here is a simple middleware that logs tool calls:
|
|
#
|
|
|
|
|
|
async def logging_middleware(
|
|
kwargs: dict,
|
|
next_handler: Callable,
|
|
) -> AsyncGenerator[ToolResponse, None]:
|
|
"""A middleware that logs tool execution."""
|
|
# Access the tool call from kwargs
|
|
tool_call = kwargs["tool_call"]
|
|
|
|
# Pre-processing: log before tool execution
|
|
print(f"[Middleware] Calling tool: {tool_call['name']}")
|
|
print(f"[Middleware] Input: {tool_call['input']}")
|
|
|
|
# Call the next handler (either another middleware or the actual tool)
|
|
async for response in await next_handler(**kwargs):
|
|
# Post-processing: log the response
|
|
print(f"[Middleware] Response: {response.content[0]['text']}")
|
|
yield response
|
|
|
|
# This will execute after all responses are yielded
|
|
print(f"[Middleware] Tool {tool_call['name']} completed")
|
|
|
|
|
|
# %%
|
|
# Let's register this middleware with a toolkit and test it:
|
|
#
|
|
|
|
|
|
async def search_tool(query: str) -> ToolResponse:
|
|
"""A simple search tool.
|
|
|
|
Args:
|
|
query (`str`):
|
|
The search query.
|
|
|
|
Returns:
|
|
`ToolResponse`:
|
|
The search result.
|
|
"""
|
|
return ToolResponse(
|
|
content=[
|
|
TextBlock(
|
|
type="text",
|
|
text=f"Search results for '{query}'",
|
|
),
|
|
],
|
|
)
|
|
|
|
|
|
async def example_logging_middleware() -> None:
|
|
"""Example of using logging middleware."""
|
|
# Create a toolkit and register the tool
|
|
toolkit = Toolkit()
|
|
toolkit.register_tool_function(search_tool)
|
|
|
|
# Register the middleware
|
|
toolkit.register_middleware(logging_middleware)
|
|
|
|
# Call the tool
|
|
result = await toolkit.call_tool_function(
|
|
ToolUseBlock(
|
|
type="tool_use",
|
|
id="1",
|
|
name="search_tool",
|
|
input={"query": "AgentScope"},
|
|
),
|
|
)
|
|
|
|
async for response in result:
|
|
print(f"\n[Final] {response.content[0]['text']}\n")
|
|
|
|
|
|
print("=" * 60)
|
|
print("Example 1: Logging Middleware")
|
|
print("=" * 60)
|
|
asyncio.run(example_logging_middleware())
|
|
|
|
# %%
|
|
# Modifying Input and Output
|
|
# ------------------------------
|
|
#
|
|
# Middleware can also modify the tool call input and the response content:
|
|
#
|
|
|
|
|
|
async def transform_middleware(
|
|
kwargs: dict,
|
|
next_handler: Callable,
|
|
) -> AsyncGenerator[ToolResponse, None]:
|
|
"""A middleware that transforms input and output."""
|
|
# Access the tool call from kwargs
|
|
tool_call = kwargs["tool_call"]
|
|
|
|
# Pre-processing: modify the input
|
|
original_query = tool_call["input"]["query"]
|
|
tool_call["input"]["query"] = f"[TRANSFORMED] {original_query}"
|
|
|
|
async for response in await next_handler(**kwargs):
|
|
# Post-processing: modify the response
|
|
original_text = response.content[0]["text"]
|
|
response.content[0]["text"] = f"{original_text} [MODIFIED]"
|
|
yield response
|
|
|
|
|
|
async def example_transform_middleware() -> None:
|
|
"""Example of transforming middleware."""
|
|
toolkit = Toolkit()
|
|
toolkit.register_tool_function(search_tool)
|
|
toolkit.register_middleware(transform_middleware)
|
|
|
|
result = await toolkit.call_tool_function(
|
|
ToolUseBlock(
|
|
type="tool_use",
|
|
id="2",
|
|
name="search_tool",
|
|
input={"query": "middleware"},
|
|
),
|
|
)
|
|
|
|
async for response in result:
|
|
print(f"Result: {response.content[0]['text']}")
|
|
|
|
|
|
print("\n" + "=" * 60)
|
|
print("Example 2: Transform Middleware")
|
|
print("=" * 60)
|
|
asyncio.run(example_transform_middleware())
|
|
|
|
# %%
|
|
# Authorization Middleware
|
|
# ------------------------------
|
|
#
|
|
# You can use middleware to implement authorization checks and skip tool execution if not authorized:
|
|
#
|
|
|
|
|
|
async def authorization_middleware(
|
|
kwargs: dict,
|
|
next_handler: Callable,
|
|
) -> AsyncGenerator[ToolResponse, None]:
|
|
"""A middleware that checks authorization."""
|
|
# Access the tool call from kwargs
|
|
tool_call = kwargs["tool_call"]
|
|
|
|
# Check if the tool is authorized (simple example)
|
|
authorized_tools = {"search_tool"}
|
|
|
|
if tool_call["name"] not in authorized_tools:
|
|
# Skip execution and return error directly
|
|
print(f"[Auth] Tool {tool_call['name']} is not authorized")
|
|
yield ToolResponse(
|
|
content=[
|
|
TextBlock(
|
|
type="text",
|
|
text=f"Error: Tool '{tool_call['name']}' is not authorized", # noqa: E501
|
|
),
|
|
],
|
|
)
|
|
return
|
|
|
|
# Tool is authorized, proceed
|
|
print(f"[Auth] Tool {tool_call['name']} is authorized")
|
|
async for response in await next_handler(**kwargs):
|
|
yield response
|
|
|
|
|
|
async def unauthorized_tool(data: str) -> ToolResponse:
|
|
"""An unauthorized tool.
|
|
|
|
Args:
|
|
data (`str`):
|
|
Some data.
|
|
|
|
Returns:
|
|
`ToolResponse`:
|
|
The result.
|
|
"""
|
|
return ToolResponse(
|
|
content=[TextBlock(type="text", text=f"Processing {data}")],
|
|
)
|
|
|
|
|
|
async def example_authorization_middleware() -> None:
|
|
"""Example of authorization middleware."""
|
|
toolkit = Toolkit()
|
|
toolkit.register_tool_function(search_tool)
|
|
toolkit.register_tool_function(unauthorized_tool)
|
|
toolkit.register_middleware(authorization_middleware)
|
|
|
|
# Try authorized tool
|
|
print("\nCalling authorized tool:")
|
|
result = await toolkit.call_tool_function(
|
|
ToolUseBlock(
|
|
type="tool_use",
|
|
id="3",
|
|
name="search_tool",
|
|
input={"query": "test"},
|
|
),
|
|
)
|
|
async for response in result:
|
|
print(f"Result: {response.content[0]['text']}")
|
|
|
|
# Try unauthorized tool
|
|
print("\nCalling unauthorized tool:")
|
|
result = await toolkit.call_tool_function(
|
|
ToolUseBlock(
|
|
type="tool_use",
|
|
id="4",
|
|
name="unauthorized_tool",
|
|
input={"data": "test"},
|
|
),
|
|
)
|
|
async for response in result:
|
|
print(f"Result: {response.content[0]['text']}")
|
|
|
|
|
|
print("\n" + "=" * 60)
|
|
print("Example 3: Authorization Middleware")
|
|
print("=" * 60)
|
|
asyncio.run(example_authorization_middleware())
|
|
|
|
# %%
|
|
# Multiple Middleware (Onion Model)
|
|
# ------------------------------
|
|
#
|
|
# When multiple middleware are registered, they form an onion-like structure.
|
|
# The execution order follows the onion model:
|
|
#
|
|
# - **Pre-processing**: Executes in the order middleware are registered
|
|
# - **Post-processing**: Executes in reverse order (inner to outer)
|
|
#
|
|
# This is because the actual tool response object is passed through the middleware chain,
|
|
# and each middleware modifies it in place.
|
|
#
|
|
|
|
|
|
async def middleware_1(
|
|
kwargs: dict,
|
|
next_handler: Callable,
|
|
) -> AsyncGenerator[ToolResponse, None]:
|
|
"""First middleware."""
|
|
# Access the tool call from kwargs
|
|
tool_call = kwargs["tool_call"]
|
|
|
|
# Pre-processing
|
|
print("[M1] Pre-processing")
|
|
tool_call["input"]["query"] += " [M1]"
|
|
|
|
async for response in await next_handler(**kwargs):
|
|
# Post-processing
|
|
response.content[0]["text"] += " [M1]"
|
|
print("[M1] Post-processing")
|
|
yield response
|
|
|
|
|
|
async def middleware_2(
|
|
kwargs: dict,
|
|
next_handler: Callable,
|
|
) -> AsyncGenerator[ToolResponse, None]:
|
|
"""Second middleware."""
|
|
# Access the tool call from kwargs
|
|
tool_call = kwargs["tool_call"]
|
|
|
|
# Pre-processing
|
|
print("[M2] Pre-processing")
|
|
tool_call["input"]["query"] += " [M2]"
|
|
|
|
async for response in await next_handler(**kwargs):
|
|
# Post-processing
|
|
response.content[0]["text"] += " [M2]"
|
|
print("[M2] Post-processing")
|
|
yield response
|
|
|
|
|
|
async def example_multiple_middleware() -> None:
|
|
"""Example of multiple middleware."""
|
|
toolkit = Toolkit()
|
|
toolkit.register_tool_function(search_tool)
|
|
|
|
# Register middleware in order
|
|
toolkit.register_middleware(middleware_1)
|
|
toolkit.register_middleware(middleware_2)
|
|
|
|
result = await toolkit.call_tool_function(
|
|
ToolUseBlock(
|
|
type="tool_use",
|
|
id="5",
|
|
name="search_tool",
|
|
input={"query": "test"},
|
|
),
|
|
)
|
|
|
|
async for response in result:
|
|
print(f"\nFinal result: {response.content[0]['text']}")
|
|
|
|
|
|
print("\n" + "=" * 60)
|
|
print("Example 4: Multiple Middleware (Onion Model)")
|
|
print("=" * 60)
|
|
print("\nExecution flow:")
|
|
print("M1 Pre → M2 Pre → Tool → M2 Post → M1 Post")
|
|
print()
|
|
asyncio.run(example_multiple_middleware())
|
|
|
|
# %%
|
|
# Use Cases
|
|
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
|
#
|
|
# The middleware system is useful for various scenarios:
|
|
#
|
|
# - **Logging and Monitoring**: Track tool usage and performance
|
|
# - **Authorization**: Control access to specific tools
|
|
# - **Rate Limiting**: Limit the frequency of tool calls
|
|
# - **Caching**: Cache tool responses for repeated calls
|
|
# - **Error Handling**: Add retry logic or graceful degradation
|
|
# - **Input Validation**: Validate and sanitize tool inputs
|
|
# - **Output Transformation**: Format or filter tool outputs
|
|
# - **Metrics Collection**: Collect statistics about tool usage
|
|
#
|
|
# .. note::
|
|
# - Middleware are applied in the order they are registered
|
|
# - The same ``ToolResponse`` object is passed through the middleware chain and modified in place
|
|
# - Middleware can completely skip tool execution by not calling ``next_handler``
|
|
# - All middleware must be async generator functions that yield ``ToolResponse`` objects
|