跳到内容

功能工具

功能工具提供了一种机制,使模型能够检索额外信息以帮助它们生成响应。

当将代理可能需要的所有上下文放入系统提示中不切实际或不可能时,或者当您希望通过将生成响应所需的部分逻辑推迟到另一个(不一定是 AI 驱动的)工具来使代理的行为更具确定性或可靠性时,它们非常有用。

功能工具 vs. RAG

功能工具基本上是 RAG(检索增强生成)的“R”——它们通过让模型请求额外信息来增强模型可以做的事情。

PydanticAI 工具和 RAG 之间的主要语义区别在于,RAG 是向量搜索的同义词,而 PydanticAI 工具用途更广泛。(注意:我们将来可能会添加对向量搜索功能的支持,特别是用于生成嵌入的 API。请参阅 #58

有多种方法可以向代理注册工具

@agent.tool 被认为是默认装饰器,因为在大多数情况下,工具将需要访问代理上下文。

这是一个同时使用两者的示例

dice_game.py
import random

from pydantic_ai import Agent, RunContext

agent = Agent(
    'google-gla:gemini-1.5-flash',  # (1)!
    deps_type=str,  # (2)!
    system_prompt=(
        "You're a dice game, you should roll the die and see if the number "
        "you get back matches the user's guess. If so, tell them they're a winner. "
        "Use the player's name in the response."
    ),
)


@agent.tool_plain  # (3)!
def roll_die() -> str:
    """Roll a six-sided die and return the result."""
    return str(random.randint(1, 6))


@agent.tool  # (4)!
def get_player_name(ctx: RunContext[str]) -> str:
    """Get the player's name."""
    return ctx.deps


dice_result = agent.run_sync('My guess is 4', deps='Anne')  # (5)!
print(dice_result.data)
#> Congratulations Anne, you guessed correctly! You're a winner!
  1. 这是一个非常简单的任务,因此我们可以使用快速且廉价的 Gemini flash 模型。
  2. 我们将用户的姓名作为依赖项传递,为了简单起见,我们仅使用姓名作为字符串作为依赖项。
  3. 此工具不需要任何上下文,它只返回一个随机数。在这种情况下,您可以使用动态系统提示。
  4. 此工具需要玩家姓名,因此它使用 RunContext 来访问依赖项,在本例中依赖项只是玩家姓名。
  5. 运行代理,将玩家姓名作为依赖项传递。

(此示例是完整的,可以“按原样”运行)

让我们打印来自该游戏的消息,看看发生了什么

dice_game_messages.py
from dice_game import dice_result

print(dice_result.all_messages())
"""
[
    ModelRequest(
        parts=[
            SystemPromptPart(
                content="You're a dice game, you should roll the die and see if the number you get back matches the user's guess. If so, tell them they're a winner. Use the player's name in the response.",
                timestamp=datetime.datetime(...),
                dynamic_ref=None,
                part_kind='system-prompt',
            ),
            UserPromptPart(
                content='My guess is 4',
                timestamp=datetime.datetime(...),
                part_kind='user-prompt',
            ),
        ],
        kind='request',
    ),
    ModelResponse(
        parts=[
            ToolCallPart(
                tool_name='roll_die', args={}, tool_call_id=None, part_kind='tool-call'
            )
        ],
        model_name='gemini-1.5-flash',
        timestamp=datetime.datetime(...),
        kind='response',
    ),
    ModelRequest(
        parts=[
            ToolReturnPart(
                tool_name='roll_die',
                content='4',
                tool_call_id=None,
                timestamp=datetime.datetime(...),
                part_kind='tool-return',
            )
        ],
        kind='request',
    ),
    ModelResponse(
        parts=[
            ToolCallPart(
                tool_name='get_player_name',
                args={},
                tool_call_id=None,
                part_kind='tool-call',
            )
        ],
        model_name='gemini-1.5-flash',
        timestamp=datetime.datetime(...),
        kind='response',
    ),
    ModelRequest(
        parts=[
            ToolReturnPart(
                tool_name='get_player_name',
                content='Anne',
                tool_call_id=None,
                timestamp=datetime.datetime(...),
                part_kind='tool-return',
            )
        ],
        kind='request',
    ),
    ModelResponse(
        parts=[
            TextPart(
                content="Congratulations Anne, you guessed correctly! You're a winner!",
                part_kind='text',
            )
        ],
        model_name='gemini-1.5-flash',
        timestamp=datetime.datetime(...),
        kind='response',
    ),
]
"""

我们可以用图表表示这一点

sequenceDiagram
    participant Agent
    participant LLM

    Note over Agent: Send prompts
    Agent ->> LLM: System: "You're a dice game..."<br>User: "My guess is 4"
    activate LLM
    Note over LLM: LLM decides to use<br>a tool

    LLM ->> Agent: Call tool<br>roll_die()
    deactivate LLM
    activate Agent
    Note over Agent: Rolls a six-sided die

    Agent -->> LLM: ToolReturn<br>"4"
    deactivate Agent
    activate LLM
    Note over LLM: LLM decides to use<br>another tool

    LLM ->> Agent: Call tool<br>get_player_name()
    deactivate LLM
    activate Agent
    Note over Agent: Retrieves player name
    Agent -->> LLM: ToolReturn<br>"Anne"
    deactivate Agent
    activate LLM
    Note over LLM: LLM constructs final response

    LLM ->> Agent: ModelResponse<br>"Congratulations Anne, ..."
    deactivate LLM
    Note over Agent: Game session complete

通过 kwarg 注册功能工具

除了使用装饰器之外,我们还可以通过 Agent 构造函数tools 参数注册工具。当您想要重用工具时,这非常有用,并且还可以对工具进行更细粒度的控制。

dice_game_tool_kwarg.py
import random

from pydantic_ai import Agent, RunContext, Tool


def roll_die() -> str:
    """Roll a six-sided die and return the result."""
    return str(random.randint(1, 6))


def get_player_name(ctx: RunContext[str]) -> str:
    """Get the player's name."""
    return ctx.deps


agent_a = Agent(
    'google-gla:gemini-1.5-flash',
    deps_type=str,
    tools=[roll_die, get_player_name],  # (1)!
)
agent_b = Agent(
    'google-gla:gemini-1.5-flash',
    deps_type=str,
    tools=[  # (2)!
        Tool(roll_die, takes_ctx=False),
        Tool(get_player_name, takes_ctx=True),
    ],
)
dice_result = agent_b.run_sync('My guess is 4', deps='Anne')
print(dice_result.data)
#> Congratulations Anne, you guessed correctly! You're a winner!
  1. 通过 Agent 构造函数注册工具的最简单方法是传递函数列表,检查函数签名以确定工具是否接受 RunContext
  2. agent_aagent_b 是相同的 — 但我们可以使用 Tool 来重用工具定义,并对工具的定义方式进行更细粒度的控制,例如设置其名称或描述,或使用自定义的 prepare 方法。

(此示例是完整的,可以“按原样”运行)

功能工具 vs. 结构化结果

顾名思义,功能工具使用模型的“工具”或“函数”API,让模型知道有哪些可调用的工具。工具或函数也用于定义结构化响应的模式,因此模型可能可以访问许多工具,其中一些工具调用功能工具,而另一些工具结束运行并返回结果。

功能工具和模式

函数参数从函数签名中提取,除了 RunContext 之外的所有参数都用于构建该工具调用的模式。

更好的是,PydanticAI 从函数中提取文档字符串,并(感谢 griffe)从文档字符串中提取参数描述,并将它们添加到模式中。

Griffe 支持googlenumpysphinx 样式文档字符串中提取参数描述。PydanticAI 将推断要使用的格式基于文档字符串,但您可以使用 docstring_format 显式设置它。您还可以通过设置 require_parameter_descriptions=True 来强制执行参数要求。如果缺少参数描述,这将引发 UserError

为了演示工具的模式,这里我们使用 FunctionModel 来打印模型将接收的模式

tool_schema.py
from pydantic_ai import Agent
from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart
from pydantic_ai.models.function import AgentInfo, FunctionModel

agent = Agent()


@agent.tool_plain(docstring_format='google', require_parameter_descriptions=True)
def foobar(a: int, b: str, c: dict[str, list[float]]) -> str:
    """Get me foobar.

    Args:
        a: apple pie
        b: banana cake
        c: carrot smoothie
    """
    return f'{a} {b} {c}'


def print_schema(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
    tool = info.function_tools[0]
    print(tool.description)
    #> Get me foobar.
    print(tool.parameters_json_schema)
    """
    {
        'additionalProperties': False,
        'properties': {
            'a': {'description': 'apple pie', 'type': 'integer'},
            'b': {'description': 'banana cake', 'type': 'string'},
            'c': {
                'additionalProperties': {'items': {'type': 'number'}, 'type': 'array'},
                'description': 'carrot smoothie',
                'type': 'object',
            },
        },
        'required': ['a', 'b', 'c'],
        'type': 'object',
    }
    """
    return ModelResponse(parts=[TextPart('foobar')])


agent.run_sync('hello', model=FunctionModel(print_schema))

(此示例是完整的,可以“按原样”运行)

工具的返回类型可以是 Pydantic 可以序列化为 JSON 的任何内容,因为某些模型(例如 Gemini)支持半结构化返回值,有些模型(OpenAI)期望文本,但似乎同样擅长从数据中提取含义。如果返回 Python 对象并且模型期望字符串,则该值将被序列化为 JSON。

如果一个工具具有可以表示为 JSON 模式中的对象的单个参数(例如,dataclass、TypedDict、pydantic 模型),则该工具的模式被简化为仅该对象。

这是一个示例,我们在其中使用 TestModel.last_model_request_parameters 来检查将传递给模型的工具模式。

single_parameter_tool.py
from pydantic import BaseModel

from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel

agent = Agent()


class Foobar(BaseModel):
    """This is a Foobar"""

    x: int
    y: str
    z: float = 3.14


@agent.tool_plain
def foobar(f: Foobar) -> str:
    return str(f)


test_model = TestModel()
result = agent.run_sync('hello', model=test_model)
print(result.data)
#> {"foobar":"x=0 y='a' z=3.14"}
print(test_model.last_model_request_parameters.function_tools)
"""
[
    ToolDefinition(
        name='foobar',
        description='This is a Foobar',
        parameters_json_schema={
            'properties': {
                'x': {'type': 'integer'},
                'y': {'type': 'string'},
                'z': {'default': 3.14, 'type': 'number'},
            },
            'required': ['x', 'y'],
            'title': 'Foobar',
            'type': 'object',
        },
        outer_typed_dict_key=None,
    )
]
"""

(此示例是完整的,可以“按原样”运行)

动态功能工具

工具可以选择使用另一个函数定义:prepare,它在运行的每个步骤中调用,以自定义传递给模型的工具的定义,或者完全从该步骤中省略该工具。

可以通过 prepare kwarg 向任何工具注册机制注册 prepare 方法

prepare 方法的类型应为 ToolPrepareFunc,一个函数,它接受 RunContext 和一个预构建的 ToolDefinition,并且应该返回该 ToolDefinition(无论是否修改),返回新的 ToolDefinition,或返回 None 以指示不应为该步骤注册此工具。

这是一个简单的 prepare 方法,仅当依赖项的值为 42 时才包含该工具。

与前面的示例一样,我们使用 TestModel 来演示行为,而无需调用真实模型。

tool_only_if_42.py
from typing import Union

from pydantic_ai import Agent, RunContext
from pydantic_ai.tools import ToolDefinition

agent = Agent('test')


async def only_if_42(
    ctx: RunContext[int], tool_def: ToolDefinition
) -> Union[ToolDefinition, None]:
    if ctx.deps == 42:
        return tool_def


@agent.tool(prepare=only_if_42)
def hitchhiker(ctx: RunContext[int], answer: str) -> str:
    return f'{ctx.deps} {answer}'


result = agent.run_sync('testing...', deps=41)
print(result.data)
#> success (no tool calls)
result = agent.run_sync('testing...', deps=42)
print(result.data)
#> {"hitchhiker":"42 a"}

(此示例是完整的,可以“按原样”运行)

这是一个更复杂的示例,我们在其中根据 deps 的值更改 name 参数的描述

为了变化,我们使用 Tool 数据类创建此工具。

customize_name.py
from __future__ import annotations

from typing import Literal

from pydantic_ai import Agent, RunContext
from pydantic_ai.models.test import TestModel
from pydantic_ai.tools import Tool, ToolDefinition


def greet(name: str) -> str:
    return f'hello {name}'


async def prepare_greet(
    ctx: RunContext[Literal['human', 'machine']], tool_def: ToolDefinition
) -> ToolDefinition | None:
    d = f'Name of the {ctx.deps} to greet.'
    tool_def.parameters_json_schema['properties']['name']['description'] = d
    return tool_def


greet_tool = Tool(greet, prepare=prepare_greet)
test_model = TestModel()
agent = Agent(test_model, tools=[greet_tool], deps_type=Literal['human', 'machine'])

result = agent.run_sync('testing...', deps='human')
print(result.data)
#> {"greet":"hello a"}
print(test_model.last_model_request_parameters.function_tools)
"""
[
    ToolDefinition(
        name='greet',
        description='',
        parameters_json_schema={
            'additionalProperties': False,
            'properties': {
                'name': {'type': 'string', 'description': 'Name of the human to greet.'}
            },
            'required': ['name'],
            'type': 'object',
        },
        outer_typed_dict_key=None,
    )
]
"""

(此示例是完整的,可以“按原样”运行)