"""AG-UI adapter for handling requests.""" from __future__ import annotations import json import uuid import warnings from base64 import b64decode from collections.abc import Sequence from dataclasses import KW_ONLY, dataclass from functools import cached_property from typing import ( TYPE_CHECKING, Any, Literal, ) from pydantic import ValidationError from typing_extensions import assert_never from ... import ExternalToolset, ToolDefinition from ..._utils import is_str_dict from ...messages import ( AudioUrl, BinaryContent, CachePoint, CompactionPart, DocumentUrl, FilePart, ForceDownloadMode, ImageUrl, ModelMessage, ModelRequest, ModelResponse, NativeToolCallPart, NativeToolReturnPart, RetryPromptPart, SpeechPart, SystemPromptPart, TextContent, TextPart, ThinkingPart, ToolAvailabilityDeltaPart, ToolCallPart, ToolPartKind, ToolReturnPart, UploadedFile, UserContent, UserPromptPart, VideoUrl, narrow_message_parts, ) from ...output import OutputDataT from ...tools import ( AgentDepsT, DeferredToolApprovalResult, DeferredToolResults, ) from ...toolsets import AbstractToolset from .._adapter import ( DEFAULT_ALLOWED_CONTENT_TYPES, compaction_part_from_payload, compaction_payload, tool_availability_delta_from_payload, ) try: from ag_ui.core import ( ActivityMessage, AssistantMessage, BaseEvent, BinaryInputContent, DeveloperMessage, FunctionCall, Message, RunAgentInput, SystemMessage, TextInputContent, Tool as AGUITool, ToolCall, ToolMessage, UserMessage, ) from .. import MessagesBuilder, UIAdapter, UIEventStream from ._event_stream import AGUIEventStream from ._forward_compat import skip_unknown_tagged_items from ._interrupt import ( HAS_INTERRUPTS, ResumeEntry, interrupt_id_to_tool_call_id, resume_entry_to_approval, ) from ._utils import ( BUILTIN_TOOL_CALL_ID_PREFIX, COMPACTION_ACTIVITY_TYPE, DEFAULT_AG_UI_VERSION, ENCRYPTED_VALUE_VERSION, FILE_ACTIVITY_TYPE, MULTIMODAL_VERSION, REASONING_VERSION, TOOL_AVAILABILITY_DELTA_ACTIVITY_TYPE, UPLOADED_FILE_ACTIVITY_TYPE, dump_tool_return_content, parse_ag_ui_version, parse_builtin_tool_call_id, parse_encrypted_outcome, parse_encrypted_tool_kind, rehydrate_tool_return_content, thinking_encrypted_metadata, tool_kind_encrypted_value_kwargs, warn_tool_kind_not_persisted, ) except ImportError as e: raise ImportError( 'Please install the `ag-ui-protocol` package to use AG-UI integration, ' 'you can use the `ag-ui` optional group — `pip install "pydantic-ai-slim[ag-ui]"`' ) from e if TYPE_CHECKING: from ag_ui.core import ( AudioInputContent, DocumentInputContent, ImageInputContent, ReasoningMessage, VideoInputContent, ) from starlette.requests import Request from ...agent import AbstractAgent else: try: from ag_ui.core import ReasoningMessage except ImportError: class ReasoningMessage: """Stub for ag-ui-protocol < 0.1.11 — no instances exist, so pattern matching is a no-op.""" try: from ag_ui.core import AudioInputContent, DocumentInputContent, ImageInputContent, VideoInputContent except ImportError: class ImageInputContent: """Stub for ag-ui-protocol < 0.1.15.""" class AudioInputContent: """Stub for ag-ui-protocol < 0.1.15.""" class VideoInputContent: """Stub for ag-ui-protocol < 0.1.15.""" class DocumentInputContent: """Stub for ag-ui-protocol < 0.1.15.""" __all__ = ['AGUIAdapter'] # Frontend toolset class _AGUIFrontendToolset(ExternalToolset[AgentDepsT]): """Toolset for AG-UI frontend tools.""" def __init__(self, tools: list[AGUITool]): """Initialize the toolset with AG-UI tools. Args: tools: List of AG-UI tool definitions. """ super().__init__( [ ToolDefinition( name=tool.name, description=tool.description, parameters_json_schema=tool.parameters or {}, ) for tool in tools ] ) @property def label(self) -> str: """Return the label for this toolset.""" return 'the AG-UI frontend tools' # pragma: no cover def _new_message_id() -> str: """Generate a new unique message ID.""" return str(uuid.uuid4()) def _user_content_to_input( item: str | TextContent | ImageUrl | VideoUrl | AudioUrl | DocumentUrl | BinaryContent | UploadedFile | CachePoint, *, use_multimodal: bool = False, ) -> ( TextInputContent | BinaryInputContent | ImageInputContent | AudioInputContent | VideoInputContent | DocumentInputContent | None ): """Convert a user content item to AG-UI input content. When `use_multimodal` is True (ag-ui >= 0.1.15), media URLs are emitted as typed multimodal input content (e.g. `ImageInputContent`) instead of generic `BinaryInputContent`. """ if isinstance(item, str): return TextInputContent(type='text', text=item) elif isinstance(item, TextContent): return TextInputContent(type='text', text=item.content) elif isinstance(item, (ImageUrl, VideoUrl, AudioUrl, DocumentUrl)): if use_multimodal: from ._multimodal import media_url_to_multimodal return media_url_to_multimodal(item) return BinaryInputContent(type='binary', url=item.url, mime_type=item.media_type or '') elif isinstance(item, BinaryContent): if use_multimodal: from ._multimodal import binary_to_multimodal return binary_to_multimodal(item) return BinaryInputContent(type='binary', data=item.base64, mime_type=item.media_type) elif isinstance(item, UploadedFile): # UploadedFile holds an opaque provider file_id (e.g. 'file-abc123'), not a URL or # binary data, so it can't be mapped to AG-UI input content. Skipped like CachePoint. return None elif isinstance(item, CachePoint): return None else: assert_never(item) @dataclass class AGUIAdapter(UIAdapter[RunAgentInput, Message, BaseEvent, AgentDepsT, OutputDataT]): """UI adapter for the Agent-User Interaction (AG-UI) protocol.""" _: KW_ONLY ag_ui_version: str = DEFAULT_AG_UI_VERSION """AG-UI protocol version controlling behavior thresholds. Accepts any version string (e.g. `'0.1.11'`). Defaults to the version detected from the installed `ag-ui-protocol` package. Known thresholds: - `< 0.1.11`: emits `THINKING_*` events during streaming, drops `ThinkingPart` from `dump_messages` output. - `>= 0.1.11`: emits `REASONING_*` events with encrypted metadata during streaming, and includes `ThinkingPart` as `ReasoningMessage` in `dump_messages` output for full round-trip fidelity of thinking signatures and provider metadata. - `>= 0.1.15`: emits typed multimodal input content (`ImageInputContent`, `AudioInputContent`, `VideoInputContent`, `DocumentInputContent`) instead of generic `BinaryInputContent`. `load_messages` always accepts `ReasoningMessage` and multimodal content types regardless of this setting, and `build_run_input` skips inbound content types the installed `ag-ui-protocol` predates rather than rejecting the request. """ preserve_file_data: bool = False """Whether to round-trip `FilePart` and `UploadedFile` through reserved `pydantic_ai_*` [activity messages](https://docs.ag-ui.com/concepts/messages). Defaults to `False`. AG-UI has no native representation for agent-generated files ([`FilePart`][pydantic_ai.messages.FilePart]) or uploaded-file references ([`UploadedFile`][pydantic_ai.messages.UploadedFile]), so when this is `True` they are serialized as sidecar activity messages on `dump_messages` and reconstructed on `load_messages`. A frontend only completes the round-trip if it echoes these activity messages back on the next request. This is a representation setting, not a security one: honoring a reconstructed inbound `UploadedFile` still requires [`allow_uploaded_files`][pydantic_ai.ui.UIAdapter.allow_uploaded_files], which the shared `sanitize_messages` step enforces regardless of this flag. Multimodal tool-return files are unaffected — they ride inline in `ToolMessage.content`. """ @classmethod def build_run_input(cls, body: bytes) -> RunAgentInput: """Build an AG-UI run input object from the request body. A message `role` or input content `type` introduced by a protocol version newer than the installed `ag-ui-protocol` is skipped with a warning rather than failing the whole request, per the backwards-compatibility policy in `pydantic_ai/ui/AGENTS.md`. Only items the installed models cannot dispatch at all are skipped: a body that is invalid for any other reason still raises, so a client bug isn't converted into silent misbehavior. """ try: return RunAgentInput.model_validate_json(body) except ValidationError: payload, skipped = skip_unknown_tagged_items(body) if not skipped: raise # Validated outside the `except` block so a body that is *also* malformed reports the # remaining errors on their own rather than chained behind the unknown-tag failure. run_input = RunAgentInput.model_validate(payload) warnings.warn( f'AG-UI content the installed ag-ui-protocol {DEFAULT_AG_UI_VERSION} does not support ' f'({", ".join(sorted(skipped))}) was skipped; upgrade `ag-ui-protocol` to accept it.', UserWarning, stacklevel=2, ) return run_input def build_event_stream(self) -> UIEventStream[RunAgentInput, BaseEvent, AgentDepsT, OutputDataT]: """Build an AG-UI event stream transformer.""" return AGUIEventStream(self.run_input, accept=self.accept, ag_ui_version=self.ag_ui_version) @classmethod async def from_request( cls, request: Request, *, agent: AbstractAgent[AgentDepsT, OutputDataT], ag_ui_version: str = DEFAULT_AG_UI_VERSION, preserve_file_data: bool = False, manage_system_prompt: Literal['server', 'client'] = 'server', allowed_file_url_schemes: frozenset[str] = frozenset({'http', 'https'}), allowed_file_url_force_download: frozenset[ForceDownloadMode] = frozenset(), allow_uploaded_files: bool = False, allowed_content_types: frozenset[str] | None = DEFAULT_ALLOWED_CONTENT_TYPES, **kwargs: Any, ) -> AGUIAdapter[AgentDepsT, OutputDataT]: """Extends [`from_request`][pydantic_ai.ui.UIAdapter.from_request] with AG-UI-specific parameters.""" return await super().from_request( request, agent=agent, ag_ui_version=ag_ui_version, preserve_file_data=preserve_file_data, manage_system_prompt=manage_system_prompt, allowed_file_url_schemes=allowed_file_url_schemes, allowed_file_url_force_download=allowed_file_url_force_download, allow_uploaded_files=allow_uploaded_files, allowed_content_types=allowed_content_types, **kwargs, ) @cached_property def messages(self) -> list[ModelMessage]: """Pydantic AI messages from the AG-UI run input.""" return self.load_messages(self.run_input.messages, preserve_file_data=self.preserve_file_data) @cached_property def toolset(self) -> AbstractToolset[AgentDepsT] | None: """Toolset representing frontend tools from the AG-UI run input.""" if self.run_input.tools: return _AGUIFrontendToolset[AgentDepsT](self.run_input.tools) return None @cached_property def state(self) -> dict[str, Any] | None: """Frontend state from the AG-UI run input.""" state = self.run_input.state if is_str_dict(state) and state: return state return None @cached_property def conversation_id(self) -> str | None: """Conversation ID from the AG-UI `RunAgentInput.threadId`.""" return self.run_input.thread_id @cached_property def deferred_tool_results(self) -> DeferredToolResults | None: """Translate AG-UI `RunAgentInput.resume[]` into Pydantic AI `DeferredToolResults`. See [docs.ag-ui.com/concepts/interrupts](https://docs.ag-ui.com/concepts/interrupts). Each `ResumeEntry` is mapped to an approval keyed by the original `tool_call_id`. The payload is validated against the same Pydantic model whose JSON schema is advertised on `Interrupt.response_schema`, and the mapping is **deny-by-default**: approval requires a payload that validates with `approved=True`. Any other shape is treated as a denial so a malformed or hostile client cannot accidentally execute a tool that requires human approval. - `status == 'cancelled'` → `ToolDenied('Cancelled by user.')` - `payload.approved is True` with a valid `payload.editedArgs` dict → `ToolApproved(override_args=...)` - `payload.approved is True` without edits → `ToolApproved()` - Anything else (`False`, missing, `null`, non-bool `approved`, non-dict payload, a non-dict `editedArgs`, or a non-string `reason`) → `ToolDenied(payload.reason)` if `reason` is a non-empty string on a payload that validated, else `ToolDenied()` (which carries the default `"The tool call was denied."` message). Returns `None` when `resume` is missing or empty, or when the installed ag-ui-protocol predates the interrupt lifecycle. """ if not HAS_INTERRUPTS: return None resume: list[ResumeEntry] | None = getattr(self.run_input, 'resume', None) if not resume: return None approvals: dict[str, DeferredToolApprovalResult | bool] = { interrupt_id_to_tool_call_id(entry.interrupt_id): resume_entry_to_approval(entry) for entry in resume } return DeferredToolResults(approvals=approvals) @classmethod def load_messages(cls, messages: Sequence[Message], *, preserve_file_data: bool = False) -> list[ModelMessage]: # noqa: C901 """Transform AG-UI messages into Pydantic AI messages.""" builder = MessagesBuilder() tool_calls: dict[str, str] = {} # Tool call ID to tool name mapping. tool_kinds: dict[str, ToolPartKind] = {} # Tool call ID to `tool_kind` claim mapping. # `ToolCall`/`ToolMessage.encrypted_value` only exists on the installed model from 0.1.11 # onward; older versions drop the client's claim, so the field is only read when present. use_encrypted_value = parse_ag_ui_version(DEFAULT_AG_UI_VERSION) >= ENCRYPTED_VALUE_VERSION for msg in messages: match msg: case UserMessage(content=content): if isinstance(content, str): builder.add(UserPromptPart(content=content)) else: user_prompt_content: list[UserContent] = [] for part in content: match part: case TextInputContent(text=text): user_prompt_content.append(text) case BinaryInputContent(): if part.url: try: binary_part = BinaryContent.from_data_uri(part.url) except ValueError: media_type_constructors = { 'image': ImageUrl, 'video': VideoUrl, 'audio': AudioUrl, } media_type_prefix = part.mime_type.split('/', 1)[0] constructor = media_type_constructors.get(media_type_prefix, DocumentUrl) binary_part = constructor(url=part.url, media_type=part.mime_type) elif part.data: binary_part = BinaryContent( data=b64decode(part.data), media_type=part.mime_type ) else: # pragma: no cover raise ValueError('BinaryInputContent must have either a `url` or `data` field.') user_prompt_content.append(binary_part) case ( ImageInputContent() | AudioInputContent() | VideoInputContent() | DocumentInputContent() ): from ._multimodal import ( multimodal_input_to_content, ) user_prompt_content.append(multimodal_input_to_content(part)) case _: assert_never(part) if user_prompt_content: content_to_add = ( user_prompt_content[0] if len(user_prompt_content) == 1 and isinstance(user_prompt_content[0], str) else user_prompt_content ) builder.add(UserPromptPart(content=content_to_add)) case SystemMessage(content=content) | DeveloperMessage(content=content): builder.add(SystemPromptPart(content=content)) case AssistantMessage(content=content, tool_calls=tool_calls_list): if content: builder.add(TextPart(content=content)) if tool_calls_list: for tool_call in tool_calls_list: tool_call_id = tool_call.id tool_name = tool_call.function.name tool_calls[tool_call_id] = tool_name # The claim is client-supplied, so it's set on the base part and promoted # best-effort by the final `narrow_message_parts` pass (which strips it if # it doesn't validate against the typed subclass). tool_kind = ( parse_encrypted_tool_kind(tool_call.encrypted_value) if use_encrypted_value else None ) if tool_kind is not None: tool_kinds[tool_call_id] = tool_kind builtin_id = parse_builtin_tool_call_id(tool_call_id) if builtin_id is not None: provider_name, original_id = builtin_id builder.add( NativeToolCallPart( tool_name=tool_name, args=tool_call.function.arguments, tool_call_id=original_id, provider_name=provider_name, tool_kind=tool_kind, ) ) else: builder.add( ToolCallPart( tool_name=tool_name, tool_call_id=tool_call_id, args=tool_call.function.arguments, tool_kind=tool_kind, ) ) case ToolMessage() as tool_msg: tool_call_id = tool_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.') # Rehydrate here (not in a later `ModelMessagesTypeAdapter` pass) so structured and # multimodal content comes back as real types; see `rehydrate_tool_return_content`. content = rehydrate_tool_return_content(tool_msg.content) # Fall back to the paired call's claim: `ToolCallResultEvent` has no metadata # slot, so client-built ToolMessages usually carry no `encrypted_value`. Error # results stay untyped — typed return parts imply success to their readers. # A non-success outcome claim (the return would otherwise reload as `'success'`, # changing how it serializes to the provider) also keeps the return untyped. tool_kind = None outcome: Literal['success', 'failed', 'denied', 'interrupted'] = 'success' encrypted_outcome = ( parse_encrypted_outcome(tool_msg.encrypted_value) if use_encrypted_value else None ) if encrypted_outcome is not None: outcome = encrypted_outcome elif tool_msg.error is not None: outcome = 'failed' else: encrypted_tool_kind = ( parse_encrypted_tool_kind(tool_msg.encrypted_value) if use_encrypted_value else None ) tool_kind = encrypted_tool_kind or tool_kinds.get(tool_call_id) builtin_id = parse_builtin_tool_call_id(tool_call_id) if builtin_id is not None: provider_name, original_id = builtin_id builder.add( NativeToolReturnPart( tool_name=tool_name, content=content, tool_call_id=original_id, provider_name=provider_name, tool_kind=tool_kind, outcome=outcome, ) ) else: # The final `narrow_message_parts` pass parses the rehydrated content into a typed # return subclass when the `tool_kind` claim validates, and leaves the base # `ToolReturnPart` (dropping the claim) when it doesn't. builder.add( ToolReturnPart( tool_name=tool_name, content=content, tool_call_id=tool_call_id, tool_kind=tool_kind, outcome=outcome, ) ) case ReasoningMessage() as reasoning_msg: try: metadata: dict[str, Any] = ( json.loads(reasoning_msg.encrypted_value) if reasoning_msg.encrypted_value else {} ) if not isinstance(metadata, dict): metadata = {} except json.JSONDecodeError: metadata = {} builder.add( ThinkingPart( content=reasoning_msg.content, id=metadata.get('id'), signature=metadata.get('signature'), provider_name=metadata.get('provider_name'), provider_details=metadata.get('provider_details'), ) ) case ActivityMessage() as activity_msg: if activity_msg.activity_type == TOOL_AVAILABILITY_DELTA_ACTIVITY_TYPE: builder.add(tool_availability_delta_from_payload(activity_msg.content)) elif activity_msg.activity_type == COMPACTION_ACTIVITY_TYPE: if (compaction_part := compaction_part_from_payload(activity_msg.content)) is not None: builder.add(compaction_part) elif activity_msg.activity_type == FILE_ACTIVITY_TYPE and preserve_file_data: activity_content = activity_msg.content url = activity_content.get('url', '') if not url: raise ValueError( f'ActivityMessage with activity_type={FILE_ACTIVITY_TYPE!r} must have a non-empty url.' ) binary_content = BinaryContent.from_data_uri(url) vendor_metadata = activity_content.get('vendor_metadata') # `vendor_metadata` is client-supplied and typed `Any`; assignment on the # (non-`validate_assignment`) `BinaryContent` dataclass bypasses validation, # so ignore anything that isn't a dict rather than let it reach the provider. if is_str_dict(vendor_metadata): binary_content.vendor_metadata = vendor_metadata builder.add( FilePart( content=binary_content, id=activity_content.get('id'), provider_name=activity_content.get('provider_name'), provider_details=activity_content.get('provider_details'), ) ) elif activity_msg.activity_type == UPLOADED_FILE_ACTIVITY_TYPE and preserve_file_data: activity_content = activity_msg.content file_id = activity_content.get('file_id', '') provider_name = activity_content.get('provider_name', '') if not file_id or not provider_name: raise ValueError( f'ActivityMessage with activity_type={UPLOADED_FILE_ACTIVITY_TYPE!r}' ' must have non-empty file_id and provider_name.' ) builder.add( UserPromptPart( content=[ UploadedFile( file_id=file_id, provider_name=provider_name, vendor_metadata=activity_content.get('vendor_metadata'), media_type=activity_content.get('media_type'), identifier=activity_content.get('identifier'), ) ] ) ) case _: if TYPE_CHECKING: assert_never(msg) warnings.warn( f'AG-UI message type {type(msg).__name__} is not yet implemented; skipping.', UserWarning, stacklevel=2, ) # Parts above are built as base `ToolCallPart`/`ToolReturnPart`/`NativeTool*Part` carrying a # `tool_kind` claim; promote them to their typed subclasses in one best-effort pass. return narrow_message_parts(builder.messages) @staticmethod def _dump_request_parts( # noqa: C901 msg: ModelRequest, *, ag_ui_version: str = DEFAULT_AG_UI_VERSION, preserve_file_data: bool = False, ) -> list[Message]: """Convert a `ModelRequest` into AG-UI messages. Uses a flush pattern to preserve part ordering: buffered user content is flushed before each tool message, so a `ToolReturnPart` that precedes a `UserPromptPart` in the original request keeps its position instead of being reordered after the user prompt. """ use_multimodal = parse_ag_ui_version(ag_ui_version) >= MULTIMODAL_VERSION # `ToolMessage.encrypted_value` (the `tool_kind` carrier here) landed in 0.1.11 — see # `tool_kind_encrypted_value`. use_encrypted_value = parse_ag_ui_version(ag_ui_version) >= ENCRYPTED_VALUE_VERSION result: list[Message] = [] system_content: list[str] = [] user_content: list[ TextInputContent | BinaryInputContent | ImageInputContent | AudioInputContent | VideoInputContent | DocumentInputContent ] = [] def flush_user_content() -> None: nonlocal user_content if not user_content: return # Simplify to plain string if only a single text item. if len(user_content) == 1 and isinstance(user_content[0], TextInputContent): result.append(UserMessage(id=_new_message_id(), content=user_content[0].text)) else: result.append(UserMessage(id=_new_message_id(), content=user_content)) user_content = [] for part in msg.parts: if isinstance(part, SystemPromptPart): system_content.append(part.content) elif isinstance(part, UserPromptPart): if isinstance(part.content, str): user_content.append(TextInputContent(type='text', text=part.content)) else: for item in part.content: if isinstance(item, UploadedFile) and preserve_file_data: # AG-UI has no native uploaded-file message type. We repurpose # ActivityMessage with a reserved `pydantic_ai_*` activity_type # for round-trip fidelity. See UploadedFileActivityContent. flush_user_content() uploaded_content: dict[str, Any] = { 'file_id': item.file_id, 'provider_name': item.provider_name, 'media_type': item.media_type, 'identifier': item.identifier, } if item.vendor_metadata is not None: uploaded_content['vendor_metadata'] = item.vendor_metadata result.append( ActivityMessage( id=_new_message_id(), activity_type=UPLOADED_FILE_ACTIVITY_TYPE, content=uploaded_content, ) ) else: converted = _user_content_to_input(item, use_multimodal=use_multimodal) if converted is not None: user_content.append(converted) elif isinstance(part, ToolReturnPart): flush_user_content() # Tool-return files ride inline in `ToolMessage.content` (see `dump_tool_return_content`). # A non-success outcome rides the `encrypted_value` carrier alongside `tool_kind`, # since a `ToolMessage` has no outcome slot. result.append( ToolMessage( id=_new_message_id(), content=dump_tool_return_content(part.content), tool_call_id=part.tool_call_id, error=part.model_response_str(wrap_if_error=False) if part.outcome in ('failed', 'denied') else None, **tool_kind_encrypted_value_kwargs( part.tool_kind, outcome=part.outcome, supported=use_encrypted_value ), ) ) elif isinstance(part, ToolAvailabilityDeltaPart): flush_user_content() result.append( ActivityMessage( id=_new_message_id(), activity_type=TOOL_AVAILABILITY_DELTA_ACTIVITY_TYPE, content={ 'added': part.tools_added, 'tool_call_id': part.tool_call_id, }, ) ) elif isinstance(part, RetryPromptPart): if part.tool_name: flush_user_content() result.append( ToolMessage( id=_new_message_id(), content=part.model_response(), tool_call_id=part.tool_call_id, error=part.model_response(), ) ) else: user_content.append(TextInputContent(type='text', text=part.model_response())) elif isinstance(part, SpeechPart): # pragma: no cover pass # Realtime audio parts are not rendered in AG-UI else: assert_never(part) messages: list[Message] = [] if system_content: messages.append(SystemMessage(id=_new_message_id(), content='\n'.join(system_content))) flush_user_content() messages.extend(result) return messages @staticmethod def _dump_response_parts( # noqa: C901 msg: ModelResponse, *, ag_ui_version: str = DEFAULT_AG_UI_VERSION, preserve_file_data: bool = False ) -> list[Message]: """Convert a `ModelResponse` into AG-UI messages. Uses a flush pattern to preserve part ordering: text that appears after tool calls gets its own AssistantMessage, and ThinkingPart/FilePart boundaries trigger a flush so content on either side doesn't get merged. """ result: list[Message] = [] text_content: list[str] = [] tool_calls_list: list[ToolCall] = [] tool_messages: list[ToolMessage] = [] version = parse_ag_ui_version(ag_ui_version) # `ReasoningMessage` and the `tool_kind` carrier `ToolCall`/`ToolMessage.encrypted_value` # both landed in 0.1.11, so these two thresholds coincide — see `tool_kind_encrypted_value`. use_reasoning = version >= REASONING_VERSION use_encrypted_value = version >= ENCRYPTED_VALUE_VERSION builtin_returns = {part.tool_call_id: part for part in msg.parts if isinstance(part, NativeToolReturnPart)} def flush() -> None: nonlocal text_content, tool_calls_list, tool_messages if not text_content and not tool_calls_list: return result.append( AssistantMessage( id=_new_message_id(), content='\n'.join(text_content) if text_content else None, tool_calls=tool_calls_list if tool_calls_list else None, ) ) result.extend(tool_messages) text_content = [] tool_calls_list = [] tool_messages = [] for part in msg.parts: if isinstance(part, TextPart): if tool_calls_list: flush() text_content.append(part.content) elif isinstance(part, ThinkingPart): if use_reasoning: from ag_ui.core import ReasoningMessage flush() encrypted = thinking_encrypted_metadata(part) result.append( ReasoningMessage( id=_new_message_id(), content=part.content, encrypted_value=json.dumps(encrypted) if encrypted else None, ) ) elif isinstance(part, ToolCallPart): tool_calls_list.append( ToolCall( id=part.tool_call_id, function=FunctionCall(name=part.tool_name, arguments=part.args_as_json_str()), **tool_kind_encrypted_value_kwargs(part.tool_kind, supported=use_encrypted_value), ) ) elif isinstance(part, NativeToolCallPart): prefixed_id = '|'.join([BUILTIN_TOOL_CALL_ID_PREFIX, part.provider_name or '', part.tool_call_id]) tool_calls_list.append( ToolCall( id=prefixed_id, function=FunctionCall(name=part.tool_name, arguments=part.args_as_json_str()), **tool_kind_encrypted_value_kwargs(part.tool_kind, supported=use_encrypted_value), ) ) if builtin_return := builtin_returns.get(part.tool_call_id): # Built-in tool-return files also ride inline in `ToolMessage.content` (see above). tool_messages.append( ToolMessage( id=_new_message_id(), content=dump_tool_return_content(builtin_return.content), tool_call_id=prefixed_id, error=builtin_return.model_response_str(wrap_if_error=False) if builtin_return.outcome in ('failed', 'denied') else None, **tool_kind_encrypted_value_kwargs( builtin_return.tool_kind, outcome=builtin_return.outcome, supported=use_encrypted_value ), ) ) elif isinstance(part, NativeToolReturnPart): # Emitted when matching NativeToolCallPart is processed above. pass elif isinstance(part, FilePart): if preserve_file_data: # AG-UI has no native file message type. We repurpose ActivityMessage # with a reserved `pydantic_ai_*` activity_type for round-trip fidelity. # See FileActivityContent. flush() file_content: dict[str, Any] = { 'url': part.content.data_uri, 'media_type': part.content.media_type, } if part.id is not None: file_content['id'] = part.id if part.provider_name is not None: file_content['provider_name'] = part.provider_name if part.provider_details is not None: file_content['provider_details'] = part.provider_details if part.content.vendor_metadata is not None: file_content['vendor_metadata'] = part.content.vendor_metadata result.append( ActivityMessage( id=_new_message_id(), activity_type=FILE_ACTIVITY_TYPE, content=file_content, ) ) elif isinstance(part, CompactionPart): flush() result.append( ActivityMessage( id=_new_message_id(), activity_type=COMPACTION_ACTIVITY_TYPE, content=compaction_payload(part), ) ) elif isinstance(part, SpeechPart): # pragma: no cover pass # Realtime audio parts are not rendered in AG-UI else: assert_never(part) flush() return result @classmethod def dump_messages( cls, messages: Sequence[ModelMessage], *, ag_ui_version: str = DEFAULT_AG_UI_VERSION, preserve_file_data: bool = False, ) -> list[Message]: """Transform Pydantic AI messages into AG-UI messages. Note: The round-trip `dump_messages` -> `load_messages` is not fully lossless: - `ModelRequest.metadata` and top-level `ModelResponse.provider_details` are lost. AG-UI has no trusted message-level carrier for framework or provider state; general client-controlled metadata must not be restored as server-side state. - `TextPart.id`, `.provider_name`, `.provider_details` are lost. - `ToolCallPart.id`, `.provider_name`, `.provider_details` are lost. - `ToolCallPart.args` and `NativeToolCallPart.args` that don't parse as a JSON object are rewritten to `'{"INVALID_JSON":""}'` (see [`args_as_json_str`][pydantic_ai.messages.BaseToolCallPart.args_as_json_str]), so the raw string is no longer recoverable as args on reload. Unlike the live event stream, which emits them verbatim so streamed fragments stay concatenable, history has to hold a sendable value. - `NativeToolCallPart.id`, `.provider_details` are lost (only `.provider_name` survives via the prefixed tool call ID). - `NativeToolReturnPart.provider_details` is lost. - `tool_kind` is lost when `ag_ui_version < '0.1.11'` (before its `encrypted_value` carrier existed), so typed tool parts reload as their base classes. - `tool_kind` is not restored on error/denied tool returns (a typed return implies success to its readers), so those reload as plain `ToolReturnPart`. - A non-`'success'` `outcome` on a (native) tool return survives via the `encrypted_value` carrier from 0.1.11 (`ToolMessage` has no outcome slot). Below that, `'failed'` survives via `ToolMessage.error`, `'denied'` reloads as `'failed'`, and `'interrupted'` reloads as `'success'`. - `RetryPromptPart` becomes `ToolReturnPart` (or `UserPromptPart`) on reload. - A `NativeToolReturnPart` is always emitted directly after its `NativeToolCallPart`, so any part that originally sat between them — e.g. a `CompactionPart` — reloads after the pair instead. Provider adapters emit compaction parts outside call/return pairs, so this only affects hand-constructed histories. - `CachePoint` and `UploadedFile` content items are dropped (unless `preserve_file_data=True`). - `FileUrl.force_download` is dropped when `ag_ui_version < '0.1.15'` (before typed multimodal content gained a metadata carrier). - `ThinkingPart` is dropped when `ag_ui_version='0.1.10'`. - `FilePart` is silently dropped unless `preserve_file_data=True`. - `UploadedFile` in a multi-item `UserPromptPart` is split into a separate activity message when `preserve_file_data=True`, which reloads as a separate `UserPromptPart`. - `MultiModalContent` items in `ToolReturnPart`/`NativeToolReturnPart.content` always round-trip, regardless of `preserve_file_data`: the full content (files as base64/URL dicts) is serialized inline into the JSON `ToolMessage.content` and rehydrated on reload through the `ToolReturnContent` union. The same serialization is used for both history (`dump_messages`) and the live event stream (`ToolCallResultEvent.content`), so files survive either round-trip. - Part ordering within a `ModelResponse` may change when text follows tool calls. Args: messages: A sequence of ModelMessage objects to convert. ag_ui_version: AG-UI protocol version controlling `ThinkingPart` emission. preserve_file_data: Whether to include `FilePart` and `UploadedFile` items as `ActivityMessage`s. (Multimodal tool-return files always ride inline in `ToolMessage.content` and are unaffected.) Returns: A list of AG-UI Message objects. """ result: list[Message] = [] if parse_ag_ui_version(ag_ui_version) < ENCRYPTED_VALUE_VERSION and any( isinstance(part, (ToolCallPart, ToolReturnPart, NativeToolCallPart, NativeToolReturnPart)) and part.tool_kind is not None for msg in messages for part in msg.parts ): warn_tool_kind_not_persisted(ag_ui_version) for msg in messages: if isinstance(msg, ModelRequest): request_messages = cls._dump_request_parts( msg, ag_ui_version=ag_ui_version, preserve_file_data=preserve_file_data ) result.extend(request_messages) elif isinstance(msg, ModelResponse): result.extend( cls._dump_response_parts(msg, ag_ui_version=ag_ui_version, preserve_file_data=preserve_file_data) ) else: assert_never(msg) return result