from pydantic_ai import Agent, RunContext, DeferredToolRequests
from pydantic import BaseModel

agent = Agent(
    'openai:gpt-4.1',
    output_type=str,
    # 通过 deferred=True 标记需要审批的工具
)

@agent.tool(deferred=True)
async def delete_user_account(ctx: RunContext, user_id: str) -> str:
    """永久删除用户账户。需要人工审批。"""
    # 这段代码只在人工审批通过后才会执行
    await db.users.delete(user_id)
    return f"用户 {user_id} 已删除"

# 运行 Agent——如果它尝试调用 delete_user_account，
# 执行会暂停并返回 DeferredToolRequests 对象
result = await agent.run("清理不活跃账户")

if isinstance(result.output, DeferredToolRequests):
    # 将待执行的工具调用展示给人类审批
    for tool_call in result.output.calls:
        print(f"Agent 想要执行: {tool_call.tool_name}({tool_call.args})")
        approved = input("是否批准？(y/n): ")
        if approved.lower() == 'y':
            # 带审批结果恢复执行
            result = await agent.run(
                "已批准",
                message_history=result.all_messages(),
                deferred_tool_requests=result.output,
            )

from dataclasses import dataclass, field
from pydantic import BaseModel
from pydantic_ai import Agent
from pydantic_graph import BaseNode, End, Graph, GraphRunContext

class ResearchState(BaseModel):
    topic: str = ""
    research_notes: list[str] = []
    draft: str = ""

@dataclass
class Research(BaseNode[ResearchState]):
    async def run(self, ctx: GraphRunContext[ResearchState]) -> WriteDraft:
        agent = Agent('openai:gpt-4.1', output_type=str)
        result = await agent.run(f"调研主题：{ctx.state.topic}")
        ctx.state.research_notes.append(result.output)
        return WriteDraft()

@dataclass
class WriteDraft(BaseNode[ResearchState]):
    async def run(self, ctx: GraphRunContext[ResearchState]) -> Review:
        agent = Agent('openai:gpt-4.1', output_type=str)
        notes = "\n".join(ctx.state.research_notes)
        result = await agent.run(f"根据以下笔记写初稿：\n{notes}")
        ctx.state.draft = result.output
        return Review()

@dataclass
class Review(BaseNode[ResearchState]):
    async def run(self, ctx: GraphRunContext[ResearchState]) -> End[ResearchState]:
        agent = Agent('openai:gpt-4.1', output_type=bool)
        result = await agent.run(f"这篇初稿可以吗？{ctx.state.draft}")
        if result.output:
            return End(ctx.state)
        else:
            return Research()  # 回到研究节点继续迭代

# 构建并运行图
graph = Graph(nodes=[Research, WriteDraft, Review])
state = ResearchState(topic="2026年AI Agent安全")
result = await graph.run(state=state)

from pydantic_ai import Agent
from pydantic_ai.mcp import MCPServerStdio

# 连接到支持完整协议的 MCP 服务器
mcp_server = MCPServerStdio(
    command="npx",
    args=["-y", "@my/mcp-server"],
)

agent = Agent(
    'openai:gpt-4.1',
    toolsets=[mcp_server],
)

# MCP 服务器现在可以：
# 1. 暴露 Agent 正常调用的工具
# 2. 请求客户端 LLM 执行采样（代表服务器推理）
# 3. 通过交互表单请求结构化用户输入（表单、确认）
# 4. 返回富内容：图片、音频、嵌入资源

result = await agent.run(
    "分析上传的图片，并在继续操作前请用户确认"
)

from pydantic_ai import Agent, RunContext
from pydantic_ai.capabilities import Capability

# 定义一个默认不加载的能力
database_capability = Capability(
    name="database_tools",
    description="用于查询和修改生产数据库的工具",
    # 这些工具只在 Agent 调用 load_capability 时才会加载
    tools=[query_db, update_db, delete_record],
)

agent = Agent(
    'openai:gpt-4.1',
    output_type=str,
    capabilities=[database_capability],
    # 注意：数据库工具不在初始 system prompt 中！
)

# Agent 以精简的上下文启动。
# 当需要数据库访问时，它调用：
#   load_capability(id="database_tools")
# 此时数据库工具才会被加载到上下文中。
# 使用完后可以卸载以节省 token。

from pydantic_ai import Agent, UsageLimits
from pydantic_ai.usage import Usage
import logfire

# 配置可观测性
logfire.configure()
logfire.instrument_pydantic_ai()

agent = Agent(
    'openai:gpt-4.1',
    output_type=str,
    # 设置硬性用量限制
    usage_limits=UsageLimits(
        request_tokens=50_000,      # 每次运行最多 50K 输入 token
        response_tokens=4_096,       # 最多 4K 输出 token
        total_tokens=54_006,         # 硬性上限
        requests=10,                 # 每次运行最多 10 次 LLM 调用
    ),
)

try:
    result = await agent.run("撰写一份关于AI安全的综合报告")
    # 检查实际用量
    print(f"输入 token：{result.usage.input_tokens}")
    print(f"输出 token：{result.usage.output_tokens}")
    print(f"总计：{result.usage.total_tokens}")
except UsageLimitExceeded as e:
    print(f"Agent 触及预算上限：{e}")
    # 优雅处理——不会有意外账单