pydantic-ai-tool-system

PydanticAI 工具系统

工具注册

根据是否需要上下文使用两种装饰器：

from pydantic_ai import Agent, RunContext

agent = Agent('openai:gpt-4o')

# @agent.tool - First param MUST be RunContext
@agent.tool
async def get_user_data(ctx: RunContext[MyDeps], user_id: int) -> str:
    """Get user data from database.

    Args:
        ctx: The run context with dependencies.
        user_id: The user's ID.
    """
    return await ctx.deps.db.get_user(user_id)

# @agent.tool_plain - NO context parameter allowed
@agent.tool_plain
def calculate_total(prices: list[float]) -> float:
    """Calculate total price.

    Args:
        prices: List of prices to sum.
    """
    return sum(prices)

关键规则

1. @agent.tool：第一个参数必须是 RunContext[DepsType]
2. @agent.tool_plain：不得包含 RunContext 参数
3. 文档字符串：必填，帮助大模型了解工具的用途
4. Google 风格文档字符串：用于参数说明

检查项（请在文件中核实，不要凭记忆判断）

1. 装饰器与签名匹配 — 如果第一个参数是 RunContext[...]，装饰器必须是 @agent.tool（不能是 @agent.tool_plain）。**通过条件：**同一 def 行的装饰器栈中包含 @agent.tool，且第一个参数的类型为 RunContext[...]。
2. 纯工具 — 使用 @agent.tool_plain 时，参数列表不得包含 RunContext。**通过条件：**快速检查签名，确认其中没有 RunContext。
3. 模型文档字符串 — 文档字符串不能为空；如果工具有参数，则需说明这些参数（使用 Google 风格的 Args:，或者在使用 docstring_format='sphinx' 时使用 Sphinx 风格的 :param）。**通过条件：**签名中的每个参数都在文档字符串正文中提及。

文档字符串格式

Google 风格（默认）：

@agent.tool_plain
async def search(query: str, limit: int = 10) -> list[str]:
    """Search for items.

    Args:
        query: The search query.
        limit: Maximum results to return.
    """

Sphinx 风格：

@agent.tool_plain(docstring_format='sphinx')
async def search(query: str) -> list[str]:
    """Search for items.

    :param query: The search query.
    """

工具返回类型

工具可以返回各种类型：

# String (direct)
@agent.tool_plain
def get_info() -> str:
    return "Some information"

# Pydantic model (serialized to JSON)
@agent.tool_plain
def get_user() -> User:
    return User(name="John", age=30)

# Dict (serialized to JSON)
@agent.tool_plain
def get_data() -> dict[str, Any]:
    return {"key": "value"}

# ToolReturn for custom content types
from pydantic_ai import ToolReturn, ImageUrl

@agent.tool_plain
def get_image() -> ToolReturn:
    return ToolReturn(content=[ImageUrl(url="https://...")])

访问上下文

RunContext 提供：

@agent.tool
async def my_tool(ctx: RunContext[MyDeps]) -> str:
    # Dependencies
    db = ctx.deps.db
    api = ctx.deps.api_client

    # Model info
    model_name = ctx.model.model_name

    # Usage tracking
    tokens_used = ctx.usage.total_tokens

    # Retry info
    attempt = ctx.retry  # Current retry attempt (0-based)
    max_retries = ctx.max_retries

    # Message history
    messages = ctx.messages

    return "result"

工具准备函数

根据每个请求动态修改工具：

from pydantic_ai.tools import ToolDefinition

async def prepare_tools(
    ctx: RunContext[MyDeps],
    tool_defs: list[ToolDefinition]
) -> list[ToolDefinition]:
    """Filter or modify tools based on context."""
    if ctx.deps.user_role != 'admin':
        # Hide admin tools from non-admins
        return [t for t in tool_defs if not t.name.startswith('admin_')]
    return tool_defs

agent = Agent('openai:gpt-4o', prepare_tools=prepare_tools)

工具集

对工具进行分组和组合：

from pydantic_ai import FunctionToolset, CombinedToolset

# Create a toolset
db_tools = FunctionToolset()

@db_tools.tool
def query_users(name: str) -> list[dict]:
    """Query users by name."""
    ...

@db_tools.tool
def update_user(id: int, data: dict) -> bool:
    """Update user data."""
    ...

# Use in agent
agent = Agent('openai:gpt-4o', toolsets=[db_tools])

# Combine toolsets
all_tools = CombinedToolset([db_tools, api_tools])

常见错误

错误：在 tool_plain 中使用上下文

@agent.tool_plain
async def bad_tool(ctx: RunContext[MyDeps]) -> str:  # ERROR!
    ...

错误：tool 中缺少上下文

@agent.tool
def bad_tool(user_id: int) -> str:  # ERROR!
    ...

错误：上下文不是第一个参数

@agent.tool
def bad_tool(user_id: int, ctx: RunContext[MyDeps]) -> str:  # ERROR!
    ...

异步与同步

两者都可用，但 I/O 操作优先使用异步：

# Async (preferred for I/O operations)
@agent.tool
async def fetch_data(ctx: RunContext[Deps]) -> str:
    return await ctx.deps.client.get('/data')

# Sync (fine for CPU-bound operations)
@agent.tool_plain
def compute(x: int, y: int) -> int:
    return x * y