Skip to content

Decorators

Decorator functions for creating middleware from simple async functions. Each decorator wraps a single hook function into a full AgentMiddleware instance, so you do not need to subclass AgentMiddleware for straightforward use cases.

The before_tool_call, after_tool_call, and on_tool_error decorators also accept an optional tools parameter to restrict the middleware to specific tool names.

before_run

pydantic_ai_middleware.before_run(func)

Create middleware from a before_run function.

Example
Python
@before_run
async def log_input(prompt: str, deps, ctx) -> str:
    print(f"Input: {prompt}")
    ctx.set("logged", True)  # Store data in context
    return prompt
Source code in src/pydantic_ai_middleware/decorators.py
Python
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def before_run(
    func: BeforeRunFunc[DepsT],
) -> AgentMiddleware[DepsT]:
    """Create middleware from a before_run function.

    Example:
        ```python
        @before_run
        async def log_input(prompt: str, deps, ctx) -> str:
            print(f"Input: {prompt}")
            ctx.set("logged", True)  # Store data in context
            return prompt
        ```
    """
    return _FunctionMiddleware(before_run_func=func)

after_run

pydantic_ai_middleware.after_run(func)

Create middleware from an after_run function.

Example
Python
@after_run
async def log_output(prompt: str, output, deps, ctx):
    print(f"Output: {output}")
    return output
Source code in src/pydantic_ai_middleware/decorators.py
Python
165
166
167
168
169
170
171
172
173
174
175
176
177
178
def after_run(
    func: AfterRunFunc[DepsT],
) -> AgentMiddleware[DepsT]:
    """Create middleware from an after_run function.

    Example:
        ```python
        @after_run
        async def log_output(prompt: str, output, deps, ctx):
            print(f"Output: {output}")
            return output
        ```
    """
    return _FunctionMiddleware(after_run_func=func)

before_model_request

pydantic_ai_middleware.before_model_request(func)

Create middleware from a before_model_request function.

Example
Python
@before_model_request
async def log_messages(messages: list, deps, ctx) -> list:
    print(f"Sending {len(messages)} messages")
    return messages
Source code in src/pydantic_ai_middleware/decorators.py
Python
181
182
183
184
185
186
187
188
189
190
191
192
193
194
def before_model_request(
    func: BeforeModelRequestFunc[DepsT],
) -> AgentMiddleware[DepsT]:
    """Create middleware from a before_model_request function.

    Example:
        ```python
        @before_model_request
        async def log_messages(messages: list, deps, ctx) -> list:
            print(f"Sending {len(messages)} messages")
            return messages
        ```
    """
    return _FunctionMiddleware(before_model_request_func=func)

before_tool_call

pydantic_ai_middleware.before_tool_call(func=None, *, tools=None)

Python
before_tool_call(
    func: BeforeToolCallFunc[DepsT],
) -> AgentMiddleware[DepsT]
Python
before_tool_call(
    func: None = None, *, tools: set[str] | None = ...
) -> Callable[
    [BeforeToolCallFunc[DepsT]], AgentMiddleware[DepsT]
]

Create middleware from a before_tool_call function.

Can be used as a plain decorator or with tools parameter:

Example
Python
@before_tool_call
async def validate_all(tool_name, tool_args, deps, ctx):
    return tool_args

@before_tool_call(tools={"send_email"})
async def validate_email(tool_name, tool_args, deps, ctx):
    return tool_args
Source code in src/pydantic_ai_middleware/decorators.py
Python
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
def before_tool_call(
    func: BeforeToolCallFunc[DepsT] | None = None,
    *,
    tools: set[str] | None = None,
) -> AgentMiddleware[DepsT] | Callable[[BeforeToolCallFunc[DepsT]], AgentMiddleware[DepsT]]:
    """Create middleware from a before_tool_call function.

    Can be used as a plain decorator or with `tools` parameter:

    Example:
        ```python
        @before_tool_call
        async def validate_all(tool_name, tool_args, deps, ctx):
            return tool_args

        @before_tool_call(tools={"send_email"})
        async def validate_email(tool_name, tool_args, deps, ctx):
            return tool_args
        ```
    """
    if func is not None:
        # Used as @before_tool_call (no parentheses)
        return _FunctionMiddleware(before_tool_call_func=func, tool_names=tools)

    # Used as @before_tool_call(tools={"send_email"})
    def decorator(f: BeforeToolCallFunc[DepsT]) -> AgentMiddleware[DepsT]:
        return _FunctionMiddleware(before_tool_call_func=f, tool_names=tools)

    return decorator

after_tool_call

pydantic_ai_middleware.after_tool_call(func=None, *, tools=None)

Python
after_tool_call(
    func: AfterToolCallFunc[DepsT],
) -> AgentMiddleware[DepsT]
Python
after_tool_call(
    func: None = None, *, tools: set[str] | None = ...
) -> Callable[
    [AfterToolCallFunc[DepsT]], AgentMiddleware[DepsT]
]

Create middleware from an after_tool_call function.

Can be used as a plain decorator or with tools parameter:

Example
Python
@after_tool_call
async def log_all(tool_name, tool_args, result, deps, ctx):
    return result

@after_tool_call(tools={"send_email"})
async def log_email(tool_name, tool_args, result, deps, ctx):
    return result
Source code in src/pydantic_ai_middleware/decorators.py
Python
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
def after_tool_call(
    func: AfterToolCallFunc[DepsT] | None = None,
    *,
    tools: set[str] | None = None,
) -> AgentMiddleware[DepsT] | Callable[[AfterToolCallFunc[DepsT]], AgentMiddleware[DepsT]]:
    """Create middleware from an after_tool_call function.

    Can be used as a plain decorator or with `tools` parameter:

    Example:
        ```python
        @after_tool_call
        async def log_all(tool_name, tool_args, result, deps, ctx):
            return result

        @after_tool_call(tools={"send_email"})
        async def log_email(tool_name, tool_args, result, deps, ctx):
            return result
        ```
    """
    if func is not None:
        return _FunctionMiddleware(after_tool_call_func=func, tool_names=tools)

    def decorator(f: AfterToolCallFunc[DepsT]) -> AgentMiddleware[DepsT]:
        return _FunctionMiddleware(after_tool_call_func=f, tool_names=tools)

    return decorator

on_tool_error

pydantic_ai_middleware.on_tool_error(func=None, *, tools=None)

Python
on_tool_error(
    func: OnToolErrorFunc[DepsT],
) -> AgentMiddleware[DepsT]
Python
on_tool_error(
    func: None = None, *, tools: set[str] | None = ...
) -> Callable[
    [OnToolErrorFunc[DepsT]], AgentMiddleware[DepsT]
]

Create middleware from an on_tool_error function.

Can be used as a plain decorator or with tools parameter:

Example
Python
@on_tool_error
async def handle_all_errors(tool_name, tool_args, error, deps, ctx):
    return None  # Re-raise original

@on_tool_error(tools={"send_email"})
async def handle_email_errors(tool_name, tool_args, error, deps, ctx):
    return RuntimeError("Email failed")
Source code in src/pydantic_ai_middleware/decorators.py
Python
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
def on_tool_error(
    func: OnToolErrorFunc[DepsT] | None = None,
    *,
    tools: set[str] | None = None,
) -> AgentMiddleware[DepsT] | Callable[[OnToolErrorFunc[DepsT]], AgentMiddleware[DepsT]]:
    """Create middleware from an on_tool_error function.

    Can be used as a plain decorator or with `tools` parameter:

    Example:
        ```python
        @on_tool_error
        async def handle_all_errors(tool_name, tool_args, error, deps, ctx):
            return None  # Re-raise original

        @on_tool_error(tools={"send_email"})
        async def handle_email_errors(tool_name, tool_args, error, deps, ctx):
            return RuntimeError("Email failed")
        ```
    """
    if func is not None:
        return _FunctionMiddleware(on_tool_error_func=func, tool_names=tools)

    def decorator(f: OnToolErrorFunc[DepsT]) -> AgentMiddleware[DepsT]:
        return _FunctionMiddleware(on_tool_error_func=f, tool_names=tools)

    return decorator

on_error

pydantic_ai_middleware.on_error(func)

Create middleware from an on_error function.

Example
Python
@on_error
async def handle_error(error: Exception, deps, ctx) -> Exception | None:
    print(f"Error occurred: {error}")
    return None  # Re-raise original
Source code in src/pydantic_ai_middleware/decorators.py
Python
328
329
330
331
332
333
334
335
336
337
338
339
340
341
def on_error(
    func: OnErrorFunc[DepsT],
) -> AgentMiddleware[DepsT]:
    """Create middleware from an on_error function.

    Example:
        ```python
        @on_error
        async def handle_error(error: Exception, deps, ctx) -> Exception | None:
            print(f"Error occurred: {error}")
            return None  # Re-raise original
        ```
    """
    return _FunctionMiddleware(on_error_func=func)