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[py]
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",
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
- Explore Scenario Basics for advanced testing patterns
- Learn about Scripted Simulations for precise control
- Check out more Agent Integration patterns for other frameworks