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