Progress

Author: DouweM

pydantic_ai_slim/pydantic_ai/__init__.py

@@ -42,6 +42,7 @@
     BinaryImage,
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
+    CustomEvent,
     DocumentFormat,
     DocumentMediaType,
     DocumentUrl,
@@ -68,6 +69,7 @@
     PartEndEvent,
     PartStartEvent,
     RetryPromptPart,
+    Return,
     SystemPromptPart,
     TextPart,
     TextPartDelta,
@@ -141,6 +143,7 @@
     'BinaryContent',
     'BuiltinToolCallPart',
     'BuiltinToolReturnPart',
+    'CustomEvent',
     'DocumentFormat',
     'DocumentMediaType',
     'DocumentUrl',
@@ -168,6 +171,7 @@
     'PartEndEvent',
     'PartStartEvent',
     'RetryPromptPart',
+    'Return',
     'SystemPromptPart',
     'TextPart',
     'TextPartDelta',

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -438,7 +438,7 @@ async def stream(
                 _output_schema=ctx.deps.output_schema,
                 _model_request_parameters=model_request_parameters,
                 _output_validators=ctx.deps.output_validators,
-                _run_ctx=build_run_context(ctx),
+                _run_ctx=run_context,
                 _usage_limits=ctx.deps.usage_limits,
                 _tool_manager=ctx.deps.tool_manager,
             )
@@ -646,6 +646,7 @@ async def _run_stream() -> AsyncIterator[_messages.HandleResponseEvent]:  # noqa
 
                     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)
                             return
                         alternatives.insert(0, 'return text')
@@ -1072,7 +1073,8 @@ async def _call_tool(
     except ToolRetryError as e:
         return e.tool_retry, None
 
-    if isinstance(tool_result, _messages.ToolReturn):
+    tool_return: _messages.Return | None = None
+    if isinstance(tool_result, _messages.Return):
         tool_return = tool_result
     else:
         result_is_list = isinstance(tool_result, list)
@@ -1083,8 +1085,8 @@ async def _call_tool(
         for content in contents:
             if isinstance(content, _messages.ToolReturn):
                 raise exceptions.UserError(
-                    f'The return value of tool {tool_call.tool_name!r} contains invalid nested `ToolReturn` objects. '
-                    f'`ToolReturn` should be used directly.'
+                    f'The return value of tool {tool_call.tool_name!r} contains invalid nested `Return` objects. '
+                    f'`Return` should be used directly.'
                 )
             elif isinstance(content, _messages.MultiModalContent):
                 identifier = content.identifier
@@ -1116,10 +1118,13 @@ async def _call_tool(
         tool_name=tool_call.tool_name,
         tool_call_id=tool_call.tool_call_id,
         content=tool_return.return_value,  # type: ignore
-        metadata=tool_return.metadata,
     )
 
-    return return_part, tool_return.content or None
+    if isinstance(tool_return, _messages.ToolReturn):
+        return_part.metadata = tool_return.metadata
+        return return_part, tool_return.content or None
+    else:
+        return return_part, None
 
 
 @dataclasses.dataclass

pydantic_ai_slim/pydantic_ai/_function_schema.py

@@ -6,7 +6,7 @@
 from __future__ import annotations as _annotations
 
 from collections.abc import Awaitable, Callable
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from inspect import Parameter, signature
 from typing import TYPE_CHECKING, Any, Concatenate, cast, get_origin
 
@@ -19,7 +19,7 @@
 from pydantic_core import SchemaValidator, core_schema
 from typing_extensions import ParamSpec, TypeIs, TypeVar
 
-from pydantic_ai.messages import CustomEvent, ToolReturn
+from pydantic_ai.messages import CustomEvent, Return
 
 from ._griffe import doc_descriptions
 from ._run_context import RunContext
@@ -61,18 +61,26 @@ async def call(self, args_dict: dict[str, Any], ctx: RunContext[Any]) -> Any:
                 'RunContext.event_stream needs to be set to use FunctionSchema.call with async iterators'
             )
 
-            async for event_payload in self.function(*args, **kwargs):
-                if isinstance(event_payload, ToolReturn):
-                    return event_payload
+            return_value: Return | None = None
+            async for event_data in self.function(*args, **kwargs):
+                if return_value is not None:
+                    from .exceptions import UserError
 
-                event = (
-                    cast(CustomEvent, event_payload)
-                    if isinstance(event_payload, CustomEvent)
-                    else CustomEvent(payload=event_payload)
-                )
+                    raise UserError('Return value must be the last value yielded by the function')
+
+                if isinstance(event_data, Return):
+                    return_value = cast(Return, event_data)
+                    continue
+
+                if isinstance(event_data, CustomEvent):
+                    event = cast(CustomEvent, event_data)
+                    if ctx.tool_call_id:
+                        event = replace(event, tool_call_id=ctx.tool_call_id)
+                else:
+                    event = CustomEvent(data=event_data, tool_call_id=ctx.tool_call_id)
                 await ctx.event_stream.send(event)
-            # TODO (DouweM): Raise if events are yielded after ToolReturn
-            return None
+
+            return return_value
         elif self.is_async:
             function = cast(Callable[[Any], Awaitable[str]], self.function)
             return await function(*args, **kwargs)

pydantic_ai_slim/pydantic_ai/_output.py

@@ -134,6 +134,8 @@ async def execute_traced_output_function(
     ) as span:
         try:
             output = await function_schema.call(args, run_context)
+            if isinstance(output, _messages.Return):
+                output = cast(_messages.Return[Any], output).return_value
         except ModelRetry as r:
             if wrap_validation_errors:
                 m = _messages.RetryPromptPart(

pydantic_ai_slim/pydantic_ai/messages.py

@@ -23,7 +23,7 @@
 if TYPE_CHECKING:
     from .models.instrumented import InstrumentationSettings
 
-EventPayloadT = TypeVar('EventPayloadT', default=Any)
+EventDataT = TypeVar('EventPayloadT', default=Any)
 
 
 AudioMediaType: TypeAlias = Literal['audio/wav', 'audio/mpeg', 'audio/ogg', 'audio/flac', 'audio/aiff', 'audio/aac']
@@ -617,9 +617,23 @@ 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
+
 
 @dataclass(repr=False)
-class ToolReturn:
+class ToolReturn(Return[ReturnValueT]):
     """A structured return value for tools that need to provide both a return value and custom content to the model.
 
     This class allows tools to return complex responses that include:
@@ -628,9 +642,6 @@ class ToolReturn:
     - Optional metadata for application use
     """
 
-    return_value: Any
-    """The return value to be used in the tool response."""
-
     _: KW_ONLY
 
     content: str | Sequence[UserContent] | None = None
@@ -1779,19 +1790,16 @@ class BuiltinToolResultEvent:
 
 
 @dataclass(repr=False)
-class CustomEvent(Generic[EventPayloadT]):
+class CustomEvent(Generic[EventDataT]):
     """An event indicating the result of a function tool call."""
 
-    payload: EventPayloadT
-    """The payload of the custom event."""
+    data: EventDataT
+    """The data of the custom event."""
 
     _: KW_ONLY
 
-    name: str | None = None
-    """The optional name of the custom event."""
-
-    id: str | None = None
-    """The optional ID of the custom event."""
+    tool_call_id: str | None = None
+    """The tool call ID, if any, that this event is associated with."""
 
     event_kind: Literal['custom'] = 'custom'
     """Event type identifier, used as a discriminator."""

pydantic_ai_slim/pydantic_ai/output.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Awaitable, Callable, Sequence
+from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
 from dataclasses import dataclass
 from typing import Any, Generic, Literal
 
@@ -11,7 +11,7 @@
 
 from . import _utils, exceptions
 from ._json_schema import InlineDefsJsonSchemaTransformer
-from .messages import ToolCallPart
+from .messages import Return, ToolCallPart
 from .tools import DeferredToolRequests, ObjectJsonSchema, RunContext, ToolDefinition
 
 __all__ = (
@@ -45,10 +45,19 @@
 StructuredOutputMode = Literal['tool', 'native', 'prompted']
 """Output modes that can be used for structured output. Used by ModelProfile.default_structured_output_mode"""
 
-
-OutputTypeOrFunction = TypeAliasType(
-    'OutputTypeOrFunction', type[T_co] | Callable[..., Awaitable[T_co] | T_co], type_params=(T_co,)
+OutputFunction = TypeAliasType(
+    'OutputFunction',
+    Callable[..., AsyncIterator[Return[T_co] | Any]] | Callable[..., Awaitable[T_co]] | Callable[..., T_co],
+    type_params=(T_co,),
 )
+"""Definition of an output function.
+
+You should not need to import or use this type directly.
+
+See [output docs](../output.md) for more information.
+"""
+
+OutputTypeOrFunction = TypeAliasType('OutputTypeOrFunction', OutputFunction[T_co] | type[T_co], type_params=(T_co,))
 """Definition of an output type or function.
 
 You should not need to import or use this type directly.

pydantic_ai_slim/pydantic_ai/result.py

@@ -170,6 +170,7 @@ 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
             )
@@ -191,6 +192,7 @@ 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/ui/_event_stream.py

@@ -15,6 +15,7 @@
     BuiltinToolCallPart,
     BuiltinToolResultEvent,  # pyright: ignore[reportDeprecated]
     BuiltinToolReturnPart,
+    CustomEvent,
     FilePart,
     FinalResultEvent,
     FunctionToolCallEvent,
@@ -229,7 +230,7 @@ async def _turn_to(self, to_turn: Literal['request', 'response'] | None) -> Asyn
             async for e in self.before_response():
                 yield e
 
-    async def handle_event(self, event: NativeEvent) -> AsyncIterator[EventT]:
+    async def handle_event(self, event: NativeEvent) -> AsyncIterator[EventT]:  # noqa: C901
         """Transform a Pydantic AI event into one or more protocol-specific events.
 
         This method dispatches to specific `handle_*` methods based on event type:
@@ -240,6 +241,7 @@ async def handle_event(self, event: NativeEvent) -> AsyncIterator[EventT]:
         - [`FinalResultEvent`][pydantic_ai.messages.FinalResultEvent] -> `handle_final_result`
         - [`FunctionToolCallEvent`][pydantic_ai.messages.FunctionToolCallEvent] -> `handle_function_tool_call`
         - [`FunctionToolResultEvent`][pydantic_ai.messages.FunctionToolResultEvent] -> `handle_function_tool_result`
+        - [`CustomEvent`][pydantic_ai.messages.CustomEvent] -> `handle_custom_event`
         - [`AgentRunResultEvent`][pydantic_ai.run.AgentRunResultEvent] -> `handle_run_result`
 
         Subclasses are encouraged to override the individual `handle_*` methods rather than this one.
@@ -264,6 +266,9 @@ async def handle_event(self, event: NativeEvent) -> AsyncIterator[EventT]:
             case FunctionToolResultEvent():
                 async for e in self.handle_function_tool_result(event):
                     yield e
+            case CustomEvent():
+                async for e in self.handle_custom_event(event):
+                    yield e
             case AgentRunResultEvent():
                 async for e in self.handle_run_result(event):
                     yield e
@@ -581,6 +586,15 @@ async def handle_function_tool_result(self, event: FunctionToolResultEvent) -> A
         return  # pragma: no cover
         yield  # Make this an async generator
 
+    async def handle_custom_event(self, event: CustomEvent) -> AsyncIterator[EventT]:
+        """Handle a `CustomEvent`.
+
+        Args:
+            event: The custom event.
+        """
+        return
+        yield  # Make this an async generator
+
     async def handle_run_result(self, event: AgentRunResultEvent) -> AsyncIterator[EventT]:
         """Handle an `AgentRunResultEvent`.
 

pydantic_ai_slim/pydantic_ai/ui/ag_ui/_event_stream.py

@@ -14,6 +14,7 @@
 from ...messages import (
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
+    CustomEvent,
     FunctionToolResultEvent,
     RetryPromptPart,
     TextPart,
@@ -234,3 +235,7 @@ async def handle_function_tool_result(self, event: FunctionToolResultEvent) -> A
                 for item in possible_event:  # type: ignore[reportUnknownMemberType]
                     if isinstance(item, BaseEvent):  # pragma: no branch
                         yield item
+
+    async def handle_custom_event(self, event: CustomEvent) -> AsyncIterator[BaseEvent]:
+        if isinstance(event.data, BaseEvent):
+            yield event.data

pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py

@@ -11,6 +11,7 @@
 from ...messages import (
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
+    CustomEvent,
     FilePart,
     FunctionToolResultEvent,
     RetryPromptPart,
@@ -185,3 +186,7 @@ async def handle_function_tool_result(self, event: FunctionToolResultEvent) -> A
             yield ToolOutputAvailableChunk(tool_call_id=result.tool_call_id, output=result.content)
 
         # ToolCallResultEvent.content may hold user parts (e.g. text, images) that Vercel AI does not currently have events for
+
+    async def handle_custom_event(self, event: CustomEvent) -> AsyncIterator[BaseChunk]:
+        if isinstance(event.data, BaseChunk):
+            yield event.data