Make DeferredToolCalls work with streaming

DouweM · DouweM · commit 0360e7791308 · 2025-06-30T15:42:28.000Z
diff --git a/pydantic_ai_slim/pydantic_ai/_agent_graph.py b/pydantic_ai_slim/pydantic_ai/_agent_graph.py
@@ -21,7 +21,7 @@
 from pydantic_graph.nodes import End, NodeRunEndT
 
 from . import _output, _system_prompt, exceptions, messages as _messages, models, result, usage as _usage
-from .output import DeferredToolCalls, OutputDataT, OutputSpec
+from .output import OutputDataT, OutputSpec
 from .settings import ModelSettings, merge_model_settings
 from .tools import RunContext, ToolDefinition, ToolKind
 
@@ -310,6 +310,7 @@ async def stream(
                 ctx.deps.output_validators,
                 build_run_context(ctx),
                 ctx.deps.usage_limits,
+                ctx.deps.toolset,
             )
             yield agent_stream
             # In case the user didn't manually consume the full stream, ensure it is fully consumed here,
@@ -497,6 +498,13 @@ async def _handle_tool_calls(
         if final_result_holder:
             final_result = final_result_holder[0]
             self._next_node = self._handle_final_result(ctx, final_result, parts)
+        elif deferred_tool_calls := ctx.deps.toolset.get_deferred_tool_calls(tool_calls):
+            if not ctx.deps.output_schema.deferred_tool_calls:
+                raise exceptions.UserError(
+                    'There are deferred tool calls but DeferredToolCalls is not among output types.'
+                )
+            final_result = result.FinalResult(cast(NodeRunEndT, deferred_tool_calls), None, None)
+            self._next_node = self._handle_final_result(ctx, final_result, parts)
         else:
             instructions = await ctx.deps.get_instructions(run_context)
             self._next_node = ModelRequestNode[DepsT, NodeRunEndT](
@@ -578,16 +586,11 @@ async def process_function_tools(  # noqa: C901
     """
     run_context = build_run_context(ctx)
 
-    unknown_calls: list[_messages.ToolCallPart] = []
-    tool_calls_by_kind: dict[ToolKind, list[_messages.ToolCallPart]] = defaultdict(list)
-    # TODO: Make Toolset.tool_defs a dict
-    tool_defs_by_name: dict[str, ToolDefinition] = {tool_def.name: tool_def for tool_def in toolset.tool_defs}
+    tool_calls_by_kind: dict[ToolKind | Literal['unknown'], list[_messages.ToolCallPart]] = defaultdict(list)
     for call in tool_calls:
-        try:
-            tool_def = tool_defs_by_name[call.tool_name]
-            tool_calls_by_kind[tool_def.kind].append(call)
-        except KeyError:
-            unknown_calls.append(call)
+        tool_def = toolset.get_tool_def(call.tool_name)
+        kind = tool_def.kind if tool_def else 'unknown'
+        tool_calls_by_kind[kind].append(call)
 
     # first, look for the output tool call
     for call in tool_calls_by_kind['output']:
@@ -642,9 +645,9 @@ async def process_function_tools(  # noqa: C901
     else:
         calls_to_run.extend(tool_calls_by_kind['function'])
 
-    if unknown_calls:
+    if tool_calls_by_kind['unknown']:
         ctx.state.increment_retries(ctx.deps.max_result_retries)
-        calls_to_run.extend(unknown_calls)
+        calls_to_run.extend(tool_calls_by_kind['unknown'])
 
     for call in calls_to_run:
         yield _messages.FunctionToolCallEvent(call)
@@ -718,7 +721,6 @@ def process_content(content: Any) -> Any:
         for k in sorted(results_by_index):
             parts.append(results_by_index[k])
 
-    deferred_calls: list[_messages.ToolCallPart] = []
     for call in tool_calls_by_kind['deferred']:
         if final_result:
             parts.append(
@@ -730,18 +732,6 @@ def process_content(content: Any) -> Any:
             )
         else:
             yield _messages.FunctionToolCallEvent(call)
-            deferred_calls.append(call)
-
-    if deferred_calls:
-        if not ctx.deps.output_schema.deferred_tool_calls:
-            raise exceptions.UserError('There are pending tool calls but DeferredToolCalls is not among output types.')
-
-        deferred_tool_names = [call.tool_name for call in deferred_calls]
-        deferred_tool_defs = {
-            tool_def.name: tool_def for tool_def in toolset.tool_defs if tool_def.name in deferred_tool_names
-        }
-        output_data = cast(NodeRunEndT, DeferredToolCalls(deferred_calls, deferred_tool_defs))
-        final_result = result.FinalResult(output_data)
 
     parts.extend(user_parts)
 
diff --git a/pydantic_ai_slim/pydantic_ai/_output.py b/pydantic_ai_slim/pydantic_ai/_output.py
@@ -3,7 +3,7 @@
 import inspect
 import json
 from abc import ABC, abstractmethod
-from collections.abc import Awaitable, Iterable, Iterator, Sequence
+from collections.abc import Awaitable, Iterable, Sequence
 from dataclasses import dataclass, field, replace
 from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Union, cast, overload
 
@@ -549,16 +549,6 @@ def find_named_tool(
                 if part.tool_name == tool_name:
                     return part, self.tools[tool_name]
 
-    def find_tool(
-        self,
-        parts: Iterable[_messages.ModelResponsePart],
-    ) -> Iterator[tuple[_messages.ToolCallPart, OutputTool[OutputDataT]]]:
-        """Find a tool that matches one of the calls."""
-        for part in parts:
-            if isinstance(part, _messages.ToolCallPart):  # pragma: no branch
-                if result := self.tools.get(part.tool_name):
-                    yield part, result
-
 
 @dataclass(init=False)
 class ToolOrTextOutputSchema(ToolOutputSchema[OutputDataT], PlainTextOutputSchema[OutputDataT]):
diff --git a/pydantic_ai_slim/pydantic_ai/agent.py b/pydantic_ai_slim/pydantic_ai/agent.py
@@ -1042,12 +1042,13 @@ async def stream_to_final(
                                         output_schema, _output.TextOutputSchema
                                     ):
                                         return FinalResult(s, None, None)
-                                    elif isinstance(new_part, _messages.ToolCallPart) and isinstance(
-                                        output_schema, _output.ToolOutputSchema
-                                    ):  # pragma: no branch
-                                        for call, _ in output_schema.find_tool([new_part]):
-                                            return FinalResult(s, call.tool_name, call.tool_call_id)
-                                    # TODO: Handle DeferredToolCalls
+                                    elif isinstance(new_part, _messages.ToolCallPart) and (
+                                        tool_def := graph_ctx.deps.toolset.get_tool_def(new_part.tool_name)
+                                    ):
+                                        if tool_def.kind == 'output':
+                                            return FinalResult(s, new_part.tool_name, new_part.tool_call_id)
+                                        elif tool_def.kind == 'deferred':
+                                            return FinalResult(s, None, None)
                             return None
 
                         final_result = await stream_to_final(streamed_response)
@@ -1072,16 +1073,20 @@ async def on_complete() -> None:
 
                                 # TODO: Should we move on to the CallToolsNode here, instead of doing this ourselves?
                                 parts: list[_messages.ModelRequestPart] = []
+                                # final_result_holder: list[result.FinalResult[models.StreamedResponse]] = []
                                 async for _event in _agent_graph.process_function_tools(
                                     graph_ctx.deps.toolset,
                                     tool_calls,
                                     final_result,
                                     graph_ctx,
                                     parts,
+                                    # final_result_holder,
                                 ):
                                     pass
                                 if parts:
                                     messages.append(_messages.ModelRequest(parts))
+                                # if final_result_holder:
+                                #     final_result = final_result_holder[0]
 
                             yield StreamedRunResult(
                                 messages,
@@ -1093,6 +1098,7 @@ async def on_complete() -> None:
                                 graph_ctx.deps.output_validators,
                                 final_result.tool_name,
                                 on_complete,
+                                graph_ctx.deps.toolset,
                             )
                             break
                         # TODO: There may be deferred tool calls, process those.
diff --git a/pydantic_ai_slim/pydantic_ai/result.py b/pydantic_ai_slim/pydantic_ai/result.py
@@ -5,11 +5,13 @@
 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.toolset import RunToolset
+
 from . import _utils, exceptions, messages as _messages, models
 from ._output import (
     OutputDataT_inv,
@@ -47,6 +49,7 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
     _output_validators: list[OutputValidator[AgentDepsT, OutputDataT]]
     _run_ctx: RunContext[AgentDepsT]
     _usage_limits: UsageLimits | None
+    _toolset: RunToolset[AgentDepsT]
 
     _agent_stream_iterator: AsyncIterator[AgentStreamEvent] | None = field(default=None, init=False)
     _final_result_event: FinalResultEvent | None = field(default=None, init=False)
@@ -102,6 +105,12 @@ async def _validate_response(
             result_data = await output_tool.process(
                 call, self._run_ctx, allow_partial=allow_partial, wrap_validation_errors=False
             )
+        elif deferred_tool_calls := self._toolset.get_deferred_tool_calls(message.parts):
+            if not self._output_schema.deferred_tool_calls:
+                raise exceptions.UserError(
+                    'There are deferred tool calls but DeferredToolCalls is not among output types.'
+                )
+            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))
 
@@ -113,8 +122,6 @@ async def _validate_response(
                 'Invalid response, unable to process text output'
             )
 
-        # TODO: Possibly return DeferredToolCalls here?
-
         for validator in self._output_validators:
             result_data = await validator.validate(result_data, call, self._run_ctx)
         return result_data
@@ -136,13 +143,19 @@ def _get_final_result_event(e: _messages.ModelResponseStreamEvent) -> _messages.
                 """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._toolset.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
@@ -177,6 +190,7 @@ class StreamedRunResult(Generic[AgentDepsT, OutputDataT]):
     _output_validators: list[OutputValidator[AgentDepsT, OutputDataT]]
     _output_tool_name: str | None
     _on_complete: Callable[[], Awaitable[None]]
+    _toolset: RunToolset[AgentDepsT]
 
     _initial_run_ctx_usage: Usage = field(init=False)
     is_complete: bool = field(default=False, init=False)
@@ -382,7 +396,6 @@ async def get_output(self) -> OutputDataT:
             pass
         message = self._stream_response.get()
         await self._marked_completed(message)
-        # TODO: Possibly return DeferredToolCalls here?
         return await self.validate_structured_output(message)
 
     @deprecated('`get_data` is deprecated, use `get_output` instead.')
@@ -423,6 +436,12 @@ async def validate_structured_output(
             result_data = await output_tool.process(
                 call, self._run_ctx, allow_partial=allow_partial, wrap_validation_errors=False
             )
+        elif deferred_tool_calls := self._toolset.get_deferred_tool_calls(message.parts):
+            if not self._output_schema.deferred_tool_calls:
+                raise exceptions.UserError(
+                    'There are deferred tool calls but DeferredToolCalls is not among output types.'
+                )
+            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))
 
diff --git a/pydantic_ai_slim/pydantic_ai/toolset.py b/pydantic_ai_slim/pydantic_ai/toolset.py
@@ -2,7 +2,7 @@
 
 import asyncio
 from abc import ABC, abstractmethod
-from collections.abc import Awaitable, Iterator, Sequence
+from collections.abc import Awaitable, Iterable, Iterator, Sequence
 from contextlib import AsyncExitStack, contextmanager
 from dataclasses import dataclass, field, replace
 from functools import partial
@@ -14,6 +14,8 @@
 from pydantic_core import SchemaValidator
 from typing_extensions import Self
 
+from pydantic_ai.output import DeferredToolCalls
+
 from . import messages as _messages
 from ._output import BaseOutputSchema, OutputValidator, ToolRetryError
 from ._run_context import AgentDepsT, RunContext
@@ -70,6 +72,9 @@ def tool_defs(self) -> list[ToolDefinition]:
     def tool_names(self) -> list[str]:
         return [tool_def.name for tool_def in self.tool_defs]
 
+    def get_tool_def(self, name: str) -> ToolDefinition | None:
+        return next((tool_def for tool_def in self.tool_defs if tool_def.name == name), None)
+
     @abstractmethod
     def _get_tool_args_validator(self, ctx: RunContext[AgentDepsT], name: str) -> SchemaValidator:
         raise NotImplementedError()
@@ -676,6 +681,25 @@ async def call_tool(
                 self._retries.pop(name, None)
                 return output
 
+    def get_deferred_tool_calls(self, parts: Iterable[_messages.ModelResponsePart]) -> DeferredToolCalls | None:
+        deferred_calls_and_defs = [
+            (part, tool_def)
+            for part in parts
+            if isinstance(part, _messages.ToolCallPart)
+            and (tool_def := self.get_tool_def(part.tool_name))
+            and tool_def.kind == 'deferred'
+        ]
+        if not deferred_calls_and_defs:
+            return None
+
+        deferred_calls: list[_messages.ToolCallPart] = []
+        deferred_tool_defs: dict[str, ToolDefinition] = {}
+        for part, tool_def in deferred_calls_and_defs:
+            deferred_calls.append(part)
+            deferred_tool_defs[part.tool_name] = tool_def
+
+        return DeferredToolCalls(deferred_calls, deferred_tool_defs)
+
     @contextmanager
     def _with_retry(self, name: str, ctx: RunContext[AgentDepsT]) -> Iterator[RunContext[AgentDepsT]]:
         try:
diff --git a/tests/test_streaming.py b/tests/test_streaming.py
