AgentMiddleware

pydantic_ai_middleware.AgentMiddleware

Bases: ABC, Generic[DepsT]

Base middleware class.

Inherit and override only the methods you need. Middleware hooks are called in order for before_* methods and in reverse order for after_* methods.

Attributes:

Name	Type	Description
tool_names	set[str] | None	Set of tool names this middleware applies to. If None (default), the middleware applies to all tools.
timeout	float | None	Maximum time in seconds for any hook call. If None (default), no timeout is enforced.

Source code in src/pydantic_ai_middleware/base.py

class AgentMiddleware(ABC, Generic[DepsT]):
    """Base middleware class.

    Inherit and override only the methods you need.
    Middleware hooks are called in order for before_* methods
    and in reverse order for after_* methods.

    Attributes:
        tool_names: Set of tool names this middleware applies to.
            If None (default), the middleware applies to all tools.
        timeout: Maximum time in seconds for any hook call.
            If None (default), no timeout is enforced.
    """

    tool_names: set[str] | None = None
    timeout: float | None = None

    def _should_handle_tool(self, tool_name: str) -> bool:
        """Check if this middleware should handle the given tool.

        Args:
            tool_name: The name of the tool to check.

        Returns:
            True if the middleware should handle this tool.
        """
        if self.tool_names is None:
            return True
        return tool_name in self.tool_names

    async def before_run(
        self,
        prompt: str | Sequence[Any],
        deps: DepsT | None,
        ctx: ScopedContext | None = None,
    ) -> str | Sequence[Any]:
        """Called before the agent runs.

        Args:
            prompt: The user prompt.
            deps: The agent dependencies.
            ctx: Scoped context for data sharing (write: before_run, read: config/metadata).

        Returns:
            The (possibly modified) prompt.

        Raises:
            InputBlocked: To block the input.
        """
        return prompt

    async def after_run(
        self,
        prompt: str | Sequence[Any],
        output: Any,
        deps: DepsT | None,
        ctx: ScopedContext | None = None,
    ) -> Any:
        """Called after the agent finishes.

        Args:
            prompt: The original user prompt.
            output: The agent output.
            deps: The agent dependencies.
            ctx: Scoped context for data sharing (write: after_run, read: all previous).

        Returns:
            The (possibly modified) output.
        """
        return output

    async def before_model_request(
        self,
        messages: list[ModelMessage],
        deps: DepsT | None,
        ctx: ScopedContext | None = None,
    ) -> list[ModelMessage]:
        """Called before each model request.

        Args:
            messages: The messages to send to the model.
            deps: The agent dependencies.
            ctx: Scoped context for data sharing (write: before_model_request, read: before_run).

        Returns:
            The (possibly modified) messages.
        """
        return messages

    async def before_tool_call(
        self,
        tool_name: str,
        tool_args: dict[str, Any],
        deps: DepsT | None,
        ctx: ScopedContext | None = None,
    ) -> dict[str, Any] | ToolPermissionResult:
        """Called before a tool is called.

        Args:
            tool_name: The name of the tool being called.
            tool_args: The arguments to the tool.
            deps: The agent dependencies.
            ctx: Scoped context for data sharing
                (write: before_tool_call, read: before_run, before_model_request).

        Returns:
            The (possibly modified) tool arguments, or a ToolPermissionResult
            for structured permission decisions.

        Raises:
            ToolBlocked: To block the tool call.
        """
        return tool_args

    async def on_tool_error(
        self,
        tool_name: str,
        tool_args: dict[str, Any],
        error: Exception,
        deps: DepsT | None,
        ctx: ScopedContext | None = None,
    ) -> Exception | None:
        """Called when a tool execution raises an error.

        This hook provides tool-specific error context (tool_name, tool_args)
        that is not available in the generic on_error hook.

        Args:
            tool_name: The name of the tool that failed.
            tool_args: The arguments that were passed to the tool.
            error: The exception raised by the tool.
            deps: The agent dependencies.
            ctx: Scoped context for data sharing (write: on_tool_error, read: all hooks).

        Returns:
            A different exception to raise, or None to re-raise the original.
        """
        return None

    async def after_tool_call(
        self,
        tool_name: str,
        tool_args: dict[str, Any],
        result: Any,
        deps: DepsT | None,
        ctx: ScopedContext | None = None,
    ) -> Any:
        """Called after a tool is called.

        Args:
            tool_name: The name of the tool that was called.
            tool_args: The arguments that were passed to the tool.
            result: The result from the tool.
            deps: The agent dependencies.
            ctx: Scoped context for data sharing (write: after_tool_call, read: all before_* hooks).

        Returns:
            The (possibly modified) result.
        """
        return result

    async def on_error(
        self,
        error: Exception,
        deps: DepsT | None,
        ctx: ScopedContext | None = None,
    ) -> Exception | None:
        """Called when an error occurs.

        Args:
            error: The exception that occurred.
            deps: The agent dependencies.
            ctx: Scoped context for data sharing (write: on_error, read: all hooks).

        Returns:
            A different exception to raise, or None to re-raise the original.
        """
        return None

tool_names = None class-attribute instance-attribute

timeout = None class-attribute instance-attribute

_should_handle_tool(tool_name)

Check if this middleware should handle the given tool.

Parameters:

Name	Type	Description	Default
tool_name	str	The name of the tool to check.	required

Returns:

Type	Description
bool	True if the middleware should handle this tool.

Source code in src/pydantic_ai_middleware/base.py

def _should_handle_tool(self, tool_name: str) -> bool:
    """Check if this middleware should handle the given tool.

    Args:
        tool_name: The name of the tool to check.

    Returns:
        True if the middleware should handle this tool.
    """
    if self.tool_names is None:
        return True
    return tool_name in self.tool_names

before_run(prompt, deps, ctx=None) async

Called before the agent runs.

Parameters:

Name	Type	Description	Default
prompt	str | Sequence[Any]	The user prompt.	required
deps	DepsT | None	The agent dependencies.	required
ctx	ScopedContext | None	Scoped context for data sharing (write: before_run, read: config/metadata).	None

Returns:

Type	Description
str | Sequence[Any]	The (possibly modified) prompt.

Raises:

Type	Description
InputBlocked	To block the input.

Source code in src/pydantic_ai_middleware/base.py

async def before_run(
    self,
    prompt: str | Sequence[Any],
    deps: DepsT | None,
    ctx: ScopedContext | None = None,
) -> str | Sequence[Any]:
    """Called before the agent runs.

    Args:
        prompt: The user prompt.
        deps: The agent dependencies.
        ctx: Scoped context for data sharing (write: before_run, read: config/metadata).

    Returns:
        The (possibly modified) prompt.

    Raises:
        InputBlocked: To block the input.
    """
    return prompt

after_run(prompt, output, deps, ctx=None) async

Called after the agent finishes.

Parameters:

Name	Type	Description	Default
prompt	str | Sequence[Any]	The original user prompt.	required
output	Any	The agent output.	required
deps	DepsT | None	The agent dependencies.	required
ctx	ScopedContext | None	Scoped context for data sharing (write: after_run, read: all previous).	None

Returns:

Type	Description
Any	The (possibly modified) output.

Source code in src/pydantic_ai_middleware/base.py

async def after_run(
    self,
    prompt: str | Sequence[Any],
    output: Any,
    deps: DepsT | None,
    ctx: ScopedContext | None = None,
) -> Any:
    """Called after the agent finishes.

    Args:
        prompt: The original user prompt.
        output: The agent output.
        deps: The agent dependencies.
        ctx: Scoped context for data sharing (write: after_run, read: all previous).

    Returns:
        The (possibly modified) output.
    """
    return output

before_model_request(messages, deps, ctx=None) async

Called before each model request.

Parameters:

Name	Type	Description	Default
messages	list[ModelMessage]	The messages to send to the model.	required
deps	DepsT | None	The agent dependencies.	required
ctx	ScopedContext | None	Scoped context for data sharing (write: before_model_request, read: before_run).	None

Returns:

Type	Description
list[ModelMessage]	The (possibly modified) messages.

Source code in src/pydantic_ai_middleware/base.py

async def before_model_request(
    self,
    messages: list[ModelMessage],
    deps: DepsT | None,
    ctx: ScopedContext | None = None,
) -> list[ModelMessage]:
    """Called before each model request.

    Args:
        messages: The messages to send to the model.
        deps: The agent dependencies.
        ctx: Scoped context for data sharing (write: before_model_request, read: before_run).

    Returns:
        The (possibly modified) messages.
    """
    return messages

before_tool_call(tool_name, tool_args, deps, ctx=None) async

Called before a tool is called.

Parameters:

Name	Type	Description	Default
tool_name	str	The name of the tool being called.	required
tool_args	dict[str, Any]	The arguments to the tool.	required
deps	DepsT | None	The agent dependencies.	required
ctx	ScopedContext | None	Scoped context for data sharing (write: before_tool_call, read: before_run, before_model_request).	None

Returns:

Type	Description
dict[str, Any] | ToolPermissionResult	The (possibly modified) tool arguments, or a ToolPermissionResult
dict[str, Any] | ToolPermissionResult	for structured permission decisions.

Raises:

Type	Description
ToolBlocked	To block the tool call.

Source code in src/pydantic_ai_middleware/base.py

async def before_tool_call(
    self,
    tool_name: str,
    tool_args: dict[str, Any],
    deps: DepsT | None,
    ctx: ScopedContext | None = None,
) -> dict[str, Any] | ToolPermissionResult:
    """Called before a tool is called.

    Args:
        tool_name: The name of the tool being called.
        tool_args: The arguments to the tool.
        deps: The agent dependencies.
        ctx: Scoped context for data sharing
            (write: before_tool_call, read: before_run, before_model_request).

    Returns:
        The (possibly modified) tool arguments, or a ToolPermissionResult
        for structured permission decisions.

    Raises:
        ToolBlocked: To block the tool call.
    """
    return tool_args

on_tool_error(tool_name, tool_args, error, deps, ctx=None) async

Called when a tool execution raises an error.

This hook provides tool-specific error context (tool_name, tool_args) that is not available in the generic on_error hook.

Parameters:

Name	Type	Description	Default
tool_name	str	The name of the tool that failed.	required
tool_args	dict[str, Any]	The arguments that were passed to the tool.	required
error	Exception	The exception raised by the tool.	required
deps	DepsT | None	The agent dependencies.	required
ctx	ScopedContext | None	Scoped context for data sharing (write: on_tool_error, read: all hooks).	None

Returns:

Type	Description
Exception | None	A different exception to raise, or None to re-raise the original.

Source code in src/pydantic_ai_middleware/base.py

async def on_tool_error(
    self,
    tool_name: str,
    tool_args: dict[str, Any],
    error: Exception,
    deps: DepsT | None,
    ctx: ScopedContext | None = None,
) -> Exception | None:
    """Called when a tool execution raises an error.

    This hook provides tool-specific error context (tool_name, tool_args)
    that is not available in the generic on_error hook.

    Args:
        tool_name: The name of the tool that failed.
        tool_args: The arguments that were passed to the tool.
        error: The exception raised by the tool.
        deps: The agent dependencies.
        ctx: Scoped context for data sharing (write: on_tool_error, read: all hooks).

    Returns:
        A different exception to raise, or None to re-raise the original.
    """
    return None

after_tool_call(tool_name, tool_args, result, deps, ctx=None) async

Called after a tool is called.

Parameters:

Name	Type	Description	Default
tool_name	str	The name of the tool that was called.	required
tool_args	dict[str, Any]	The arguments that were passed to the tool.	required
result	Any	The result from the tool.	required
deps	DepsT | None	The agent dependencies.	required
ctx	ScopedContext | None	Scoped context for data sharing (write: after_tool_call, read: all before_* hooks).	None

Returns:

Type	Description
Any	The (possibly modified) result.

Source code in src/pydantic_ai_middleware/base.py

async def after_tool_call(
    self,
    tool_name: str,
    tool_args: dict[str, Any],
    result: Any,
    deps: DepsT | None,
    ctx: ScopedContext | None = None,
) -> Any:
    """Called after a tool is called.

    Args:
        tool_name: The name of the tool that was called.
        tool_args: The arguments that were passed to the tool.
        result: The result from the tool.
        deps: The agent dependencies.
        ctx: Scoped context for data sharing (write: after_tool_call, read: all before_* hooks).

    Returns:
        The (possibly modified) result.
    """
    return result

on_error(error, deps, ctx=None) async

Called when an error occurs.

Parameters:

Name	Type	Description	Default
error	Exception	The exception that occurred.	required
deps	DepsT | None	The agent dependencies.	required
ctx	ScopedContext | None	Scoped context for data sharing (write: on_error, read: all hooks).	None

Returns:

Type	Description
Exception | None	A different exception to raise, or None to re-raise the original.

Source code in src/pydantic_ai_middleware/base.py

async def on_error(
    self,
    error: Exception,
    deps: DepsT | None,
    ctx: ScopedContext | None = None,
) -> Exception | None:
    """Called when an error occurs.

    Args:
        error: The exception that occurred.
        deps: The agent dependencies.
        ctx: Scoped context for data sharing (write: on_error, read: all hooks).

    Returns:
        A different exception to raise, or None to re-raise the original.
    """
    return None