mkdir pydantic-agent-demo && cd pydantic-agent-demo
python -m venv .venv && source .venv/bin/activate

# Install PydanticAI with your preferred model backend
pip install 'pydantic-ai[openai]' python-dotenv

# Or for Anthropic:
# pip install 'pydantic-ai[anthropic]'
# Or for Gemini:
# pip install 'pydantic-ai[google-generativeai]'

OPENAI_API_KEY=sk-...your-key...

# schemas.py
from pydantic import BaseModel, Field
from typing import Optional


class SkillRequirement(BaseModel):
    skill: str = Field(description="Name of the required skill or technology")
    level: str = Field(description="Required proficiency level: beginner, intermediate, or expert")
    is_required: bool = Field(description="True if mandatory, False if nice-to-have")


class SalaryRange(BaseModel):
    min_usd: Optional[int] = Field(None, description="Minimum annual salary in USD, or None if not specified")
    max_usd: Optional[int] = Field(None, description="Maximum annual salary in USD, or None if not specified")
    currency: str = Field(default="USD", description="Currency code")


class JobAnalysis(BaseModel):
    """Structured analysis of a job listing."""

    job_title: str = Field(description="Extracted job title")
    company: Optional[str] = Field(None, description="Company name if mentioned")
    experience_years_min: int = Field(description="Minimum years of experience required")
    experience_years_max: Optional[int] = Field(None, description="Maximum years or None if open-ended")
    skills: list[SkillRequirement] = Field(description="List of required and nice-to-have skills")
    salary: SalaryRange = Field(description="Extracted salary range")
    is_remote: bool = Field(description="True if remote work is explicitly mentioned")
    seniority_level: str = Field(description="One of: junior, mid, senior, staff, principal, director")
    summary: str = Field(description="One-paragraph summary of the role in plain language")
    suitability_score: float = Field(
        ge=0.0,
        le=10.0,
        description="Overall suitability score from 0-10 based on clarity and completeness of the listing",
    )

# agent.py
import os
from dataclasses import dataclass
from dotenv import load_dotenv
from pydantic_ai import Agent, RunContext
from pydantic_ai.models.openai import OpenAIModel
from schemas import JobAnalysis

load_dotenv()


@dataclass
class AgentDeps:
    """Dependencies injected into the agent at runtime."""
    user_profile: dict  # The candidate's own profile for suitability scoring
    min_salary_threshold: int = 80_000  # User's minimum acceptable salary


# Create the agent: model + output type + dependency type
job_agent = Agent(
    model=OpenAIModel("gpt-4o"),
    result_type=JobAnalysis,  # The agent MUST return this type
    deps_type=AgentDeps,
    system_prompt="""You are an expert technical recruiter who analyzes job listings.
    Extract all structured information from the job description provided.
    For suitability_score, rate how well-written and complete the listing is (not candidate fit).
    Be precise with salary extraction — only fill in values explicitly stated in the text.""",
)


@job_agent.tool
async def get_market_salary(
    ctx: RunContext[AgentDeps],
    job_title: str,
    location: str = "Remote",
) -> dict:
    """Look up typical market salary ranges for a given job title and location.

    Args:
        ctx: Run context with injected dependencies.
        job_title: The job title to look up (e.g., 'Senior Python Developer').
        location: City or region for location-adjusted data.

    Returns:
        A dict with median, p25, and p75 annual salary in USD.
    """
    # In production, call a salary API like Levels.fyi, Glassdoor, or Payscale
    # Mock implementation for tutorial
    salary_data = {
        "senior python developer": {"median": 155_000, "p25": 130_000, "p75": 185_000},
        "machine learning engineer": {"median": 175_000, "p25": 145_000, "p75": 210_000},
        "ai engineer": {"median": 180_000, "p25": 150_000, "p75": 220_000},
        "data scientist": {"median": 140_000, "p25": 115_000, "p75": 170_000},
    }
    key = job_title.lower()
    data = salary_data.get(key, {"median": 120_000, "p25": 95_000, "p75": 150_000})
    return {
        "job_title": job_title,
        "location": location,
        "market_salary_usd": data,
        "user_minimum_usd": ctx.deps.min_salary_threshold,
        "meets_minimum": data["median"] >= ctx.deps.min_salary_threshold,
    }

# run.py
import asyncio
from agent import job_agent, AgentDeps

JOB_LISTING = """
Senior AI Engineer at TechCorp (Remote)

We're looking for a Senior AI Engineer to join our growing team. You'll be building
production ML systems and LLM-powered applications.

Requirements:
- 5+ years of Python experience
- 3+ years working with machine learning frameworks (PyTorch, TensorFlow)
- Experience with LLM APIs (OpenAI, Anthropic)
- Strong knowledge of MLOps practices (Docker, Kubernetes, CI/CD)
- Experience with vector databases (Pinecone, Weaviate) is a plus

Compensation: $160,000 - $220,000 + equity + benefits

This is a fully remote position. We offer flexible hours and a async-first culture.
"""


async def main():
    deps = AgentDeps(
        user_profile={"name": "Alice", "years_experience": 6},
        min_salary_threshold=150_000,
    )

    result = await job_agent.run(
        user_prompt=f"Analyze this job listing:\n\n{JOB_LISTING}",
        deps=deps,
    )

    analysis: JobAnalysis = result.data  # Fully typed, validated output

    print(f"Job Title: {analysis.job_title}")
    print(f"Seniority: {analysis.seniority_level}")
    print(f"Experience: {analysis.experience_years_min}+ years")
    print(f"Remote: {analysis.is_remote}")
    print(f"Salary: ${analysis.salary.min_usd:,} - ${analysis.salary.max_usd:,}")
    print(f"\nRequired Skills:")
    for skill in analysis.skills:
        req_marker = "[REQUIRED]" if skill.is_required else "[NICE TO HAVE]"
        print(f"  {req_marker} {skill.skill} ({skill.level})")
    print(f"\nSuitability Score: {analysis.suitability_score}/10")
    print(f"\nSummary: {analysis.summary}")


if __name__ == "__main__":
    asyncio.run(main())

# test_agent.py
import pytest
from pydantic_ai import models
from pydantic_ai.models.test import TestModel
from agent import job_agent, AgentDeps
from schemas import JobAnalysis, SalaryRange


@pytest.mark.asyncio
async def test_job_agent_returns_valid_schema():
    """Test that the agent returns a valid JobAnalysis even with a test model."""

    # TestModel returns minimal valid data matching the result_type schema
    with models.override(job_agent, TestModel()):
        deps = AgentDeps(user_profile={}, min_salary_threshold=100_000)
        result = await job_agent.run("Analyze this listing: Senior Dev, 3+ years Python.", deps=deps)

    # PydanticAI guarantees result.data matches JobAnalysis
    assert isinstance(result.data, JobAnalysis)
    assert isinstance(result.data.salary, SalaryRange)
    assert 0.0 <= result.data.suitability_score <= 10.0


@pytest.mark.asyncio
async def test_tool_receives_deps():
    """Test that tools can access injected dependencies."""
    received_threshold = None

    @job_agent.tool
    async def capture_threshold(ctx, job_title: str) -> str:
        nonlocal received_threshold
        received_threshold = ctx.deps.min_salary_threshold
        return "captured"

    with models.override(job_agent, TestModel(custom_result_text='{"job_title":"test"}')):
        deps = AgentDeps(user_profile={}, min_salary_threshold=200_000)
        await job_agent.run("analyze", deps=deps)

    assert received_threshold == 200_000

pip install pytest pytest-asyncio
pytest test_agent.py -v

async def stream_example():
    deps = AgentDeps(user_profile={}, min_salary_threshold=100_000)

    async with job_agent.run_stream(
        f"Analyze: {JOB_LISTING}",
        deps=deps,
    ) as stream:
        # Stream partial text as it arrives
        async for text in stream.stream_text():
            print(text, end="", flush=True)

        # Get the fully validated final result
        final: JobAnalysis = await stream.get_data()
        print(f"\n\nFinal validated result: {final.job_title}")