1
from pydantic import BaseModel, Field
2
from typing import List, Optional
3
from pydantic_ai import Agent
4

5
class ChainStep(BaseModel):
6
    """Represents a single step in our prompt chain."""
7
    prompt_template: str = Field(..., description="Template for the prompt, with {input} placeholder")
8
    name: str = Field(..., description="Name of this step for debugging")
9
    validation_rules: Optional[List[str]] = Field(None, description="Rules to validate the output")
10

11
class PromptChain:
12
    """Executes a sequence of prompts, passing output from each to the next."""
13

14
    def __init__(self, steps: List[ChainStep], model_type: str = "openai:gpt-4o"):
15
        self.steps = steps
16
        self.agent = Agent(model_type)
17

18
    async def execute(self, input_data: str) -> str:
19
        """Execute the chain on the input data."""
20
        result = input_data
21

22
        for step in self.steps:
23
            # Format the prompt with the previous result
24
            prompt = step.prompt_template.format(input=result)
25

26
            # Call the LLM using Pydantic AI
27
            response = await self.agent.run(prompt)
28
            result = response.data
29

30
            # Validate if rules are specified
31
            if step.validation_rules:
32
                self._validate_output(result, step.validation_rules)
33

34
        return result
35

36
    def _validate_output(self, output: str, rules: List[str]):
37
        """Apply validation rules to the output."""
38
        # Implementation depends on your specific rules
39
        pass
40

41
# Example: Building a research summarization chain
42
research_chain = PromptChain([
43
    ChainStep(
44
        name="extract_key_points",
45
        prompt_template="Extract the key points from this text: {input}"
46
    ),
47
    ChainStep(
48
        name="organize_themes",
49
        prompt_template="Organize these key points into themes: {input}"
50
    ),
51
    ChainStep(
52
        name="write_summary",
53
        prompt_template="Write a concise summary based on these themes: {input}"
54
    )
55
])

1
from pydantic import BaseModel, Field
2
from typing import Dict, Optional
3
from pydantic_ai import Agent
4

5
class RouteRequest(BaseModel):
6
    """Input request that needs routing."""
7
    query: str = Field(..., description="The user query to classify and route")
8
    context: Optional[Dict] = Field(None, description="Additional context for routing decisions")
9

10
class RouteResponse(BaseModel):
11
    """Routing decision from our classifier."""
12
    route: str = Field(..., description="Selected route identifier")
13
    confidence: float = Field(..., ge=0, le=1, description="Confidence in this routing decision")
14
    reasoning: str = Field(..., description="Explanation for the routing choice")
15

16
class Router:
17
    """Intelligently routes queries to specialized handlers."""
18

19
    def __init__(self, classifier_prompt: str, handlers: Dict[str, Agent]):
20
        self.classifier_prompt = classifier_prompt
21
        self.handlers = handlers
22
        # Use a structured output type for reliable routing decisions
23
        self.classifier = Agent(
24
            'anthropic:claude-3-sonnet-20240229',
25
            result_type=RouteResponse
26
        )
27

28
    async def route_and_process(self, request: RouteRequest) -> str:
29
        """Classify the request and route to appropriate handler."""
30
        # Build a description of available routes
31
        routes_desc = "\n".join([
32
            f"- {route_id}: {handler.system_prompt[:50]}..."
33
            for route_id, handler in self.handlers.items()
34
        ])
35

36
        # Classify the request
37
        prompt = self.classifier_prompt.format(
38
            query=request.query,
39
            available_routes=routes_desc,
40
            **request.context or {}
41
        )
42

43
        route_response = await self.classifier.run(prompt)
44

45
        # Handle unknown routes gracefully
46
        handler = self.handlers.get(route_response.data.route)
47
        if not handler:
48
            return f"I'm not sure how to handle that. Routing confidence was {route_response.data.confidence}"
49

50
        # Process with the specialized handler
51
        result = await handler.run(request.query)
52
        return result.data
53

54
# Example: Customer support router
55
support_router = Router(
56
    classifier_prompt="""Classify this customer query and select the best handler:
57
    Query: {query}
58

59
    Available handlers:
60
    {available_routes}
61

62
    Select the most appropriate handler based on the query type.""",
63
    handlers={
64
        "technical": Agent("openai:gpt-4o", system_prompt="You are a technical support specialist..."),
65
        "billing": Agent("openai:gpt-4o", system_prompt="You are a billing support agent..."),
66
        "general": Agent("openai:gpt-4o", system_prompt="You are a friendly customer service agent...")
67
    }
68
)

1
from pydantic import BaseModel, Field
2
from typing import List, Dict, Any
3
from pydantic_ai import Agent
4
import asyncio
5

6
class SectioningTask(BaseModel):
7
    """A subtask that can be processed independently."""
8
    section_id: str = Field(..., description="Unique identifier for this section")
9
    prompt_template: str = Field(..., description="Prompt template for processing this section")
10
    data: Dict[str, Any] = Field(default_factory=dict, description="Section-specific data")
11

12
class ParallelExecutor:
13
    """Execute tasks in parallel for improved performance."""
14

15
    def __init__(self, model_type: str = "openai:gpt-4o"):
16
        self.model_type = model_type
17

18
    async def execute_sectioning(self, tasks: List[SectioningTask], shared_data: Dict[str, Any] = None) -> Dict[str, str]:
19
        """Process multiple sections in parallel."""
20
        shared_data = shared_data or {}
21

22
        async def process_section(task):
23
            # Create an agent for this section
24
            agent = Agent(self.model_type)
25

26
            # Combine shared and section-specific data
27
            prompt = task.prompt_template.format(**{**shared_data, **task.data})
28

29
            # Process asynchronously
30
            result = await agent.run(prompt)
31
            return task.section_id, result.data
32

33
        # Execute all sections concurrently
34
        coroutines = [process_section(task) for task in tasks]
35
        results_list = await asyncio.gather(*coroutines)
36

37
        # Convert to dictionary for easy access
38
        return {section_id: result for section_id, result in results_list}
39

40
    async def execute_voting(self, prompt: str, num_votes: int = 3) -> Dict[str, Any]:
41
        """Get multiple perspectives on the same prompt."""
42
        async def get_vote(vote_num):
43
            agent = Agent(self.model_type, system_prompt=f"You are assistant {vote_num}. Provide your perspective.")
44
            result = await agent.run(prompt)
45
            return result.data
46

47
        # Collect all votes
48
        votes = await asyncio.gather(*[get_vote(i) for i in range(num_votes)])
49

50
        # Analyze consensus (simplified example)
51
        return {
52
            "votes": votes,
53
            "consensus": self._find_consensus(votes)
54
        }
55

56
    def _find_consensus(self, votes: List[str]) -> str:
57
        """Analyze votes to find consensus (simplified)."""
58
        # In practice, you might use another LLM call or more sophisticated analysis
59
        return "Majority opinion: " + votes[0]  # Placeholder
60

61
# Example: Parallel document analysis
62
async def analyze_document(document: str):
63
    executor = ParallelExecutor()
64

65
    # Split into sections for parallel processing
66
    tasks = [
67
        SectioningTask(
68
            section_id="summary",
69
            prompt_template="Summarize this document: {document}",
70
            data={"document": document}
71
        ),
72
        SectioningTask(
73
            section_id="key_points",
74
            prompt_template="Extract key points from: {document}",
75
            data={"document": document}
76
        ),
77
        SectioningTask(
78
            section_id="sentiment",
79
            prompt_template="Analyze the sentiment of: {document}",
80
            data={"document": document}
81
        )
82
    ]
83

84
    results = await executor.execute_sectioning(tasks)
85
    return results

1
from pydantic import BaseModel, Field
2
from typing import List, Dict, Any
3
from pydantic_ai import Agent
4

5
class SubTask(BaseModel):
6
    """A subtask created by the orchestrator."""
7
    task_id: str = Field(..., description="Unique identifier")
8
    type: str = Field(..., description="Type of subtask - determines which worker to use")
9
    description: str = Field(..., description="What needs to be done")
10
    dependencies: List[str] = Field(default_factory=list, description="IDs of tasks this depends on")
11
    context: Dict[str, Any] = Field(default_factory=dict, description="Additional context for the worker")
12

13
class TaskDecomposition(BaseModel):
14
    """The orchestrator's plan for solving a complex task."""
15
    subtasks: List[SubTask]
16
    execution_order: List[str] = Field(..., description="Suggested order of execution")
17

18
class OrchestratorSystem:
19
    """Dynamic task decomposition and execution system."""
20

21
    def __init__(
22
        self,
23
        orchestrator_prompt: str,
24
        worker_prompts: Dict[str, str],
25
        orchestrator_model: str = "anthropic:claude-3-opus-20240229",
26
        worker_model: str = "openai:gpt-4o"
27
    ):
28
        self.orchestrator_prompt = orchestrator_prompt
29
        self.worker_prompts = worker_prompts
30

31
        # Orchestrator with structured output for reliable task decomposition
32
        self.orchestrator = Agent(
33
            orchestrator_model,
34
            result_type=TaskDecomposition,
35
            system_prompt="You are a task orchestrator. Break down complex tasks into manageable subtasks."
36
        )
37

38
        self.worker_model = worker_model
39
        self.completed_tasks = {}
40

41
    async def process_task(self, task: str, context: Dict[str, Any] = None) -> Dict[str, Any]:
42
        """Process a complex task through orchestration."""
43
        context = context or {}
44

45
        # Step 1: Decompose the task
46
        decompose_prompt = self.orchestrator_prompt.format(
47
            task=task,
48
            available_workers=list(self.worker_prompts.keys()),
49
            **context
50
        )
51

52
        decomposition = await self.orchestrator.run(decompose_prompt)
53
        subtasks = decomposition.data.subtasks
54

55
        # Step 2: Execute subtasks respecting dependencies
56
        results = {}
57
        for subtask in self._order_by_dependencies(subtasks):
58
            # Wait for dependencies
59
            await self._wait_for_dependencies(subtask, results)
60

61
            # Execute subtask
62
            worker_result = await self._execute_subtask(subtask, results)
63
            results[subtask.task_id] = worker_result
64

65
        # Step 3: Synthesize results
66
        synthesis_result = await self._synthesize_results(task, subtasks, results)
67

68
        return {
69
            "task": task,
70
            "subtasks": [st.dict() for st in subtasks],
71
            "results": results,
72
            "final_result": synthesis_result
73
        }
74

75
    async def _execute_subtask(self, subtask: SubTask, completed_results: Dict[str, Any]) -> str:
76
        """Execute a single subtask with the appropriate worker."""
77
        worker_prompt = self.worker_prompts.get(subtask.type)
78
        if not worker_prompt:
79
            return f"No worker available for task type: {subtask.type}"
80

81
        # Create context including dependency results
82
        context = {
83
            "description": subtask.description,
84
            "dependencies": {dep_id: completed_results.get(dep_id) for dep_id in subtask.dependencies},
85
            **subtask.context
86
        }
87

88
        worker = Agent(
89
            self.worker_model,
90
            system_prompt=worker_prompt
91
        )
92

93
        result = await worker.run(str(context))
94
        return result.data
95

96
    def _order_by_dependencies(self, subtasks: List[SubTask]) -> List[SubTask]:
97
        """Order subtasks respecting dependencies (simplified topological sort)."""
98
        # In practice, implement proper topological sorting
99
        return sorted(subtasks, key=lambda x: len(x.dependencies))
100

101
    async def _wait_for_dependencies(self, subtask: SubTask, results: Dict[str, Any]):
102
        """Wait for all dependencies to complete."""
103
        # In a real implementation, this would handle async coordination
104
        pass
105

106
    async def _synthesize_results(self, original_task: str, subtasks: List[SubTask], results: Dict[str, Any]) -> str:
107
        """Combine all results into a final answer."""
108
        synthesis_agent = Agent(
109
            self.orchestrator.model_name,
110
            system_prompt="You are a synthesis expert. Combine subtask results into a coherent final answer."
111
        )
112

113
        synthesis_prompt = f"""
114
        Original task: {original_task}
115

116
        Completed subtasks and results:
117
        {self._format_results_for_synthesis(subtasks, results)}
118

119
        Synthesize these results into a comprehensive final solution.
120
        """
121

122
        final_result = await synthesis_agent.run(synthesis_prompt)
123
        return final_result.data
124

125
    def _format_results_for_synthesis(self, subtasks: List[SubTask], results: Dict[str, Any]) -> str:
126
        """Format results for the synthesis step."""
127
        formatted = []
128
        for subtask in subtasks:
129
            result = results.get(subtask.task_id, "No result")
130
            formatted.append(f"- {subtask.description}: {result}")
131
        return "\n".join(formatted)
132

133
# Example: Research orchestrator
134
research_orchestrator = OrchestratorSystem(
135
    orchestrator_prompt="""Break down this research task into subtasks:
136
    Task: {task}
137

138
    Available workers: {available_workers}
139

140
    Create a plan with specific subtasks that can be executed by the available workers.""",
141
    worker_prompts={
142
        "search": "You are a search specialist. Find relevant information based on the given query.",
143
        "analyze": "You are an analysis expert. Analyze the provided information and extract insights.",
144
        "synthesize": "You are a synthesis specialist. Combine multiple pieces of information coherently.",
145
        "fact_check": "You are a fact checker. Verify the accuracy of the provided claims."
146
    }
147
)

1
from pydantic import BaseModel, Field
2
from typing import List, Dict, Literal
3
from pydantic_ai import Agent
4

5
class EvaluationCriteria(BaseModel):
6
    """Criteria for evaluating generated content."""
7
    name: str = Field(..., description="Name of this criterion")
8
    description: str = Field(..., description="What this criterion measures")
9
    threshold: float = Field(..., ge=0, le=1, description="Minimum score to pass")
10
    weight: float = Field(1.0, description="Importance weight for this criterion")
11

12
class Evaluation(BaseModel):
13
    """Structured evaluation of generated content."""
14
    status: Literal["PASS", "NEEDS_IMPROVEMENT", "FAIL"] = Field(..., description="Overall status")
15
    criteria_scores: Dict[str, float] = Field(..., description="Individual criterion scores")
16
    feedback: str = Field(..., description="Specific, actionable feedback for improvement")
17
    strengths: List[str] = Field(default_factory=list, description="What worked well")
18
    improvements: List[str] = Field(default_factory=list, description="Specific improvements needed")
19

20
    @property
21
    def passed(self) -> bool:
22
        """Check if the evaluation passed all criteria."""
23
        return self.status == "PASS"
24

25
    @property
26
    def overall_score(self) -> float:
27
        """Calculate weighted overall score."""
28
        if not self.criteria_scores:
29
            return 0.0
30
        return sum(score for score in self.criteria_scores.values()) / len(self.criteria_scores)
31

32
class EvaluatorOptimizerSystem:
33
    """Iterative improvement through evaluation and optimization."""
34

35
    def __init__(
36
        self,
37
        optimizer_prompt: str,
38
        evaluator_prompt: str,
39
        criteria: List[EvaluationCriteria],
40
        max_iterations: int = 5,
41
        optimizer_model: str = "openai:gpt-4o",
42
        evaluator_model: str = "anthropic:claude-3-sonnet-20240229"
43
    ):
44
        self.optimizer_prompt = optimizer_prompt
45
        self.evaluator_prompt = evaluator_prompt
46
        self.criteria = criteria
47
        self.max_iterations = max_iterations
48

49
        # Different models for different strengths
50
        self.optimizer = Agent(optimizer_model, system_prompt="You are a content creator focused on quality.")
51
        self.evaluator = Agent(
52
            evaluator_model,
53
            result_type=Evaluation,
54
            system_prompt="You are a critical evaluator. Provide honest, constructive feedback."
55
        )
56

57
    async def optimize(self, task: str, context: Dict[str, str] = None) -> Dict[str, Any]:
58
        """Run the optimization loop."""
59
        context = context or {}
60
        history = []
61

62
        # Initial generation
63
        content = await self._generate_initial(task, context)
64

65
        for iteration in range(self.max_iterations):
66
            # Evaluate current content
67
            evaluation = await self._evaluate_content(task, content, context)
68

69
            # Track history
70
            history.append({
71
                "iteration": iteration,
72
                "content": content,
73
                "evaluation": evaluation.dict(),
74
                "score": evaluation.overall_score
75
            })
76

77
            # Check if we're done
78
            if evaluation.passed:
79
                break
80

81
            # Generate improved version
82
            content = await self._improve_content(task, content, evaluation, context)
83

84
        return {
85
            "task": task,
86
            "final_content": content,
87
            "iterations": len(history),
88
            "passed": evaluation.passed,
89
            "final_score": evaluation.overall_score,
90
            "history": history
91
        }
92

93
    async def _generate_initial(self, task: str, context: Dict[str, str]) -> str:
94
        """Generate the initial content."""
95
        prompt = self.optimizer_prompt.format(
96
            task=task,
97
            **context
98
        )
99
        response = await self.optimizer.run(prompt)
100
        return response.data
101

102
    async def _evaluate_content(self, task: str, content: str, context: Dict[str, str]) -> Evaluation:
103
        """Evaluate content against criteria."""
104
        criteria_text = "\n".join([
105
            f"- {c.name}: {c.description} (minimum score: {c.threshold})"
106
            for c in self.criteria
107
        ])
108

109
        prompt = self.evaluator_prompt.format(
110
            task=task,
111
            content=content,
112
            criteria=criteria_text,
113
            **context
114
        )
115

116
        response = await self.evaluator.run(prompt)
117
        return response.data
118

119
    async def _improve_content(self, task: str, content: str, evaluation: Evaluation, context: Dict[str, str]) -> str:
120
        """Generate improved content based on feedback."""
121
        improvement_prompt = f"""
122
        Task: {task}
123

124
        Previous attempt:
125
        {content}
126

127
        Evaluation feedback:
128
        {evaluation.feedback}
129

130
        Specific improvements needed:
131
        {chr(10).join(f'- {imp}' for imp in evaluation.improvements)}
132

133
        Generate an improved version that addresses all the feedback while maintaining the strengths.
134
        """
135

136
        response = await self.optimizer.run(improvement_prompt)
137
        return response.data
138

139
# Example: Blog post optimizer
140
blog_optimizer = EvaluatorOptimizerSystem(
141
    optimizer_prompt="Write a blog post about: {task}\n\nTone: {tone}\nAudience: {audience}",
142
    evaluator_prompt="""Evaluate this blog post:
143
    Task: {task}
144
    Content: {content}
145

146
    Criteria:
147
    {criteria}
148

149
    Provide specific, actionable feedback for improvement.""",
150
    criteria=[
151
        EvaluationCriteria(
152
            name="clarity",
153
            description="Ideas are expressed clearly and logically",
154
            threshold=0.8
155
        ),
156
        EvaluationCriteria(
157
            name="engagement",
158
            description="Content is engaging and holds reader attention",
159
            threshold=0.7
160
        ),
161
        EvaluationCriteria(
162
            name="accuracy",
163
            description="Information is accurate and well-researched",
164
            threshold=0.9
165
        )
166
    ]
167
)

1
from pydantic import BaseModel, Field, validator
2
from datetime import datetime
3
from typing import List, Optional
4
from pydantic_ai import Agent
5

6
# Define what we expect from our agent
7
class CustomerInquiry(BaseModel):
8
    """Structured representation of a customer inquiry."""
9
    category: str = Field(description="Type of inquiry: technical, billing, or general")
10
    urgency: int = Field(ge=1, le=5, description="Urgency level from 1 (low) to 5 (critical)")
11
    summary: str = Field(description="Brief summary of the issue")
12
    customer_sentiment: float = Field(ge=-1, le=1, description="Sentiment score from -1 (angry) to 1 (happy)")
13
    requires_human: bool = Field(description="Whether this needs human intervention")
14
    suggested_actions: List[str] = Field(description="Recommended next steps")
15

16
    @validator('category')
17
    def validate_category(cls, v):
18
        valid_categories = ['technical', 'billing', 'general']
19
        if v.lower() not in valid_categories:
20
            raise ValueError(f"Category must be one of {valid_categories}")
21
        return v.lower()
22

23
    @validator('suggested_actions')
24
    def validate_actions(cls, v):
25
        if not v:
26
            raise ValueError("At least one suggested action is required")
27
        return v
28

29
# Create an agent that outputs structured data
30
support_classifier = Agent(
31
    'openai:gpt-4o',
32
    result_type=CustomerInquiry,
33
    system_prompt="""You are a customer support classifier. Analyze customer messages and
34
    extract structured information to help route and prioritize support tickets."""
35
)
36

37
# Use it with confidence
38
async def process_customer_message(message: str) -> CustomerInquiry:
39
    result = await support_classifier.run(
40
        f"Analyze this customer message: {message}"
41
    )
42
    return result.data  # This is guaranteed to be a valid CustomerInquiry
43

44
# Example usage
45
inquiry = await process_customer_message("My internet has been down for 3 days and I'm furious!")
46
print(f"Category: {inquiry.category}")
47
print(f"Urgency: {inquiry.urgency}/5")
48
print(f"Needs human: {inquiry.requires_human}")

1
from dataclasses import dataclass
2
from pydantic_ai import Agent, RunContext
3
from typing import Dict, List
4
import asyncio
5

6
@dataclass
7
class Dependencies:
8
    """Dependencies that will be injected into tool calls."""
9
    database: object  # Your database connection
10
    api_client: object  # External API client
11
    user_id: str  # Current user context
12

13
# Create an agent with dependencies
14
agent = Agent(
15
    'openai:gpt-4o',
16
    deps_type=Dependencies,
17
    system_prompt='You are a helpful assistant with access to user data and external services.'
18
)
19

20
@agent.tool
21
async def get_user_orders(ctx: RunContext[Dependencies]) -> List[Dict]:
22
    """Fetch user's order history from the database."""
23
    # Note how we access injected dependencies via ctx.deps
24
    orders = await ctx.deps.database.get_orders(ctx.deps.user_id)
25
    return [
26
        {
27
            "order_id": order.id,
28
            "date": order.date.isoformat(),
29
            "total": float(order.total),
30
            "status": order.status
31
        }
32
        for order in orders
33
    ]
34

35
@agent.tool
36
async def check_shipping_status(ctx: RunContext[Dependencies], order_id: str) -> Dict:
37
    """Check shipping status with external shipping API."""
38
    # Tools can take parameters and access dependencies
39
    tracking = await ctx.deps.api_client.get_tracking(order_id)
40
    return {
41
        "order_id": order_id,
42
        "status": tracking.status,
43
        "location": tracking.current_location,
44
        "estimated_delivery": tracking.eta.isoformat() if tracking.eta else None
45
    }
46

47
@agent.tool
48
def calculate_loyalty_points(ctx: RunContext[Dependencies], order_total: float) -> int:
49
    """Calculate loyalty points for an order (synchronous tools work too!)."""
50
    # Business logic can be encapsulated in tools
51
    points_rate = 10  # 10 points per dollar
52
    bonus_multiplier = 2 if order_total > 100 else 1
53
    return int(order_total * points_rate * bonus_multiplier)
54

55
# Use the agent with injected dependencies
56
async def handle_customer_query(query: str, user_id: str):
57
    deps = Dependencies(
58
        database=db_connection,
59
        api_client=shipping_api,
60
        user_id=user_id
61
    )
62

63
    result = await agent.run(query, deps=deps)
64
    return result.data
65

66
# Example: The agent can now use tools intelligently
67
response = await handle_customer_query(
68
    "What's the status of my recent orders and how many points did I earn?",
69
    user_id="user123"
70
)

1
from pydantic_ai.models.test import TestModel
2
from pydantic_ai.models.function import FunctionModel, AgentInfo
3
from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart
4

5
# Production code stays the same
6
order_agent = Agent(
7
    'openai:gpt-4o',
8
    deps_type=Dependencies,
9
    system_prompt='You help customers with their orders.'
10
)
11

12
@order_agent.tool
13
async def get_order_details(ctx: RunContext[Dependencies], order_id: str) -> Dict:
14
    """Fetch order details from database."""
15
    return await ctx.deps.database.get_order(order_id)
16

17
# For testing, we can inject mock dependencies
18
class MockDatabase:
19
    async def get_order(self, order_id: str) -> Dict:
20
        # Return test data instead of hitting real database
21
        return {
22
            "order_id": order_id,
23
            "status": "shipped",
24
            "items": ["Test Item 1", "Test Item 2"]
25
        }
26

27
# Test with mocked dependencies and model
28
async def test_order_lookup():
29
    test_deps = Dependencies(
30
        database=MockDatabase(),
31
        api_client=None,  # Not needed for this test
32
        user_id="test_user"
33
    )
34

35
    # Use TestModel to avoid API calls
36
    with order_agent.override(model=TestModel()):
37
        result = await order_agent.run(
38
            "What's the status of order ABC123?",
39
            deps=test_deps
40
        )
41

42
        # Assertions about the result
43
        assert "shipped" in result.data.lower()
44

45
# For more complex testing scenarios
46
async def custom_model_function(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
47
    """Custom function that simulates model responses based on input."""
48
    user_message = messages[-1].content
49

50
    if "order" in user_message.lower():
51
        # Simulate the model calling our tool
52
        return ModelResponse(
53
            parts=[TextPart("I'll check that order for you.")],
54
            tool_calls=[{
55
                "tool_name": "get_order_details",
56
                "args": {"order_id": "ABC123"}
57
            }]
58
        )
59

60
    return ModelResponse(parts=[TextPart("How can I help you?")])
61

62
# Test with custom model behavior
63
async def test_complex_interaction():
64
    with order_agent.override(model=FunctionModel(custom_model_function)):
65
        result = await order_agent.run(
66
            "Check order ABC123",
67
            deps=test_deps
68
        )
69
        # Now we can test the full flow including tool calls

1
from pydantic_ai import Agent
2
from pydantic_ai.models.test import TestModel
3
from pydantic_ai.models.function import FunctionModel, AgentInfo
4
import pytest
5
from datetime import datetime
6

7
# 1. Unit Testing with TestModel
8
class TestCustomerSupportAgent:
9
    def setup_method(self):
10
        """Set up test fixtures."""
11
        self.agent = Agent(
12
            'openai:gpt-4o',
13
            system_prompt="You are a helpful customer support agent."
14
        )
15

16
    def test_basic_response(self):
17
        """Test that agent responds appropriately to basic queries."""
18
        # TestModel returns predictable responses
19
        with self.agent.override(model=TestModel()):
20
            result = self.agent.run_sync("Hello, I need help")
21

22
            # Check that we got a response
23
            assert result.data is not None
24
            assert isinstance(result.data, str)
25
            assert len(result.data) > 0
26

27
    def test_tool_calling(self):
28
        """Test that agent calls tools correctly."""
29
        tool_called = False
30

31
        @self.agent.tool
32
        def check_order_status(order_id: str) -> str:
33
            nonlocal tool_called
34
            tool_called = True
35
            return f"Order {order_id} is shipped"
36

37
        # Custom model that always calls our tool
38
        async def model_function(messages, info):
39
            return ModelResponse(
40
                parts=[TextPart("Let me check that order")],
41
                tool_calls=[{"tool_name": "check_order_status", "args": {"order_id": "123"}}]
42
            )
43

44
        with self.agent.override(model=FunctionModel(model_function)):
45
            result = self.agent.run_sync("Check order 123")
46
            assert tool_called
47
            assert "shipped" in result.data
48

49
# 2. Integration Testing with Mocked Services
50
class TestIntegrationScenarios:
51
    @pytest.mark.asyncio
52
    async def test_multi_tool_workflow(self):
53
        """Test complex workflows involving multiple tools."""
54
        agent = Agent(
55
            'openai:gpt-4o',
56
            deps_type=Dependencies
57
        )
58

59
        calls_made = []
60

61
        @agent.tool
62
        async def search_products(ctx: RunContext[Dependencies], query: str) -> List[Dict]:
63
            calls_made.append(('search', query))
64
            return [
65
                {"id": "1", "name": "Product A", "price": 99.99},
66
                {"id": "2", "name": "Product B", "price": 149.99}
67
            ]
68

69
        @agent.tool
70
        async def check_inventory(ctx: RunContext[Dependencies], product_id: str) -> bool:
71
            calls_made.append(('inventory', product_id))
72
            return True
73

74
        @agent.tool
75
        async def calculate_shipping(ctx: RunContext[Dependencies], product_id: str, zip_code: str) -> float:
76
            calls_made.append(('shipping', product_id, zip_code))
77
            return 9.99
78

79
        # Mock the model to execute a specific workflow
80
        async def workflow_model(messages, info):
81
            # This simulates the LLM orchestrating multiple tool calls
82
            return ModelResponse(
83
                parts=[TextPart("I'll help you find products and check shipping")],
84
                tool_calls=[
85
                    {"tool_name": "search_products", "args": {"query": "laptop"}},
86
                    {"tool_name": "check_inventory", "args": {"product_id": "1"}},
87
                    {"tool_name": "calculate_shipping", "args": {"product_id": "1", "zip_code": "10001"}}
88
                ]
89
            )
90

91
        test_deps = Dependencies(
92
            database=None,
93
            api_client=None,
94
            user_id="test"
95
        )
96

97
        with agent.override(model=FunctionModel(workflow_model)):
98
            result = await agent.run(
99
                "Find laptops and calculate shipping to 10001",
100
                deps=test_deps
101
            )
102

103
        # Verify the workflow executed correctly
104
        assert len(calls_made) == 3
105
        assert calls_made[0][0] == 'search'
106
        assert calls_made[1][0] == 'inventory'
107
        assert calls_made[2][0] == 'shipping'
108

109
# 3. End-to-End Testing with Recorded Responses
110
class TestEndToEnd:
111
    def test_customer_journey(self):
112
        """Test a complete customer interaction journey."""
113
        # For E2E tests, you might use recorded real LLM responses
114
        recorded_responses = {
115
            "greeting": "Hello! How can I help you today?",
116
            "order_query": "I'll check your order status right away.",
117
            "followup": "Is there anything else I can help you with?"
118
        }
119

120
        agent = Agent('openai:gpt-4o')
121

122
        # Override with recorded responses
123
        response_index = 0
124
        def get_next_response(messages, info):
125
            nonlocal response_index
126
            responses = list(recorded_responses.values())
127
            response = responses[response_index % len(responses)]
128
            response_index += 1
129
            return ModelResponse(parts=[TextPart(response)])
130

131
        with agent.override(model=FunctionModel(get_next_response)):
132
            # Simulate customer journey
133
            response1 = agent.run_sync("Hi")
134
            assert "Hello" in response1.data
135

136
            response2 = agent.run_sync("What's my order status?")
137
            assert "check" in response2.data
138

139
            response3 = agent.run_sync("Thanks!")
140
            assert "else" in response3.data

1
import logfire
2
from pydantic_ai import Agent
3
from datetime import datetime
4
import json
5

6
# Configure Logfire for your application
7
logfire.configure()
8

9
# Create an instrumented agent
10
agent = Agent(
11
    'openai:gpt-4o',
12
    system_prompt='You are a helpful assistant.',
13
    instrument=True  # Enable automatic instrumentation
14
)
15

16
# Custom metrics tracking
17
class AgentMetrics:
18
    def __init__(self):
19
        self.reset_daily_metrics()
20

21
    def reset_daily_metrics(self):
22
        self.metrics = {
23
            "total_requests": 0,
24
            "successful_requests": 0,
25
            "failed_requests": 0,
26
            "tool_calls": {},
27
            "response_times": [],
28
            "token_usage": {
29
                "prompt_tokens": 0,
30
                "completion_tokens": 0
31
            }
32
        }
33

34
    def track_request(self, duration: float, success: bool, tokens: Dict):
35
        self.metrics["total_requests"] += 1
36
        if success:
37
            self.metrics["successful_requests"] += 1
38
        else:
39
            self.metrics["failed_requests"] += 1
40

41
        self.metrics["response_times"].append(duration)
42
        self.metrics["token_usage"]["prompt_tokens"] += tokens.get("prompt_tokens", 0)
43
        self.metrics["token_usage"]["completion_tokens"] += tokens.get("completion_tokens", 0)
44

45
    def track_tool_call(self, tool_name: str):
46
        if tool_name not in self.metrics["tool_calls"]:
47
            self.metrics["tool_calls"][tool_name] = 0
48
        self.metrics["tool_calls"][tool_name] += 1
49

50
    def get_summary(self) -> Dict:
51
        response_times = self.metrics["response_times"]
52
        return {
53
            "total_requests": self.metrics["total_requests"],
54
            "success_rate": self.metrics["successful_requests"] / max(self.metrics["total_requests"], 1),
55
            "avg_response_time": sum(response_times) / len(response_times) if response_times else 0,
56
            "p95_response_time": sorted(response_times)[int(len(response_times) * 0.95)] if response_times else 0,
57
            "tool_usage": self.metrics["tool_calls"],
58
            "token_usage": self.metrics["token_usage"],
59
            "estimated_cost": self._estimate_cost()
60
        }
61

62
    def _estimate_cost(self) -> float:
63
        # Rough cost estimation (adjust based on your model)
64
        prompt_cost = 0.01 / 1000  # $0.01 per 1K tokens
65
        completion_cost = 0.03 / 1000  # $0.03 per 1K tokens
66

67
        return (
68
            self.metrics["token_usage"]["prompt_tokens"] * prompt_cost +
69
            self.metrics["token_usage"]["completion_tokens"] * completion_cost
70
        )
71

72
# Use in production with monitoring
73
metrics = AgentMetrics()
74

75
async def monitored_agent_call(query: str) -> Dict:
76
    start_time = datetime.now()
77

78
    try:
79
        # Log the request
80
        logfire.info("Agent request started", query=query)
81

82
        # Execute the agent
83
        result = await agent.run(query)
84

85
        # Track success
86
        duration = (datetime.now() - start_time).total_seconds()
87
        metrics.track_request(duration, True, result.usage)
88

89
        # Log successful completion
90
        logfire.info(
91
            "Agent request completed",
92
            duration=duration,
93
            tokens_used=result.usage
94
        )
95

96
        return {
97
            "success": True,
98
            "data": result.data,
99
            "duration": duration
100
        }
101

102
    except Exception as e:
103
        # Track failure
104
        duration = (datetime.now() - start_time).total_seconds()
105
        metrics.track_request(duration, False, {})
106

107
        # Log error with context
108
        logfire.error(
109
            "Agent request failed",
110
            error=str(e),
111
            query=query,
112
            duration=duration
113
        )
114

115
        return {
116
            "success": False,
117
            "error": str(e),
118
            "duration": duration
119
        }
120

121
# Periodic metrics reporting
122
async def report_metrics():
123
    summary = metrics.get_summary()
124
    logfire.info("Agent metrics summary", **summary)
125

126
    # Alert on concerning metrics
127
    if summary["success_rate"] < 0.95:
128
        logfire.warning("Low success rate detected", success_rate=summary["success_rate"])
129

130
    if summary["estimated_cost"] > 100:  # $100
131
        logfire.warning("High token usage cost", cost=summary["estimated_cost"])

1
# Customer Support RAG Agent
2
class ProductKnowledge(BaseModel):
3
    product_id: str
4
    features: List[str]
5
    price: float
6
    availability: bool
7
    similar_products: List[str]
8

9
support_agent = Agent(
10
    'openai:gpt-4o',
11
    result_type=ProductKnowledge,
12
    system_prompt="""You are an e-commerce support specialist. Use the product database
13
    to answer customer questions accurately and suggest alternatives when needed."""
14
)
15

16
@support_agent.tool
17
async def search_products(query: str) -> List[Dict]:
18
    # RAG implementation to search product database
19
    results = await vector_store.search(query, top_k=5)
20
    return [doc.to_dict() for doc in results]
21

22
# Order Management Agent
23
order_agent = Agent(
24
    'openai:gpt-4o',
25
    deps_type=OrderSystemDeps,
26
    system_prompt="You help customers manage their orders, including updates and returns."
27
)
28

29
@order_agent.tool
30
async def update_shipping_address(ctx: RunContext[OrderSystemDeps], order_id: str, new_address: str) -> bool:
31
    # Validate order status allows address change
32
    order = await ctx.deps.db.get_order(order_id)
33
    if order.status not in ['pending', 'processing']:
34
        raise ValueError("Cannot update address after order ships")
35

36
    # Update address
37
    return await ctx.deps.db.update_order_address(order_id, new_address)

1
# Multi-stage research workflow
2
research_orchestrator = OrchestratorSystem(
3
    orchestrator_prompt="""Break down this research question into specific sub-questions
4
    that can be investigated independently. Consider:
5
    - What information needs to be gathered?
6
    - What sources should be consulted?
7
    - What analysis is required?
8
    - How should findings be synthesized?""",
9
    worker_prompts={
10
        "literature_search": "Search academic literature for relevant papers on the given topic.",
11
        "data_extraction": "Extract key findings and data from the provided sources.",
12
        "statistical_analysis": "Perform statistical analysis on the extracted data.",
13
        "synthesis": "Synthesize findings into a coherent narrative with citations.",
14
        "fact_checking": "Verify claims and check for contradictions in the findings."
15
    }
16
)
17

18
# Example usage for research task
19
async def conduct_research(topic: str) -> ResearchReport:
20
    # The orchestrator dynamically creates a research plan
21
    result = await research_orchestrator.process_task(
22
        f"Research the effectiveness of {topic} including recent studies and meta-analyses",
23
        context={
24
            "output_format": "academic_paper",
25
            "citation_style": "APA",
26
            "max_sources": 50
27
        }
28
    )
29

30
    return ResearchReport(
31
        topic=topic,
32
        sections=result["results"],
33
        synthesis=result["final_result"],
34
        sources=extract_sources(result)
35
    )