PRP Template for Pydantic AI Agents

This commit is contained in:
Cole Medin
2025-07-20 08:01:14 -05:00
parent 84d49cf30a
commit 1bcba59231
30 changed files with 6134 additions and 88 deletions

View File

@@ -0,0 +1,55 @@
# Execute Pydantic AI Agent PRP
Implement a Pydantic AI agent using the PRP file.
## PRP File: $ARGUMENTS
## Execution Process
1. **Load PRP**
- Read the specified Pydantic AI PRP file
- Understand all agent requirements and research findings
- Follow all instructions in the PRP and extend research if needed
- Review main_agent_reference patterns for implementation guidance
- Do more web searches and Pydantic AI documentation review as needed
2. **ULTRATHINK**
- Think hard before executing the agent implementation plan
- Break down agent development into smaller steps using your todos tools
- Use the TodoWrite tool to create and track your agent implementation plan
- Follow main_agent_reference patterns for configuration and structure
- Plan agent.py, tools.py, dependencies.py, and testing approach
3. **Execute the plan**
- Implement the Pydantic AI agent following the PRP
- Create agent with environment-based configuration (settings.py, providers.py)
- Use string output by default (no result_type unless structured output needed)
- Implement tools with @agent.tool decorators and proper error handling
- Add comprehensive testing with TestModel and FunctionModel
4. **Validate**
- Test agent import and instantiation
- Run TestModel validation for rapid development testing
- Test tool registration and functionality
- Run pytest test suite if created
- Verify agent follows main_agent_reference patterns
5. **Complete**
- Ensure all PRP checklist items done
- Test agent with example queries
- Verify security patterns (environment variables, error handling)
- Report completion status
- Read the PRP again to ensure complete implementation
6. **Reference the PRP**
- You can always reference the PRP again if needed
## Pydantic AI-Specific Patterns to Follow
- **Configuration**: Use environment-based setup like main_agent_reference
- **Output**: Default to string output, only use result_type when validation needed
- **Tools**: Use @agent.tool with RunContext for dependency injection
- **Testing**: Include TestModel validation for development
- **Security**: Environment variables for API keys, proper error handling
Note: If validation fails, use error patterns in PRP to fix and retry. Follow main_agent_reference for proven Pydantic AI implementation patterns.

View File

@@ -0,0 +1,94 @@
# Create PRP
## Feature file: $ARGUMENTS
Generate a complete PRP for general feature implementation with thorough research. Ensure context is passed to the AI agent to enable self-validation and iterative refinement. Read the feature file first to understand what needs to be created, how the examples provided help, and any other considerations.
The AI agent only gets the context you are appending to the PRP and training data. Assuma the AI agent has access to the codebase and the same knowledge cutoff as you, so its important that your research findings are included or referenced in the PRP. The Agent has Websearch capabilities, so pass urls to documentation and examples.
## Research Process
1. **Codebase Analysis**
- Search for similar features/patterns in the codebase
- Identify files to reference in PRP
- Note existing conventions to follow
- Check test patterns for validation approach
2. **External Research**
- Search for similar features/patterns online
- Library documentation (include specific URLs)
- Implementation examples (GitHub/StackOverflow/blogs)
- Best practices and common pitfalls
- Use Archon MCP server to gather latest Pydantic AI documentation
- Web search for specific patterns and best practices relevant to the agent type
- Research model provider capabilities and limitations
- Investigate tool integration patterns and security considerations
- Document async/sync patterns and testing strategies
3. **User Clarification** (if needed)
- Specific patterns to mirror and where to find them?
- Integration requirements and where to find them?
4. **Analyzing Initial Requirements**
- Read and understand the agent feature requirements
- Identify the type of agent needed (chat, tool-enabled, workflow, structured output)
- Determine required model providers and external integrations
- Assess complexity and scope of the agent implementation
5. **Agent Architecture Planning**
- Design agent structure (agent.py, tools.py, models.py, dependencies.py)
- Plan dependency injection patterns and external service integrations
- Design structured output models using Pydantic validation
- Plan tool registration and parameter validation strategies
- Design testing approach with TestModel/FunctionModel patterns
6. **Implementation Blueprint Creation**
- Create detailed agent implementation steps
- Plan model provider configuration and fallback strategies
- Design tool error handling and retry mechanisms
- Plan security implementation (API keys, input validation, rate limiting)
- Design validation loops with agent behavior testing
## PRP Generation
Using PRPs/templates/prp_pydantic_aibase.md as template:
### Critical Context to Include and pass to the AI agent as part of the PRP
- **Documentation**: URLs with specific sections
- **Code Examples**: Real snippets from codebase
- **Gotchas**: Library quirks, version issues
- **Patterns**: Existing approaches to follow
### Implementation Blueprint
- Start with pseudocode showing approach
- Reference real files for patterns
- Include error handling strategy
- list tasks to be completed to fullfill the PRP in the order they should be completed
### Validation Gates (Must be Executable) eg for python
```bash
# Syntax/Style
ruff check --fix && mypy .
# Unit Tests
uv run pytest tests/ -v
```
*** CRITICAL AFTER YOU ARE DONE RESEARCHING AND EXPLORING THE CODEBASE BEFORE YOU START WRITING THE PRP ***
*** ULTRATHINK ABOUT THE PRP AND PLAN YOUR APPROACH THEN START WRITING THE PRP ***
## Output
Save as: `PRPs/{feature-name}.md`
## Quality Checklist
- [ ] All necessary context included
- [ ] Validation gates are executable by AI
- [ ] References existing patterns
- [ ] Clear implementation path
- [ ] Error handling documented
Score the PRP on a scale of 1-10 (confidence level to succeed in one-pass implementation using claude codes)
Remember: The goal is one-pass implementation success through comprehensive context.

View File

@@ -0,0 +1,176 @@
# PydanticAI Context Engineering - Global Rules for AI Agent Development
This file contains the global rules and principles that apply to ALL PydanticAI agent development work. These rules are specialized for building production-grade AI agents with tools, memory, and structured outputs.
## 🔄 PydanticAI Core Principles
**IMPORTANT: These principles apply to ALL PydanticAI agent development:**
### Agent Development Workflow
- **Always start with INITIAL.md** - Define agent requirements before generating PRPs
- **Use the PRP pattern**: INITIAL.md → `/generate-pydantic-ai-prp INITIAL.md``/execute-pydantic-ai-prp PRPs/filename.md`
- **Follow validation loops** - Each PRP must include agent testing with TestModel/FunctionModel
- **Context is King** - Include ALL necessary PydanticAI patterns, examples, and documentation
### Research Methodology for AI Agents
- **Web search extensively** - Always research PydanticAI patterns and best practices
- **Study official documentation** - ai.pydantic.dev is the authoritative source
- **Pattern extraction** - Identify reusable agent architectures and tool patterns
- **Gotcha documentation** - Document async patterns, model limits, and context management issues
## 📚 Project Awareness & Context
- **Use consistent PydanticAI naming conventions** and agent structure patterns
- **Follow established agent directory organization** patterns (agent.py, tools.py, models.py)
- **Leverage PydanticAI examples extensively** - Study existing patterns before creating new agents
## 🧱 Agent Structure & Modularity
- **Never create files longer than 500 lines** - Split into modules when approaching limit
- **Organize agent code into clearly separated modules** grouped by responsibility:
- `agent.py` - Main agent definition and execution logic
- `tools.py` - Tool functions used by the agent
- `models.py` - Pydantic output models and dependency classes
- `dependencies.py` - Context dependencies and external service integrations
- **Use clear, consistent imports** - Import from pydantic_ai package appropriately
- **Use environment variables for API keys** - Never hardcode sensitive information
## 🤖 PydanticAI Development Standards
### Agent Creation Patterns
- **Use model-agnostic design** - Support multiple providers (OpenAI, Anthropic, Gemini)
- **Implement dependency injection** - Use deps_type for external services and context
- **Define structured outputs** - Use Pydantic models for result validation
- **Include comprehensive system prompts** - Both static and dynamic instructions
### Tool Integration Standards
- **Use @agent.tool decorator** for context-aware tools with RunContext[DepsType]
- **Use @agent.tool_plain decorator** for simple tools without context dependencies
- **Implement proper parameter validation** - Use Pydantic models for tool parameters
- **Handle tool errors gracefully** - Implement retry mechanisms and error recovery
### Model Provider Configuration
```python
# Use environment-based configuration, never hardcode model strings
from pydantic_settings import BaseSettings
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.models.openai import OpenAIModel
class Settings(BaseSettings):
# LLM Configuration
llm_provider: str = "openai"
llm_api_key: str
llm_model: str = "gpt-4"
llm_base_url: str = "https://api.openai.com/v1"
class Config:
env_file = ".env"
def get_llm_model():
settings = Settings()
provider = OpenAIProvider(
base_url=settings.llm_base_url,
api_key=settings.llm_api_key
)
return OpenAIModel(settings.llm_model, provider=provider)
```
### Testing Standards for AI Agents
- **Use TestModel for development** - Fast validation without API calls
- **Use FunctionModel for custom behavior** - Control agent responses in tests
- **Use Agent.override() for testing** - Replace models in test contexts
- **Test both sync and async patterns** - Ensure compatibility with different execution modes
- **Test tool validation** - Verify tool parameter schemas and error handling
## ✅ Task Management for AI Development
- **Break agent development into clear steps** with specific completion criteria
- **Mark tasks complete immediately** after finishing agent implementations
- **Update task status in real-time** as agent development progresses
- **Test agent behavior** before marking implementation tasks complete
## 📎 PydanticAI Coding Standards
### Agent Architecture
```python
# Follow main_agent_reference patterns - no result_type unless structured output needed
from pydantic_ai import Agent, RunContext
from dataclasses import dataclass
@dataclass
class AgentDependencies:
"""Dependencies for agent execution"""
api_key: str
session_id: str = None
# Simple agent with string output (default)
agent = Agent(
get_llm_model(), # Use environment-based configuration
deps_type=AgentDependencies,
system_prompt="You are a helpful assistant..."
)
@agent.tool
async def example_tool(
ctx: RunContext[AgentDependencies],
query: str
) -> str:
"""Tool with proper context access"""
return await external_api_call(ctx.deps.api_key, query)
```
### Security Best Practices
- **API key management** - Use environment variables, never commit keys
- **Input validation** - Use Pydantic models for all tool parameters
- **Rate limiting** - Implement proper request throttling for external APIs
- **Prompt injection prevention** - Validate and sanitize user inputs
- **Error handling** - Never expose sensitive information in error messages
### Common PydanticAI Gotchas
- **Async/sync mixing issues** - Be consistent with async/await patterns throughout
- **Model token limits** - Different models have different context limits, plan accordingly
- **Dependency injection complexity** - Keep dependency graphs simple and well-typed
- **Tool error handling failures** - Always implement proper retry and fallback mechanisms
- **Context state management** - Design stateless tools when possible for reliability
## 🔍 Research Standards for AI Agents
- **Use Archon MCP server** - Leverage available PydanticAI documentation via RAG
- **Study official examples** - ai.pydantic.dev/examples has working implementations
- **Research model capabilities** - Understand provider-specific features and limitations
- **Document integration patterns** - Include external service integration examples
## 🎯 Implementation Standards for AI Agents
- **Follow the PRP workflow religiously** - Don't skip agent validation steps
- **Always test with TestModel first** - Validate agent logic before using real models
- **Use existing agent patterns** rather than creating from scratch
- **Include comprehensive error handling** for tool failures and model errors
- **Test streaming patterns** when implementing real-time agent interactions
## 🚫 Anti-Patterns to Always Avoid
- ❌ Don't skip agent testing - Always use TestModel/FunctionModel for validation
- ❌ Don't hardcode model strings - Use environment-based configuration like main_agent_reference
- ❌ Don't use result_type unless structured output is specifically needed - default to string
- ❌ Don't ignore async patterns - PydanticAI has specific async/sync considerations
- ❌ Don't create complex dependency graphs - Keep dependencies simple and testable
- ❌ Don't forget tool error handling - Implement proper retry and graceful degradation
- ❌ Don't skip input validation - Use Pydantic models for all external inputs
## 🔧 Tool Usage Standards for AI Development
- **Use web search extensively** for PydanticAI research and documentation
- **Follow PydanticAI command patterns** for slash commands and agent workflows
- **Use agent validation loops** to ensure quality at each development step
- **Test with multiple model providers** to ensure agent compatibility
## 🧪 Testing & Reliability for AI Agents
- **Always create comprehensive agent tests** for tools, outputs, and error handling
- **Test agent behavior with TestModel** before using real model providers
- **Include edge case testing** for tool failures and model provider issues
- **Test both structured and unstructured outputs** to ensure agent flexibility
- **Validate dependency injection** works correctly in test environments
These global rules apply specifically to PydanticAI agent development and ensure production-ready AI applications with proper error handling, testing, and security practices.

View File

@@ -0,0 +1,25 @@
## FEATURE:
Build a simple customer support chatbot using PydanticAI that can answer basic questions and escalate complex issues to human agents.
## EXAMPLES:
- Basic chat agent with conversation memory
- Tool-enabled agent with web search capabilities
- Structured output agent for data validation
- Testing examples with TestModel and FunctionModel
## DOCUMENTATION:
- PydanticAI Official Documentation: https://ai.pydantic.dev/
- Agent Creation Guide: https://ai.pydantic.dev/agents/
- Tool Integration: https://ai.pydantic.dev/tools/
- Testing Patterns: https://ai.pydantic.dev/testing/
- Model Providers: https://ai.pydantic.dev/models/
## OTHER CONSIDERATIONS:
- Use environment variables for API key configuration instead of hardcoded model strings
- Keep agents simple - default to string output unless structured output is specifically needed
- Follow the main_agent_reference patterns for configuration and providers
- Always include comprehensive testing with TestModel for development

View File

@@ -0,0 +1,410 @@
---
name: "PydanticAI Agent PRP Template"
description: "Template for generating comprehensive PRPs for PydanticAI agent development projects"
---
## Purpose
[Brief description of the PydanticAI agent to be built and its main purpose]
## Core Principles
1. **PydanticAI Best Practices**: Deep integration with PydanticAI patterns for agent creation, tools, and structured outputs
2. **Production Ready**: Include security, testing, and monitoring for production deployments
3. **Type Safety First**: Leverage PydanticAI's type-safe design and Pydantic validation throughout
4. **Context Engineering Integration**: Apply proven context engineering workflows to AI agent development
5. **Comprehensive Testing**: Use TestModel and FunctionModel for thorough agent validation
## ⚠️ Implementation Guidelines: Don't Over-Engineer
**IMPORTANT**: Keep your agent implementation focused and practical. Don't build unnecessary complexity.
### What NOT to do:
-**Don't create dozens of tools** - Build only the tools your agent actually needs
-**Don't over-complicate dependencies** - Keep dependency injection simple and focused
-**Don't add unnecessary abstractions** - Follow main_agent_reference patterns directly
-**Don't build complex workflows** unless specifically required
-**Don't add structured output** unless validation is specifically needed (default to string)
### What TO do:
-**Start simple** - Build the minimum viable agent that meets requirements
-**Add tools incrementally** - Implement only what the agent needs to function
-**Follow main_agent_reference** - Use proven patterns, don't reinvent
-**Use string output by default** - Only add result_type when validation is required
-**Test early and often** - Use TestModel to validate as you build
### Key Question:
**"Does this agent really need this feature to accomplish its core purpose?"**
If the answer is no, don't build it. Keep it simple, focused, and functional.
---
## Goal
[Detailed description of what the agent should accomplish]
## Why
[Explanation of why this agent is needed and what problem it solves]
## What
### Agent Type Classification
- [ ] **Chat Agent**: Conversational interface with memory and context
- [ ] **Tool-Enabled Agent**: Agent with external tool integration capabilities
- [ ] **Workflow Agent**: Multi-step task processing and orchestration
- [ ] **Structured Output Agent**: Complex data validation and formatting
### Model Provider Requirements
- [ ] **OpenAI**: `openai:gpt-4o` or `openai:gpt-4o-mini`
- [ ] **Anthropic**: `anthropic:claude-3-5-sonnet-20241022` or `anthropic:claude-3-5-haiku-20241022`
- [ ] **Google**: `gemini-1.5-flash` or `gemini-1.5-pro`
- [ ] **Fallback Strategy**: Multiple provider support with automatic failover
### External Integrations
- [ ] Database connections (specify type: PostgreSQL, MongoDB, etc.)
- [ ] REST API integrations (list required services)
- [ ] File system operations
- [ ] Web scraping or search capabilities
- [ ] Real-time data sources
### Success Criteria
- [ ] Agent successfully handles specified use cases
- [ ] All tools work correctly with proper error handling
- [ ] Structured outputs validate according to Pydantic models
- [ ] Comprehensive test coverage with TestModel and FunctionModel
- [ ] Security measures implemented (API keys, input validation, rate limiting)
- [ ] Performance meets requirements (response time, throughput)
## All Needed Context
### PydanticAI Documentation & Research
```yaml
# MCP servers
- mcp: Archon
query: "PydanticAI agent creation model providers tools dependencies"
why: Core framework understanding and latest patterns
# ESSENTIAL PYDANTIC AI DOCUMENTATION - Must be researched
- url: https://ai.pydantic.dev/
why: Official PydanticAI documentation with getting started guide
content: Agent creation, model providers, dependency injection patterns
- url: https://ai.pydantic.dev/agents/
why: Comprehensive agent architecture and configuration patterns
content: System prompts, output types, execution methods, agent composition
- url: https://ai.pydantic.dev/tools/
why: Tool integration patterns and function registration
content: @agent.tool decorators, RunContext usage, parameter validation
- url: https://ai.pydantic.dev/testing/
why: Testing strategies specific to PydanticAI agents
content: TestModel, FunctionModel, Agent.override(), pytest patterns
- url: https://ai.pydantic.dev/models/
why: Model provider configuration and authentication
content: OpenAI, Anthropic, Gemini setup, API key management, fallback models
# Prebuilt examples
- path: examples/
why: Reference implementations for Pydantic AI agents
content: A bunch of already built simple Pydantic AI examples to reference including how to set up models and providers
- path: examples/cli.py
why: Shows real-world interaction with Pydantic AI agents
content: Conversational CLI with streaming, tool call visibility, and conversation handling - demonstrates how users actually interact with agents
```
### Agent Architecture Research
```yaml
# PydanticAI Architecture Patterns (follow main_agent_reference)
agent_structure:
configuration:
- settings.py: Environment-based configuration with pydantic-settings
- providers.py: Model provider abstraction with get_llm_model()
- Environment variables for API keys and model selection
- Never hardcode model strings like "openai:gpt-4o"
agent_definition:
- Default to string output (no result_type unless structured output needed)
- Use get_llm_model() from providers.py for model configuration
- System prompts as string constants or functions
- Dataclass dependencies for external services
tool_integration:
- @agent.tool for context-aware tools with RunContext[DepsType]
- Tool functions as pure functions that can be called independently
- Proper error handling and logging in tool implementations
- Dependency injection through RunContext.deps
testing_strategy:
- TestModel for rapid development validation
- FunctionModel for custom behavior testing
- Agent.override() for test isolation
- Comprehensive tool testing with mocks
```
### Security and Production Considerations
```yaml
# PydanticAI Security Patterns (research required)
security_requirements:
api_management:
environment_variables: ["OPENAI_API_KEY", "ANTHROPIC_API_KEY", "GEMINI_API_KEY"]
secure_storage: "Never commit API keys to version control"
rotation_strategy: "Plan for key rotation and management"
input_validation:
sanitization: "Validate all user inputs with Pydantic models"
prompt_injection: "Implement prompt injection prevention strategies"
rate_limiting: "Prevent abuse with proper throttling"
output_security:
data_filtering: "Ensure no sensitive data in agent responses"
content_validation: "Validate output structure and content"
logging_safety: "Safe logging without exposing secrets"
```
### Common PydanticAI Gotchas (research and document)
```yaml
# Agent-specific gotchas to research and address
implementation_gotchas:
async_patterns:
issue: "Mixing sync and async agent calls inconsistently"
research: "PydanticAI async/await best practices"
solution: "[To be documented based on research]"
model_limits:
issue: "Different models have different capabilities and token limits"
research: "Model provider comparison and capabilities"
solution: "[To be documented based on research]"
dependency_complexity:
issue: "Complex dependency graphs can be hard to debug"
research: "Dependency injection best practices in PydanticAI"
solution: "[To be documented based on research]"
tool_error_handling:
issue: "Tool failures can crash entire agent runs"
research: "Error handling and retry patterns for tools"
solution: "[To be documented based on research]"
```
## Implementation Blueprint
### Technology Research Phase
**RESEARCH REQUIRED - Complete before implementation:**
**PydanticAI Framework Deep Dive:**
- [ ] Agent creation patterns and best practices
- [ ] Model provider configuration and fallback strategies
- [ ] Tool integration patterns (@agent.tool vs @agent.tool_plain)
- [ ] Dependency injection system and type safety
- [ ] Testing strategies with TestModel and FunctionModel
**Agent Architecture Investigation:**
- [ ] Project structure conventions (agent.py, tools.py, models.py, dependencies.py)
- [ ] System prompt design (static vs dynamic)
- [ ] Structured output validation with Pydantic models
- [ ] Async/sync patterns and streaming support
- [ ] Error handling and retry mechanisms
**Security and Production Patterns:**
- [ ] API key management and secure configuration
- [ ] Input validation and prompt injection prevention
- [ ] Rate limiting and monitoring strategies
- [ ] Logging and observability patterns
- [ ] Deployment and scaling considerations
### Agent Implementation Plan
```yaml
Implementation Task 1 - Agent Architecture Setup (Follow main_agent_reference):
CREATE agent project structure:
- settings.py: Environment-based configuration with pydantic-settings
- providers.py: Model provider abstraction with get_llm_model()
- agent.py: Main agent definition (default string output)
- tools.py: Tool functions with proper decorators
- dependencies.py: External service integrations (dataclasses)
- tests/: Comprehensive test suite
Implementation Task 2 - Core Agent Development:
IMPLEMENT agent.py following main_agent_reference patterns:
- Use get_llm_model() from providers.py for model configuration
- System prompt as string constant or function
- Dependency injection with dataclass
- NO result_type unless structured output specifically needed
- Error handling and logging
Implementation Task 3 - Tool Integration:
DEVELOP tools.py:
- Tool functions with @agent.tool decorators
- RunContext[DepsType] integration for dependency access
- Parameter validation with proper type hints
- Error handling and retry mechanisms
- Tool documentation and schema generation
Implementation Task 4 - Data Models and Dependencies:
CREATE models.py and dependencies.py:
- Pydantic models for structured outputs
- Dependency classes for external services
- Input validation models for tools
- Custom validators and constraints
Implementation Task 5 - Comprehensive Testing:
IMPLEMENT testing suite:
- TestModel integration for rapid development
- FunctionModel tests for custom behavior
- Agent.override() patterns for isolation
- Integration tests with real providers
- Tool validation and error scenario testing
Implementation Task 6 - Security and Configuration:
SETUP security patterns:
- Environment variable management for API keys
- Input sanitization and validation
- Rate limiting implementation
- Secure logging and monitoring
- Production deployment configuration
```
## Validation Loop
### Level 1: Agent Structure Validation
```bash
# Verify complete agent project structure
find agent_project -name "*.py" | sort
test -f agent_project/agent.py && echo "Agent definition present"
test -f agent_project/tools.py && echo "Tools module present"
test -f agent_project/models.py && echo "Models module present"
test -f agent_project/dependencies.py && echo "Dependencies module present"
# Verify proper PydanticAI imports
grep -q "from pydantic_ai import Agent" agent_project/agent.py
grep -q "@agent.tool" agent_project/tools.py
grep -q "from pydantic import BaseModel" agent_project/models.py
# Expected: All required files with proper PydanticAI patterns
# If missing: Generate missing components with correct patterns
```
### Level 2: Agent Functionality Validation
```bash
# Test agent can be imported and instantiated
python -c "
from agent_project.agent import agent
print('Agent created successfully')
print(f'Model: {agent.model}')
print(f'Tools: {len(agent.tools)}')
"
# Test with TestModel for validation
python -c "
from pydantic_ai.models.test import TestModel
from agent_project.agent import agent
test_model = TestModel()
with agent.override(model=test_model):
result = agent.run_sync('Test message')
print(f'Agent response: {result.output}')
"
# Expected: Agent instantiation works, tools registered, TestModel validation passes
# If failing: Debug agent configuration and tool registration
```
### Level 3: Comprehensive Testing Validation
```bash
# Run complete test suite
cd agent_project
python -m pytest tests/ -v
# Test specific agent behavior
python -m pytest tests/test_agent.py::test_agent_response -v
python -m pytest tests/test_tools.py::test_tool_validation -v
python -m pytest tests/test_models.py::test_output_validation -v
# Expected: All tests pass, comprehensive coverage achieved
# If failing: Fix implementation based on test failures
```
### Level 4: Production Readiness Validation
```bash
# Verify security patterns
grep -r "API_KEY" agent_project/ | grep -v ".py:" # Should not expose keys
test -f agent_project/.env.example && echo "Environment template present"
# Check error handling
grep -r "try:" agent_project/ | wc -l # Should have error handling
grep -r "except" agent_project/ | wc -l # Should have exception handling
# Verify logging setup
grep -r "logging\|logger" agent_project/ | wc -l # Should have logging
# Expected: Security measures in place, error handling comprehensive, logging configured
# If issues: Implement missing security and production patterns
```
## Final Validation Checklist
### Agent Implementation Completeness
- [ ] Complete agent project structure: `agent.py`, `tools.py`, `models.py`, `dependencies.py`
- [ ] Agent instantiation with proper model provider configuration
- [ ] Tool registration with @agent.tool decorators and RunContext integration
- [ ] Structured outputs with Pydantic model validation
- [ ] Dependency injection properly configured and tested
- [ ] Comprehensive test suite with TestModel and FunctionModel
### PydanticAI Best Practices
- [ ] Type safety throughout with proper type hints and validation
- [ ] Security patterns implemented (API keys, input validation, rate limiting)
- [ ] Error handling and retry mechanisms for robust operation
- [ ] Async/sync patterns consistent and appropriate
- [ ] Documentation and code comments for maintainability
### Production Readiness
- [ ] Environment configuration with .env files and validation
- [ ] Logging and monitoring setup for observability
- [ ] Performance optimization and resource management
- [ ] Deployment readiness with proper configuration management
- [ ] Maintenance and update strategies documented
---
## Anti-Patterns to Avoid
### PydanticAI Agent Development
- ❌ Don't skip TestModel validation - always test with TestModel during development
- ❌ Don't hardcode API keys - use environment variables for all credentials
- ❌ Don't ignore async patterns - PydanticAI has specific async/sync requirements
- ❌ Don't create complex tool chains - keep tools focused and composable
- ❌ Don't skip error handling - implement comprehensive retry and fallback mechanisms
### Agent Architecture
- ❌ Don't mix agent types - clearly separate chat, tool, workflow, and structured output patterns
- ❌ Don't ignore dependency injection - use proper type-safe dependency management
- ❌ Don't skip output validation - always use Pydantic models for structured responses
- ❌ Don't forget tool documentation - ensure all tools have proper descriptions and schemas
### Security and Production
- ❌ Don't expose sensitive data - validate all outputs and logs for security
- ❌ Don't skip input validation - sanitize and validate all user inputs
- ❌ Don't ignore rate limiting - implement proper throttling for external services
- ❌ Don't deploy without monitoring - include proper observability from the start
**RESEARCH STATUS: [TO BE COMPLETED]** - Complete comprehensive PydanticAI research before implementation begins.

View File

@@ -0,0 +1,532 @@
# Pydantic AI Context Engineering Template
A comprehensive template for building production-grade AI agents using Pydantic AI with context engineering best practices, tools integration, structured outputs, and comprehensive testing patterns.
## 🚀 Quick Start - Copy Template
**Get started in 2 minutes:**
```bash
# Clone the context engineering repository
git clone https://github.com/coleam00/Context-Engineering-Intro.git
cd Context-Engineering-Intro/use-cases/pydantic-ai
# 1. Copy this template to your new project
python copy_template.py /path/to/my-agent-project
# 2. Navigate to your project
cd /path/to/my-agent-project
# 3. Set up environment
python -m venv venv
source venv/bin/activate # On Windows: venv\Scripts\activate
# 4. Start building with the PRP workflow
# Edit PRPs/INITIAL.md with your requirements, then:
/generate-pydantic-ai-prp PRPs/INITIAL.md
/execute-pydantic-ai-prp PRPs/generated_prp.md
```
## 📖 What is This Template?
This template provides everything you need to build sophisticated Pydantic AI agents using proven context engineering workflows. It combines:
- **Pydantic AI Best Practices**: Type-safe agents with tools, structured outputs, and dependency injection
- **Context Engineering Workflows**: Proven PRP (Problem Requirements Planning) methodology
- **Production Patterns**: Security, testing, monitoring, and deployment-ready code
- **Working Examples**: Complete agent implementations you can learn from and extend
## 🎯 PRP Framework Workflow
This template uses a 3-step context engineering workflow for building AI agents:
### 1. **Define Requirements** (`PRPs/INITIAL.md`)
Start by clearly defining what your agent needs to do:
```markdown
# Customer Support Agent - Initial Requirements
## Overview
Build an intelligent customer support agent that can handle inquiries,
access customer data, and escalate issues appropriately.
## Core Requirements
- Multi-turn conversations with context and memory
- Customer authentication and account access
- Account balance and transaction queries
- Payment processing and refund handling
...
```
### 2. **Generate Implementation Plan**
```bash
/generate-pydantic-ai-prp PRPs/INITIAL.md
```
This creates a comprehensive Problem Requirements Planning document that includes:
- Pydantic AI technology research and best practices
- Agent architecture design with tools and dependencies
- Implementation roadmap with validation loops
- Security patterns and production considerations
### 3. **Execute Implementation**
```bash
/execute-pydantic-ai-prp PRPs/your_agent.md
```
This implements the complete agent based on the PRP, including:
- Agent creation with proper model provider configuration
- Tool integration with error handling and validation
- Structured output models with Pydantic validation
- Comprehensive testing with TestModel and FunctionModel
- Security patterns and production deployment setup
## 📂 Template Structure
```
pydantic-ai/
├── CLAUDE.md # Pydantic AI global development rules
├── copy_template.py # Template deployment script
├── .claude/commands/
│ ├── generate-pydantic-ai-prp.md # PRP generation for agents
│ └── execute-pydantic-ai-prp.md # PRP execution for agents
├── PRPs/
│ ├── templates/
│ │ └── prp_pydantic_ai_base.md # Base PRP template for agents
│ └── INITIAL.md # Example agent requirements
├── examples/
│ ├── basic_chat_agent/ # Simple conversational agent
│ │ ├── agent.py # Agent with memory and context
│ │ └── README.md # Usage guide
│ ├── tool_enabled_agent/ # Agent with external tools
│ │ ├── agent.py # Web search + calculator tools
│ │ └── requirements.txt # Dependencies
│ └── testing_examples/ # Comprehensive testing patterns
│ ├── test_agent_patterns.py # TestModel, FunctionModel examples
│ └── pytest.ini # Test configuration
└── README.md # This file
```
## 🤖 Agent Examples Included
### 1. Main Agent Reference (`examples/main_agent_reference/`)
**The canonical reference implementation** showing proper Pydantic AI patterns:
- Environment-based configuration with `settings.py` and `providers.py`
- Clean separation of concerns between email and research agents
- Tool integration with external APIs (Gmail, Brave Search)
- Production-ready error handling and logging
**Key Files:**
- `settings.py`: Environment configuration with pydantic-settings
- `providers.py`: Model provider abstraction with `get_llm_model()`
- `research_agent.py`: Multi-tool agent with web search and email integration
- `email_agent.py`: Specialized agent for Gmail draft creation
### 2. Basic Chat Agent (`examples/basic_chat_agent/`)
A simple conversational agent demonstrating core patterns:
- **Environment-based model configuration** (follows main_agent_reference)
- **String output by default** (no `result_type` unless needed)
- System prompts (static and dynamic)
- Conversation memory with dependency injection
**Key Features:**
- Simple string responses (not structured output)
- Settings-based configuration pattern
- Conversation context tracking
- Clean, minimal implementation
### 3. Tool-Enabled Agent (`examples/tool_enabled_agent/`)
An agent with tool integration capabilities:
- **Environment-based configuration** (follows main_agent_reference)
- **String output by default** (no unnecessary structure)
- Web search and calculation tools
- Error handling and retry mechanisms
**Key Features:**
- `@agent.tool` decorator patterns
- RunContext for dependency injection
- Tool error handling and recovery
- Simple string responses from tools
### 4. Structured Output Agent (`examples/structured_output_agent/`)
**NEW**: Shows when to use `result_type` for data validation:
- **Environment-based configuration** (follows main_agent_reference)
- **Structured output with Pydantic validation** (when specifically needed)
- Data analysis with statistical tools
- Professional report generation
**Key Features:**
- Demonstrates proper use of `result_type`
- Pydantic validation for business reports
- Data analysis tools with numerical statistics
- Clear documentation on when to use structured vs string output
### 5. Testing Examples (`examples/testing_examples/`)
Comprehensive testing patterns for Pydantic AI agents:
- TestModel for rapid development validation
- FunctionModel for custom behavior testing
- Agent.override() for test isolation
- Pytest fixtures and async testing
**Key Features:**
- Unit testing without API costs
- Mock dependency injection
- Tool validation and error scenario testing
- Integration testing patterns
## 🛠️ Core Pydantic AI Patterns
### Environment-Based Configuration (from main_agent_reference)
```python
# settings.py - Environment configuration
from pydantic_settings import BaseSettings
from pydantic import Field
class Settings(BaseSettings):
llm_provider: str = Field(default="openai")
llm_api_key: str = Field(...)
llm_model: str = Field(default="gpt-4")
llm_base_url: str = Field(default="https://api.openai.com/v1")
class Config:
env_file = ".env"
# providers.py - Model provider abstraction
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.models.openai import OpenAIModel
def get_llm_model() -> OpenAIModel:
settings = Settings()
provider = OpenAIProvider(
base_url=settings.llm_base_url,
api_key=settings.llm_api_key
)
return OpenAIModel(settings.llm_model, provider=provider)
```
### Simple Agent (String Output - Default)
```python
from pydantic_ai import Agent, RunContext
from dataclasses import dataclass
@dataclass
class AgentDependencies:
"""Dependencies for agent execution"""
api_key: str
session_id: str = None
# Simple agent - no result_type, defaults to string
agent = Agent(
get_llm_model(), # Environment-based configuration
deps_type=AgentDependencies,
system_prompt="You are a helpful assistant..."
)
```
### Structured Output Agent (When Validation Needed)
```python
from pydantic import BaseModel, Field
class AnalysisReport(BaseModel):
"""Use result_type ONLY when you need validation"""
summary: str
confidence: float = Field(ge=0.0, le=1.0)
insights: list[str] = Field(min_items=1)
# Structured agent - result_type specified for validation
structured_agent = Agent(
get_llm_model(),
deps_type=AgentDependencies,
result_type=AnalysisReport, # Only when structure is required
system_prompt="You are a data analyst..."
)
```
### Tool Integration
```python
@agent.tool
async def example_tool(
ctx: RunContext[AgentDependencies],
query: str
) -> str:
"""Tool with proper error handling - returns string."""
try:
result = await external_api_call(ctx.deps.api_key, query)
return f"API result: {result}"
except Exception as e:
logger.error(f"Tool error: {e}")
return f"Tool temporarily unavailable: {str(e)}"
```
### Testing with TestModel
```python
from pydantic_ai.models.test import TestModel
def test_simple_agent():
"""Test simple agent with string output."""
test_model = TestModel()
with agent.override(model=test_model):
result = agent.run_sync("Test message")
assert isinstance(result.data, str) # String output
def test_structured_agent():
"""Test structured agent with validation."""
test_model = TestModel(
custom_output_text='{"summary": "Test", "confidence": 0.8, "insights": ["insight1"]}'
)
with structured_agent.override(model=test_model):
result = structured_agent.run_sync("Analyze this data")
assert isinstance(result.data, AnalysisReport) # Validated object
assert 0.0 <= result.data.confidence <= 1.0
```
## 🎯 When to Use String vs Structured Output
### Use String Output (Default) ✅
**Most agents should use string output** - don't specify `result_type`:
```python
# ✅ Simple chat agent
chat_agent = Agent(get_llm_model(), system_prompt="You are helpful...")
# ✅ Tool-enabled agent
tool_agent = Agent(get_llm_model(), tools=[search_tool], system_prompt="...")
# Result: agent.run() returns string
result = agent.run_sync("Hello")
print(result.data) # "Hello! How can I help you today?"
```
**When to use string output:**
- Conversational agents
- Creative writing
- Flexible responses
- Human-readable output
- Simple tool responses
### Use Structured Output (Specific Cases) 🎯
**Only use `result_type` when you need validation:**
```python
# ✅ Data analysis requiring validation
analysis_agent = Agent(
get_llm_model(),
result_type=AnalysisReport, # Pydantic model with validation
system_prompt="You are a data analyst..."
)
# Result: agent.run() returns validated Pydantic object
result = analysis_agent.run_sync("Analyze sales data")
print(result.data.confidence) # 0.85 (validated 0.0-1.0)
```
**When to use structured output:**
- Data validation required
- API integrations needing specific schemas
- Business reports with consistent formatting
- Downstream processing requiring type safety
- Database insertion with validated fields
### Key Rule 📏
**Start with string output. Only add `result_type` when you specifically need validation or structure.**
## 🔒 Security Best Practices
This template includes production-ready security patterns:
### API Key Management
```bash
# Environment variables (never commit to code)
export LLM_API_KEY="your-api-key-here"
export LLM_MODEL="gpt-4"
export LLM_BASE_URL="https://api.openai.com/v1"
# Or use .env file (git-ignored)
echo "LLM_API_KEY=your-api-key-here" > .env
echo "LLM_MODEL=gpt-4" >> .env
```
### Input Validation
```python
from pydantic import BaseModel, Field
class ToolInput(BaseModel):
query: str = Field(max_length=1000, description="Search query")
max_results: int = Field(ge=1, le=10, default=5)
```
### Error Handling
```python
@agent.tool
async def secure_tool(ctx: RunContext[Deps], input_data: str) -> str:
try:
# Validate and sanitize input
cleaned_input = sanitize_input(input_data)
result = await process_safely(cleaned_input)
return result
except Exception as e:
# Log error without exposing sensitive data
logger.error(f"Tool error: {type(e).__name__}")
return "An error occurred. Please try again."
```
## 🧪 Testing Your Agents
### Development Testing (Fast, No API Costs)
```python
from pydantic_ai.models.test import TestModel
# Test with TestModel for rapid iteration
test_model = TestModel()
with agent.override(model=test_model):
result = agent.run_sync("Test input")
print(result.data)
```
### Custom Behavior Testing
```python
from pydantic_ai.models.test import FunctionModel
def custom_response(messages, tools):
"""Custom function to control agent responses."""
return '{"response": "Custom test response", "confidence": 0.9}'
function_model = FunctionModel(function=custom_response)
with agent.override(model=function_model):
result = agent.run_sync("Test input")
```
### Integration Testing
```python
# Test with real models (use sparingly due to costs)
@pytest.mark.integration
async def test_agent_integration():
result = await agent.run("Real test message")
assert result.data.confidence > 0.5
```
## 🚀 Deployment Patterns
### Environment Configuration
```python
# settings.py - Production configuration
from pydantic_settings import BaseSettings
class Settings(BaseSettings):
# LLM Configuration
llm_api_key: str
llm_model: str = "gpt-4"
llm_base_url: str = "https://api.openai.com/v1"
# Production settings
app_env: str = "production"
log_level: str = "INFO"
retries: int = 3
class Config:
env_file = ".env"
# agent.py - Use environment settings
agent = Agent(
get_llm_model(), # From providers.py
retries=settings.retries,
system_prompt="Production agent..."
)
```
### Docker Deployment
```dockerfile
FROM python:3.11-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install -r requirements.txt
COPY . .
EXPOSE 8000
CMD ["python", "-m", "uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
```
## 🎓 Learning Path
### 1. Start with Examples
- Run `examples/basic_chat_agent/agent.py` to see a simple agent
- Explore `examples/tool_enabled_agent/` for tool integration
- Study `examples/testing_examples/` for testing patterns
### 2. Use the PRP Workflow
- Edit `PRPs/INITIAL.md` with your agent requirements
- Generate a PRP: `/generate-pydantic-ai-prp PRPs/INITIAL.md`
- Execute the PRP: `/execute-pydantic-ai-prp PRPs/generated_file.md`
### 3. Build Your Own Agent
- Start with the basic chat agent pattern
- Add tools for external capabilities
- Implement structured outputs for your use case
- Add comprehensive testing and error handling
### 4. Production Deployment
- Implement security patterns from `CLAUDE.md`
- Add monitoring and logging
- Set up CI/CD with automated testing
- Deploy with proper scaling and availability
## 🤝 Common Gotchas & Solutions
Based on extensive Pydantic AI research, here are common issues and solutions:
### Async/Sync Patterns
```python
# ❌ Don't mix sync and async inconsistently
def bad_tool(ctx):
return asyncio.run(some_async_function()) # Anti-pattern
# ✅ Be consistent with async patterns
@agent.tool
async def good_tool(ctx: RunContext[Deps]) -> str:
result = await some_async_function()
return result
```
### Model Token Limits
```python
# ✅ Handle different model capabilities
from pydantic_ai.models import FallbackModel
model = FallbackModel([
"openai:gpt-4o", # High capability, higher cost
"openai:gpt-4o-mini", # Fallback option
])
```
### Tool Error Handling
```python
# ✅ Implement proper retry and fallback
@agent.tool
async def resilient_tool(ctx: RunContext[Deps], query: str) -> str:
for attempt in range(3):
try:
return await external_api_call(query)
except TemporaryError:
if attempt == 2:
return "Service temporarily unavailable"
await asyncio.sleep(2 ** attempt)
```
## 📚 Additional Resources
- **Official Pydantic AI Documentation**: https://ai.pydantic.dev/
- **Model Provider Guides**: https://ai.pydantic.dev/models/
- **Tool Integration Patterns**: https://ai.pydantic.dev/tools/
- **Testing Strategies**: https://ai.pydantic.dev/testing/
- **Context Engineering Methodology**: See main repository README
## 🆘 Support & Contributing
- **Issues**: Report problems with the template or examples
- **Improvements**: Contribute additional examples or patterns
- **Questions**: Ask about Pydantic AI integration or context engineering
This template is part of the larger Context Engineering framework. See the main repository for more context engineering templates and methodologies.
---
**Ready to build production-grade AI agents?** Start with `python copy_template.py my-agent-project` and follow the PRP workflow! 🚀

View File

@@ -0,0 +1,291 @@
#!/usr/bin/env python3
"""
PydanticAI Template Copy Script
Copies the complete PydanticAI context engineering template to a target directory
for starting new PydanticAI agent development projects.
Usage:
python copy_template.py <target_directory>
Example:
python copy_template.py my-agent-project
python copy_template.py /path/to/my-new-agent
"""
import os
import sys
import shutil
import argparse
from pathlib import Path
from typing import List, Tuple
def get_template_files() -> List[Tuple[str, str]]:
"""
Get list of template files to copy with their relative paths.
Returns:
List of (source_path, relative_path) tuples
"""
template_root = Path(__file__).parent
files_to_copy = []
# Core template files
core_files = [
"CLAUDE.md",
"README.md",
]
for file in core_files:
source_path = template_root / file
if source_path.exists():
# Rename README.md to readme_template.md in target
target_name = "README_TEMPLATE.md" if file == "README.md" else file
files_to_copy.append((str(source_path), target_name))
# Claude commands directory
commands_dir = template_root / ".claude" / "commands"
if commands_dir.exists():
for file in commands_dir.glob("*.md"):
rel_path = f".claude/commands/{file.name}"
files_to_copy.append((str(file), rel_path))
# PRPs directory
prps_dir = template_root / "PRPs"
if prps_dir.exists():
# Copy templates subdirectory
templates_dir = prps_dir / "templates"
if templates_dir.exists():
for file in templates_dir.glob("*.md"):
rel_path = f"PRPs/templates/{file.name}"
files_to_copy.append((str(file), rel_path))
# Copy INITIAL.md example
initial_file = prps_dir / "INITIAL.md"
if initial_file.exists():
files_to_copy.append((str(initial_file), "PRPs/INITIAL.md"))
# Examples directory - copy all examples
examples_dir = template_root / "examples"
if examples_dir.exists():
for example_dir in examples_dir.iterdir():
if example_dir.is_dir():
# Copy all files in each example directory
for file in example_dir.rglob("*"):
if file.is_file():
rel_path = file.relative_to(template_root)
files_to_copy.append((str(file), str(rel_path)))
return files_to_copy
def create_directory_structure(target_dir: Path, files: List[Tuple[str, str]]) -> None:
"""
Create directory structure for all files.
Args:
target_dir: Target directory path
files: List of (source_path, relative_path) tuples
"""
directories = set()
for _, rel_path in files:
dir_path = target_dir / Path(rel_path).parent
directories.add(dir_path)
for directory in directories:
directory.mkdir(parents=True, exist_ok=True)
def copy_template_files(target_dir: Path, files: List[Tuple[str, str]]) -> int:
"""
Copy all template files to target directory.
Args:
target_dir: Target directory path
files: List of (source_path, relative_path) tuples
Returns:
Number of files copied successfully
"""
copied_count = 0
for source_path, rel_path in files:
target_path = target_dir / rel_path
try:
shutil.copy2(source_path, target_path)
copied_count += 1
print(f"{rel_path}")
except Exception as e:
print(f"{rel_path} - Error: {e}")
return copied_count
def validate_template_integrity(target_dir: Path) -> bool:
"""
Validate that essential template files were copied correctly.
Args:
target_dir: Target directory path
Returns:
True if template appears complete, False otherwise
"""
essential_files = [
"CLAUDE.md",
"README_TEMPLATE.md",
".claude/commands/generate-pydantic-ai-prp.md",
".claude/commands/execute-pydantic-ai-prp.md",
"PRPs/templates/prp_pydantic_ai_base.md",
"PRPs/INITIAL.md",
"examples/basic_chat_agent/agent.py",
"examples/testing_examples/test_agent_patterns.py"
]
missing_files = []
for file_path in essential_files:
if not (target_dir / file_path).exists():
missing_files.append(file_path)
if missing_files:
print(f"\n⚠️ Warning: Some essential files are missing:")
for file in missing_files:
print(f" - {file}")
return False
return True
def print_next_steps(target_dir: Path) -> None:
"""
Print helpful next steps for using the template.
Args:
target_dir: Target directory path
"""
print(f"""
🎉 PydanticAI template successfully copied to: {target_dir}
📋 Next Steps:
1. Navigate to your new project:
cd {target_dir}
2. Set up your environment:
# Create virtual environment
python -m venv venv
source venv/bin/activate # On Windows: venv\\Scripts\\activate
# Install packages ahead of time or let your AI coding assistant handle taht
3. Start building your agent:
# 1. Edit PRPs/INITIAL.md with your agent requirements
# 2. Generate PRP: /generate-pydantic-ai-prp PRPs/INITIAL.md
# 3. Execute PRP: /execute-pydantic-ai-prp PRPs/generated_prp.md
5. Read the documentation:
# Check README.md for complete usage guide
# Check CLAUDE.md for PydanticAI development rules
🔗 Useful Resources:
- PydanticAI Docs: https://ai.pydantic.dev/
- Examples: See examples/ directory
- Testing: See examples/testing_examples/
Happy agent building! 🤖
""")
def main():
"""Main function for the copy template script."""
parser = argparse.ArgumentParser(
description="Copy PydanticAI context engineering template to a new project directory",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
python copy_template.py my-agent-project
python copy_template.py /path/to/my-new-agent
python copy_template.py ../customer-support-agent
"""
)
parser.add_argument(
"target_directory",
help="Target directory for the new PydanticAI project"
)
parser.add_argument(
"--force",
action="store_true",
help="Overwrite target directory if it exists"
)
parser.add_argument(
"--dry-run",
action="store_true",
help="Show what would be copied without actually copying"
)
if len(sys.argv) == 1:
parser.print_help()
return
args = parser.parse_args()
# Convert target directory to Path object
target_dir = Path(args.target_directory).resolve()
# Check if target directory exists
if target_dir.exists():
if target_dir.is_file():
print(f"❌ Error: {target_dir} is a file, not a directory")
return
if list(target_dir.iterdir()) and not args.force:
print(f"❌ Error: {target_dir} is not empty")
print("Use --force to overwrite existing directory")
return
if args.force and not args.dry_run:
print(f"⚠️ Overwriting existing directory: {target_dir}")
# Get list of files to copy
print("📂 Scanning PydanticAI template files...")
files_to_copy = get_template_files()
if not files_to_copy:
print("❌ Error: No template files found. Make sure you're running this from the template directory.")
return
print(f"Found {len(files_to_copy)} files to copy")
if args.dry_run:
print(f"\n🔍 Dry run - would copy to: {target_dir}")
for _, rel_path in files_to_copy:
print(f"{rel_path}")
return
# Create target directory and structure
print(f"\n📁 Creating directory structure in: {target_dir}")
target_dir.mkdir(parents=True, exist_ok=True)
create_directory_structure(target_dir, files_to_copy)
# Copy files
print(f"\n📋 Copying template files:")
copied_count = copy_template_files(target_dir, files_to_copy)
# Validate template integrity
print(f"\n✅ Copied {copied_count}/{len(files_to_copy)} files successfully")
if validate_template_integrity(target_dir):
print("✅ Template integrity check passed")
print_next_steps(target_dir)
else:
print("⚠️ Template may be incomplete. Check for missing files.")
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,191 @@
"""
Basic Chat Agent with Memory and Context
A simple conversational agent that demonstrates core PydanticAI patterns:
- Environment-based model configuration
- System prompts for personality and behavior
- Basic conversation handling with memory
- String output (default, no result_type needed)
"""
import logging
from dataclasses import dataclass
from typing import Optional
from pydantic_settings import BaseSettings
from pydantic import Field
from pydantic_ai import Agent, RunContext
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.models.openai import OpenAIModel
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
logger = logging.getLogger(__name__)
class Settings(BaseSettings):
"""Configuration settings for the chat agent."""
# LLM Configuration
llm_provider: str = Field(default="openai")
llm_api_key: str = Field(...)
llm_model: str = Field(default="gpt-4")
llm_base_url: str = Field(default="https://api.openai.com/v1")
class Config:
env_file = ".env"
case_sensitive = False
def get_llm_model() -> OpenAIModel:
"""Get configured LLM model from environment settings."""
try:
settings = Settings()
provider = OpenAIProvider(
base_url=settings.llm_base_url,
api_key=settings.llm_api_key
)
return OpenAIModel(settings.llm_model, provider=provider)
except Exception:
# For testing without env vars
import os
os.environ.setdefault("LLM_API_KEY", "test-key")
settings = Settings()
provider = OpenAIProvider(
base_url=settings.llm_base_url,
api_key="test-key"
)
return OpenAIModel(settings.llm_model, provider=provider)
@dataclass
class ConversationContext:
"""Simple context for conversation state management."""
user_name: Optional[str] = None
conversation_count: int = 0
preferred_language: str = "English"
session_id: Optional[str] = None
SYSTEM_PROMPT = """
You are a friendly and helpful AI assistant.
Your personality:
- Warm and approachable
- Knowledgeable but humble
- Patient and understanding
- Encouraging and supportive
Guidelines:
- Keep responses conversational and natural
- Be helpful without being overwhelming
- Ask follow-up questions when appropriate
- Remember context from the conversation
- Adapt your tone to match the user's needs
"""
# Create the basic chat agent - note: no result_type, defaults to string
chat_agent = Agent(
get_llm_model(),
deps_type=ConversationContext,
system_prompt=SYSTEM_PROMPT
)
@chat_agent.system_prompt
def dynamic_context_prompt(ctx) -> str:
"""Dynamic system prompt that includes conversation context."""
prompt_parts = []
if ctx.deps.user_name:
prompt_parts.append(f"The user's name is {ctx.deps.user_name}.")
if ctx.deps.conversation_count > 0:
prompt_parts.append(f"This is message #{ctx.deps.conversation_count + 1} in your conversation.")
if ctx.deps.preferred_language != "English":
prompt_parts.append(f"The user prefers to communicate in {ctx.deps.preferred_language}.")
return " ".join(prompt_parts) if prompt_parts else ""
async def chat_with_agent(message: str, context: Optional[ConversationContext] = None) -> str:
"""
Main function to chat with the agent.
Args:
message: User's message to the agent
context: Optional conversation context for memory
Returns:
String response from the agent
"""
if context is None:
context = ConversationContext()
# Increment conversation count
context.conversation_count += 1
# Run the agent with the message and context
result = await chat_agent.run(message, deps=context)
return result.data
def chat_with_agent_sync(message: str, context: Optional[ConversationContext] = None) -> str:
"""
Synchronous version of chat_with_agent for simple use cases.
Args:
message: User's message to the agent
context: Optional conversation context for memory
Returns:
String response from the agent
"""
if context is None:
context = ConversationContext()
# Increment conversation count
context.conversation_count += 1
# Run the agent synchronously
result = chat_agent.run_sync(message, deps=context)
return result.data
# Example usage and demonstration
if __name__ == "__main__":
import asyncio
async def demo_conversation():
"""Demonstrate the basic chat agent with a simple conversation."""
print("=== Basic Chat Agent Demo ===\n")
# Create conversation context
context = ConversationContext(
user_name="Alex",
preferred_language="English"
)
# Sample conversation
messages = [
"Hello! My name is Alex, nice to meet you.",
"Can you help me understand what PydanticAI is?",
"That's interesting! What makes it different from other AI frameworks?",
"Thanks for the explanation. Can you recommend some good resources to learn more?"
]
for message in messages:
print(f"User: {message}")
response = await chat_with_agent(message, context)
print(f"Agent: {response}")
print("-" * 50)
# Run the demo
asyncio.run(demo_conversation())

View File

@@ -0,0 +1,214 @@
#!/usr/bin/env python3
"""Conversational CLI with real-time streaming and tool call visibility for Pydantic AI agents."""
import asyncio
import sys
import os
from typing import List
# Add parent directory to Python path for imports
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from rich.console import Console
from rich.panel import Panel
from rich.prompt import Prompt
from rich.live import Live
from rich.text import Text
from pydantic_ai import Agent
from agents.research_agent import research_agent
from agents.dependencies import ResearchAgentDependencies
from agents.settings import settings
console = Console()
async def stream_agent_interaction(user_input: str, conversation_history: List[str]) -> tuple[str, str]:
"""Stream agent interaction with real-time tool call display."""
try:
# Set up dependencies
research_deps = ResearchAgentDependencies(brave_api_key=settings.brave_api_key)
# Build context with conversation history
context = "\n".join(conversation_history[-6:]) if conversation_history else ""
prompt = f"""Previous conversation:
{context}
User: {user_input}
Respond naturally and helpfully."""
# Stream the agent execution
async with research_agent.iter(prompt, deps=research_deps) as run:
async for node in run:
# Handle user prompt node
if Agent.is_user_prompt_node(node):
pass # Clean start - no processing messages
# Handle model request node - stream the thinking process
elif Agent.is_model_request_node(node):
# Show assistant prefix at the start
console.print("[bold blue]Assistant:[/bold blue] ", end="")
# Stream model request events for real-time text
response_text = ""
async with node.stream(run.ctx) as request_stream:
async for event in request_stream:
# Handle different event types based on their type
event_type = type(event).__name__
if event_type == "PartDeltaEvent":
# Extract content from delta
if hasattr(event, 'delta') and hasattr(event.delta, 'content_delta'):
delta_text = event.delta.content_delta
if delta_text:
console.print(delta_text, end="")
response_text += delta_text
elif event_type == "FinalResultEvent":
console.print() # New line after streaming
# Handle tool calls - this is the key part
elif Agent.is_call_tools_node(node):
# Stream tool execution events
async with node.stream(run.ctx) as tool_stream:
async for event in tool_stream:
event_type = type(event).__name__
if event_type == "FunctionToolCallEvent":
# Extract tool name from the part attribute
tool_name = "Unknown Tool"
args = None
# Check if the part attribute contains the tool call
if hasattr(event, 'part'):
part = event.part
# Check if part has tool_name directly
if hasattr(part, 'tool_name'):
tool_name = part.tool_name
elif hasattr(part, 'function_name'):
tool_name = part.function_name
elif hasattr(part, 'name'):
tool_name = part.name
# Check for arguments in part
if hasattr(part, 'args'):
args = part.args
elif hasattr(part, 'arguments'):
args = part.arguments
# Debug: print part attributes to understand structure
if tool_name == "Unknown Tool" and hasattr(event, 'part'):
part_attrs = [attr for attr in dir(event.part) if not attr.startswith('_')]
console.print(f" [dim red]Debug - Part attributes: {part_attrs}[/dim red]")
# Try to get more details about the part
if hasattr(event.part, '__dict__'):
console.print(f" [dim red]Part dict: {event.part.__dict__}[/dim red]")
console.print(f" 🔹 [cyan]Calling tool:[/cyan] [bold]{tool_name}[/bold]")
# Show tool args if available
if args and isinstance(args, dict):
# Show first few characters of each arg
arg_preview = []
for key, value in list(args.items())[:3]:
val_str = str(value)
if len(val_str) > 50:
val_str = val_str[:47] + "..."
arg_preview.append(f"{key}={val_str}")
console.print(f" [dim]Args: {', '.join(arg_preview)}[/dim]")
elif args:
args_str = str(args)
if len(args_str) > 100:
args_str = args_str[:97] + "..."
console.print(f" [dim]Args: {args_str}[/dim]")
elif event_type == "FunctionToolResultEvent":
# Display tool result
result = str(event.tool_return) if hasattr(event, 'tool_return') else "No result"
if len(result) > 100:
result = result[:97] + "..."
console.print(f" ✅ [green]Tool result:[/green] [dim]{result}[/dim]")
# Handle end node
elif Agent.is_end_node(node):
# Don't show "Processing complete" - keep it clean
pass
# Get final result
final_result = run.result
final_output = final_result.output if hasattr(final_result, 'output') else str(final_result)
# Return both streamed and final content
return (response_text.strip(), final_output)
except Exception as e:
console.print(f"[red]❌ Error: {e}[/red]")
return ("", f"Error: {e}")
async def main():
"""Main conversation loop."""
# Show welcome
welcome = Panel(
"[bold blue]🤖 Pydantic AI Research Assistant[/bold blue]\n\n"
"[green]Real-time tool execution visibility[/green]\n"
"[dim]Type 'exit' to quit[/dim]",
style="blue",
padding=(1, 2)
)
console.print(welcome)
console.print()
conversation_history = []
while True:
try:
# Get user input
user_input = Prompt.ask("[bold green]You").strip()
# Handle exit
if user_input.lower() in ['exit', 'quit']:
console.print("\n[yellow]👋 Goodbye![/yellow]")
break
if not user_input:
continue
# Add to history
conversation_history.append(f"User: {user_input}")
# Stream the interaction and get response
streamed_text, final_response = await stream_agent_interaction(user_input, conversation_history)
# Handle the response display
if streamed_text:
# Response was streamed, just add spacing
console.print()
conversation_history.append(f"Assistant: {streamed_text}")
elif final_response and final_response.strip():
# Response wasn't streamed, display with proper formatting
console.print(f"[bold blue]Assistant:[/bold blue] {final_response}")
console.print()
conversation_history.append(f"Assistant: {final_response}")
else:
# No response
console.print()
except KeyboardInterrupt:
console.print("\n[yellow]Use 'exit' to quit[/yellow]")
continue
except Exception as e:
console.print(f"[red]Error: {e}[/red]")
continue
if __name__ == "__main__":
asyncio.run(main())

View File

@@ -0,0 +1,9 @@
# ===== LLM Configuration =====
# Provider: openai, anthropic, gemini, ollama, etc.
LLM_PROVIDER=openai
# Your LLM API key
LLM_API_KEY=sk-your-openai-api-key-here
# LLM to use for the agents (e.g., gpt-4.1-mini, gpt-4.1, claude-4-sonnet)
LLM_CHOICE=gpt-4.1-mini
# Base URL for the LLM API (change for Ollama or other providers)
LLM_BASE_URL=https://api.openai.com/v1

View File

@@ -0,0 +1,103 @@
"""
Core data models for the multi-agent system.
"""
from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Any
from datetime import datetime
class ResearchQuery(BaseModel):
"""Model for research query requests."""
query: str = Field(..., description="Research topic to investigate")
max_results: int = Field(10, ge=1, le=50, description="Maximum number of results to return")
include_summary: bool = Field(True, description="Whether to include AI-generated summary")
class BraveSearchResult(BaseModel):
"""Model for individual Brave search results."""
title: str = Field(..., description="Title of the search result")
url: str = Field(..., description="URL of the search result")
description: str = Field(..., description="Description/snippet from the search result")
score: float = Field(0.0, ge=0.0, le=1.0, description="Relevance score")
class Config:
"""Pydantic configuration."""
json_schema_extra = {
"example": {
"title": "Understanding AI Safety",
"url": "https://example.com/ai-safety",
"description": "A comprehensive guide to AI safety principles...",
"score": 0.95
}
}
class EmailDraft(BaseModel):
"""Model for email draft creation."""
to: List[str] = Field(..., min_length=1, description="List of recipient email addresses")
subject: str = Field(..., min_length=1, description="Email subject line")
body: str = Field(..., min_length=1, description="Email body content")
cc: Optional[List[str]] = Field(None, description="List of CC recipients")
bcc: Optional[List[str]] = Field(None, description="List of BCC recipients")
class Config:
"""Pydantic configuration."""
json_schema_extra = {
"example": {
"to": ["john@example.com"],
"subject": "AI Research Summary",
"body": "Dear John,\n\nHere's the latest research on AI safety...",
"cc": ["team@example.com"]
}
}
class EmailDraftResponse(BaseModel):
"""Response model for email draft creation."""
draft_id: str = Field(..., description="Gmail draft ID")
message_id: str = Field(..., description="Message ID")
thread_id: Optional[str] = Field(None, description="Thread ID if part of a thread")
created_at: datetime = Field(default_factory=datetime.now, description="Draft creation timestamp")
class ResearchEmailRequest(BaseModel):
"""Model for research + email draft request."""
research_query: str = Field(..., description="Topic to research")
email_context: str = Field(..., description="Context for email generation")
recipient_email: str = Field(..., description="Email recipient")
email_subject: Optional[str] = Field(None, description="Optional email subject")
class ResearchResponse(BaseModel):
"""Response model for research queries."""
query: str = Field(..., description="Original research query")
results: List[BraveSearchResult] = Field(..., description="Search results")
summary: Optional[str] = Field(None, description="AI-generated summary of results")
total_results: int = Field(..., description="Total number of results found")
timestamp: datetime = Field(default_factory=datetime.now, description="Query timestamp")
class AgentResponse(BaseModel):
"""Generic agent response model."""
success: bool = Field(..., description="Whether the operation was successful")
data: Optional[Dict[str, Any]] = Field(None, description="Response data")
error: Optional[str] = Field(None, description="Error message if failed")
tools_used: List[str] = Field(default_factory=list, description="List of tools used")
class ChatMessage(BaseModel):
"""Model for chat messages in the CLI."""
role: str = Field(..., description="Message role (user/assistant)")
content: str = Field(..., description="Message content")
timestamp: datetime = Field(default_factory=datetime.now, description="Message timestamp")
tools_used: Optional[List[Dict[str, Any]]] = Field(None, description="Tools used in response")
class SessionState(BaseModel):
"""Model for maintaining session state."""
session_id: str = Field(..., description="Unique session identifier")
user_id: Optional[str] = Field(None, description="User identifier")
messages: List[ChatMessage] = Field(default_factory=list, description="Conversation history")
created_at: datetime = Field(default_factory=datetime.now, description="Session creation time")
last_activity: datetime = Field(default_factory=datetime.now, description="Last activity timestamp")

View File

@@ -0,0 +1,61 @@
"""
Flexible provider configuration for LLM models.
Based on examples/agent/providers.py pattern.
"""
from typing import Optional
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.models.openai import OpenAIModel
from .settings import settings
def get_llm_model(model_choice: Optional[str] = None) -> OpenAIModel:
"""
Get LLM model configuration based on environment variables.
Args:
model_choice: Optional override for model choice
Returns:
Configured OpenAI-compatible model
"""
llm_choice = model_choice or settings.llm_model
base_url = settings.llm_base_url
api_key = settings.llm_api_key
# Create provider based on configuration
provider = OpenAIProvider(base_url=base_url, api_key=api_key)
return OpenAIModel(llm_choice, provider=provider)
def get_model_info() -> dict:
"""
Get information about current model configuration.
Returns:
Dictionary with model configuration info
"""
return {
"llm_provider": settings.llm_provider,
"llm_model": settings.llm_model,
"llm_base_url": settings.llm_base_url,
"app_env": settings.app_env,
"debug": settings.debug,
}
def validate_llm_configuration() -> bool:
"""
Validate that LLM configuration is properly set.
Returns:
True if configuration is valid
"""
try:
# Check if we can create a model instance
get_llm_model()
return True
except Exception as e:
print(f"LLM configuration validation failed: {e}")
return False

View File

@@ -0,0 +1,263 @@
"""
Research Agent that uses Brave Search and can invoke Email Agent.
"""
import logging
from typing import Dict, Any, List, Optional
from dataclasses import dataclass
from pydantic_ai import Agent, RunContext
from .providers import get_llm_model
from .email_agent import email_agent, EmailAgentDependencies
from .tools import search_web_tool
logger = logging.getLogger(__name__)
SYSTEM_PROMPT = """
You are an expert research assistant with the ability to search the web and create email drafts. Your primary goal is to help users find relevant information and communicate findings effectively.
Your capabilities:
1. **Web Search**: Use Brave Search to find current, relevant information on any topic
2. **Email Creation**: Create professional email drafts through Gmail when requested
When conducting research:
- Use specific, targeted search queries
- Analyze search results for relevance and credibility
- Synthesize information from multiple sources
- Provide clear, well-organized summaries
- Include source URLs for reference
When creating emails:
- Use research findings to create informed, professional content
- Adapt tone and detail level to the intended recipient
- Include relevant sources and citations when appropriate
- Ensure emails are clear, concise, and actionable
Always strive to provide accurate, helpful, and actionable information.
"""
@dataclass
class ResearchAgentDependencies:
"""Dependencies for the research agent - only configuration, no tool instances."""
brave_api_key: str
gmail_credentials_path: str
gmail_token_path: str
session_id: Optional[str] = None
# Initialize the research agent
research_agent = Agent(
get_llm_model(),
deps_type=ResearchAgentDependencies,
system_prompt=SYSTEM_PROMPT
)
@research_agent.tool
async def search_web(
ctx: RunContext[ResearchAgentDependencies],
query: str,
max_results: int = 10
) -> List[Dict[str, Any]]:
"""
Search the web using Brave Search API.
Args:
query: Search query
max_results: Maximum number of results to return (1-20)
Returns:
List of search results with title, URL, description, and score
"""
try:
# Ensure max_results is within valid range
max_results = min(max(max_results, 1), 20)
results = await search_web_tool(
api_key=ctx.deps.brave_api_key,
query=query,
count=max_results
)
logger.info(f"Found {len(results)} results for query: {query}")
return results
except Exception as e:
logger.error(f"Web search failed: {e}")
return [{"error": f"Search failed: {str(e)}"}]
@research_agent.tool
async def create_email_draft(
ctx: RunContext[ResearchAgentDependencies],
recipient_email: str,
subject: str,
context: str,
research_summary: Optional[str] = None
) -> Dict[str, Any]:
"""
Create an email draft based on research context using the Email Agent.
Args:
recipient_email: Email address of the recipient
subject: Email subject line
context: Context or purpose for the email
research_summary: Optional research findings to include
Returns:
Dictionary with draft creation results
"""
try:
# Prepare the email content prompt
if research_summary:
email_prompt = f"""
Create a professional email to {recipient_email} with the subject "{subject}".
Context: {context}
Research Summary:
{research_summary}
Please create a well-structured email that:
1. Has an appropriate greeting
2. Provides clear context
3. Summarizes the key research findings professionally
4. Includes actionable next steps if appropriate
5. Ends with a professional closing
The email should be informative but concise, and maintain a professional yet friendly tone.
"""
else:
email_prompt = f"""
Create a professional email to {recipient_email} with the subject "{subject}".
Context: {context}
Please create a well-structured email that addresses the context provided.
"""
# Create dependencies for email agent
email_deps = EmailAgentDependencies(
gmail_credentials_path=ctx.deps.gmail_credentials_path,
gmail_token_path=ctx.deps.gmail_token_path,
session_id=ctx.deps.session_id
)
# Run the email agent
result = await email_agent.run(
email_prompt,
deps=email_deps,
usage=ctx.usage # Pass usage for token tracking
)
logger.info(f"Email agent invoked for recipient: {recipient_email}")
return {
"success": True,
"agent_response": result.data,
"recipient": recipient_email,
"subject": subject,
"context": context
}
except Exception as e:
logger.error(f"Failed to create email draft via Email Agent: {e}")
return {
"success": False,
"error": str(e),
"recipient": recipient_email,
"subject": subject
}
@research_agent.tool
async def summarize_research(
ctx: RunContext[ResearchAgentDependencies],
search_results: List[Dict[str, Any]],
topic: str,
focus_areas: Optional[str] = None
) -> Dict[str, Any]:
"""
Create a comprehensive summary of research findings.
Args:
search_results: List of search result dictionaries
topic: Main research topic
focus_areas: Optional specific areas to focus on
Returns:
Dictionary with research summary
"""
try:
if not search_results:
return {
"summary": "No search results provided for summarization.",
"key_points": [],
"sources": []
}
# Extract key information
sources = []
descriptions = []
for result in search_results:
if "title" in result and "url" in result:
sources.append(f"- {result['title']}: {result['url']}")
if "description" in result:
descriptions.append(result["description"])
# Create summary content
content_summary = "\n".join(descriptions[:5]) # Limit to top 5 descriptions
sources_list = "\n".join(sources[:10]) # Limit to top 10 sources
focus_text = f"\nSpecific focus areas: {focus_areas}" if focus_areas else ""
summary = f"""
Research Summary: {topic}{focus_text}
Key Findings:
{content_summary}
Sources:
{sources_list}
"""
return {
"summary": summary,
"topic": topic,
"sources_count": len(sources),
"key_points": descriptions[:5]
}
except Exception as e:
logger.error(f"Failed to summarize research: {e}")
return {
"summary": f"Failed to summarize research: {str(e)}",
"key_points": [],
"sources": []
}
# Convenience function to create research agent with dependencies
def create_research_agent(
brave_api_key: str,
gmail_credentials_path: str,
gmail_token_path: str,
session_id: Optional[str] = None
) -> Agent:
"""
Create a research agent with specified dependencies.
Args:
brave_api_key: Brave Search API key
gmail_credentials_path: Path to Gmail credentials.json
gmail_token_path: Path to Gmail token.json
session_id: Optional session identifier
Returns:
Configured research agent
"""
return research_agent

View File

@@ -0,0 +1,58 @@
"""
Configuration management using pydantic-settings.
"""
import os
from typing import Optional
from pydantic_settings import BaseSettings
from pydantic import Field, field_validator, ConfigDict
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
class Settings(BaseSettings):
"""Application settings with environment variable support."""
model_config = ConfigDict(
env_file=".env",
env_file_encoding="utf-8",
case_sensitive=False
)
# LLM Configuration
llm_provider: str = Field(default="openai")
llm_api_key: str = Field(...)
llm_model: str = Field(default="gpt-4")
llm_base_url: Optional[str] = Field(default="https://api.openai.com/v1")
# Brave Search Configuration
brave_api_key: str = Field(...)
brave_search_url: str = Field(
default="https://api.search.brave.com/res/v1/web/search"
)
# Application Configuration
app_env: str = Field(default="development")
log_level: str = Field(default="INFO")
debug: bool = Field(default=False)
@field_validator("llm_api_key", "brave_api_key")
@classmethod
def validate_api_keys(cls, v):
"""Ensure API keys are not empty."""
if not v or v.strip() == "":
raise ValueError("API key cannot be empty")
return v
# Global settings instance
try:
settings = Settings()
except Exception:
# For testing, create settings with dummy values
import os
os.environ.setdefault("LLM_API_KEY", "test_key")
os.environ.setdefault("BRAVE_API_KEY", "test_key")
settings = Settings()

View File

@@ -0,0 +1,120 @@
"""
Pure tool functions for multi-agent system.
These are standalone functions that can be imported and used by any agent.
"""
import os
import base64
import logging
import httpx
from typing import List, Dict, Any, Optional
from datetime import datetime
from agents.models import BraveSearchResult
logger = logging.getLogger(__name__)
# Brave Search Tool Function
async def search_web_tool(
api_key: str,
query: str,
count: int = 10,
offset: int = 0,
country: Optional[str] = None,
lang: Optional[str] = None
) -> List[Dict[str, Any]]:
"""
Pure function to search the web using Brave Search API.
Args:
api_key: Brave Search API key
query: Search query
count: Number of results to return (1-20)
offset: Offset for pagination
country: Country code for localized results
lang: Language code for results
Returns:
List of search results as dictionaries
Raises:
ValueError: If query is empty or API key missing
Exception: If API request fails
"""
if not api_key or not api_key.strip():
raise ValueError("Brave API key is required")
if not query or not query.strip():
raise ValueError("Query cannot be empty")
# Ensure count is within valid range
count = min(max(count, 1), 20)
headers = {
"X-Subscription-Token": api_key,
"Accept": "application/json"
}
params = {
"q": query,
"count": count,
"offset": offset
}
if country:
params["country"] = country
if lang:
params["lang"] = lang
logger.info(f"Searching Brave for: {query}")
async with httpx.AsyncClient() as client:
try:
response = await client.get(
"https://api.search.brave.com/res/v1/web/search",
headers=headers,
params=params,
timeout=30.0
)
# Handle rate limiting
if response.status_code == 429:
raise Exception("Rate limit exceeded. Check your Brave API quota.")
# Handle authentication errors
if response.status_code == 401:
raise Exception("Invalid Brave API key")
# Handle other errors
if response.status_code != 200:
raise Exception(f"Brave API returned {response.status_code}: {response.text}")
data = response.json()
# Extract web results
web_results = data.get("web", {}).get("results", [])
# Convert to our format
results = []
for idx, result in enumerate(web_results):
# Calculate a simple relevance score based on position
score = 1.0 - (idx * 0.05) # Decrease by 0.05 for each position
score = max(score, 0.1) # Minimum score of 0.1
results.append({
"title": result.get("title", ""),
"url": result.get("url", ""),
"description": result.get("description", ""),
"score": score
})
logger.info(f"Found {len(results)} results for query: {query}")
return results
except httpx.RequestError as e:
logger.error(f"Request error during Brave search: {e}")
raise Exception(f"Request failed: {str(e)}")
except Exception as e:
logger.error(f"Error during Brave search: {e}")
raise

View File

@@ -0,0 +1,303 @@
"""
Structured Output Agent for Data Validation
Demonstrates when to use structured outputs with PydanticAI:
- Environment-based model configuration (following main_agent_reference)
- Structured output validation with Pydantic models (result_type specified)
- Data extraction and validation use case
- Professional report generation with consistent formatting
"""
import logging
from dataclasses import dataclass
from typing import Optional, List
from pydantic_settings import BaseSettings
from pydantic import BaseModel, Field
from pydantic_ai import Agent, RunContext
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.models.openai import OpenAIModel
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
logger = logging.getLogger(__name__)
class Settings(BaseSettings):
"""Configuration settings for the structured output agent."""
# LLM Configuration
llm_provider: str = Field(default="openai")
llm_api_key: str = Field(...)
llm_model: str = Field(default="gpt-4")
llm_base_url: str = Field(default="https://api.openai.com/v1")
class Config:
env_file = ".env"
case_sensitive = False
def get_llm_model() -> OpenAIModel:
"""Get configured LLM model from environment settings."""
try:
settings = Settings()
provider = OpenAIProvider(
base_url=settings.llm_base_url,
api_key=settings.llm_api_key
)
return OpenAIModel(settings.llm_model, provider=provider)
except Exception:
# For testing without env vars
import os
os.environ.setdefault("LLM_API_KEY", "test-key")
settings = Settings()
provider = OpenAIProvider(
base_url=settings.llm_base_url,
api_key="test-key"
)
return OpenAIModel(settings.llm_model, provider=provider)
@dataclass
class AnalysisDependencies:
"""Dependencies for the analysis agent."""
report_format: str = "business" # business, technical, academic
include_recommendations: bool = True
session_id: Optional[str] = None
class DataInsight(BaseModel):
"""Individual insight extracted from data."""
insight: str = Field(description="The key insight or finding")
confidence: float = Field(ge=0.0, le=1.0, description="Confidence level in this insight")
data_points: List[str] = Field(description="Supporting data points")
class DataAnalysisReport(BaseModel):
"""Structured output for data analysis with validation."""
# Required fields
summary: str = Field(description="Executive summary of the analysis")
key_insights: List[DataInsight] = Field(
min_items=1,
max_items=10,
description="Key insights discovered in the data"
)
# Validated fields
confidence_score: float = Field(
ge=0.0, le=1.0,
description="Overall confidence in the analysis"
)
data_quality: str = Field(
pattern="^(excellent|good|fair|poor)$",
description="Assessment of data quality"
)
# Optional structured fields
recommendations: Optional[List[str]] = Field(
default=None,
description="Actionable recommendations based on findings"
)
limitations: Optional[List[str]] = Field(
default=None,
description="Limitations or caveats in the analysis"
)
# Metadata
analysis_type: str = Field(description="Type of analysis performed")
data_sources: List[str] = Field(description="Sources of data analyzed")
SYSTEM_PROMPT = """
You are an expert data analyst specializing in extracting structured insights from various data sources.
Your role:
- Analyze provided data with statistical rigor
- Extract meaningful insights and patterns
- Assess data quality and reliability
- Provide actionable recommendations
- Structure findings in a consistent, professional format
Guidelines:
- Be objective and evidence-based in your analysis
- Clearly distinguish between facts and interpretations
- Provide confidence levels for your insights
- Highlight both strengths and limitations of the data
- Ensure all outputs follow the required structured format
"""
# Create structured output agent - NOTE: result_type specified for data validation
structured_agent = Agent(
get_llm_model(),
deps_type=AnalysisDependencies,
result_type=DataAnalysisReport, # This is when we DO want structured output
system_prompt=SYSTEM_PROMPT
)
@structured_agent.tool
def analyze_numerical_data(
ctx: RunContext[AnalysisDependencies],
data_description: str,
numbers: List[float]
) -> str:
"""
Analyze numerical data and provide statistical insights.
Args:
data_description: Description of what the numbers represent
numbers: List of numerical values to analyze
Returns:
Statistical analysis summary
"""
try:
if not numbers:
return "No numerical data provided for analysis."
# Basic statistical calculations
count = len(numbers)
total = sum(numbers)
average = total / count
minimum = min(numbers)
maximum = max(numbers)
# Calculate variance and standard deviation
variance = sum((x - average) ** 2 for x in numbers) / count
std_dev = variance ** 0.5
# Simple trend analysis
if count > 1:
trend = "increasing" if numbers[-1] > numbers[0] else "decreasing"
else:
trend = "insufficient data"
analysis = f"""
Statistical Analysis of {data_description}:
- Count: {count} data points
- Average: {average:.2f}
- Range: {minimum:.2f} to {maximum:.2f}
- Standard Deviation: {std_dev:.2f}
- Overall Trend: {trend}
- Data Quality: {'good' if std_dev < average * 0.5 else 'variable'}
"""
logger.info(f"Analyzed {count} data points for: {data_description}")
return analysis.strip()
except Exception as e:
logger.error(f"Error in numerical analysis: {e}")
return f"Error analyzing numerical data: {str(e)}"
async def analyze_data(
data_input: str,
dependencies: Optional[AnalysisDependencies] = None
) -> DataAnalysisReport:
"""
Analyze data and return structured report.
Args:
data_input: Raw data or description to analyze
dependencies: Optional analysis configuration
Returns:
Structured DataAnalysisReport with validation
"""
if dependencies is None:
dependencies = AnalysisDependencies()
result = await structured_agent.run(data_input, deps=dependencies)
return result.data
def analyze_data_sync(
data_input: str,
dependencies: Optional[AnalysisDependencies] = None
) -> DataAnalysisReport:
"""
Synchronous version of analyze_data.
Args:
data_input: Raw data or description to analyze
dependencies: Optional analysis configuration
Returns:
Structured DataAnalysisReport with validation
"""
import asyncio
return asyncio.run(analyze_data(data_input, dependencies))
# Example usage and demonstration
if __name__ == "__main__":
import asyncio
async def demo_structured_output():
"""Demonstrate structured output validation."""
print("=== Structured Output Agent Demo ===\n")
# Sample data scenarios
scenarios = [
{
"title": "Sales Performance Data",
"data": """
Monthly sales data for Q4 2024:
October: $125,000
November: $142,000
December: $158,000
Customer satisfaction scores: 4.2, 4.5, 4.1, 4.6, 4.3
Return rate: 3.2%
"""
},
{
"title": "Website Analytics",
"data": """
Website traffic analysis:
- Daily visitors: 5,200 average
- Bounce rate: 35%
- Page load time: 2.1 seconds
- Conversion rate: 3.8%
- Mobile traffic: 68%
"""
}
]
for scenario in scenarios:
print(f"Analysis: {scenario['title']}")
print(f"Input Data: {scenario['data'][:100]}...")
# Configure for business report
deps = AnalysisDependencies(
report_format="business",
include_recommendations=True
)
try:
report = await analyze_data(scenario['data'], deps)
print(f"Summary: {report.summary}")
print(f"Confidence: {report.confidence_score}")
print(f"Data Quality: {report.data_quality}")
print(f"Key Insights: {len(report.key_insights)} found")
for i, insight in enumerate(report.key_insights, 1):
print(f" {i}. {insight.insight} (confidence: {insight.confidence})")
if report.recommendations:
print(f"Recommendations: {len(report.recommendations)}")
for i, rec in enumerate(report.recommendations, 1):
print(f" {i}. {rec}")
print("=" * 60)
except Exception as e:
print(f"Analysis failed: {e}")
print("=" * 60)
# Run the demo
asyncio.run(demo_structured_output())

View File

@@ -0,0 +1,18 @@
[tool:pytest]
testpaths = .
python_files = test_*.py
python_classes = Test*
python_functions = test_*
addopts =
-v
--tb=short
--strict-markers
--disable-warnings
markers =
integration: Integration tests
slow: Slow running tests
asyncio: Async tests
filterwarnings =
ignore::DeprecationWarning
ignore::PendingDeprecationWarning
asyncio_mode = auto

View File

@@ -0,0 +1,399 @@
"""
Comprehensive PydanticAI Testing Examples
Demonstrates testing patterns and best practices for PydanticAI agents:
- TestModel for fast development validation
- FunctionModel for custom behavior testing
- Agent.override() for test isolation
- Pytest fixtures and async testing
- Tool validation and error handling tests
"""
import pytest
import asyncio
from unittest.mock import Mock, AsyncMock
from dataclasses import dataclass
from typing import Optional, List
from pydantic import BaseModel
from pydantic_ai import Agent, RunContext
from pydantic_ai.models.test import TestModel, FunctionModel
@dataclass
class TestDependencies:
"""Test dependencies for agent testing."""
database: Mock
api_client: Mock
user_id: str = "test_user_123"
class TestResponse(BaseModel):
"""Test response model for validation."""
message: str
confidence: float = 0.8
actions: List[str] = []
# Create test agent for demonstrations
test_agent = Agent(
model="openai:gpt-4o-mini", # Will be overridden in tests
deps_type=TestDependencies,
result_type=TestResponse,
system_prompt="You are a helpful test assistant."
)
@test_agent.tool
async def mock_database_query(
ctx: RunContext[TestDependencies],
query: str
) -> str:
"""Mock database query tool for testing."""
try:
# Simulate database call
result = await ctx.deps.database.execute_query(query)
return f"Database result: {result}"
except Exception as e:
return f"Database error: {str(e)}"
@test_agent.tool
def mock_api_call(
ctx: RunContext[TestDependencies],
endpoint: str,
data: Optional[dict] = None
) -> str:
"""Mock API call tool for testing."""
try:
# Simulate API call
response = ctx.deps.api_client.post(endpoint, json=data)
return f"API response: {response}"
except Exception as e:
return f"API error: {str(e)}"
class TestAgentBasics:
"""Test basic agent functionality with TestModel."""
@pytest.fixture
def test_dependencies(self):
"""Create mock dependencies for testing."""
return TestDependencies(
database=AsyncMock(),
api_client=Mock(),
user_id="test_user_123"
)
def test_agent_with_test_model(self, test_dependencies):
"""Test agent behavior with TestModel."""
test_model = TestModel()
with test_agent.override(model=test_model):
result = test_agent.run_sync(
"Hello, please help me with a simple task.",
deps=test_dependencies
)
# TestModel returns a JSON summary by default
assert result.data.message is not None
assert isinstance(result.data.confidence, float)
assert isinstance(result.data.actions, list)
def test_agent_custom_test_model_output(self, test_dependencies):
"""Test agent with custom TestModel output."""
test_model = TestModel(
custom_output_text='{"message": "Custom test response", "confidence": 0.9, "actions": ["test_action"]}'
)
with test_agent.override(model=test_model):
result = test_agent.run_sync(
"Test message",
deps=test_dependencies
)
assert result.data.message == "Custom test response"
assert result.data.confidence == 0.9
assert result.data.actions == ["test_action"]
@pytest.mark.asyncio
async def test_agent_async_with_test_model(self, test_dependencies):
"""Test async agent behavior with TestModel."""
test_model = TestModel()
with test_agent.override(model=test_model):
result = await test_agent.run(
"Async test message",
deps=test_dependencies
)
assert result.data.message is not None
assert result.data.confidence >= 0.0
class TestAgentTools:
"""Test agent tool functionality."""
@pytest.fixture
def mock_dependencies(self):
"""Create mock dependencies with configured responses."""
database_mock = AsyncMock()
database_mock.execute_query.return_value = "Test data from database"
api_mock = Mock()
api_mock.post.return_value = {"status": "success", "data": "test_data"}
return TestDependencies(
database=database_mock,
api_client=api_mock,
user_id="test_user_456"
)
@pytest.mark.asyncio
async def test_database_tool_success(self, mock_dependencies):
"""Test database tool with successful response."""
test_model = TestModel(call_tools=['mock_database_query'])
with test_agent.override(model=test_model):
result = await test_agent.run(
"Please query the database for user data",
deps=mock_dependencies
)
# Verify database was called
mock_dependencies.database.execute_query.assert_called()
# TestModel should include tool results
assert "mock_database_query" in result.data.message
@pytest.mark.asyncio
async def test_database_tool_error(self, mock_dependencies):
"""Test database tool with error handling."""
# Configure mock to raise exception
mock_dependencies.database.execute_query.side_effect = Exception("Connection failed")
test_model = TestModel(call_tools=['mock_database_query'])
with test_agent.override(model=test_model):
result = await test_agent.run(
"Query the database",
deps=mock_dependencies
)
# Tool should handle the error gracefully
assert "mock_database_query" in result.data.message
def test_api_tool_with_data(self, mock_dependencies):
"""Test API tool with POST data."""
test_model = TestModel(call_tools=['mock_api_call'])
with test_agent.override(model=test_model):
result = test_agent.run_sync(
"Make an API call to create a new record",
deps=mock_dependencies
)
# Verify API was called
mock_dependencies.api_client.post.assert_called()
# Check tool execution in response
assert "mock_api_call" in result.data.message
class TestAgentWithFunctionModel:
"""Test agent behavior with FunctionModel for custom responses."""
@pytest.fixture
def test_dependencies(self):
"""Create basic test dependencies."""
return TestDependencies(
database=AsyncMock(),
api_client=Mock()
)
def test_function_model_custom_behavior(self, test_dependencies):
"""Test agent with FunctionModel for custom behavior."""
def custom_response_func(messages, tools):
"""Custom function to generate specific responses."""
last_message = messages[-1].content if messages else ""
if "error" in last_message.lower():
return '{"message": "Error detected and handled", "confidence": 0.6, "actions": ["error_handling"]}'
else:
return '{"message": "Normal operation", "confidence": 0.9, "actions": ["standard_response"]}'
function_model = FunctionModel(function=custom_response_func)
with test_agent.override(model=function_model):
# Test normal case
result1 = test_agent.run_sync(
"Please help me with a normal request",
deps=test_dependencies
)
assert result1.data.message == "Normal operation"
assert result1.data.confidence == 0.9
# Test error case
result2 = test_agent.run_sync(
"There's an error in the system",
deps=test_dependencies
)
assert result2.data.message == "Error detected and handled"
assert result2.data.confidence == 0.6
assert "error_handling" in result2.data.actions
class TestAgentValidation:
"""Test agent output validation and error scenarios."""
@pytest.fixture
def test_dependencies(self):
"""Create test dependencies."""
return TestDependencies(
database=AsyncMock(),
api_client=Mock()
)
def test_invalid_output_handling(self, test_dependencies):
"""Test how agent handles invalid output format."""
# TestModel with invalid JSON output
test_model = TestModel(
custom_output_text='{"message": "test", "invalid_field": "should_not_exist"}'
)
with test_agent.override(model=test_model):
# This should either succeed with validation or raise appropriate error
try:
result = test_agent.run_sync(
"Test invalid output",
deps=test_dependencies
)
# If it succeeds, Pydantic should filter out invalid fields
assert hasattr(result.data, 'message')
assert not hasattr(result.data, 'invalid_field')
except Exception as e:
# Or it might raise a validation error, which is also acceptable
assert "validation" in str(e).lower() or "error" in str(e).lower()
def test_missing_required_fields(self, test_dependencies):
"""Test handling of missing required fields in output."""
# TestModel with missing required message field
test_model = TestModel(
custom_output_text='{"confidence": 0.8}'
)
with test_agent.override(model=test_model):
try:
result = test_agent.run_sync(
"Test missing fields",
deps=test_dependencies
)
# Should either provide default or raise validation error
if hasattr(result.data, 'message'):
assert result.data.message is not None
except Exception as e:
# Validation error is expected for missing required fields
assert any(keyword in str(e).lower() for keyword in ['validation', 'required', 'missing'])
class TestAgentIntegration:
"""Integration tests for complete agent workflows."""
@pytest.fixture
def full_mock_dependencies(self):
"""Create fully configured mock dependencies."""
database_mock = AsyncMock()
database_mock.execute_query.return_value = {
"user_id": "123",
"name": "Test User",
"status": "active"
}
api_mock = Mock()
api_mock.post.return_value = {
"status": "success",
"transaction_id": "txn_123456"
}
return TestDependencies(
database=database_mock,
api_client=api_mock,
user_id="test_integration_user"
)
@pytest.mark.asyncio
async def test_complete_workflow(self, full_mock_dependencies):
"""Test complete agent workflow with multiple tools."""
test_model = TestModel(call_tools='all') # Call all available tools
with test_agent.override(model=test_model):
result = await test_agent.run(
"Please look up user information and create a new transaction",
deps=full_mock_dependencies
)
# Verify both tools were potentially called
assert result.data.message is not None
assert isinstance(result.data.actions, list)
# Verify mocks were called
full_mock_dependencies.database.execute_query.assert_called()
full_mock_dependencies.api_client.post.assert_called()
class TestAgentErrorRecovery:
"""Test agent error handling and recovery patterns."""
@pytest.fixture
def failing_dependencies(self):
"""Create dependencies that will fail for testing error handling."""
database_mock = AsyncMock()
database_mock.execute_query.side_effect = Exception("Database connection failed")
api_mock = Mock()
api_mock.post.side_effect = Exception("API service unavailable")
return TestDependencies(
database=database_mock,
api_client=api_mock,
user_id="failing_test_user"
)
@pytest.mark.asyncio
async def test_tool_error_recovery(self, failing_dependencies):
"""Test agent behavior when tools fail."""
test_model = TestModel(call_tools='all')
with test_agent.override(model=test_model):
# Agent should handle tool failures gracefully
result = await test_agent.run(
"Try to access database and API",
deps=failing_dependencies
)
# Even with tool failures, agent should return a valid response
assert result.data.message is not None
assert isinstance(result.data.confidence, float)
# Pytest configuration and utilities
@pytest.fixture(scope="session")
def event_loop():
"""Create an instance of the default event loop for the test session."""
loop = asyncio.get_event_loop_policy().new_event_loop()
yield loop
loop.close()
def pytest_configure(config):
"""Configure pytest with custom markers."""
config.addinivalue_line(
"markers", "integration: mark test as integration test"
)
config.addinivalue_line(
"markers", "slow: mark test as slow running"
)
if __name__ == "__main__":
# Run tests directly
pytest.main([__file__, "-v"])

View File

@@ -0,0 +1,374 @@
"""
Tool-Enabled Agent with Web Search and Calculator
Demonstrates PydanticAI tool integration patterns:
- Environment-based model configuration
- Tool registration with @agent.tool decorator
- RunContext for dependency injection
- Parameter validation with type hints
- Error handling and retry mechanisms
- String output (default, no result_type needed)
"""
import logging
import math
import json
import asyncio
from dataclasses import dataclass
from typing import Optional, List, Dict, Any
from datetime import datetime
import aiohttp
from pydantic_settings import BaseSettings
from pydantic import Field
from pydantic_ai import Agent, RunContext
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.models.openai import OpenAIModel
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
logger = logging.getLogger(__name__)
class Settings(BaseSettings):
"""Configuration settings for the tool-enabled agent."""
# LLM Configuration
llm_provider: str = Field(default="openai")
llm_api_key: str = Field(...)
llm_model: str = Field(default="gpt-4")
llm_base_url: str = Field(default="https://api.openai.com/v1")
class Config:
env_file = ".env"
case_sensitive = False
def get_llm_model() -> OpenAIModel:
"""Get configured LLM model from environment settings."""
try:
settings = Settings()
provider = OpenAIProvider(
base_url=settings.llm_base_url,
api_key=settings.llm_api_key
)
return OpenAIModel(settings.llm_model, provider=provider)
except Exception:
# For testing without env vars
import os
os.environ.setdefault("LLM_API_KEY", "test-key")
settings = Settings()
provider = OpenAIProvider(
base_url=settings.llm_base_url,
api_key="test-key"
)
return OpenAIModel(settings.llm_model, provider=provider)
@dataclass
class ToolDependencies:
"""Dependencies for tool-enabled agent."""
session: Optional[aiohttp.ClientSession] = None
api_timeout: int = 10
max_search_results: int = 5
calculation_precision: int = 6
session_id: Optional[str] = None
SYSTEM_PROMPT = """
You are a helpful research assistant with access to web search and calculation tools.
Your capabilities:
- Web search for current information and facts
- Mathematical calculations and data analysis
- Data processing and formatting
- Source verification and citation
Guidelines:
- Always use tools when you need current information or calculations
- Cite sources when providing factual information
- Show your work for mathematical calculations
- Be precise and accurate in your responses
- If tools fail, explain the limitation and provide what you can
"""
# Create the tool-enabled agent - note: no result_type, defaults to string
tool_agent = Agent(
get_llm_model(),
deps_type=ToolDependencies,
system_prompt=SYSTEM_PROMPT
)
@tool_agent.tool
async def web_search(
ctx: RunContext[ToolDependencies],
query: str,
max_results: Optional[int] = None
) -> str:
"""
Search the web for current information.
Args:
query: Search query string
max_results: Maximum number of results to return (default: 5)
Returns:
Formatted search results with titles, snippets, and URLs
"""
if not ctx.deps.session:
return "Web search unavailable: No HTTP session configured"
max_results = max_results or ctx.deps.max_search_results
try:
# Using DuckDuckGo Instant Answer API as a simple example
# In production, use proper search APIs like Google, Bing, or DuckDuckGo
search_url = "https://api.duckduckgo.com/"
params = {
"q": query,
"format": "json",
"pretty": "1",
"no_redirect": "1"
}
async with ctx.deps.session.get(
search_url,
params=params,
timeout=ctx.deps.api_timeout
) as response:
if response.status == 200:
data = await response.json()
results = []
# Process instant answer if available
if data.get("AbstractText"):
results.append({
"title": "Instant Answer",
"snippet": data["AbstractText"],
"url": data.get("AbstractURL", "")
})
# Process related topics
for topic in data.get("RelatedTopics", [])[:max_results-len(results)]:
if isinstance(topic, dict) and "Text" in topic:
results.append({
"title": topic.get("FirstURL", "").split("/")[-1].replace("_", " "),
"snippet": topic["Text"],
"url": topic.get("FirstURL", "")
})
if not results:
return f"No results found for query: {query}"
# Format results
formatted_results = []
for i, result in enumerate(results, 1):
formatted_results.append(
f"{i}. **{result['title']}**\n"
f" {result['snippet']}\n"
f" Source: {result['url']}"
)
return "\n\n".join(formatted_results)
else:
return f"Search failed with status: {response.status}"
except asyncio.TimeoutError:
return f"Search timed out after {ctx.deps.api_timeout} seconds"
except Exception as e:
return f"Search error: {str(e)}"
@tool_agent.tool
def calculate(
ctx: RunContext[ToolDependencies],
expression: str,
description: Optional[str] = None
) -> str:
"""
Perform mathematical calculations safely.
Args:
expression: Mathematical expression to evaluate
description: Optional description of what's being calculated
Returns:
Calculation result with formatted output
"""
try:
# Safe evaluation - only allow mathematical operations
allowed_names = {
"abs": abs, "round": round, "min": min, "max": max,
"sum": sum, "pow": pow, "sqrt": math.sqrt,
"sin": math.sin, "cos": math.cos, "tan": math.tan,
"log": math.log, "log10": math.log10, "exp": math.exp,
"pi": math.pi, "e": math.e
}
# Remove any potentially dangerous operations
safe_expression = expression.replace("__", "").replace("import", "")
# Evaluate the expression
result = eval(safe_expression, {"__builtins__": {}}, allowed_names)
# Format result with appropriate precision
if isinstance(result, float):
result = round(result, ctx.deps.calculation_precision)
output = f"Calculation: {expression} = {result}"
if description:
output = f"{description}\n{output}"
return output
except Exception as e:
return f"Calculation error: {str(e)}\nExpression: {expression}"
@tool_agent.tool
def format_data(
ctx: RunContext[ToolDependencies],
data: str,
format_type: str = "table"
) -> str:
"""
Format data into structured output.
Args:
data: Raw data to format
format_type: Type of formatting (table, list, json)
Returns:
Formatted data string
"""
try:
lines = data.strip().split('\n')
if format_type == "table":
# Simple table formatting
if len(lines) > 1:
header = lines[0]
rows = lines[1:]
# Basic table formatting
formatted = f"| {header} |\n"
formatted += f"|{'-' * (len(header) + 2)}|\n"
for row in rows[:10]: # Limit to 10 rows
formatted += f"| {row} |\n"
return formatted
else:
return data
elif format_type == "list":
# Bullet point list
formatted_lines = [f"{line.strip()}" for line in lines if line.strip()]
return "\n".join(formatted_lines)
elif format_type == "json":
# Try to parse and format as JSON
try:
parsed = json.loads(data)
return json.dumps(parsed, indent=2)
except json.JSONDecodeError:
# If not valid JSON, create simple key-value structure
items = {}
for i, line in enumerate(lines):
items[f"item_{i+1}"] = line.strip()
return json.dumps(items, indent=2)
return data
except Exception as e:
return f"Formatting error: {str(e)}"
@tool_agent.tool
def get_current_time(ctx: RunContext[ToolDependencies]) -> str:
"""
Get the current date and time.
Returns:
Current timestamp in a readable format
"""
now = datetime.now()
return now.strftime("%Y-%m-%d %H:%M:%S UTC")
async def ask_agent(
question: str,
dependencies: Optional[ToolDependencies] = None
) -> str:
"""
Ask the tool-enabled agent a question.
Args:
question: Question or request for the agent
dependencies: Optional tool dependencies
Returns:
String response from the agent
"""
if dependencies is None:
# Create HTTP session for web search
session = aiohttp.ClientSession()
dependencies = ToolDependencies(session=session)
try:
result = await tool_agent.run(question, deps=dependencies)
return result.data
finally:
# Clean up session if we created it
if dependencies.session and not dependencies.session.closed:
await dependencies.session.close()
def ask_agent_sync(question: str) -> str:
"""
Synchronous version of ask_agent.
Args:
question: Question or request for the agent
Returns:
String response from the agent
"""
return asyncio.run(ask_agent(question))
# Example usage and demonstration
if __name__ == "__main__":
async def demo_tools():
"""Demonstrate the tool-enabled agent capabilities."""
print("=== Tool-Enabled Agent Demo ===\n")
# Create dependencies with HTTP session
session = aiohttp.ClientSession()
dependencies = ToolDependencies(session=session)
try:
# Sample questions that exercise different tools
questions = [
"What's the current time?",
"Calculate the square root of 144 plus 25% of 200",
"Search for recent news about artificial intelligence",
"Format this data as a table: Name,Age\nAlice,25\nBob,30\nCharlie,35"
]
for question in questions:
print(f"Question: {question}")
response = await ask_agent(question, dependencies)
print(f"Answer: {response}")
print("-" * 60)
finally:
await session.close()
# Run the demo
asyncio.run(demo_tools())