pydantic-ai-slim 0.6.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

pydantic_ai/models/__init__.py

@@ -13,19 +13,32 @@ from contextlib import asynccontextmanager, contextmanager
 from dataclasses import dataclass, field, replace
 from datetime import datetime
 from functools import cache, cached_property
-from typing import Generic, TypeVar, overload
+from typing import Any, Generic, TypeVar, overload
 
 import httpx
 from typing_extensions import Literal, TypeAliasType, TypedDict
 
-from pydantic_ai.profiles import DEFAULT_PROFILE, ModelProfile, ModelProfileSpec
-
 from .. import _utils
 from .._output import OutputObjectDefinition
 from .._parts_manager import ModelResponsePartsManager
+from .._run_context import RunContext
+from ..builtin_tools import AbstractBuiltinTool
 from ..exceptions import UserError
-from ..messages import FileUrl, ModelMessage, ModelRequest, ModelResponse, ModelResponseStreamEvent, VideoUrl
+from ..messages import (
+    AgentStreamEvent,
+    FileUrl,
+    FinalResultEvent,
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    ModelResponseStreamEvent,
+    PartStartEvent,
+    TextPart,
+    ToolCallPart,
+    VideoUrl,
+)
 from ..output import OutputMode
+from ..profiles import DEFAULT_PROFILE, ModelProfile, ModelProfileSpec
 from ..profiles._json_schema import JsonSchemaTransformer
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
@@ -336,12 +349,17 @@ class ModelRequestParameters:
     """Configuration for an agent's request to a model, specifically related to tools and output handling."""
 
     function_tools: list[ToolDefinition] = field(default_factory=list)
+    builtin_tools: list[AbstractBuiltinTool] = field(default_factory=list)
 
     output_mode: OutputMode = 'text'
     output_object: OutputObjectDefinition | None = None
     output_tools: list[ToolDefinition] = field(default_factory=list)
     allow_text_output: bool = True
 
+    @cached_property
+    def tool_defs(self) -> dict[str, ToolDefinition]:
+        return {tool_def.name: tool_def for tool_def in [*self.function_tools, *self.output_tools]}
+
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
@@ -387,6 +405,7 @@ class Model(ABC):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         """Make a request to the model and return a streaming response."""
         # This method is not required, but you need to implement it if you want to support streamed responses
@@ -499,14 +518,40 @@ class Model(ABC):
 class StreamedResponse(ABC):
     """Streamed response from an LLM when calling a tool."""
 
+    model_request_parameters: ModelRequestParameters
+    final_result_event: FinalResultEvent | None = field(default=None, init=False)
+
     _parts_manager: ModelResponsePartsManager = field(default_factory=ModelResponsePartsManager, init=False)
-    _event_iterator: AsyncIterator[ModelResponseStreamEvent] | None = field(default=None, init=False)
+    _event_iterator: AsyncIterator[AgentStreamEvent] | None = field(default=None, init=False)
     _usage: Usage = field(default_factory=Usage, init=False)
 
-    def __aiter__(self) -> AsyncIterator[ModelResponseStreamEvent]:
-        """Stream the response as an async iterable of [`ModelResponseStreamEvent`][pydantic_ai.messages.ModelResponseStreamEvent]s."""
+    def __aiter__(self) -> AsyncIterator[AgentStreamEvent]:
+        """Stream the response as an async iterable of [`AgentStreamEvent`][pydantic_ai.messages.AgentStreamEvent]s.
+
+        This proxies the `_event_iterator()` and emits all events, while also checking for matches
+        on the result schema and emitting a [`FinalResultEvent`][pydantic_ai.messages.FinalResultEvent] if/when the
+        first match is found.
+        """
         if self._event_iterator is None:
-            self._event_iterator = self._get_event_iterator()
+
+            async def iterator_with_final_event(
+                iterator: AsyncIterator[ModelResponseStreamEvent],
+            ) -> AsyncIterator[AgentStreamEvent]:
+                async for event in iterator:
+                    yield event
+                    if (
+                        final_result_event := _get_final_result_event(event, self.model_request_parameters)
+                    ) is not None:
+                        self.final_result_event = final_result_event
+                        yield final_result_event
+                        break
+
+                # If we broke out of the above loop, we need to yield the rest of the events
+                # If we didn't, this will just be a no-op
+                async for event in iterator:
+                    yield event
+
+            self._event_iterator = iterator_with_final_event(self._get_event_iterator())
         return self._event_iterator
 
     @abstractmethod
@@ -808,3 +853,16 @@ def _customize_output_object(transformer: type[JsonSchemaTransformer], o: Output
         json_schema=json_schema,
         strict=schema_transformer.is_strict_compatible if o.strict is None else o.strict,
     )
+
+
+def _get_final_result_event(e: ModelResponseStreamEvent, params: ModelRequestParameters) -> FinalResultEvent | None:
+    """Return an appropriate FinalResultEvent if `e` corresponds to a part that will produce a final result."""
+    if isinstance(e, PartStartEvent):
+        new_part = e.part
+        if isinstance(new_part, TextPart) and params.allow_text_output:  # pragma: no branch
+            return FinalResultEvent(tool_name=None, tool_call_id=None)
+        elif isinstance(new_part, ToolCallPart) and (tool_def := params.tool_defs.get(new_part.tool_name)):
+            if tool_def.kind == 'output':
+                return FinalResultEvent(tool_name=new_part.tool_name, tool_call_id=new_part.tool_call_id)
+            elif tool_def.kind == 'deferred':
+                return FinalResultEvent(tool_name=None, tool_call_id=None)

pydantic_ai/models/anthropic.py

@@ -8,12 +8,26 @@ from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from typing import Any, Literal, Union, cast, overload
 
+from anthropic.types.beta import (
+    BetaCitationsDelta,
+    BetaCodeExecutionToolResultBlock,
+    BetaCodeExecutionToolResultBlockParam,
+    BetaInputJSONDelta,
+    BetaServerToolUseBlockParam,
+    BetaWebSearchToolResultBlockParam,
+)
 from typing_extensions import assert_never
 
+from pydantic_ai.builtin_tools import CodeExecutionTool, WebSearchTool
+
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
+from .._run_context import RunContext
 from .._utils import guard_tool_call_id as _guard_tool_call_id
+from ..exceptions import UserError
 from ..messages import (
     BinaryContent,
+    BuiltinToolCallPart,
+    BuiltinToolReturnPart,
     DocumentUrl,
     ImageUrl,
     ModelMessage,
@@ -33,13 +47,21 @@ from ..profiles import ModelProfileSpec
 from ..providers import Provider, infer_provider
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
-from . import Model, ModelRequestParameters, StreamedResponse, check_allow_model_requests, download_item, get_user_agent
+from . import (
+    Model,
+    ModelRequestParameters,
+    StreamedResponse,
+    check_allow_model_requests,
+    download_item,
+    get_user_agent,
+)
 
 try:
     from anthropic import NOT_GIVEN, APIStatusError, AsyncAnthropic, AsyncStream
     from anthropic.types.beta import (
         BetaBase64PDFBlockParam,
         BetaBase64PDFSourceParam,
+        BetaCodeExecutionTool20250522Param,
         BetaContentBlock,
         BetaContentBlockParam,
         BetaImageBlockParam,
@@ -55,6 +77,7 @@ try:
         BetaRawMessageStopEvent,
         BetaRawMessageStreamEvent,
         BetaRedactedThinkingBlock,
+        BetaServerToolUseBlock,
         BetaSignatureDelta,
         BetaTextBlock,
         BetaTextBlockParam,
@@ -66,9 +89,13 @@ try:
         BetaToolChoiceParam,
         BetaToolParam,
         BetaToolResultBlockParam,
+        BetaToolUnionParam,
         BetaToolUseBlock,
         BetaToolUseBlockParam,
+        BetaWebSearchTool20250305Param,
+        BetaWebSearchToolResultBlock,
     )
+    from anthropic.types.beta.beta_web_search_tool_20250305_param import UserLocation
     from anthropic.types.model_param import ModelParam
 
 except ImportError as _import_error:
@@ -171,13 +198,14 @@ class AnthropicModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         check_allow_model_requests()
         response = await self._messages_create(
             messages, True, cast(AnthropicModelSettings, model_settings or {}), model_request_parameters
         )
         async with response:
-            yield await self._process_streamed_response(response)
+            yield await self._process_streamed_response(response, model_request_parameters)
 
     @property
     def model_name(self) -> AnthropicModelName:
@@ -218,6 +246,7 @@ class AnthropicModel(Model):
     ) -> BetaMessage | AsyncStream[BetaRawMessageStreamEvent]:
         # standalone function to make it easier to override
         tools = self._get_tools(model_request_parameters)
+        tools += self._get_builtin_tools(model_request_parameters)
         tool_choice: BetaToolChoiceParam | None
 
         if not tools:
@@ -236,6 +265,7 @@ class AnthropicModel(Model):
         try:
             extra_headers = model_settings.get('extra_headers', {})
             extra_headers.setdefault('User-Agent', get_user_agent())
+            extra_headers.setdefault('anthropic-beta', 'code-execution-2025-05-22')
             return await self.client.beta.messages.create(
                 max_tokens=model_settings.get('max_tokens', 4096),
                 system=system_prompt or NOT_GIVEN,
@@ -264,6 +294,24 @@ class AnthropicModel(Model):
         for item in response.content:
             if isinstance(item, BetaTextBlock):
                 items.append(TextPart(content=item.text))
+            elif isinstance(item, (BetaWebSearchToolResultBlock, BetaCodeExecutionToolResultBlock)):
+                items.append(
+                    BuiltinToolReturnPart(
+                        provider_name='anthropic',
+                        tool_name=item.type,
+                        content=item.content,
+                        tool_call_id=item.tool_use_id,
+                    )
+                )
+            elif isinstance(item, BetaServerToolUseBlock):
+                items.append(
+                    BuiltinToolCallPart(
+                        provider_name='anthropic',
+                        tool_name=item.name,
+                        args=cast(dict[str, Any], item.input),
+                        tool_call_id=item.id,
+                    )
+                )
             elif isinstance(item, BetaRedactedThinkingBlock):  # pragma: no cover
                 warnings.warn(
                     'Pydantic AI currently does not handle redacted thinking blocks. '
@@ -284,7 +332,9 @@ class AnthropicModel(Model):
 
         return ModelResponse(items, usage=_map_usage(response), model_name=response.model, vendor_id=response.id)
 
-    async def _process_streamed_response(self, response: AsyncStream[BetaRawMessageStreamEvent]) -> StreamedResponse:
+    async def _process_streamed_response(
+        self, response: AsyncStream[BetaRawMessageStreamEvent], model_request_parameters: ModelRequestParameters
+    ) -> StreamedResponse:
         peekable_response = _utils.PeekableAsyncStream(response)
         first_chunk = await peekable_response.peek()
         if isinstance(first_chunk, _utils.Unset):
@@ -293,13 +343,35 @@ class AnthropicModel(Model):
         # Since Anthropic doesn't provide a timestamp in the message, we'll use the current time
         timestamp = datetime.now(tz=timezone.utc)
         return AnthropicStreamedResponse(
-            _model_name=self._model_name, _response=peekable_response, _timestamp=timestamp
+            model_request_parameters=model_request_parameters,
+            _model_name=self._model_name,
+            _response=peekable_response,
+            _timestamp=timestamp,
         )
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[BetaToolParam]:
-        tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools]
-        if model_request_parameters.output_tools:
-            tools += [self._map_tool_definition(r) for r in model_request_parameters.output_tools]
+        return [self._map_tool_definition(r) for r in model_request_parameters.tool_defs.values()]
+
+    def _get_builtin_tools(self, model_request_parameters: ModelRequestParameters) -> list[BetaToolUnionParam]:
+        tools: list[BetaToolUnionParam] = []
+        for tool in model_request_parameters.builtin_tools:
+            if isinstance(tool, WebSearchTool):
+                user_location = UserLocation(type='approximate', **tool.user_location) if tool.user_location else None
+                tools.append(
+                    BetaWebSearchTool20250305Param(
+                        name='web_search',
+                        type='web_search_20250305',
+                        allowed_domains=tool.allowed_domains,
+                        blocked_domains=tool.blocked_domains,
+                        user_location=user_location,
+                    )
+                )
+            elif isinstance(tool, CodeExecutionTool):  # pragma: no branch
+                tools.append(BetaCodeExecutionTool20250522Param(name='code_execution', type='code_execution_20250522'))
+            else:  # pragma: no cover
+                raise UserError(
+                    f'`{tool.__class__.__name__}` is not supported by `AnthropicModel`. If it should be, please file an issue.'
+                )
         return tools
 
     async def _map_message(self, messages: list[ModelMessage]) -> tuple[str, list[BetaMessageParam]]:  # noqa: C901
@@ -338,11 +410,26 @@ class AnthropicModel(Model):
                 if len(user_content_params) > 0:
                     anthropic_messages.append(BetaMessageParam(role='user', content=user_content_params))
             elif isinstance(m, ModelResponse):
-                assistant_content_params: list[BetaTextBlockParam | BetaToolUseBlockParam | BetaThinkingBlockParam] = []
+                assistant_content_params: list[
+                    BetaTextBlockParam
+                    | BetaToolUseBlockParam
+                    | BetaServerToolUseBlockParam
+                    | BetaWebSearchToolResultBlockParam
+                    | BetaCodeExecutionToolResultBlockParam
+                    | BetaThinkingBlockParam
+                ] = []
                 for response_part in m.parts:
                     if isinstance(response_part, TextPart):
-                        if response_part.content:  # Only add non-empty text
+                        if response_part.content:
                             assistant_content_params.append(BetaTextBlockParam(text=response_part.content, type='text'))
+                    elif isinstance(response_part, ToolCallPart):
+                        tool_use_block_param = BetaToolUseBlockParam(
+                            id=_guard_tool_call_id(t=response_part),
+                            type='tool_use',
+                            name=response_part.tool_name,
+                            input=response_part.args_as_dict(),
+                        )
+                        assistant_content_params.append(tool_use_block_param)
                     elif isinstance(response_part, ThinkingPart):
                         # NOTE: We only send thinking part back for Anthropic, otherwise they raise an error.
                         if response_part.signature is not None:  # pragma: no branch
@@ -351,14 +438,31 @@ class AnthropicModel(Model):
                                     thinking=response_part.content, signature=response_part.signature, type='thinking'
                                 )
                             )
+                    elif isinstance(response_part, BuiltinToolCallPart):
+                        if response_part.provider_name == 'anthropic':
+                            server_tool_use_block_param = BetaServerToolUseBlockParam(
+                                id=_guard_tool_call_id(t=response_part),
+                                type='server_tool_use',
+                                name=cast(Literal['web_search', 'code_execution'], response_part.tool_name),
+                                input=response_part.args_as_dict(),
+                            )
+                            assistant_content_params.append(server_tool_use_block_param)
+                    elif isinstance(response_part, BuiltinToolReturnPart):
+                        if response_part.provider_name == 'anthropic':
+                            tool_use_id = _guard_tool_call_id(t=response_part)
+                            if response_part.tool_name == 'web_search_tool_result':
+                                server_tool_result_block_param = BetaWebSearchToolResultBlockParam(
+                                    tool_use_id=tool_use_id, type=response_part.tool_name, content=response_part.content
+                                )
+                            elif response_part.tool_name == 'code_execution_tool_result':
+                                server_tool_result_block_param = BetaCodeExecutionToolResultBlockParam(
+                                    tool_use_id=tool_use_id, type=response_part.tool_name, content=response_part.content
+                                )
+                            else:
+                                raise ValueError(f'Unsupported tool name: {response_part.tool_name}')
+                            assistant_content_params.append(server_tool_result_block_param)
                     else:
-                        tool_use_block_param = BetaToolUseBlockParam(
-                            id=_guard_tool_call_id(t=response_part),
-                            type='tool_use',
-                            name=response_part.tool_name,
-                            input=response_part.args_as_dict(),
-                        )
-                        assistant_content_params.append(tool_use_block_param)
+                        assert_never(response_part)
                 if len(assistant_content_params) > 0:
                     anthropic_messages.append(BetaMessageParam(role='assistant', content=assistant_content_params))
             else:
@@ -476,7 +580,10 @@ class AnthropicStreamedResponse(StreamedResponse):
         async for event in self._response:
             self._usage += _map_usage(event)
 
-            if isinstance(event, BetaRawContentBlockStartEvent):
+            if isinstance(event, BetaRawMessageStartEvent):
+                pass
+
+            elif isinstance(event, BetaRawContentBlockStartEvent):
                 current_block = event.content_block
                 if isinstance(current_block, BetaTextBlock) and current_block.text:
                     maybe_event = self._parts_manager.handle_text_delta(
@@ -499,6 +606,8 @@ class AnthropicStreamedResponse(StreamedResponse):
                     )
                     if maybe_event is not None:  # pragma: no branch
                         yield maybe_event
+                elif isinstance(current_block, BetaServerToolUseBlock):
+                    pass
 
             elif isinstance(event, BetaRawContentBlockDeltaEvent):
                 if isinstance(event.delta, BetaTextDelta):
@@ -528,8 +637,16 @@ class AnthropicStreamedResponse(StreamedResponse):
                     )
                     if maybe_event is not None:  # pragma: no branch
                         yield maybe_event
+                elif isinstance(event.delta, BetaInputJSONDelta):
+                    pass
+                # TODO(Marcelo): We need to handle citations.
+                elif isinstance(event.delta, BetaCitationsDelta):
+                    pass
+
+            elif isinstance(event, BetaRawMessageDeltaEvent):
+                pass
 
-            elif isinstance(event, (BetaRawContentBlockStopEvent, BetaRawMessageStopEvent)):
+            elif isinstance(event, (BetaRawContentBlockStopEvent, BetaRawMessageStopEvent)):  # pragma: no branch
                 current_block = None
 
     @property

pydantic_ai/models/bedrock.py

@@ -15,9 +15,13 @@ import anyio.to_thread
 from typing_extensions import ParamSpec, assert_never
 
 from pydantic_ai import _utils, usage
+from pydantic_ai._run_context import RunContext
+from pydantic_ai.exceptions import UserError
 from pydantic_ai.messages import (
     AudioUrl,
     BinaryContent,
+    BuiltinToolCallPart,
+    BuiltinToolReturnPart,
     DocumentUrl,
     ImageUrl,
     ModelMessage,
@@ -227,10 +231,7 @@ class BedrockConverseModel(Model):
         super().__init__(settings=settings, profile=profile or provider.model_profile)
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[ToolTypeDef]:
-        tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools]
-        if model_request_parameters.output_tools:
-            tools += [self._map_tool_definition(r) for r in model_request_parameters.output_tools]
-        return tools
+        return [self._map_tool_definition(r) for r in model_request_parameters.tool_defs.values()]
 
     @staticmethod
     def _map_tool_definition(f: ToolDefinition) -> ToolTypeDef:
@@ -266,10 +267,15 @@ class BedrockConverseModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         settings = cast(BedrockModelSettings, model_settings or {})
         response = await self._messages_create(messages, True, settings, model_request_parameters)
-        yield BedrockStreamedResponse(_model_name=self.model_name, _event_stream=response)
+        yield BedrockStreamedResponse(
+            model_request_parameters=model_request_parameters,
+            _model_name=self.model_name,
+            _event_stream=response,
+        )
 
     async def _process_response(self, response: ConverseResponseTypeDef) -> ModelResponse:
         items: list[ModelResponsePart] = []
@@ -342,6 +348,9 @@ class BedrockConverseModel(Model):
         if tool_config:
             params['toolConfig'] = tool_config
 
+        if model_request_parameters.builtin_tools:
+            raise UserError('Bedrock does not support built-in tools')
+
         # Bedrock supports a set of specific extra parameters
         if model_settings:
             if guardrail_config := model_settings.get('bedrock_guardrail_config', None):
@@ -478,6 +487,8 @@ class BedrockConverseModel(Model):
                         else:
                             # NOTE: We don't pass the thinking part to Bedrock for models other than Claude since it raises an error.
                             pass
+                    elif isinstance(item, (BuiltinToolCallPart, BuiltinToolReturnPart)):
+                        pass
                     else:
                         assert isinstance(item, ToolCallPart)
                         content.append(self._map_tool_call(item))

pydantic_ai/models/cohere.py

@@ -7,10 +7,13 @@ from typing import Literal, Union, cast
 from typing_extensions import assert_never
 
 from pydantic_ai._thinking_part import split_content_into_text_and_thinking
+from pydantic_ai.exceptions import UserError
 
 from .. import ModelHTTPError, usage
 from .._utils import generate_tool_call_id as _generate_tool_call_id, guard_tool_call_id as _guard_tool_call_id
 from ..messages import (
+    BuiltinToolCallPart,
+    BuiltinToolReturnPart,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -166,6 +169,10 @@ class CohereModel(Model):
         model_request_parameters: ModelRequestParameters,
     ) -> V2ChatResponse:
         tools = self._get_tools(model_request_parameters)
+
+        if model_request_parameters.builtin_tools:
+            raise UserError('Cohere does not support built-in tools')
+
         cohere_messages = self._map_messages(messages)
         try:
             return await self.client.chat(
@@ -223,6 +230,9 @@ class CohereModel(Model):
                         pass
                     elif isinstance(item, ToolCallPart):
                         tool_calls.append(self._map_tool_call(item))
+                    elif isinstance(item, (BuiltinToolCallPart, BuiltinToolReturnPart)):  # pragma: no cover
+                        # This is currently never returned from cohere
+                        pass
                     else:
                         assert_never(item)
                 message_param = AssistantChatMessageV2(role='assistant')
@@ -238,10 +248,7 @@ class CohereModel(Model):
         return cohere_messages
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[ToolV2]:
-        tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools]
-        if model_request_parameters.output_tools:
-            tools += [self._map_tool_definition(r) for r in model_request_parameters.output_tools]
-        return tools
+        return [self._map_tool_definition(r) for r in model_request_parameters.tool_defs.values()]
 
     @staticmethod
     def _map_tool_call(t: ToolCallPart) -> ToolCallV2:

pydantic_ai/models/fallback.py

@@ -3,10 +3,11 @@ from __future__ import annotations as _annotations
 from collections.abc import AsyncIterator
 from contextlib import AsyncExitStack, asynccontextmanager, suppress
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING, Any, Callable
 
 from opentelemetry.trace import get_current_span
 
+from pydantic_ai._run_context import RunContext
 from pydantic_ai.models.instrumented import InstrumentedModel
 
 from ..exceptions import FallbackExceptionGroup, ModelHTTPError
@@ -83,6 +84,7 @@ class FallbackModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         """Try each model in sequence until one succeeds."""
         exceptions: list[Exception] = []
@@ -92,7 +94,7 @@ class FallbackModel(Model):
             async with AsyncExitStack() as stack:
                 try:
                     response = await stack.enter_async_context(
-                        model.request_stream(messages, model_settings, customized_model_request_parameters)
+                        model.request_stream(messages, model_settings, customized_model_request_parameters, run_context)
                     )
                 except Exception as exc:
                     if self._fallback_on(exc):

pydantic_ai/models/function.py

@@ -7,16 +7,17 @@ from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime
 from itertools import chain
-from typing import Callable, Union
+from typing import Any, Callable, Union
 
 from typing_extensions import TypeAlias, assert_never, overload
 
-from pydantic_ai.profiles import ModelProfileSpec
-
 from .. import _utils, usage
+from .._run_context import RunContext
 from .._utils import PeekableAsyncStream
 from ..messages import (
     BinaryContent,
+    BuiltinToolCallPart,
+    BuiltinToolReturnPart,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -30,6 +31,7 @@ from ..messages import (
     UserContent,
     UserPromptPart,
 )
+from ..profiles import ModelProfileSpec
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import Model, ModelRequestParameters, StreamedResponse
@@ -145,6 +147,7 @@ class FunctionModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         agent_info = AgentInfo(
             model_request_parameters.function_tools,
@@ -163,7 +166,11 @@ class FunctionModel(Model):
         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)
+        yield FunctionStreamedResponse(
+            model_request_parameters=model_request_parameters,
+            _model_name=self._model_name,
+            _iter=response_stream,
+        )
 
     @property
     def model_name(self) -> str:
@@ -331,6 +338,13 @@ def _estimate_usage(messages: Iterable[ModelMessage]) -> usage.Usage:
                     response_tokens += _estimate_string_tokens(part.content)
                 elif isinstance(part, ToolCallPart):
                     response_tokens += 1 + _estimate_string_tokens(part.args_as_json_str())
+                # TODO(Marcelo): We need to add coverage here.
+                elif isinstance(part, BuiltinToolCallPart):  # pragma: no cover
+                    call = part
+                    response_tokens += 1 + _estimate_string_tokens(call.args_as_json_str())
+                # TODO(Marcelo): We need to add coverage here.
+                elif isinstance(part, BuiltinToolReturnPart):  # pragma: no cover
+                    response_tokens += _estimate_string_tokens(part.model_response_str())
                 else:
                     assert_never(part)
         else: