Progress

Author: DouweM

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -13,6 +13,7 @@
 from typing import TYPE_CHECKING, Any, Generic, Literal, TypeGuard, cast
 
 import anyio
+from anyio.streams.memory import MemoryObjectSendStream
 from opentelemetry.trace import Tracer
 from typing_extensions import TypeVar, assert_never
 
@@ -599,7 +600,7 @@ async def _run():  # noqa: C901
                                         if text:
                                             try:
                                                 self._next_node = await self._handle_text_response(
-                                                    ctx, text, text_processor
+                                                    ctx, text, text_processor, send_stream
                                                 )
                                                 return
                                             except ToolRetryError:
@@ -668,8 +669,9 @@ async def _run():  # noqa: C901
 
                             if text_processor := output_schema.text_processor:
                                 if text:
-                                    # TODO (DouweM): This could call an output function that yields custom events, but we're not in an event stream here?
-                                    self._next_node = await self._handle_text_response(ctx, text, text_processor)
+                                    self._next_node = await self._handle_text_response(
+                                        ctx, text, text_processor, send_stream
+                                    )
                                     return
                                 alternatives.insert(0, 'return text')
 
@@ -737,8 +739,10 @@ async def _handle_text_response(
         ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
         text: str,
         text_processor: _output.BaseOutputProcessor[NodeRunEndT],
+        event_stream: MemoryObjectSendStream[_messages.CustomEvent],
     ) -> ModelRequestNode[DepsT, NodeRunEndT] | End[result.FinalResult[NodeRunEndT]]:
         run_context = build_run_context(ctx)
+        run_context = replace(run_context, event_stream=event_stream)
 
         result_data = await text_processor.process(text, run_context)
 

pydantic_ai_slim/pydantic_ai/_function_schema.py

@@ -68,9 +68,9 @@ async def call(self, args_dict: dict[str, Any], ctx: RunContext[Any]) -> Any:
                     return_value = cast(Return, event_data)
                     continue
 
-                # If there's no event stream, we're being called inside `agent.run_stream`, after event streaming has completed
-                # and final result streaming has begun, so there's no need to yield custom events.
-                # TODO (DouweM): Should we store events on the ToolReturnPart/RetryPromptPart or something, like we allow `metadata`?
+                # If there's no event stream, we're being called from inside `agent.run_stream()` or `AgentStream.get_output()`,
+                # after event streaming has completed and final result streaming has begun, so there's nowhere to yield custom events to.
+                # We could consider storing the yielded events somewhere and letting them be accessed after the fact as a list.
                 if ctx.event_stream is not None:
                     if isinstance(event_data, CustomEvent):
                         event = cast(CustomEvent, event_data)

pydantic_ai_slim/pydantic_ai/_tool_manager.py

@@ -50,10 +50,6 @@ async def for_run_step(self, ctx: RunContext[AgentDepsT]) -> ToolManager[AgentDe
         """Build a new tool manager for the next run step, carrying over the retries from the current run step."""
         if self.ctx is not None:
             if ctx.run_step == self.ctx.run_step:
-                # TODO (DouweM): Refactor to make sure it's always set
-
-                if ctx.event_stream and not self.ctx.event_stream:
-                    self.ctx.event_stream = ctx.event_stream
                 return self
 
             retries = {

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_toolset.py

@@ -11,6 +11,7 @@
 
 from pydantic_ai import AbstractToolset, FunctionToolset, WrapperToolset
 from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry
+from pydantic_ai.messages import Return, ToolReturn
 from pydantic_ai.tools import AgentDepsT, ToolDefinition
 
 from ._run_context import TemporalRunContext
@@ -43,7 +44,7 @@ class _ModelRetry:
 
 @dataclass
 class _ToolReturn:
-    result: Any
+    result: ToolReturn[Any] | Any
     kind: Literal['tool_return'] = 'tool_return'
 
 
@@ -74,6 +75,9 @@ def visit_and_replace(
     async def _wrap_call_tool_result(self, coro: Awaitable[Any]) -> CallToolResult:
         try:
             result = await coro
+            if type(result) is Return:
+                # We don't use `isinstance` because `ToolReturn` is a subclass of `Return` with additional fields, which should be returned in full.
+                result = result.return_value
             return _ToolReturn(result=result)
         except ApprovalRequired:
             return _ApprovalRequired()

pydantic_ai_slim/pydantic_ai/messages.py

@@ -617,16 +617,13 @@ def __init__(
 MultiModalContent = ImageUrl | AudioUrl | DocumentUrl | VideoUrl | BinaryContent
 UserContent: TypeAlias = str | MultiModalContent
 
-# TODO (DouweM): thinking about variance
 ReturnValueT = TypeVar('ReturnValueT', default=Any, covariant=True)
 
 
 @dataclass(repr=False)
 class Return(Generic[ReturnValueT]):
     """TODO (DouweM): Docstring."""
 
-    # TODO (DouweM): Find a better name, or get rid of this entirely?
-
     return_value: ReturnValueT
 
     __repr__ = _utils.dataclasses_no_defaults_repr
@@ -1793,11 +1790,14 @@ class BuiltinToolResultEvent:
 class CustomEvent(Generic[EventDataT]):
     """An event indicating the result of a function tool call."""
 
-    data: EventDataT
+    data: EventDataT = None
     """The data of the custom event."""
 
     _: KW_ONLY
 
+    name: str | None = None
+    """The name of the custom event."""
+
     tool_call_id: str | None = None
     """The tool call ID, if any, that this event is associated with."""
 

pydantic_ai_slim/pydantic_ai/output.py

@@ -69,7 +69,12 @@
 
 TextOutputFunc = TypeAliasType(
     'TextOutputFunc',
-    Callable[[RunContext, str], Awaitable[T_co] | T_co] | Callable[[str], Awaitable[T_co] | T_co],
+    Callable[[RunContext, str], AsyncIterator[Return[T_co] | Any]]
+    | Callable[[RunContext, str], Awaitable[T_co]]
+    | Callable[[RunContext, str], T_co]
+    | Callable[[str], AsyncIterator[Return[T_co] | Any]]
+    | Callable[[str], Awaitable[T_co]]
+    | Callable[[str], T_co],
     type_params=(T_co,),
 )
 """Definition of a function that will be called to process the model's plain text output. The function must take a single string argument.

pydantic_ai_slim/pydantic_ai/result.py

@@ -116,7 +116,6 @@ async def stream_text(self, *, delta: bool = False, debounce_by: float | None =
                 yield text
         else:
             async for text in self._stream_response_text(delta=False, debounce_by=debounce_by):
-                # TODO (DouweM): What if there's an output function?
                 for validator in self._output_validators:
                     text = await validator.validate(text, replace(self._run_ctx, partial_output=True))
                 yield text
@@ -171,7 +170,6 @@ async def validate_response_output(
                 raise exceptions.UnexpectedModelBehavior(  # pragma: no cover
                     f'Invalid response, unable to find tool call for {output_tool_name!r}'
                 )
-            # TODO (DouweM): This could call an output function that yields custom events, but we're not in an event stream here?
             return await self._tool_manager.handle_call(
                 tool_call, allow_partial=allow_partial, wrap_validation_errors=False
             )
@@ -193,7 +191,6 @@ async def validate_response_output(
                     # not part of the final result output, so we reset the accumulated text
                     text = ''
 
-            # TODO (DouweM): This could call an output function that yields custom events, but we're not in an event stream here?
             result_data = await text_processor.process(
                 text, self._run_ctx, allow_partial=allow_partial, wrap_validation_errors=False
             )

pydantic_ai_slim/pydantic_ai/toolsets/function.py

@@ -339,7 +339,7 @@ async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[
                 max_retries=max_retries,
                 args_validator=tool.function_schema.validator,
                 call_func=tool.function_schema.call,
-                is_async=tool.function_schema.is_async,
+                is_async=tool.function_schema.is_async or tool.function_schema.is_async_iterator,
             )
         return tools
 

pydantic_ai_slim/pydantic_ai/ui/ag_ui/_event_stream.py

@@ -32,6 +32,7 @@
 try:
     from ag_ui.core import (
         BaseEvent,
+        CustomEvent as AGUICustomEvent,
         EventType,
         RunAgentInput,
         RunErrorEvent,
@@ -239,3 +240,8 @@ async def handle_function_tool_result(self, event: FunctionToolResultEvent) -> A
     async def handle_custom_event(self, event: CustomEvent) -> AsyncIterator[BaseEvent]:
         if isinstance(event.data, BaseEvent):
             yield event.data
+        elif event.name:
+            data = event.data
+            if event.tool_call_id:
+                data = {'tool_call_id': event.tool_call_id, 'data': data}
+            yield AGUICustomEvent(name=event.name, value=data)

pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py

@@ -28,6 +28,7 @@
 from .request_types import RequestData
 from .response_types import (
     BaseChunk,
+    DataChunk,
     DoneChunk,
     ErrorChunk,
     FileChunk,
@@ -190,3 +191,8 @@ async def handle_function_tool_result(self, event: FunctionToolResultEvent) -> A
     async def handle_custom_event(self, event: CustomEvent) -> AsyncIterator[BaseChunk]:
         if isinstance(event.data, BaseChunk):
             yield event.data
+        elif event.name:
+            data = event.data
+            if event.tool_call_id:
+                data = {'tool_call_id': event.tool_call_id, 'data': data}
+            yield DataChunk(type=f'data-{event.name}', data=data)