pydantic_ai.models.function

由本地函数控制的模型。

FunctionModel 类似于 TestModel，但允许对模型的行为进行更精细的控制。

其主要用例是进行比 TestModel 更高级的单元测试。

这是一个最小的示例

function_model_usage.py

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

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


async def model_function(
    messages: list[ModelMessage], info: AgentInfo
) -> ModelResponse:
    print(messages)
    """
    [
        ModelRequest(
            parts=[
                UserPromptPart(
                    content='Testing my agent...',
                    timestamp=datetime.datetime(...),
                    part_kind='user-prompt',
                )
            ],
            kind='request',
        )
    ]
    """
    print(info)
    """
    AgentInfo(
        function_tools=[], allow_text_result=True, result_tools=[], model_settings=None
    )
    """
    return ModelResponse(parts=[TextPart('hello world')])


async def test_my_agent():
    """Unit test for my_agent, to be run by pytest."""
    with my_agent.override(model=FunctionModel(model_function)):
        result = await my_agent.run('Testing my agent...')
        assert result.data == 'hello world'

请参阅 使用 FunctionModel 进行单元测试 以获取详细文档。

FunctionModel dataclass

基类: Model

由本地函数控制的模型。

除了 __init__ 之外，所有方法都是私有的或与基类的方法匹配。

源代码位于 pydantic_ai_slim/pydantic_ai/models/function.py

@dataclass(init=False)
class FunctionModel(Model):
    """A model controlled by a local function.

    Apart from `__init__`, all methods are private or match those of the base class.
    """

    function: FunctionDef | None = None
    stream_function: StreamFunctionDef | None = None

    _model_name: str = field(repr=False)
    _system: str = field(default='function', repr=False)

    @overload
    def __init__(self, function: FunctionDef, *, model_name: str | None = None) -> None: ...

    @overload
    def __init__(self, *, stream_function: StreamFunctionDef, model_name: str | None = None) -> None: ...

    @overload
    def __init__(
        self, function: FunctionDef, *, stream_function: StreamFunctionDef, model_name: str | None = None
    ) -> None: ...

    def __init__(
        self,
        function: FunctionDef | None = None,
        *,
        stream_function: StreamFunctionDef | None = None,
        model_name: str | None = None,
    ):
        """Initialize a `FunctionModel`.

        Either `function` or `stream_function` must be provided, providing both is allowed.

        Args:
            function: The function to call for non-streamed requests.
            stream_function: The function to call for streamed requests.
            model_name: The name of the model. If not provided, a name is generated from the function names.
        """
        if function is None and stream_function is None:
            raise TypeError('Either `function` or `stream_function` must be provided')
        self.function = function
        self.stream_function = stream_function

        function_name = self.function.__name__ if self.function is not None else ''
        stream_function_name = self.stream_function.__name__ if self.stream_function is not None else ''
        self._model_name = model_name or f'function:{function_name}:{stream_function_name}'

    async def request(
        self,
        messages: list[ModelMessage],
        model_settings: ModelSettings | None,
        model_request_parameters: ModelRequestParameters,
    ) -> tuple[ModelResponse, usage.Usage]:
        agent_info = AgentInfo(
            model_request_parameters.function_tools,
            model_request_parameters.allow_text_result,
            model_request_parameters.result_tools,
            model_settings,
        )

        assert self.function is not None, 'FunctionModel must receive a `function` to support non-streamed requests'

        if inspect.iscoroutinefunction(self.function):
            response = await self.function(messages, agent_info)
        else:
            response_ = await _utils.run_in_executor(self.function, messages, agent_info)
            assert isinstance(response_, ModelResponse), response_
            response = response_
        response.model_name = self._model_name
        # TODO is `messages` right here? Should it just be new messages?
        return response, _estimate_usage(chain(messages, [response]))

    @asynccontextmanager
    async def request_stream(
        self,
        messages: list[ModelMessage],
        model_settings: ModelSettings | None,
        model_request_parameters: ModelRequestParameters,
    ) -> AsyncIterator[StreamedResponse]:
        agent_info = AgentInfo(
            model_request_parameters.function_tools,
            model_request_parameters.allow_text_result,
            model_request_parameters.result_tools,
            model_settings,
        )

        assert self.stream_function is not None, (
            'FunctionModel must receive a `stream_function` to support streamed requests'
        )

        response_stream = PeekableAsyncStream(self.stream_function(messages, agent_info))

        first = await response_stream.peek()
        if isinstance(first, _utils.Unset):
            raise ValueError('Stream function must return at least one item')

        yield FunctionStreamedResponse(_model_name=self._model_name, _iter=response_stream)

    @property
    def model_name(self) -> str:
        """The model name."""
        return self._model_name

    @property
    def system(self) -> str:
        """The system / model provider."""
        return self._system

__init__

__init__(
    function: FunctionDef, *, model_name: str | None = None
) -> None

__init__(
    *,
    stream_function: StreamFunctionDef,
    model_name: str | None = None
) -> None

__init__(
    function: FunctionDef,
    *,
    stream_function: StreamFunctionDef,
    model_name: str | None = None
) -> None

__init__(
    function: FunctionDef | None = None,
    *,
    stream_function: StreamFunctionDef | None = None,
    model_name: str | None = None
)

初始化 FunctionModel。

必须提供 function 或 stream_function 之一，允许同时提供两者。

参数

名称	类型	描述	默认值
function	FunctionDef | None	用于非流式请求的函数调用。	None
stream_function	StreamFunctionDef | None	用于流式请求的函数调用。	None
model_name	str | None	模型的名称。如果未提供，则会从函数名称生成名称。	None

源代码位于 pydantic_ai_slim/pydantic_ai/models/function.py

def __init__(
    self,
    function: FunctionDef | None = None,
    *,
    stream_function: StreamFunctionDef | None = None,
    model_name: str | None = None,
):
    """Initialize a `FunctionModel`.

    Either `function` or `stream_function` must be provided, providing both is allowed.

    Args:
        function: The function to call for non-streamed requests.
        stream_function: The function to call for streamed requests.
        model_name: The name of the model. If not provided, a name is generated from the function names.
    """
    if function is None and stream_function is None:
        raise TypeError('Either `function` or `stream_function` must be provided')
    self.function = function
    self.stream_function = stream_function

    function_name = self.function.__name__ if self.function is not None else ''
    stream_function_name = self.stream_function.__name__ if self.stream_function is not None else ''
    self._model_name = model_name or f'function:{function_name}:{stream_function_name}'

model_name property

model_name: str

模型名称。

system property

system: str

系统 / 模型提供商。

AgentInfo dataclass

关于代理的信息。

这作为第二个参数传递给 FunctionModel 中使用的函数。

源代码位于 pydantic_ai_slim/pydantic_ai/models/function.py

@dataclass(frozen=True)
class AgentInfo:
    """Information about an agent.

    This is passed as the second to functions used within [`FunctionModel`][pydantic_ai.models.function.FunctionModel].
    """

    function_tools: list[ToolDefinition]
    """The function tools available on this agent.

    These are the tools registered via the [`tool`][pydantic_ai.Agent.tool] and
    [`tool_plain`][pydantic_ai.Agent.tool_plain] decorators.
    """
    allow_text_result: bool
    """Whether a plain text result is allowed."""
    result_tools: list[ToolDefinition]
    """The tools that can called as the final result of the run."""
    model_settings: ModelSettings | None
    """The model settings passed to the run call."""

function_tools instance-attribute

function_tools: list[ToolDefinition]

此代理上可用的函数工具。

这些工具是通过 tool 和 tool_plain 装饰器注册的。

allow_text_result instance-attribute

allow_text_result: bool

是否允许纯文本结果。

result_tools instance-attribute

result_tools: list[ToolDefinition]

可以作为运行的最终结果调用的工具。

model_settings instance-attribute

model_settings: ModelSettings | None

传递给运行调用的模型设置。

DeltaToolCall dataclass

工具调用的增量更改。

用于描述流式结构化响应时的块。

源代码位于 pydantic_ai_slim/pydantic_ai/models/function.py

@dataclass
class DeltaToolCall:
    """Incremental change to a tool call.

    Used to describe a chunk when streaming structured responses.
    """

    name: str | None = None
    """Incremental change to the name of the tool."""
    json_args: str | None = None
    """Incremental change to the arguments as JSON"""
    tool_call_id: str | None = None
    """Incremental change to the tool call ID."""

name class-attribute instance-attribute

name: str | None = None

工具名称的增量更改。

json_args class-attribute instance-attribute

json_args: str | None = None

参数作为 JSON 的增量更改

tool_call_id class-attribute instance-attribute

tool_call_id: str | None = None

工具调用 ID 的增量更改。

DeltaToolCalls module-attribute

DeltaToolCalls: TypeAlias = dict[int, DeltaToolCall]

工具调用 ID 到增量更改的映射。

FunctionDef module-attribute

FunctionDef: TypeAlias = Callable[
    [list[ModelMessage], AgentInfo],
    Union[ModelResponse, Awaitable[ModelResponse]],
]

用于生成非流式响应的函数。

StreamFunctionDef module-attribute

StreamFunctionDef: TypeAlias = Callable[
    [list[ModelMessage], AgentInfo],
    AsyncIterator[Union[str, DeltaToolCalls]],
]

用于生成流式响应的函数。

虽然这被定义为具有 AsyncIterator[Union[str, DeltaToolCalls]] 的返回类型，但实际上应该将其视为 Union[AsyncIterator[str], AsyncIterator[DeltaToolCalls]，

例如，你需要生成所有文本或所有 DeltaToolCalls，而不是将它们混合。

FunctionStreamedResponse dataclass

基类: StreamedResponse

FunctionModel 的 StreamedResponse 实现。

源代码位于 pydantic_ai_slim/pydantic_ai/models/function.py

@dataclass
class FunctionStreamedResponse(StreamedResponse):
    """Implementation of `StreamedResponse` for [FunctionModel][pydantic_ai.models.function.FunctionModel]."""

    _model_name: str
    _iter: AsyncIterator[str | DeltaToolCalls]
    _timestamp: datetime = field(default_factory=_utils.now_utc)

    def __post_init__(self):
        self._usage += _estimate_usage([])

    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
        async for item in self._iter:
            if isinstance(item, str):
                response_tokens = _estimate_string_tokens(item)
                self._usage += usage.Usage(response_tokens=response_tokens, total_tokens=response_tokens)
                yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=item)
            else:
                delta_tool_calls = item
                for dtc_index, delta_tool_call in delta_tool_calls.items():
                    if delta_tool_call.json_args:
                        response_tokens = _estimate_string_tokens(delta_tool_call.json_args)
                        self._usage += usage.Usage(response_tokens=response_tokens, total_tokens=response_tokens)
                    maybe_event = self._parts_manager.handle_tool_call_delta(
                        vendor_part_id=dtc_index,
                        tool_name=delta_tool_call.name,
                        args=delta_tool_call.json_args,
                        tool_call_id=delta_tool_call.tool_call_id,
                    )
                    if maybe_event is not None:
                        yield maybe_event

    @property
    def model_name(self) -> str:
        """Get the model name of the response."""
        return self._model_name

    @property
    def timestamp(self) -> datetime:
        """Get the timestamp of the response."""
        return self._timestamp

model_name property

model_name: str

获取响应的模型名称。

timestamp property

timestamp: datetime

获取响应的时间戳。