Type-Safe Agent Development with PydanticAI Patterns

from dotenv import load_dotenv
import nest_asyncio

load_dotenv()
nest_asyncio.apply()

from pydantic import BaseModel, Field
from typing import Optional

class Paper(BaseModel):
    """Academic paper with validated fields."""
    title: str = Field(min_length=1)
    authors: list[str] = Field(min_length=1)
    year: int = Field(ge=1800, le=2025)
    abstract: str
    doi: Optional[str] = None

class ResearchSummary(BaseModel):
    """Final research output with validation."""
    query: str = Field(min_length=1)
    summary: str = Field(min_length=1, description="Natural language summary")
    key_papers: list[Paper] = Field(min_length=1, description="Most relevant papers found")
    confidence: float = Field(ge=0.0, le=1.0, description="Agent's confidence that an answer to the user's query has been found")

    def _repr_html_(self) -> str:
        """Rich HTML display for Jupyter notebooks."""
        papers_html = ""
        for i, paper in enumerate(self.key_papers, 1):
            authors = ", ".join(paper.authors)
            doi_link = f'<a href="https://doi.org/{paper.doi}" target="_blank">DOI</a>' if paper.doi else ""
            papers_html += f"""
            <div style="margin: 10px 0; padding: 10px; background: #f8f9fa; border-left: 3px solid #007bff;">
                <strong>{i}. {paper.title}</strong> ({paper.year}) {doi_link}<br>
                <em>{authors}</em><br>
                <small style="color: #666;">{paper.abstract[:200]}...</small>
            </div>
            """

        confidence_color = "#28a745" if self.confidence >= 0.7 else "#ffc107" if self.confidence >= 0.4 else "#dc3545"
        confidence_pct = f"{self.confidence:.0%}"

        return f"""
        <div style="font-family: sans-serif; max-width: 800px;">
            <h3 style="color: #333; border-bottom: 2px solid #007bff; padding-bottom: 10px;">
                📚 Research Results: "{self.query}"
            </h3>
            <div style="margin: 15px 0; padding: 15px; background: #e7f3ff; border-radius: 5px;">
                <strong>Summary:</strong><br>
                {self.summary}
            </div>
            <div style="margin: 10px 0;">
                <strong>Confidence:</strong>
                <span style="color: {confidence_color}; font-weight: bold;">{confidence_pct}</span>
            </div>
            <div style="margin-top: 20px;">
                <strong>Key Papers ({len(self.key_papers)}):</strong>
                {papers_html}
            </div>
        </div>
        """

    def __str__(self) -> str:
        """Plain text representation."""
        lines = [
            f"Research Results: {self.query}",
            "=" * 60,
            f"\nSummary:\n{self.summary}",
            f"\nConfidence: {self.confidence:.0%}",
            f"\nKey Papers ({len(self.key_papers)}):",
        ]

        for i, paper in enumerate(self.key_papers, 1):
            lines.append(f"\n{i}. {paper.title} ({paper.year})")
            lines.append(f"   Authors: {', '.join(paper.authors)}")
            if paper.doi:
                lines.append(f"   DOI: {paper.doi}")

        return "\n".join(lines)

# Test that validation works
try:
    invalid = Paper(title="", authors=[], year=3000, abstract="test")
except Exception as e:
    print(f"✓ Validation caught error: {type(e).__name__}")

    # print the error
    print(e)

valid = Paper(
    title="Attention Is All You Need",
    authors=["Vaswani et al."],
    year=2017,
    abstract="We propose the Transformer..."
)
print(f"✓ Valid paper created: {valid.title}")

from dataclasses import dataclass

@dataclass
class ResearchDeps:
    """Dependencies for the research agent."""
    database_url: str
    api_key: str
    max_papers_per_search: int = 10

    def search_database(self, query: str, limit: int) -> list[Paper]:
        """Simulate database search (in production, would query real DB)."""
        print(f"  🔍 Searching database for: '{query}' (limit: {limit})")

        # In production, this would be:
        # conn = connect(self.database_url, self.api_key)
        # results = conn.execute(search_query)
        # return [Paper(**row) for row in results]

        # For demo, return mock results
        return [
            Paper(
                title="Attention Is All You Need",
                authors=["Vaswani, A.", "Shazeer, N.", "Parmar, N."],
                year=2017,
                abstract="We propose a new simple network architecture, the Transformer...",
                doi="10.5555/3295222.3295349"
            ),
            Paper(
                title="BERT: Pre-training of Deep Bidirectional Transformers",
                authors=["Devlin, J.", "Chang, M.", "Lee, K."],
                year=2019,
                abstract="We introduce BERT, a method for pre-training language representations...",
                doi="10.18653/v1/N19-1423"
            ),
        ]

    def get_paper_by_id(self, paper_id: str) -> Paper:
        """Retrieve specific paper by ID."""
        print(f"  📄 Fetching paper: {paper_id}")
        return Paper(
            title="Example Paper",
            authors=["Author, A."],
            year=2023,
            abstract="This is an example paper abstract."
        )

# Create dependencies
deps = ResearchDeps(
    database_url="sqlite3:///papers.db",
    api_key="demo-api-key-12345",
    max_papers_per_search=5
)

print("✓ Dependencies configured")

from pydantic_ai import Agent, RunContext

# Create agent with system prompt and structured output
research_agent = Agent(
    'anthropic:claude-sonnet-4-5',
    deps_type=ResearchDeps,
    output_type=ResearchSummary,
    system_prompt="""You are a helpful research assistant.

    When asked to research a topic:
    1. Search for relevant papers using the search_papers tool
    2. Analyze the results carefully
    3. Provide a clear summary with the most important papers
    4. Rate your confidence based on result quality and relevance

    Be thorough but concise. Focus on the most impactful papers.
    IMPORTANT: never, under any circumstances, make up paper titles, authors, years, abstracts, or DOIs.

    IMPORTANT: If you cannot find relevant papers, state that clearly in your summary and set confidence to 0.0. Also don't return any key papers.

    IMPORTANT: only use details from the papers you find when generating the summary. do not invent information or use general knowledge.
               again if you cannot find relevant papers, state that clearly in your summary and set confidence to 0.0.

    """
)

# Register tool 1: Search for papers
@research_agent.tool
def search_papers(
    ctx: RunContext[ResearchDeps],
    query: str,
    max_results: int = 5
) -> list[Paper]:
    """
    Search for academic papers by query string.

    Args:
        query: Search query (keywords, topics, authors)
        max_results: Maximum number of papers to return (default 5)

    Returns:
        List of relevant papers
    """
    # Access dependencies via ctx.deps
    limit = min(max_results, ctx.deps.max_papers_per_search)
    papers = ctx.deps.search_database(query, limit)
    return papers

# Register tool 2: Get specific paper details
@research_agent.tool
def get_paper_details(
    ctx: RunContext[ResearchDeps],
    paper_id: str
) -> Paper:
    """
    Get detailed information about a specific paper by its ID.

    Args:
        paper_id: Unique identifier for the paper

    Returns:
        Paper object with full details
    """
    return ctx.deps.get_paper_by_id(paper_id)

print("✓ Agent created with 2 tools registered")

from pydantic_ai import UsageLimits

async def run_research(question: str):
    deps = ResearchDeps(
        database_url="sqlite3:///papers.db",
        api_key="demo-api-key-12345",
        max_papers_per_search=5
    )
    result = await research_agent.run(
        question,
        deps=deps,
        usage_limits=UsageLimits(tool_calls_limit=5),
    )
    return result

research_result = await run_research("How many species of monkeys live in north america")

research_result.all_messages()

research_result.output

@research_agent.tool
def compare_papers(
    ctx: RunContext[ResearchDeps],
    paper1_title: str,
    paper2_title: str
) -> str:
    """
    Compare two papers by their titles.

    Args:
        paper1_title: Title of first paper
        paper2_title: Title of second paper

    Returns:
        Comparison summary
    """
    # TODO: Implement comparison logic
    # In production, would search for both papers and compare:
    # - Publication dates
    # - Citation counts
    # - Research methods
    # - Key contributions

    return f"Comparison between '{paper1_title}' and '{paper2_title}': [TODO: implement]"

from pydantic_ai import RunContext

# Create agent without static system prompt
adaptive_agent = Agent(
    'anthropic:claude-haiku-4-5',
    deps_type=ResearchDeps,
    output_type=ResearchSummary
)

# Dynamic system prompt based on context
@adaptive_agent.system_prompt(dynamic=True)
def get_system_prompt(ctx: RunContext[ResearchDeps]) -> str:
    """Generate system prompt based on current context."""
    max_papers = ctx.deps.max_papers_per_search

    return f"""You are a research assistant with access to {max_papers} papers per search.

Focus on quality over quantity. When analyzing papers:
- Prioritize recent publications (last 5 years)
- Look for highly-cited works
- Consider methodological rigor
- Provide balanced summaries

Your database: {ctx.deps.database_url}
    """

print("✓ Adaptive agent created with dynamic system prompt")