Spaces:
Paused
Paused
File size: 8,771 Bytes
9792ea7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 | # -*- coding: utf-8 -*-
"""Base middleware class for AgentScope middleware system."""
from typing import AsyncGenerator, Awaitable, Callable, TYPE_CHECKING
from ..tool import ToolBase
if TYPE_CHECKING:
from ..agent import Agent
from ..model import ChatResponse
class MiddlewareBase: # pylint: disable=unused-argument
"""Base class for all middleware implementations.
Middleware provides interception mechanisms at 5 key execution points
in the Agent lifecycle:
**Onion Pattern Hooks** (with before/after logic):
- `on_reply`: Intercepts the entire reply process
- `on_reasoning`: Intercepts the reasoning/model call phase
- `on_acting`: Intercepts individual tool call execution
- `on_model_call`: Intercepts the raw model API call
**Transformer Pattern Hook** (sequential pipeline):
- `on_system_prompt`: Transforms the system prompt string
Each hook is optional - only implement the ones you need. The middleware
system will automatically detect which hooks are implemented at runtime.
Example:
```python
class LoggingMiddleware(MiddlewareBase):
async def on_reasoning(
self,
agent: Agent,
input_kwargs: dict,
next_handler: Callable[[], AsyncGenerator],
) -> AsyncGenerator:
print(f"Before reasoning for agent {agent.name}")
async for event in next_handler():
yield event
print(f"After reasoning for agent {agent.name}")
agent = Agent(
...
middlewares=[LoggingMiddleware()],
...
)
```
"""
def is_implemented(self, hook_name: str) -> bool:
"""Check if a hook method is implemented in the subclass.
Args:
hook_name: Name of the hook method to check
Returns:
True if the hook is implemented (overridden), False otherwise
"""
base_method = getattr(MiddlewareBase, hook_name, None)
sub_method = getattr(type(self), hook_name, None)
return base_method is not sub_method
async def on_reply(
self,
agent: "Agent",
input_kwargs: dict,
next_handler: Callable[..., AsyncGenerator],
) -> AsyncGenerator:
"""Hook for intercepting the reply process.
Args:
agent: The Agent instance executing this middleware
input_kwargs: Dictionary containing:
- inputs: Msg | list[Msg] | UserConfirmResultEvent |
ExternalExecutionResultEvent | None — the unified inputs
that trigger this reply (new message(s), a resumption
event from a previous outside interaction, or None).
next_handler: Callable that executes the next middleware or
original method
Yields:
AgentEvent | Msg: Events from the reply process
"""
raise RuntimeError(
f"{type(self).__name__} does not implement on_reply",
)
yield # pylint: disable=unreachable
async def on_reasoning(
self,
agent: "Agent",
input_kwargs: dict,
next_handler: Callable[..., AsyncGenerator],
) -> AsyncGenerator:
"""Hook for intercepting the reasoning process.
Args:
agent: The Agent instance executing this middleware
input_kwargs: Dictionary containing:
- tool_choice: ToolChoice (default None)
next_handler: Callable that executes the next middleware or
original method
Yields:
Various events from the reasoning process
"""
raise RuntimeError(
f"{type(self).__name__} does not implement on_reasoning",
)
yield # pylint: disable=unreachable
async def on_acting(
self,
agent: "Agent",
input_kwargs: dict,
next_handler: Callable[..., AsyncGenerator],
) -> AsyncGenerator:
"""Hook for intercepting the raw tool execution.
This hook wraps **only** the ``toolkit.call_tool`` call — i.e. the
pure I/O execution layer. Permission checking, input validation, and
context writes are handled by the agent **outside** this hook and are
therefore not visible here.
This separation makes it safe to offload the ``next_handler``
coroutine to a background task: it will never mutate agent context
on its own.
.. note::
Tools with ``is_state_injected=True`` receive the live
``agent.state`` object. Offloading such tools to a background
task may cause concurrent state mutations — guard against this
in your middleware implementation.
Args:
agent (`Agent`):
The Agent instance executing this middleware.
input_kwargs (`dict`):
Dictionary containing:
- ``tool_call`` (``ToolCallBlock``): the tool call to execute.
By the time this hook is invoked the tool call has already
been validated and permitted.
next_handler (`Callable[..., AsyncGenerator]`):
Callable that executes the next middleware or
``_acting_impl``.
Yields:
`ToolChunk | ToolResponse`:
Intermediate ``ToolChunk`` objects followed by a final
``ToolResponse`` produced by the tool.
"""
raise RuntimeError(
f"{type(self).__name__} does not implement on_acting",
)
yield # pylint: disable=unreachable
async def on_model_call(
self,
agent: "Agent",
input_kwargs: dict,
next_handler: Callable[
...,
Awaitable["ChatResponse" | AsyncGenerator["ChatResponse", None]],
],
) -> "ChatResponse" | AsyncGenerator["ChatResponse", None]:
"""Hook for intercepting the model API call.
Args:
agent: The Agent instance executing this middleware
input_kwargs: Dictionary containing:
- messages: list[Msg]
- tools: list[dict]
- tool_choice: ToolChoice
- current_model: The model instance used for this call
next_handler: Callable that executes the next middleware or
original method
Returns:
ChatResponse or AsyncGenerator[ChatResponse, None]
"""
raise RuntimeError(
f"{type(self).__name__} does not implement on_model_call",
)
async def on_compress_context(
self,
agent: "Agent",
input_kwargs: dict,
next_handler: Callable[..., Awaitable[None]],
) -> None:
"""Onion hook for `compress_context` function in `Agent` class
Args:
agent (`Agent`):
The Agent instance executing this middleware
input_kwargs (`dict`):
Dictionary containing:
- context_config: ContextConfig | None
- instructions: HintBlock | None
next_handler (`Callable[..., Awaitable[None]]`):
Callable that executes the next middleware or
original method
"""
raise RuntimeError(
f"{type(self).__name__} does not implement on_compress_context",
)
async def on_system_prompt(
self,
agent: "Agent",
current_prompt: str,
) -> str:
"""Transform the system prompt string.
This uses a transformer/pipeline pattern rather than onion pattern.
Multiple middlewares are applied sequentially, each receiving the
output of the previous one.
Args:
agent: The Agent instance executing this middleware
current_prompt: The current system prompt string
Returns:
str: The transformed system prompt
"""
raise RuntimeError(
f"{type(self).__name__} does not implement on_system_prompt",
)
async def list_tools(self) -> list[ToolBase]:
"""List available tools provided by this middleware. Optional to
implement.
Returns:
`list[ToolBase]`:
A list of tools provided by this middleware.
"""
return []
async def get_middleware_key(self) -> str:
"""Get the unique key for this middleware, used to save middleware
states in `AgentState` instances.
Optionally, middleware classes can override this method to
provide their own unique key.
"""
return self.__class__.__name__
|