Skip to content

Pydantic AI Integration

Learn how to integrate Pydantic AI agents with the Scenario testing framework

Pydantic AI agents work seamlessly with Scenario through the AgentAdapter interface. The key is managing conversation history and handling the iterative agent execution.

Basic Integration Pattern

Given this Pydantic AI agent example:

from pydantic_ai import Agent
 
agent = Agent(
    "openai:gpt-4.1-mini",
    system_prompt="""Your system prompt here...""",
)
 
@agent.tool_plain
def get_customer_order_history() -> List[OrderSummaryResponse]:
    """Get the current customer order history"""
    return http_GET_customer_order_history()
 
@agent.tool_plain
def get_order_status(order_id: str) -> OrderStatusResponse:
    """Get the status of a specific order"""
    return http_GET_order_status(order_id)
 
@agent.tool_plain
def escalate_to_human() -> dict[str, str]:
    """Escalate to human support"""
    return {"url": "https://support.example.com/tickets", "type": "escalation"}

You can integrate it with Scenario like this:

import scenario
from pydantic_ai.messages import ModelMessage
from pydantic_graph import End
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.openai import OpenAIProvider
 
class PydanticAIAgentAdapter(scenario.AgentAdapter):
    def __init__(self):
        # In-Memory History - Pydantic AI doesn't have built-in memory yet: https://github.com/pydantic/pydantic-ai/issues/530
        self.history: dict[str, List[ModelMessage]] = {}
 
    async def call(self, input: scenario.AgentInput) -> scenario.AgentReturnTypes:
        if input.thread_id not in self.history:
            self.history[input.thread_id] = []
 
        user_message = input.last_new_user_message_str()
 
        # Run the agent iteratively
        async with agent.iter(
            user_message,
            message_history=self.history[input.thread_id],
        ) as agent_run:
            next_node = agent_run.next_node
            nodes = [next_node]
 
            # Process all nodes until completion
            while not isinstance(next_node, End):
                next_node = await agent_run.next(next_node)
                nodes.append(next_node)
 
            if not agent_run.result:
                raise Exception("No result from agent")
 
            # Get new messages and update history
            new_messages = agent_run.result.new_messages()
            self.history[input.thread_id] += new_messages
 
            # Convert to OpenAI format for Scenario
            new_messages_openai_format = await OpenAIModel(
                "any", provider=OpenAIProvider(api_key="bogus")
            )._map_messages(new_messages)
 
            return new_messages_openai_format

Full Example Project

For a complete working example with a customer support agent, including tools, system prompts, and comprehensive tests, check out the create-agent-app project.

Next Steps