Build agents that go the distance.

from pydantic import BaseModel
from pydantic_ai import (
    Agent, ModelRetry, RunContext,
)
import logfire

logfire.configure(
    send_to_logfire='if-token-present',
)
logfire.instrument_pydantic_ai()

class ResearchBrief(BaseModel):
    summary: str
    sources: list[str]
    follow_up: list[str]

agent = Agent(
    'openai:gpt-5.6-sol',
    deps_type=list[str],
    output_type=ResearchBrief,
    defer_model_check=True,
)

@agent.tool
def read_notes(
    ctx: RunContext[list[str]],
) -> list[str]:
    """Read this customer's notes."""
    return ctx.deps

@agent.output_validator
def require_sources(
    brief: ResearchBrief,
) -> ResearchBrief:
    if not brief.sources:
        raise ModelRetry(
            'Cite your sources.'
        )
    return brief

if __name__ == '__main__':
    result = agent.run_sync(
        'Read notes. Write a brief.',
        deps=['We need a Python SDK.'],
    )
    print(result.output.summary)

import pytest
from pydantic_ai import (
    UnexpectedModelBehavior,
)
from pydantic_ai.messages import (
    ToolReturnPart as ToolResult,
)
from pydantic_ai.models.test import (
    TestModel,
)
from research_agent import agent

def test_research_brief():
    notes = ['We need a Python SDK.']
    output = {
        'summary': 'Ship a Python SDK.',
        'sources': ['Customer notes'],
        'follow_up': ['Prototype SDK'],
    }
    model = TestModel(
        call_tools=['read_notes'],
        custom_output_args=output,
    )
    with agent.override(model=model):
        result = agent.run_sync(
            'Write a brief from notes.',
            deps=notes,
        )

    tool_name = 'read_notes'
    tool_results = [
        part.content
        for msg in result.all_messages()
        for part in msg.parts
        if isinstance(part, ToolResult)
        and part.tool_name == tool_name
    ]
    assert tool_results == [notes]
    assert result.output.follow_up == [
        'Prototype SDK'
    ]

def test_rejects_missing_sources():
    model = TestModel(
        call_tools=[],
        custom_output_args={
            'summary': 'Ship a SDK.',
            'sources': [],
            'follow_up': [],
        },
    )
    with agent.override(model=model):
        with pytest.raises(
            UnexpectedModelBehavior
        ):
            agent.run_sync(
                'Write a brief.',
                deps=[],
            )