# -*- 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__