Pydantic AI：类型安全的 Agent 开发框架

from pydantic_ai import Agent, RunContext
from pydantic import BaseModel, Field
from typing import Literal
from dataclasses import dataclass

# 定义结构化输出
class ResearchReport(BaseModel):
    title: str = Field(description="报告标题")
    summary: str = Field(description="执行摘要，不超过200字")
    key_findings: list[str] = Field(description="关键发现列表")
    confidence: float = Field(ge=0, le=1, description="置信度 0-1")
    sources: list[str] = Field(description="信息来源 URL 列表")
    recommendation: Literal["buy", "hold", "sell"] = Field(description="投资建议")

# 定义依赖类型
@dataclass
class ResearchDeps:
    api_keys: dict[str, str]
    database_url: str
    max_sources: int = 10

# 创建 Agent
research_agent = Agent(
    model="openai:gpt-4o",
    deps_type=ResearchDeps,
    output_type=ResearchReport,  # 类型安全的输出
    system_prompt="你是一名专业的研究分析师..."
)

# 类型安全的工具定义
@research_agent.tool
async def search_database(ctx: RunContext[ResearchDeps], query: str, limit: int = 10) -> list[dict]:
    """
    搜索内部数据库。
    
    Args:
        query: 搜索查询
        limit: 返回结果数量上限
    """
    # ctx.deps 是 ResearchDeps 类型，IDE 自动补全
    async with httpx.AsyncClient() as client:
        response = await client.post(
            f"{ctx.deps.database_url}/search",
            json={"query": query, "limit": min(limit, ctx.deps.max_sources)},
            headers={"Authorization": f"Bearer {ctx.deps.api_keys['database']}"}
        )
        return response.json()["results"]

@research_agent.tool
async def fetch_web_page(ctx: RunContext[ResearchDeps], url: str) -> str:
    """获取网页内容。"""
    async with httpx.AsyncClient() as client:
        response = await client.get(url, timeout=10)
        return response.text[:5000]  # 限制内容长度

# 运行 Agent —— 输出类型自动验证
result = await research_agent.run(
    "分析 2026 年 AGI 芯片市场",
    deps=ResearchDeps(
        api_keys={"database": "xxx"},
        database_url="https://api.example.com"
    )
)

# result.output 是 ResearchReport 类型，IDE 完整支持
print(result.output.title)
print(result.output.confidence)
print(result.output.recommendation)
# 如果 LLM 输出不符合 Schema，Pydantic 会抛出
# ValidationError，而非静默返回错误数据

from pydantic_ai import Agent
from pydantic import BaseModel, validator
from typing import Optional

class CodeReview(BaseModel):
    file_name: str
    issues: list["CodeIssue"]
    overall_rating: int = Field(ge=1, le=10, description="整体评分 1-10")
    approved: bool
    
    @validator("file_name")
    def validate_filename(cls, v):
        if not v.endswith((".py", ".js", ".ts")):
            raise ValueError("仅支持 .py/.js/.ts 文件")
        return v

class CodeIssue(BaseModel):
    line_number: int = Field(ge=1)
    severity: Literal["error", "warning", "info"]
    message: str
    suggestion: Optional[str] = None

review_agent = Agent(
    model="anthropic:claude-4-sonnet",
    output_type=CodeReview,
    system_prompt="你是代码审查专家，分析代码质量问题。"
)

# Agent 输出保证符合 CodeReview Schema
result = await review_agent.run("审查以下代码: ...")
review: CodeReview = result.output  # 类型保证

from pydantic_ai.models import Model
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.models.anthropic import AnthropicModel
from pydantic_ai.models.groq import GroqModel

# 环境感知模型选择
def get_model() -> Model:
    env = os.getenv("ENVIRONMENT", "dev")
    
    if env == "production":
        return OpenAIModel("gpt-4o", api_key=os.getenv("OPENAI_API_KEY"))
    elif env == "staging":
        return AnthropicModel("claude-4-sonnet", api_key=os.getenv("ANTHROPIC_API_KEY"))
    elif env == "dev":
        return GroqModel("llama-4-70b", api_key=os.getenv("GROQ_API_KEY"))
    else:
        # 使用 Ollama 本地模型
        return OpenAIModel(
            "qwen3:72b",
            base_url="http://localhost:11434/v1",
            api_key="ollama"
        )

agent = Agent(
    model=get_model(),
    output_type=ResearchReport,
    system_prompt="..."
)

# 运行时切换模型
result = await agent.run("...", model=AnthropicModel("claude-4-opus"))

from pydantic_ai import Agent
from pydantic_ai.messages import Part

streaming_agent = Agent(
    model="openai:gpt-4o",
    system_prompt="你是技术写作助手..."
)

# 流式输出
async with streaming_agent.run_stream("写一篇关于 AGI 的短文") as result:
    async for message in result.stream_text(delta=True):
        print(message, end="", flush=True)

# 流式结构化输出（2026 新特性）
class ArticleOutline(BaseModel):
    title: str
    sections: list[str]
    estimated_words: int

async with streaming_agent.run_stream(
    "生成文章大纲",
    output_type=ArticleOutline
) as result:
    # 部分结果流式返回
    async for partial in result.stream_structured():
        print(f"当前大纲: {partial.title} ({len(partial.sections)} 节)")
    
    final = await result.get_output()
    print(f"最终大纲: {final}")

from pydantic_ai import Agent, RunContext
from dataclasses import dataclass
import httpx
import asyncio

@dataclass
class AppDeps:
    http_client: httpx.AsyncClient
    db_pool: asyncpg.Pool
    cache: redis.Redis
    config: dict

# 所有工具共享同一个依赖上下文
agent = Agent("openai:gpt-4o", deps_type=AppDeps)

@agent.tool
async def get_user_profile(ctx: RunContext[AppDeps], user_id: int) -> dict:
    # 从缓存获取
    cached = await ctx.deps.cache.get(f"user:{user_id}")
    if cached:
        return json.loads(cached)
    
    # 从数据库获取
    async with ctx.deps.db_pool.acquire() as conn:
        user = await conn.fetchrow(
            "SELECT * FROM users WHERE id = $1", user_id
        )
    
    # 写入缓存
    await ctx.deps.cache.setex(f"user:{user_id}", 3600, json.dumps(dict(user)))
    return dict(user)

@agent.tool
async def call_external_api(ctx: RunContext[AppDeps], url: str) -> dict:
    response = await ctx.deps.http_client.get(url)
    return response.json()

# 运行时注入依赖
async def main():
    async with httpx.AsyncClient() as http_client:
        db_pool = await asyncpg.create_pool(DATABASE_URL)
        cache = redis.Redis()
        
        deps = AppDeps(
            http_client=http_client,
            db_pool=db_pool,
            cache=cache,
            config={"max_results": 50}
        )
        
        result = await agent.run("查询用户 123 的信息", deps=deps)
        print(result.output)

from pydantic_ai import Agent

# 子 Agent：各自有独立的输出类型
research_agent = Agent(
    model="openai:gpt-4o",
    output_type=ResearchReport,
    system_prompt="你是研究员..."
)

writing_agent = Agent(
    model="anthropic:claude-4-sonnet",
    output_type=Article,
    system_prompt="你是技术写手..."
)

review_agent = Agent(
    model="openai:gpt-4o",
    output_type=ReviewResult,
    system_prompt="你是审稿人..."
)

# 主 Agent 协调子 Agent
class PipelineResult(BaseModel):
    research: ResearchReport
    article: Article
    review: ReviewResult
    final_status: str

pipeline_agent = Agent(
    model="openai:gpt-4o",
    output_type=PipelineResult,
    system_prompt="你是项目经理，协调多个子任务..."
)

# 主 Agent 可以委派给子 Agent
@pipeline_agent.tool
async def run_research(ctx: RunContext, topic: str) -> ResearchReport:
    result = await research_agent.run(topic, deps=ctx.deps)
    return result.output

@pipeline_agent.tool
async def run_writing(ctx: RunContext, research: ResearchReport) -> Article:
    result = await writing_agent.run(
        f"基于以下研究撰写文章: {research.model_dump_json()}",
        deps=ctx.deps
    )
    return result.output

from pydantic_ai.evals import TestCase, Evaluator

# 定义测试用例
test_cases = [
    TestCase(
        input="分析苹果公司 2026 年 Q1 财报",
        expected_output=ResearchReport(
            title="苹果 2026 Q1 财报分析",
            summary="...",
            key_findings=["营收增长 15%", "服务业务创新高"],
            confidence=0.9,
            sources=["https://investor.apple.com"],
            recommendation="buy"
        )
    ),
    # 更多测试用例...
]

# 评估器
class ResearchEvaluator(Evaluator):
    async def evaluate(self, output: ResearchReport, expected: ResearchReport) -> float:
        score = 0.0
        
        # 检查结构化字段
        if output.recommendation == expected.recommendation:
            score += 0.3
        
        # 检查发现重叠度
        overlap = len(set(output.key_findings) & set(expected.key_findings))
        score += 0.4 * (overlap / max(len(expected.key_findings), 1))
        
        # 检查置信度合理性
        if abs(output.confidence - expected.confidence) < 0.2:
            score += 0.3
        
        return score

# 运行评估
results = await Evaluator.run(
    agent=research_agent,
    test_cases=test_cases,
    evaluator=ResearchEvaluator()
)

print(f"平均得分: {results.avg_score}")
print(f"通过率: {results.pass_rate}")