Tools can now return AG-UI events separate from result sent to model (#2922)

Author: DouweM

docs/ag-ui.md

@@ -204,18 +204,18 @@ user experiences with frontend user interfaces.
 
 ### Events
 
-Pydantic AI tools can send
-[AG-UI events](https://docs.ag-ui.com/concepts/events) simply by defining a tool
-which returns a (subclass of)
-[`BaseEvent`](https://docs.ag-ui.com/sdk/python/core/events#baseevent), which allows
-for custom events and state updates.
+Pydantic AI tools can send [AG-UI events](https://docs.ag-ui.com/concepts/events) simply by returning a
+[`ToolReturn`](tools-advanced.md#advanced-tool-returns) object with a
+[`BaseEvent`](https://docs.ag-ui.com/sdk/python/core/events#baseevent) (or a list of events) as `metadata`,
+which allows for custom events and state updates.
 
 ```python {title="ag_ui_tool_events.py"}
 from ag_ui.core import CustomEvent, EventType, StateSnapshotEvent
 from pydantic import BaseModel
 
 from pydantic_ai import Agent, RunContext
 from pydantic_ai.ag_ui import StateDeps
+from pydantic_ai.messages import ToolReturn
 
 
 class DocumentState(BaseModel):
@@ -233,27 +233,35 @@ app = agent.to_ag_ui(deps=StateDeps(DocumentState()))
 
 
 @agent.tool
-async def update_state(ctx: RunContext[StateDeps[DocumentState]]) -> StateSnapshotEvent:
-    return StateSnapshotEvent(
-        type=EventType.STATE_SNAPSHOT,
-        snapshot=ctx.deps.state,
+async def update_state(ctx: RunContext[StateDeps[DocumentState]]) -> ToolReturn:
+    return ToolReturn(
+        return_value='State updated',
+        metadata=[
+            StateSnapshotEvent(
+                type=EventType.STATE_SNAPSHOT,
+                snapshot=ctx.deps.state,
+            ),
+        ],
     )
 
 
 @agent.tool_plain
-async def custom_events() -> list[CustomEvent]:
-    return [
-        CustomEvent(
-            type=EventType.CUSTOM,
-            name='count',
-            value=1,
-        ),
-        CustomEvent(
-            type=EventType.CUSTOM,
-            name='count',
-            value=2,
-        ),
-    ]
+async def custom_events() -> ToolReturn:
+    return ToolReturn(
+        return_value='Count events sent',
+        metadata=[
+            CustomEvent(
+                type=EventType.CUSTOM,
+                name='count',
+                value=1,
+            ),
+            CustomEvent(
+                type=EventType.CUSTOM,
+                name='count',
+                value=2,
+            ),
+        ]
+    )
 ```
 
 Since `app` is an ASGI application, it can be used with any ASGI server:

pydantic_ai_slim/pydantic_ai/ag_ui.py

@@ -559,15 +559,15 @@ async def _handle_tool_result_event(
         content=result.model_response_str(),
     )
 
-    # Now check for  AG-UI events returned by the tool calls.
-    content = result.content
-    if isinstance(content, BaseEvent):
-        yield content
-    elif isinstance(content, str | bytes):  # pragma: no branch
+    # Now check for AG-UI events returned by the tool calls.
+    possible_event = result.metadata or result.content
+    if isinstance(possible_event, BaseEvent):
+        yield possible_event
+    elif isinstance(possible_event, str | bytes):  # pragma: no branch
         # Avoid iterable check for strings and bytes.
         pass
-    elif isinstance(content, Iterable):  # pragma: no branch
-        for item in content:  # type: ignore[reportUnknownMemberType]
+    elif isinstance(possible_event, Iterable):  # pragma: no branch
+        for item in possible_event:  # type: ignore[reportUnknownMemberType]
             if isinstance(item, BaseEvent):  # pragma: no branch
                 yield item
 

tests/test_ag_ui.py

@@ -28,6 +28,7 @@
     SystemPromptPart,
     TextPart,
     ToolCallPart,
+    ToolReturn,
     ToolReturnPart,
     UserPromptPart,
 )
@@ -163,24 +164,22 @@ async def send_snapshot() -> StateSnapshotEvent:
     )
 
 
-async def send_custom() -> list[CustomEvent]:
-    """Display the recipe to the user.
-
-    Returns:
-        StateSnapshotEvent.
-    """
-    return [
-        CustomEvent(
-            type=EventType.CUSTOM,
-            name='custom_event1',
-            value={'key1': 'value1'},
-        ),
-        CustomEvent(
-            type=EventType.CUSTOM,
-            name='custom_event2',
-            value={'key2': 'value2'},
-        ),
-    ]
+async def send_custom() -> ToolReturn:
+    return ToolReturn(
+        return_value='Done',
+        metadata=[
+            CustomEvent(
+                type=EventType.CUSTOM,
+                name='custom_event1',
+                value={'key1': 'value1'},
+            ),
+            CustomEvent(
+                type=EventType.CUSTOM,
+                name='custom_event2',
+                value={'key2': 'value2'},
+            ),
+        ],
+    )
 
 
 def uuid_str() -> str:
@@ -785,7 +784,7 @@ async def stream_function(
                 'type': 'TOOL_CALL_RESULT',
                 'messageId': IsStr(),
                 'toolCallId': tool_call_id,
-                'content': '[{"type":"CUSTOM","timestamp":null,"raw_event":null,"name":"custom_event1","value":{"key1":"value1"}},{"type":"CUSTOM","timestamp":null,"raw_event":null,"name":"custom_event2","value":{"key2":"value2"}}]',
+                'content': 'Done',
                 'role': 'tool',
             },
             {'type': 'CUSTOM', 'name': 'custom_event1', 'value': {'key1': 'value1'}},