pydantic-ai-slim 0.4.3__py3-none-any.whl → 0.4.5__py3-none-any.whl

pydantic_ai/models/anthropic.py

@@ -266,7 +266,7 @@ class AnthropicModel(Model):
                 items.append(TextPart(content=item.text))
             elif isinstance(item, BetaRedactedThinkingBlock):  # pragma: no cover
                 warnings.warn(
-                    'PydanticAI currently does not handle redacted thinking blocks. '
+                    'Pydantic AI currently does not handle redacted thinking blocks. '
                     'If you have a suggestion on how we should handle them, please open an issue.',
                     UserWarning,
                 )

pydantic_ai/models/function.py

@@ -214,21 +214,39 @@ class DeltaToolCall:
     """Incremental change to the tool call ID."""
 
 
+@dataclass
+class DeltaThinkingPart:
+    """Incremental change to a thinking part.
+
+    Used to describe a chunk when streaming thinking responses.
+    """
+
+    content: str | None = None
+    """Incremental change to the thinking content."""
+    signature: str | None = None
+    """Incremental change to the thinking signature."""
+
+
 DeltaToolCalls: TypeAlias = dict[int, DeltaToolCall]
 """A mapping of tool call IDs to incremental changes."""
 
+DeltaThinkingCalls: TypeAlias = dict[int, DeltaThinkingPart]
+"""A mapping of thinking call IDs to incremental changes."""
+
 # TODO: Change the signature to Callable[[list[ModelMessage], ModelSettings, ModelRequestParameters], ...]
 FunctionDef: TypeAlias = Callable[[list[ModelMessage], AgentInfo], Union[ModelResponse, Awaitable[ModelResponse]]]
 """A function used to generate a non-streamed response."""
 
 # TODO: Change signature as indicated above
-StreamFunctionDef: TypeAlias = Callable[[list[ModelMessage], AgentInfo], AsyncIterator[Union[str, DeltaToolCalls]]]
+StreamFunctionDef: TypeAlias = Callable[
+    [list[ModelMessage], AgentInfo], AsyncIterator[Union[str, DeltaToolCalls, DeltaThinkingCalls]]
+]
 """A function used to generate a streamed response.
 
-While this is defined as having return type of `AsyncIterator[Union[str, DeltaToolCalls]]`, it should
-really be considered as `Union[AsyncIterator[str], AsyncIterator[DeltaToolCalls]`,
+While this is defined as having return type of `AsyncIterator[Union[str, DeltaToolCalls, DeltaThinkingCalls]]`, it should
+really be considered as `Union[AsyncIterator[str], AsyncIterator[DeltaToolCalls], AsyncIterator[DeltaThinkingCalls]]`,
 
-E.g. you need to yield all text or all `DeltaToolCalls`, not mix them.
+E.g. you need to yield all text, all `DeltaToolCalls`, or all `DeltaThinkingCalls`, not mix them.
 """
 
 
@@ -237,7 +255,7 @@ class FunctionStreamedResponse(StreamedResponse):
     """Implementation of `StreamedResponse` for [FunctionModel][pydantic_ai.models.function.FunctionModel]."""
 
     _model_name: str
-    _iter: AsyncIterator[str | DeltaToolCalls]
+    _iter: AsyncIterator[str | DeltaToolCalls | DeltaThinkingCalls]
     _timestamp: datetime = field(default_factory=_utils.now_utc)
 
     def __post_init__(self):
@@ -249,20 +267,31 @@ class FunctionStreamedResponse(StreamedResponse):
                 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
+            elif isinstance(item, dict) and item:
+                for dtc_index, delta in item.items():
+                    if isinstance(delta, DeltaThinkingPart):
+                        if delta.content:  # pragma: no branch
+                            response_tokens = _estimate_string_tokens(delta.content)
+                            self._usage += usage.Usage(response_tokens=response_tokens, total_tokens=response_tokens)
+                        yield self._parts_manager.handle_thinking_delta(
+                            vendor_part_id=dtc_index,
+                            content=delta.content,
+                            signature=delta.signature,
+                        )
+                    elif isinstance(delta, DeltaToolCall):
+                        if delta.json_args:
+                            response_tokens = _estimate_string_tokens(delta.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.name,
+                            args=delta.json_args,
+                            tool_call_id=delta.tool_call_id,
+                        )
+                        if maybe_event is not None:
+                            yield maybe_event
+                    else:
+                        assert_never(delta)
 
     @property
     def model_name(self) -> str:
@@ -299,12 +328,9 @@ def _estimate_usage(messages: Iterable[ModelMessage]) -> usage.Usage:
                 if isinstance(part, TextPart):
                     response_tokens += _estimate_string_tokens(part.content)
                 elif isinstance(part, ThinkingPart):
-                    # NOTE: We don't send ThinkingPart to the providers yet.
-                    # If you are unsatisfied with this, please open an issue.
-                    pass
+                    response_tokens += _estimate_string_tokens(part.content)
                 elif isinstance(part, ToolCallPart):
-                    call = part
-                    response_tokens += 1 + _estimate_string_tokens(call.args_as_json_str())
+                    response_tokens += 1 + _estimate_string_tokens(part.args_as_json_str())
                 else:
                     assert_never(part)
         else:

pydantic_ai/models/gemini.py

@@ -48,18 +48,10 @@ from . import (
 )
 
 LatestGeminiModelNames = Literal[
-    'gemini-1.5-flash',
-    'gemini-1.5-flash-8b',
-    'gemini-1.5-pro',
-    'gemini-1.0-pro',
     'gemini-2.0-flash',
-    'gemini-2.0-flash-lite-preview-02-05',
-    'gemini-2.0-pro-exp-02-05',
-    'gemini-2.5-flash-preview-05-20',
+    'gemini-2.0-flash-lite',
     'gemini-2.5-flash',
     'gemini-2.5-flash-lite-preview-06-17',
-    'gemini-2.5-pro-exp-03-25',
-    'gemini-2.5-pro-preview-05-06',
     'gemini-2.5-pro',
 ]
 """Latest Gemini models."""
@@ -99,15 +91,6 @@ class GeminiModelSettings(ModelSettings, total=False):
     See the [Gemini API docs](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/add-labels-to-api-calls) for use cases and limitations.
     """
 
-    gemini_thinking_config: ThinkingConfig
-    """Thinking is on by default in both the API and AI Studio.
-
-    Being on by default doesn't mean the model will send back thoughts. For that, you need to set `include_thoughts`
-    to `True`. If you want to turn it off, set `thinking_budget` to `0`.
-
-    See more about it on <https://ai.google.dev/gemini-api/docs/thinking>.
-    """
-
 
 @dataclass(init=False)
 class GeminiModel(Model):

pydantic_ai/models/google.py

@@ -73,18 +73,10 @@ except ImportError as _import_error:
     ) from _import_error
 
 LatestGoogleModelNames = Literal[
-    'gemini-1.5-flash',
-    'gemini-1.5-flash-8b',
-    'gemini-1.5-pro',
-    'gemini-1.0-pro',
     'gemini-2.0-flash',
-    'gemini-2.0-flash-lite-preview-02-05',
-    'gemini-2.0-pro-exp-02-05',
-    'gemini-2.5-flash-preview-05-20',
+    'gemini-2.0-flash-lite',
     'gemini-2.5-flash',
     'gemini-2.5-flash-lite-preview-06-17',
-    'gemini-2.5-pro-exp-03-25',
-    'gemini-2.5-pro-preview-05-06',
     'gemini-2.5-pro',
 ]
 """Latest Gemini models."""
@@ -492,8 +484,7 @@ def _content_model_response(m: ModelResponse) -> ContentDict:
             function_call = FunctionCallDict(name=item.tool_name, args=item.args_as_dict(), id=item.tool_call_id)
             parts.append({'function_call': function_call})
         elif isinstance(item, TextPart):
-            if item.content:  # pragma: no branch
-                parts.append({'text': item.content})
+            parts.append({'text': item.content})
         elif isinstance(item, ThinkingPart):  # pragma: no cover
             # NOTE: We don't send ThinkingPart to the providers yet. If you are unsatisfied with this,
             # please open an issue. The below code is the code to send thinking to the provider.

pydantic_ai/models/groq.py

@@ -79,6 +79,7 @@ PreviewGroqModelNames = Literal[
     'llama-3.2-3b-preview',
     'llama-3.2-11b-vision-preview',
     'llama-3.2-90b-vision-preview',
+    'moonshotai/kimi-k2-instruct',
 ]
 """Preview Groq models from <https://console.groq.com/docs/models#preview-models>."""
 

pydantic_ai/models/instrumented.py

@@ -156,7 +156,12 @@ class InstrumentationSettings:
         events: list[Event] = []
         instructions = InstrumentedModel._get_instructions(messages)  # pyright: ignore [reportPrivateUsage]
         if instructions is not None:
-            events.append(Event('gen_ai.system.message', body={'content': instructions, 'role': 'system'}))
+            events.append(
+                Event(
+                    'gen_ai.system.message',
+                    body={**({'content': instructions} if self.include_content else {}), 'role': 'system'},
+                )
+            )
 
         for message_index, message in enumerate(messages):
             message_events: list[Event] = []

pydantic_ai/models/mistral.py

@@ -428,7 +428,7 @@ class MistralModel(Model):
         if value_type == 'object':
             additional_properties = value.get('additionalProperties', {})
             if isinstance(additional_properties, bool):
-                return 'bool'  # pragma: no cover
+                return 'bool'  # pragma: lax no cover
             additional_properties_type = additional_properties.get('type')
             if (
                 additional_properties_type in SIMPLE_JSON_TYPE_MAPPING

pydantic_ai/models/openai.py

@@ -8,6 +8,7 @@ from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Any, Literal, Union, cast, overload
 
+from pydantic import ValidationError
 from typing_extensions import assert_never
 
 from pydantic_ai._thinking_part import split_content_into_text_and_thinking
@@ -50,7 +51,7 @@ from . import (
 
 try:
     from openai import NOT_GIVEN, APIStatusError, AsyncOpenAI, AsyncStream, NotGiven
-    from openai.types import ChatModel, chat, responses
+    from openai.types import AllModels, chat, responses
     from openai.types.chat import (
         ChatCompletionChunk,
         ChatCompletionContentPartImageParam,
@@ -80,7 +81,7 @@ __all__ = (
     'OpenAIModelName',
 )
 
-OpenAIModelName = Union[str, ChatModel]
+OpenAIModelName = Union[str, AllModels]
 """
 Possible OpenAI model names.
 
@@ -347,8 +348,19 @@ class OpenAIModel(Model):
                 raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
             raise  # pragma: no cover
 
-    def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
+    def _process_response(self, response: chat.ChatCompletion | str) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
+        # Although the OpenAI SDK claims to return a Pydantic model (`ChatCompletion`) from the chat completions function:
+        # * it hasn't actually performed validation (presumably they're creating the model with `model_construct` or something?!)
+        # * if the endpoint returns plain text, the return type is a string
+        # Thus we validate it fully here.
+        if not isinstance(response, chat.ChatCompletion):
+            raise UnexpectedModelBehavior('Invalid response from OpenAI chat completions endpoint, expected JSON data')
+
+        try:
+            response = chat.ChatCompletion.model_validate(response.model_dump())
+        except ValidationError as e:
+            raise UnexpectedModelBehavior(f'Invalid response from OpenAI chat completions endpoint: {e}') from e
         timestamp = number_to_datetime(response.created)
         choice = response.choices[0]
         items: list[ModelResponsePart] = []
@@ -1051,7 +1063,7 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
                     vendor_part_id=chunk.item_id,
                     tool_name=None,
                     args=chunk.delta,
-                    tool_call_id=chunk.item_id,
+                    tool_call_id=None,
                 )
                 if maybe_event is not None:  # pragma: no branch
                     yield maybe_event

pydantic_ai/output.py

@@ -10,7 +10,8 @@ from pydantic_core import core_schema
 from typing_extensions import TypeAliasType, TypeVar
 
 from . import _utils
-from .tools import RunContext
+from .messages import ToolCallPart
+from .tools import RunContext, ToolDefinition
 
 __all__ = (
     # classes
@@ -330,15 +331,17 @@ def StructuredDict(
     return _StructuredDict
 
 
+_OutputSpecItem = TypeAliasType(
+    '_OutputSpecItem',
+    Union[OutputTypeOrFunction[T_co], ToolOutput[T_co], NativeOutput[T_co], PromptedOutput[T_co], TextOutput[T_co]],
+    type_params=(T_co,),
+)
+
 OutputSpec = TypeAliasType(
     'OutputSpec',
     Union[
-        OutputTypeOrFunction[T_co],
-        ToolOutput[T_co],
-        NativeOutput[T_co],
-        PromptedOutput[T_co],
-        TextOutput[T_co],
-        Sequence[Union[OutputTypeOrFunction[T_co], ToolOutput[T_co], TextOutput[T_co]]],
+        _OutputSpecItem[T_co],
+        Sequence['OutputSpec[T_co]'],
     ],
     type_params=(T_co,),
 )
@@ -354,3 +357,14 @@ You should not need to import or use this type directly.
 
 See [output docs](../output.md) for more information.
 """
+
+
+@dataclass
+class DeferredToolCalls:
+    """Container for calls of deferred tools. This can be used as an agent's `output_type` and will be used as the output of the agent run if the model called any deferred tools.
+
+    See [deferred toolset docs](../toolsets.md#deferred-toolset) for more information.
+    """
+
+    tool_calls: list[ToolCallPart]
+    tool_defs: dict[str, ToolDefinition]

pydantic_ai/profiles/google.py

@@ -43,7 +43,7 @@ class GoogleJsonSchemaTransformer(JsonSchemaTransformer):
                 f' Full schema: {self.schema}\n\n'
                 f'Source of additionalProperties within the full schema: {original_schema}\n\n'
                 'If this came from a field with a type like `dict[str, MyType]`, that field will always be empty.\n\n'
-                "If Google's APIs are updated to support this properly, please create an issue on the PydanticAI GitHub"
+                "If Google's APIs are updated to support this properly, please create an issue on the Pydantic AI GitHub"
                 ' and we will fix this behavior.',
                 UserWarning,
             )

pydantic_ai/profiles/moonshotai.py

@@ -0,0 +1,8 @@
+from __future__ import annotations as _annotations
+
+from . import ModelProfile
+
+
+def moonshotai_model_profile(model_name: str) -> ModelProfile | None:
+    """Get the model profile for a MoonshotAI model."""
+    return None

pydantic_ai/providers/grok.py

@@ -1,7 +1,7 @@
 from __future__ import annotations as _annotations
 
 import os
-from typing import overload
+from typing import Literal, overload
 
 from httpx import AsyncClient as AsyncHTTPClient
 from openai import AsyncOpenAI
@@ -21,6 +21,18 @@ except ImportError as _import_error:  # pragma: no cover
         'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
     ) from _import_error
 
+# https://docs.x.ai/docs/models
+GrokModelName = Literal[
+    'grok-4',
+    'grok-4-0709',
+    'grok-3',
+    'grok-3-mini',
+    'grok-3-fast',
+    'grok-3-mini-fast',
+    'grok-2-vision-1212',
+    'grok-2-image-1212',
+]
+
 
 class GrokProvider(Provider[AsyncOpenAI]):
     """Provider for Grok API."""

pydantic_ai/providers/groq.py

@@ -12,6 +12,7 @@ from pydantic_ai.profiles.deepseek import deepseek_model_profile
 from pydantic_ai.profiles.google import google_model_profile
 from pydantic_ai.profiles.meta import meta_model_profile
 from pydantic_ai.profiles.mistral import mistral_model_profile
+from pydantic_ai.profiles.moonshotai import moonshotai_model_profile
 from pydantic_ai.profiles.qwen import qwen_model_profile
 from pydantic_ai.providers import Provider
 
@@ -47,6 +48,7 @@ class GroqProvider(Provider[AsyncGroq]):
             'qwen': qwen_model_profile,
             'deepseek': deepseek_model_profile,
             'mistral': mistral_model_profile,
+            'moonshotai/': moonshotai_model_profile,
         }
 
         for prefix, profile_func in prefix_to_profile.items():

pydantic_ai/result.py

@@ -5,11 +5,13 @@ from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
 from copy import copy
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Generic
+from typing import Generic, cast
 
 from pydantic import ValidationError
 from typing_extensions import TypeVar, deprecated, overload
 
+from pydantic_ai._tool_manager import ToolManager
+
 from . import _utils, exceptions, messages as _messages, models
 from ._output import (
     OutputDataT_inv,
@@ -19,7 +21,6 @@ from ._output import (
     PlainTextOutputSchema,
     TextOutputSchema,
     ToolOutputSchema,
-    TraceContext,
 )
 from ._run_context import AgentDepsT, RunContext
 from .messages import AgentStreamEvent, FinalResultEvent
@@ -47,8 +48,8 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
     _output_schema: OutputSchema[OutputDataT]
     _output_validators: list[OutputValidator[AgentDepsT, OutputDataT]]
     _run_ctx: RunContext[AgentDepsT]
-    _trace_ctx: TraceContext
     _usage_limits: UsageLimits | None
+    _tool_manager: ToolManager[AgentDepsT]
 
     _agent_stream_iterator: AsyncIterator[AgentStreamEvent] | None = field(default=None, init=False)
     _final_result_event: FinalResultEvent | None = field(default=None, init=False)
@@ -97,37 +98,40 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
         self, message: _messages.ModelResponse, output_tool_name: str | None, *, allow_partial: bool = False
     ) -> OutputDataT:
         """Validate a structured result message."""
-        call = None
         if isinstance(self._output_schema, ToolOutputSchema) and output_tool_name is not None:
-            match = self._output_schema.find_named_tool(message.parts, output_tool_name)
-            if match is None:
+            tool_call = next(
+                (
+                    part
+                    for part in message.parts
+                    if isinstance(part, _messages.ToolCallPart) and part.tool_name == output_tool_name
+                ),
+                None,
+            )
+            if tool_call is None:
                 raise exceptions.UnexpectedModelBehavior(  # pragma: no cover
-                    f'Invalid response, unable to find tool: {self._output_schema.tool_names()}'
+                    f'Invalid response, unable to find tool call for {output_tool_name!r}'
                 )
-
-            call, output_tool = match
-            result_data = await output_tool.process(
-                call,
-                self._run_ctx,
-                self._trace_ctx,
-                allow_partial=allow_partial,
-                wrap_validation_errors=False,
-            )
+            return await self._tool_manager.handle_call(tool_call, allow_partial=allow_partial)
+        elif deferred_tool_calls := self._tool_manager.get_deferred_tool_calls(message.parts):
+            if not self._output_schema.allows_deferred_tool_calls:
+                raise exceptions.UserError(  # pragma: no cover
+                    'A deferred tool call was present, but `DeferredToolCalls` is not among output types. To resolve this, add `DeferredToolCalls` to the list of output types for this agent.'
+                )
+            return cast(OutputDataT, deferred_tool_calls)
         elif isinstance(self._output_schema, TextOutputSchema):
             text = '\n\n'.join(x.content for x in message.parts if isinstance(x, _messages.TextPart))
 
             result_data = await self._output_schema.process(
-                text, self._run_ctx, self._trace_ctx, allow_partial=allow_partial, wrap_validation_errors=False
+                text, self._run_ctx, allow_partial=allow_partial, wrap_validation_errors=False
             )
+            for validator in self._output_validators:
+                result_data = await validator.validate(result_data, self._run_ctx)
+            return result_data
         else:
             raise exceptions.UnexpectedModelBehavior(  # pragma: no cover
                 'Invalid response, unable to process text output'
             )
 
-        for validator in self._output_validators:
-            result_data = await validator.validate(result_data, call, self._run_ctx)
-        return result_data
-
     def __aiter__(self) -> AsyncIterator[AgentStreamEvent]:
         """Stream [`AgentStreamEvent`][pydantic_ai.messages.AgentStreamEvent]s.
 
@@ -145,13 +149,19 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
                 """Return an appropriate FinalResultEvent if `e` corresponds to a part that will produce a final result."""
                 if isinstance(e, _messages.PartStartEvent):
                     new_part = e.part
-                    if isinstance(new_part, _messages.ToolCallPart) and isinstance(output_schema, ToolOutputSchema):
-                        for call, _ in output_schema.find_tool([new_part]):  # pragma: no branch
-                            return _messages.FinalResultEvent(tool_name=call.tool_name, tool_call_id=call.tool_call_id)
-                    elif isinstance(new_part, _messages.TextPart) and isinstance(
+                    if isinstance(new_part, _messages.TextPart) and isinstance(
                         output_schema, TextOutputSchema
                     ):  # pragma: no branch
                         return _messages.FinalResultEvent(tool_name=None, tool_call_id=None)
+                    elif isinstance(new_part, _messages.ToolCallPart) and (
+                        tool_def := self._tool_manager.get_tool_def(new_part.tool_name)
+                    ):
+                        if tool_def.kind == 'output':
+                            return _messages.FinalResultEvent(
+                                tool_name=new_part.tool_name, tool_call_id=new_part.tool_call_id
+                            )
+                        elif tool_def.kind == 'deferred':
+                            return _messages.FinalResultEvent(tool_name=None, tool_call_id=None)
 
             usage_checking_stream = _get_usage_checking_stream_response(
                 self._raw_stream_response, self._usage_limits, self.usage
@@ -183,10 +193,10 @@ class StreamedRunResult(Generic[AgentDepsT, OutputDataT]):
     _stream_response: models.StreamedResponse
     _output_schema: OutputSchema[OutputDataT]
     _run_ctx: RunContext[AgentDepsT]
-    _trace_ctx: TraceContext
     _output_validators: list[OutputValidator[AgentDepsT, OutputDataT]]
     _output_tool_name: str | None
     _on_complete: Callable[[], Awaitable[None]]
+    _tool_manager: ToolManager[AgentDepsT]
 
     _initial_run_ctx_usage: Usage = field(init=False)
     is_complete: bool = field(default=False, init=False)
@@ -420,40 +430,43 @@ class StreamedRunResult(Generic[AgentDepsT, OutputDataT]):
         self, message: _messages.ModelResponse, *, allow_partial: bool = False
     ) -> OutputDataT:
         """Validate a structured result message."""
-        call = None
         if isinstance(self._output_schema, ToolOutputSchema) and self._output_tool_name is not None:
-            match = self._output_schema.find_named_tool(message.parts, self._output_tool_name)
-            if match is None:
+            tool_call = next(
+                (
+                    part
+                    for part in message.parts
+                    if isinstance(part, _messages.ToolCallPart) and part.tool_name == self._output_tool_name
+                ),
+                None,
+            )
+            if tool_call is None:
                 raise exceptions.UnexpectedModelBehavior(  # pragma: no cover
-                    f'Invalid response, unable to find tool: {self._output_schema.tool_names()}'
+                    f'Invalid response, unable to find tool call for {self._output_tool_name!r}'
                 )
-
-            call, output_tool = match
-            result_data = await output_tool.process(
-                call,
-                self._run_ctx,
-                self._trace_ctx,
-                allow_partial=allow_partial,
-                wrap_validation_errors=False,
-            )
+            return await self._tool_manager.handle_call(tool_call, allow_partial=allow_partial)
+        elif deferred_tool_calls := self._tool_manager.get_deferred_tool_calls(message.parts):
+            if not self._output_schema.allows_deferred_tool_calls:
+                raise exceptions.UserError(
+                    'A deferred tool call was present, but `DeferredToolCalls` is not among output types. To resolve this, add `DeferredToolCalls` to the list of output types for this agent.'
+                )
+            return cast(OutputDataT, deferred_tool_calls)
         elif isinstance(self._output_schema, TextOutputSchema):
             text = '\n\n'.join(x.content for x in message.parts if isinstance(x, _messages.TextPart))
 
             result_data = await self._output_schema.process(
-                text, self._run_ctx, self._trace_ctx, allow_partial=allow_partial, wrap_validation_errors=False
+                text, self._run_ctx, allow_partial=allow_partial, wrap_validation_errors=False
             )
+            for validator in self._output_validators:
+                result_data = await validator.validate(result_data, self._run_ctx)  # pragma: no cover
+            return result_data
         else:
             raise exceptions.UnexpectedModelBehavior(  # pragma: no cover
                 'Invalid response, unable to process text output'
             )
 
-        for validator in self._output_validators:
-            result_data = await validator.validate(result_data, call, self._run_ctx)  # pragma: no cover
-        return result_data
-
     async def _validate_text_output(self, text: str) -> str:
         for validator in self._output_validators:
-            text = await validator.validate(text, None, self._run_ctx)  # pragma: no cover
+            text = await validator.validate(text, self._run_ctx)  # pragma: no cover
         return text
 
     async def _marked_completed(self, message: _messages.ModelResponse) -> None: