feat(ag-ui): preserve thinking signatures, files and tool returns in roundtrip#3971
feat(ag-ui): preserve thinking signatures, files and tool returns in roundtrip#3971dsfaccini wants to merge 35 commits intopydantic:mainfrom
Conversation
Add support for preserving thinking metadata (signature, provider_name, etc.) through AG-UI round-trips, enabling multi-turn conversations with Anthropic's extended thinking models. Fixes pydantic#3911 Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
e62256a to
46be3c4
Compare
|
@dsfaccini FYI I'd wait to do much more here until #3754 lands |
|
noted, there's honestly not much more to do here, I rewrote the logic a bit more cleanly as well as the tests, but this is now good to go from my side. |
| yield ThinkingEndEvent( | ||
| type=EventType.THINKING_END, | ||
| raw_event={'pydantic_ai': pydantic_ai_meta} if pydantic_ai_meta else None, | ||
| encryptedContent=part.signature, # pyright: ignore[reportCallIssue] |
There was a problem hiding this comment.
Are we allowed to add custom fields?
There was a problem hiding this comment.
yeahp, I've added explanations on the new code blocks https://docs.ag-ui.com/concepts/events#activity-events
|
|
||
| yield ThinkingEndEvent( | ||
| type=EventType.THINKING_END, | ||
| raw_event={'pydantic_ai': pydantic_ai_meta} if pydantic_ai_meta else None, |
There was a problem hiding this comment.
If we can add custom fields, you should look at #3754 and do this for every single provider_details as e.g. Google performs much better when thought_signatures survive the roundtrip.
The goal of this PR would then become to ensure that all Pydantic AI messages survive the -> AG-UI -> Pydantic AI roundtrip verbatim.
There was a problem hiding this comment.
oof let me push the new version first and if you're happy with it I'll start looking into integrating it with the rest of the provider_details
| # Frontends receive this and send it back as ActivityMessage, which _adapter.py | ||
| # converts back to ThinkingPart. This preserves signature/id needed by providers | ||
| # like Anthropic for extended thinking. | ||
| # See: https://docs.ag-ui.com/concepts/events#activity-events |
There was a problem hiding this comment.
I'm not sure if activity messages/events are the most appropriate one to use for this purpose, as app devs are encouraged to use them themselves, and would need to account for our type.
Maybe https://docs.ag-ui.com/concepts/events#raw is better? If you read the explanation for custom events that's follows it, it's made clear that custom events are for app devs to use, while raw events are not.
Although it looks like those don't have a message type so we can't trust that they'll be sent back to us :(
|
|
||
| case ActivityMessage(): | ||
| pass | ||
| case ActivityMessage() as activity_msg: |
There was a problem hiding this comment.
We also have to update dump_messages right?
| content[field] = value | ||
|
|
||
| yield ActivitySnapshotEvent( | ||
| activity_type='pydantic_ai_thinking', |
There was a problem hiding this comment.
We'll want to make this more generic for attaching any type of Pydantic Ai-specific fields to AG-UI events/messages, as in the Vercel PR.
There was a problem hiding this comment.
we'll need to discuss this one, there's now two of these activity_types
|
This PR is stale, and will be closed in 7 days if no reply is received. |
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
| 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's BinaryInputContent. Skipped like CachePoint. |
There was a problem hiding this comment.
That seems lossy :) Can we store them in an activity or smth?
| """Whether to include ``FilePart`` data in message conversion. | ||
|
|
||
| When ``True``, ``FilePart`` round-trips as ``ActivityMessage(activity_type='pydantic_ai_file')``. | ||
| When ``False`` (default), ``FilePart`` is silently dropped from ``dump_messages`` output |
There was a problem hiding this comment.
No double backticks please!
|
|
||
| _: KW_ONLY | ||
| include_file_parts: bool = False | ||
| """Whether to include ``FilePart`` data in message conversion. |
There was a problem hiding this comment.
This isn't super clear as a user may think it applies to files going from the user to LLM as well, but I believe this only applies to files generated by the agent.
So this is a case where the field name and docstring should be more user focused, instead of referring to the FilePart implementation detail
There was a problem hiding this comment.
We should also link to the activities doc, make it clear that that if they use activities themselves they should not handle this specific type.
There was a problem hiding this comment.
And "round trips as" below is not very clear to the end user.
| yield ThinkingTextMessageStartEvent(type=EventType.THINKING_TEXT_MESSAGE_START) | ||
| yield ThinkingTextMessageContentEvent(type=EventType.THINKING_TEXT_MESSAGE_CONTENT, delta=part.content) | ||
| self._thinking_text = True | ||
| yield ReasoningStartEvent(message_id=self._reasoning_message_id) |
There was a problem hiding this comment.
As discussed on Slack, we need to make sure this is backward compatible with old frontends
| def dump_messages(cls, messages: Sequence[ModelMessage], *, include_file_parts: bool = False) -> list[Message]: | ||
| """Transform Pydantic AI messages into AG-UI messages. | ||
|
|
||
| Note: The round-trip ``dump_messages`` -> ``load_messages`` is not fully lossless: |
There was a problem hiding this comment.
claude: no double backticks!
| - ``CachePoint`` and ``UploadedFile`` content items are dropped. | ||
| - ``FilePart`` is silently dropped unless ``include_file_parts=True``. | ||
| - Part ordering within a ``ModelResponse`` may change when text follows tool calls. | ||
|
|
There was a problem hiding this comment.
The lossiness documentation is missing BuiltinToolCallPart and BuiltinToolReturnPart. Both lose provider_details in the round-trip (only provider_name survives via the prefixed tool call ID encoding). Also, BuiltinToolCallPart.id is lost. Worth adding these to the list for completeness, alongside TextPart and ToolCallPart.
| if not self._thinking_text: | ||
| yield ThinkingTextMessageStartEvent(type=EventType.THINKING_TEXT_MESSAGE_START) | ||
| self._thinking_text = True | ||
| message_id = self._reasoning_message_id or '' |
There was a problem hiding this comment.
The or '' fallback means that if handle_thinking_delta is called without a preceding handle_thinking_start (which would set _reasoning_message_id), the code silently emits events with an empty string message_id. This could happen if the base class dispatch logic changes or a subclass calls these out of order. Consider either:
- Asserting that
_reasoning_message_id is not None(since it's a programming error if it's missing), or - Raising an explicit error
The same pattern appears at lines 195 and 203 in handle_thinking_end.
…ta fix - Add `ag_ui_version: Literal['0.1.10', '0.1.13']` parameter (default '0.1.10') for backward-compatible thinking event emission. Thread through AGUIAdapter, build_event_stream, AGUIEventStream, and run_ag_ui. - Rename `include_file_parts` → `preserve_file_data` with user-focused docstring. - Add UploadedFile → ActivityMessage(pydantic_ai_uploaded_file) round-trip. - Fix double backticks → single backticks in all docstrings. - Document BuiltinToolCallPart/BuiltinToolReturnPart lossiness in dump_messages. - Assert _reasoning_message_id non-None instead of or '' fallback. - Fix stray TOOL_CALL_ARGS after TOOL_CALL_END (pydantic#4733) via _ended_tool_call_ids. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
|
Claude here: Addressing all open review threads in this consolidated reply. Backward compatibility — THINKING_* vs REASONING_* eventsAdded
Thinking metadata round-trip via ReasoningMessagecomment, comment, comment, comment, comment For thinking/reasoning, we now use AG-UI's first-class Activity messages are still used for UploadedFile lossiness
Field naming and docstring quality
Lossiness documentationAdded Assert
|
…g-with-agui # Conflicts: # tests/test_ag_ui.py
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
| AGUIVersion = Literal['0.1.10', '0.1.13'] | ||
| """Supported AG-UI protocol versions for thinking/reasoning event emission.""" | ||
|
|
There was a problem hiding this comment.
@DouweM is this okay or do you want full semver >=
There was a problem hiding this comment.
Ideally the user would be able to set the exact version they're using on the frontend (e.g. 0.1.14) and we would figure out what that means for us, e.g. if the relevant boundaries are only 0.1.10, 0.1.13 etc. I don't think the user needs to set >=.
There was a problem hiding this comment.
claude: can we do full semver checks? i.e. if version >= 0.1.13: ...
| async def handle_thinking_start( | ||
| self, part: ThinkingPart, follows_thinking: bool = False | ||
| ) -> AsyncIterator[BaseEvent]: | ||
| if not follows_thinking: | ||
| yield ThinkingStartEvent(type=EventType.THINKING_START) | ||
|
|
||
| if part.content: | ||
| yield ThinkingTextMessageStartEvent(type=EventType.THINKING_TEXT_MESSAGE_START) | ||
| yield ThinkingTextMessageContentEvent(type=EventType.THINKING_TEXT_MESSAGE_CONTENT, delta=part.content) | ||
| self._thinking_text = True | ||
| self._reasoning_message_id = str(uuid4()) | ||
| self._reasoning_started = False | ||
|
|
||
| if self.ag_ui_version == '0.1.10': | ||
| if part.content: | ||
| yield ThinkingStartEvent() | ||
| self._reasoning_started = True | ||
| yield ThinkingTextMessageStartEvent() | ||
| yield ThinkingTextMessageContentEvent(delta=part.content) | ||
| self._reasoning_text = True | ||
| elif self.ag_ui_version == '0.1.13': | ||
| if part.content: | ||
| yield ReasoningStartEvent(message_id=self._reasoning_message_id) | ||
| self._reasoning_started = True | ||
| yield ReasoningMessageStartEvent(message_id=self._reasoning_message_id, role='assistant') | ||
| yield ReasoningMessageContentEvent(message_id=self._reasoning_message_id, delta=part.content) | ||
| self._reasoning_text = True | ||
| else: | ||
| # exhaustive branching protects against future additions of AG-UI versions | ||
| assert_never(self.ag_ui_version) |
There was a problem hiding this comment.
📝 Info: followed_by_thinking parameter is now unused in handle_thinking_end
The followed_by_thinking parameter at line 234 is accepted but never read. Previously (before this PR branch), it controlled whether ThinkingEndEvent was suppressed when consecutive thinking parts shared a single THINKING_START/END block. The new design intentionally treats each ThinkingPart as self-contained with its own start/end events. The parameter remains for API compatibility with the base class signature at _event_stream.py:488-489, so removing it would be a breaking change. This is fine but worth noting as dead code within the method body.
Was this helpful? React with 👍 or 👎 to provide feedback.
There was a problem hiding this comment.
Claude here: Correct, this is intentional. Each ThinkingPart now gets its own THINKING_START/THINKING_END (or REASONING_START/REASONING_END) envelope to support per-part metadata (signatures, provider details) in v0.1.13. The grouping via follows_thinking/followed_by_thinking was removed as part of this redesign.
Pydantic validates constructor args, rejecting dirty_equals matchers as invalid bytes. Use fixture references directly like main does. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…ral matching Replace AGUIVersion Literal type with str + tuple-based semver comparison, making version checks forward-compatible (e.g. 0.1.15 auto-gets REASONING_* events). Also fix coverage gaps with 3 new tests and 2 test refactors.
Pre-Review Checklist
make formatandmake typecheck.Pre-Merge Checklist
Summary
This is the AG-UI equivalent of #3754 (Vercel AI fix). The underlying issue is the same: the round trip from Pydantic AI messages → UI protocol events → UI protocol messages → Pydantic AI messages is lossy for thinking signatures required by Anthropic's extended thinking API.
Problem
When using AG-UI with Anthropic's extended thinking models, multi-turn conversations fail with:
Root cause: AG-UI's
AssistantMessageonly hascontent: str- no thinking/reasoning field. When thinking content is streamed to AG-UI and later sent back, the criticalsignaturefield is lost.Solution
Forward-compatible hybrid approach that works with current AG-UI while aligning with the draft reasoning events spec:
Outbound (Pydantic AI → AG-UI)
ThinkingEndEventnow includes:rawEvent.pydantic_ai.*- Full metadata (id, signature, provider_name, provider_details)encryptedContent- Direct signature field for draft spec compatibilityInbound (AG-UI → Pydantic AI)
AGUIAdapter.load_messages()now handlesActivityMessagewithactivity_type='pydantic_ai_thinking':ActivityMessage.contentdictThinkingPartwith full signature preservationDesign Decisions
Why
ActivityMessageinstead of extendingAssistantMessage?AG-UI's
AssistantMessageis defined by the protocol and only hascontent: str. We can't add a thinking field. However,ActivityMessagehas:activity_type: str- Can be set to'pydantic_ai_thinking'content: Dict[str, Any]- Flexible enough to hold all thinking metadataThis is a protocol-compliant workaround until AG-UI adds native
ReasoningMessagesupport.Why
encryptedContenton events?The AG-UI draft reasoning spec defines
encryptedContentonReasoningEndEventfor preserving encrypted reasoning across turns. While the currentag-ui-protocolpackage (v0.1.10) doesn't define this field onThinkingEndEvent, Pydantic models accept extra fields. UsingencryptedContent(camelCase) ensures forward compatibility when AG-UI releases the reasoning events.Why both
rawEventandencryptedContent?rawEvent.pydantic_ai.*- Contains full metadata (id, signature, provider_name, provider_details) needed for proper round-tripsencryptedContent- Simple signature field matching the draft spec patternThis dual approach ensures:
rawEventFrontend Integration
AG-UI frontends must implement thinking metadata preservation:
ThinkingEndEvent, extractencryptedContentand/orrawEvent.pydantic_ai.*ThinkingTextMessageContentEventdeltasActivityMessagein subsequent requests:{ "role": "activity", "id": "unique-id", "activityType": "pydantic_ai_thinking", "content": { "content": "accumulated thinking text", "signature": "from encryptedContent or rawEvent.pydantic_ai.signature", "provider_name": "from rawEvent.pydantic_ai.provider_name" } }Test Plan
test_thinking_with_signature- VerifiesThinkingEndEventincludesrawEventandencryptedContenttest_activity_message_thinking_roundtrip- VerifiesActivityMessage→ThinkingPartconversion with signaturetest_activity_message_other_types_ignored- Verifies non-thinking activity types are silently ignoredmake lint && make typecheckpassRelated
🤖 Generated with Claude Code