pydantic-ai-middleware

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())