Improved support for converting to and from AG_UI message format

Author: ChuckJonas

pydantic_ai_slim/pydantic_ai/ui/ag_ui/__init__.py

@@ -2,8 +2,11 @@
 
 from ._adapter import AGUIAdapter
 from ._event_stream import AGUIEventStream
+from ._messages import messages_from_ag_ui, messages_to_ag_ui
 
 __all__ = [
     'AGUIAdapter',
     'AGUIEventStream',
+    'messages_from_ag_ui',
+    'messages_to_ag_ui',
 ]

pydantic_ai_slim/pydantic_ai/ui/ag_ui/_adapter.py

@@ -11,35 +11,22 @@
 )
 
 from ... import ExternalToolset, ToolDefinition
-from ...messages import (
-    BuiltinToolCallPart,
-    BuiltinToolReturnPart,
-    ModelMessage,
-    SystemPromptPart,
-    TextPart,
-    ToolCallPart,
-    ToolReturnPart,
-    UserPromptPart,
-)
+from ...messages import ModelMessage
 from ...output import OutputDataT
 from ...tools import AgentDepsT
 from ...toolsets import AbstractToolset
 
 try:
     from ag_ui.core import (
-        AssistantMessage,
         BaseEvent,
-        DeveloperMessage,
         Message,
         RunAgentInput,
-        SystemMessage,
         Tool as AGUITool,
-        ToolMessage,
-        UserMessage,
     )
 
-    from .. import MessagesBuilder, UIAdapter, UIEventStream
-    from ._event_stream import BUILTIN_TOOL_CALL_ID_PREFIX, AGUIEventStream
+    from .. import UIAdapter, UIEventStream
+    from ._event_stream import AGUIEventStream
+    from ._messages import messages_from_ag_ui
 except ImportError as e:  # pragma: no cover
     raise ImportError(
         'Please install the `ag-ui-protocol` package to use AG-UI integration, '
@@ -119,77 +106,9 @@ def state(self) -> dict[str, Any] | None:
 
     @classmethod
     def load_messages(cls, messages: Sequence[Message]) -> list[ModelMessage]:
-        """Transform AG-UI messages into Pydantic AI messages."""
-        builder = MessagesBuilder()
-        tool_calls: dict[str, str] = {}  # Tool call ID to tool name mapping.
-
-        for msg in messages:
-            if isinstance(msg, UserMessage | SystemMessage | DeveloperMessage) or (
-                isinstance(msg, ToolMessage) and not msg.tool_call_id.startswith(BUILTIN_TOOL_CALL_ID_PREFIX)
-            ):
-                if isinstance(msg, UserMessage):
-                    builder.add(UserPromptPart(content=msg.content))
-                elif isinstance(msg, SystemMessage | DeveloperMessage):
-                    builder.add(SystemPromptPart(content=msg.content))
-                else:
-                    tool_call_id = msg.tool_call_id
-                    tool_name = tool_calls.get(tool_call_id)
-                    if tool_name is None:  # pragma: no cover
-                        raise ValueError(f'Tool call with ID {tool_call_id} not found in the history.')
-
-                    builder.add(
-                        ToolReturnPart(
-                            tool_name=tool_name,
-                            content=msg.content,
-                            tool_call_id=tool_call_id,
-                        )
-                    )
-
-            elif isinstance(msg, AssistantMessage) or (  # pragma: no branch
-                isinstance(msg, ToolMessage) and msg.tool_call_id.startswith(BUILTIN_TOOL_CALL_ID_PREFIX)
-            ):
-                if isinstance(msg, AssistantMessage):
-                    if msg.content:
-                        builder.add(TextPart(content=msg.content))
-
-                    if msg.tool_calls:
-                        for tool_call in msg.tool_calls:
-                            tool_call_id = tool_call.id
-                            tool_name = tool_call.function.name
-                            tool_calls[tool_call_id] = tool_name
-
-                            if tool_call_id.startswith(BUILTIN_TOOL_CALL_ID_PREFIX):
-                                _, provider_name, tool_call_id = tool_call_id.split('|', 2)
-                                builder.add(
-                                    BuiltinToolCallPart(
-                                        tool_name=tool_name,
-                                        args=tool_call.function.arguments,
-                                        tool_call_id=tool_call_id,
-                                        provider_name=provider_name,
-                                    )
-                                )
-                            else:
-                                builder.add(
-                                    ToolCallPart(
-                                        tool_name=tool_name,
-                                        tool_call_id=tool_call_id,
-                                        args=tool_call.function.arguments,
-                                    )
-                                )
-                else:
-                    tool_call_id = msg.tool_call_id
-                    tool_name = tool_calls.get(tool_call_id)
-                    if tool_name is None:  # pragma: no cover
-                        raise ValueError(f'Tool call with ID {tool_call_id} not found in the history.')
-                    _, provider_name, tool_call_id = tool_call_id.split('|', 2)
-
-                    builder.add(
-                        BuiltinToolReturnPart(
-                            tool_name=tool_name,
-                            content=msg.content,
-                            tool_call_id=tool_call_id,
-                            provider_name=provider_name,
-                        )
-                    )
-
-        return builder.messages
+        """Transform AG-UI messages into Pydantic AI messages.
+
+        This is a convenience method that delegates to [`messages_from_ag_ui()`][pydantic_ai.ui.ag_ui.messages_from_ag_ui].
+        You can use that function directly if you need to convert messages without creating an adapter.
+        """
+        return messages_from_ag_ui(messages)