pydantic-ai-middleware (Abandoned: Renamed to pydantic-ai-shields)

Common errors

• ModuleNotFoundError: No module named 'pydantic_ai_middleware.agent'

  cause Attempting to import from `pydantic-ai-middleware` after version `0.3.0`, which was renamed to `pydantic-ai-shields` and had its core middleware removed.
  fix Install `pydantic-ai-middleware==0.2.4` for the old functionality, or migrate to `pydantic-ai-shields` (for guardrails) or `pydantic-ai` v1.71+ (for native agent middleware).

• AttributeError: Can't instantiate abstract class MiddlewareAgent with abstract method description

  cause Using `MiddlewareAgent` from `pydantic-ai-middleware` versions `0.2.0-0.2.3` with a `pydantic-ai-slim` version (e.g., >=1.73.0) that made `description` an abstract property on `AbstractAgent`.
  fix Upgrade `pydantic-ai-middleware` to `0.2.4`. If this is not possible, downgrade `pydantic-ai-slim` to a version where `description` is not an abstract method on `AbstractAgent`.

• TypeError: 'NoneType' object is not callable (or similar error during streaming/iteration)

  cause Prior to version `0.2.3`, `AgentMiddleware`'s `after_run` and `on_error` hooks were not correctly invoked when using `MiddlewareAgent.iter()` or `run_stream()`.
  fix Upgrade `pydantic-ai-middleware` to version `0.2.3` or higher.

Warnings

• breaking The `pydantic-ai-middleware` package has been renamed to `pydantic-ai-shields` as of version `0.3.0`. The core middleware functionality (`MiddlewareAgent`, `AgentMiddleware`) has been removed from this package and integrated directly into `pydantic-ai` v1.71+.
  Fix: For new projects, use `pydantic-ai-shields` for guardrails or `pydantic-ai` v1.71+ for native agent middleware. For existing projects, downgrade to `pydantic-ai-middleware<0.3.0` or migrate to the new `pydantic-ai` native middleware.
• deprecated The `pydantic-ai` dependency was replaced with `pydantic-ai-slim` in version `0.2.2` to reduce dependency bloat. While `pydantic-ai` still works with older versions of `pydantic-ai-middleware`, using `pydantic-ai-slim` is recommended.
  Fix: Ensure `pydantic-ai-slim` is installed instead of `pydantic-ai`. Imports for `pydantic_ai.*` remain compatible.
• gotcha Instances of `MiddlewareAgent` will fail to initialize with `AttributeError: Can't instantiate abstract class MiddlewareAgent with abstract method description` when used with `pydantic-ai` versions that introduced `description` as an abstract property on `AbstractAgent` (e.g., `pydantic-ai-slim` >= 1.73.0).
  Fix: Upgrade `pydantic-ai-middleware` to version `0.2.4` or higher. If unable to upgrade, you may need to downgrade `pydantic-ai-slim` to a version that predates the `description` abstract property.
• gotcha When using `MiddlewareAgent.iter()` or `run_stream()`, `after_run` and `on_error` lifecycle hooks were not invoked in `AgentMiddleware`, making post-processing or error handling impossible for streaming operations.
  Fix: Upgrade `pydantic-ai-middleware` to version `0.2.3` or higher to ensure all middleware hooks are correctly triggered during streaming operations.

Install

• pip install pydantic-ai-middleware==0.2.4 Install last version of pydantic-ai-middleware

Imports

• MiddlewareAgent
  wrong:

  from pydantic_ai_shields.agent import MiddlewareAgent

  correct:

  from pydantic_ai_middleware.agent import MiddlewareAgent

  This class was removed from `pydantic-ai-middleware` v0.3.0+ (now `pydantic-ai-shields`) and native support was added to `pydantic-ai` v1.71+
• AgentMiddleware
  wrong:

  from pydantic_ai_shields.middleware import AgentMiddleware

  correct:

  from pydantic_ai_middleware.middleware import AgentMiddleware

  This base class was removed from `pydantic-ai-middleware` v0.3.0+ (now `pydantic-ai-shields`) and native support was added to `pydantic-ai` v1.71+
• CostTrackingMiddleware
  wrong:

  from pydantic_ai_shields.middleware import CostTrackingMiddleware

  correct:

  from pydantic_ai_middleware.middleware import CostTrackingMiddleware

  This specific middleware was moved to `pydantic-ai-shields` as of `pydantic-ai-middleware` v0.3.0
• MiddlewareChain
  wrong:

  from pydantic_ai_shields.chain import MiddlewareChain

  correct:

  from pydantic_ai_middleware.chain import MiddlewareChain

  This class was moved to `pydantic-ai-shields` as of `pydantic-ai-middleware` v0.3.0

Quickstart

This quickstart demonstrates how to wrap a `pydantic-ai` agent with custom middleware using `MiddlewareAgent` and `AgentMiddleware` from `pydantic-ai-middleware` v0.2.4. It defines a simple middleware to log before and after agent execution.

import os
from pydantic_ai import PydanticAI
from pydantic_ai.agents import Agent
from pydantic_ai_middleware.agent import MiddlewareAgent
from pydantic_ai_middleware.middleware import AgentMiddleware
import asyncio

class MyCustomMiddleware(AgentMiddleware):
    async def before_run(self, agent: Agent, input_data: str, **kwargs):
        print(f"Middleware: Before run with input: {input_data}")
        return input_data

    async def after_run(self, agent: Agent, output_data: str, **kwargs):
        print(f"Middleware: After run with output: {output_data}")
        return output_data

class MyAgent(Agent):
    def run(self, topic: str) -> str:
        """Generate a short story about the given topic."""
        return f"A story about {topic}..."

ai = PydanticAI(
    api_key=os.environ.get("OPENAI_API_KEY", "sk-DUMMY"), # Use a dummy key if not set
    model="gpt-3.5-turbo",
)
my_agent = MyAgent(ai=ai)

middleware_agent = MiddlewareAgent(
    wrapped_agent=my_agent,
    middleware=[MyCustomMiddleware()]
)

async def main():
    print("\n--- Running Middleware Agent ---")
    result = await middleware_agent.run("a brave knight")
    print(f"Agent final result: {result}")
    print("--------------------------------")

if __name__ == "__main__":
    asyncio.run(main())