用 Pydantic AI 构建类型安全的 LLM Agent：结构化输出、工具与依赖注入实战

from pydantic import BaseModel, Field
from pydantic_ai import Agent

class Diagnosis(BaseModel):
    disease: str = Field(description="最可能的疾病名称")
    confidence: float = Field(description="置信度，0-1之间")
    symptoms_matched: list[str] = Field(description="匹配到的症状列表")

doctor_agent = Agent(
    model="openai:gpt-4o",
    result_type=Diagnosis,          # 关键：声明输出类型
    system_prompt="你是一个医学诊断助手，根据症状给出结构化诊断。",
)

result = doctor_agent.run_sync("患者发热3天，咳嗽，乏力")
# result.data 是一个 Diagnosis 实例，字段类型已校验
print(result.data.disease)        # str
print(result.data.confidence)     # float，保证在 0-1
print(result.data.symptoms_matched)  # list[str]，不会出现 None 或字符串

from pydantic_ai import Agent, RunContext

agent = Agent("openai:gpt-4o", system_prompt="你是一个客服助手，可以查询订单状态。")

@agent.tool
async def order_status(ctx: RunContext[None], order_id: str) -> str:
    """查询订单状态。order_id 为订单编号。"""
    # 实际场景中这里会查数据库或调用内部 API
    mock_db = {
        "ORD-001": "已发货，预计明天送达",
        "ORD-002": "仓库备货中",
    }
    return mock_db.get(order_id, "未找到该订单")

result = agent.run_sync("我的订单 ORD-001 到哪了？")
# Agent 会自动调用 order_status("ORD-001")，然后把结果融入回答
print(result.data)

from dataclasses import dataclass
from pydantic_ai import Agent, RunContext

@dataclass
class Deps:
    db_conn: str          # 模拟数据库连接
    api_key: str          # 外部服务认证

agent = Agent(
    "openai:gpt-4o",
    deps_type=Deps,
    system_prompt="你是一个财务助手，可以查询账户余额。",
)

@agent.tool
async def balance(ctx: RunContext[Deps], account: str) -> str:
    """查询账户余额。account 为账户名。"""
    # ctx.deps 携带注入的依赖
    print(f"使用 db_conn={ctx.deps.db_conn} 查询")
    return f"账户 {account} 余额：¥12,350"

# 运行时传入依赖实例
deps = Deps(db_conn="postgres://prod-db:5432", api_key="sk-live-xxx")
result = agent.run_sync("查一下张三的余额", deps=deps)
print(result.data)

import asyncio
from dataclasses import dataclass
from pydantic import BaseModel, Field
from pydantic_ai import Agent, RunContext

# ---- 1. 结构化输出模型 ----
class TravelPlan(BaseModel):
    destination: str = Field(description="推荐目的地城市")
    budget_estimate: float = Field(description="预估人均费用，单位元")
    activities: list[str] = Field(description="推荐活动列表，至少3项")
    risk_note: str = Field(description="需要注意的风险或限制")

# ---- 2. 依赖声明 ----
@dataclass
class Deps:
    weather_api: str   # 天气服务端点（模拟）
    user_locale: str   # 用户所在城市

# ---- 3. Agent + 工具 ----
travel_agent = Agent(
    model="openai:gpt-4o",
    deps_type=Deps,
    result_type=TravelPlan,
    system_prompt=(
        "你是一个旅行规划助手。根据用户偏好和实时天气，"
        "给出结构化的旅行方案。务必调用天气工具确认目的地天气。"
    ),
)

@travel_agent.tool
async def get_weather(ctx: RunContext[Deps], city: str) -> str:
    """查询指定城市的当前天气。city 为城市名。"""
    # 模拟调用——实际替换为真实天气 API
    mock = {
        "三亚": "晴天，32°C，适合海滩活动",
        "成都": "多云，22°C，适合室内和美食",
        "哈尔滨": "小雪，-8°C，注意保暖",
    }
    return mock.get(city, "暂无该城市天气数据")

# ---- 4. 运行 ----
async def main():
    deps = Deps(weather_api="https://api.weather.example", user_locale="北京")
    result = await travel_agent.run("想去暖和的地方度假，3天左右", deps=deps)

    plan = result.data  # 类型：TravelPlan，字段已校验
    print(f"目的地：{plan.destination}")
    print(f"预算：¥{plan.budget_estimate:.0f}")
    print(f"活动：{', '.join(plan.activities)}")
    print(f"风险提示：{plan.risk_note}")

asyncio.run(main())

pip install pydantic-ai
export OPENAI_API_KEY=sk-your-key