Add AGUIAdapter.dump_messages to convert Pydantic AI messages to AG-UI messages
#3474
Conversation
|
|
||
| id_index = 0 | ||
|
|
||
| def get_next_id() -> str: |
There was a problem hiding this comment.
This is unfortunately more awkward than it would ideally be because metadata isn't universally available on all "MessageParts" and a ModelRequst/Response may contain multiple AG_UI Messages
@DouweM feels like metadata should maybe just be a base property on MessagePart?
There was a problem hiding this comment.
@ChuckJonas Yeah I think we'll need to have some dict for arbitrary values on each part. There's some related discussion in #3453 (if you go down a few comments) as well, although that focuses on developer-provided fields to be sent to provider APIs, whereas here the purpose is to store some information we receive from outside of Pydantic AI (whether that's a provider API or frontend) so that we can use it to losslessly send data back to that same provider/frontend. So we may need 2 dicts 😅
In any case, we're not going to be able to implement that in this PR, but I think we'll need it sooner rather than later.
There was a problem hiding this comment.
Note that we now have provider_details field in a bunch of places that we can use (and add where we don't have it yet).
| __all__ = ['messages_from_ag_ui', 'messages_to_ag_ui'] | ||
|
|
||
|
|
||
| def messages_from_ag_ui(messages: Sequence[Message]) -> list[ModelMessage]: |
There was a problem hiding this comment.
Feels like there could be a more "standardized" ModelMessage translator interface with to & from functions. This would be useful for converting between all the different model formats in general.
For now, I just kept it simple for now with the messages_from_ag_ui & messages_to_ag_ui functions.
There was a problem hiding this comment.
My comment above could've been here -- that's pretty much what I was thinking with load/dump methods on the adapter.
ChuckJonas
force-pushed
the
messages_from_ag_ui
branch
from
055a098 to
1080d91
Compare
| 'AGUIAdapter', | ||
| 'AGUIEventStream', | ||
| 'messages_from_ag_ui', | ||
| 'messages_to_ag_ui', |
There was a problem hiding this comment.
Quoting from Slack:
The way I imagined it, we'd have
load_messagesanddump_messagesclassmethods on the adapters, which is being implemented for Vercel in #3392. Some of the comments on that PR about lossless roundtrip may be relevant to AG-UI as well, although I'm not sure they support arbitrary metadata, and I'd be OK with an initial version that's best-effort
| __all__ = ['messages_from_ag_ui', 'messages_to_ag_ui'] | ||
|
|
||
|
|
||
| def messages_from_ag_ui(messages: Sequence[Message]) -> list[ModelMessage]: |
There was a problem hiding this comment.
My comment above could've been here -- that's pretty much what I was thinking with load/dump methods on the adapter.
|
|
||
| id_index = 0 | ||
|
|
||
| def get_next_id() -> str: |
There was a problem hiding this comment.
@ChuckJonas Yeah I think we'll need to have some dict for arbitrary values on each part. There's some related discussion in #3453 (if you go down a few comments) as well, although that focuses on developer-provided fields to be sent to provider APIs, whereas here the purpose is to store some information we receive from outside of Pydantic AI (whether that's a provider API or frontend) so that we can use it to losslessly send data back to that same provider/frontend. So we may need 2 dicts 😅
In any case, we're not going to be able to implement that in this PR, but I think we'll need it sooner rather than later.
|
@ChuckJonas FYI there's an older PR at #3068; I recommend checking out the code and comments to see if they apply here |
ChuckJonas
force-pushed
the
messages_from_ag_ui
branch
from
1080d91 to
257e448
Compare
|
This PR is stale, and will be closed in 3 days if no reply is received. |
|
@ChuckJonas Hey Charlie, are you still interested in working on this? |
|
@ChuckJonas Ah you're totally right, I had missed that most recent push. Let me get the other PR in first then, and then I'll come back to this one next week. |
ChuckJonas
force-pushed
the
messages_from_ag_ui
branch
from
257e448 to
ac319e6
Compare
DouweM
left a comment
DouweM
left a comment
There was a problem hiding this comment.
@ChuckJonas Thanks Charlies; #3392 has now been merged so you can rebase! Also please look at the implementation there and see if there are any other details we should copy. I commented on the things I noticed but there may be more.
| content = part.content if isinstance(part.content, str) else str(part.content) | ||
| ag_ui_messages.append(UserMessage(id=msg_id, content=content)) | ||
|
|
||
| elif isinstance(part, ToolReturnPart): |
There was a problem hiding this comment.
We need to handle RetryPromptPart as well. Let's add else: assert_none(part) at the end so that the type checker will ensure we're exhaustive in terms of all types part could have.
There was a problem hiding this comment.
#3392 has now been merged; you can check out the implementation there
| @staticmethod | ||
| def _convert_request_parts( | ||
| parts: Sequence[ModelRequestPart], | ||
| ag_ui_messages: list[Message], |
There was a problem hiding this comment.
Do we have to modify the list in place, or could we just return a new list?
| """Convert ModelResponse parts to AG-UI messages.""" | ||
|
|
||
| # Group consecutive assistant parts (text, tool calls) together | ||
| def is_assistant_part(part: ModelResponsePart) -> bool: |
There was a problem hiding this comment.
Isn't everything in ModelResponsePart an assistant part?
Edit: Ah this is specifically because AG-UI doesn't understand builtin tool returns that come from the assistant. Let's make that explicit.
| else: | ||
| # Each non-assistant part becomes its own message | ||
| for part in parts_list: | ||
| if isinstance(part, BuiltinToolReturnPart): |
There was a problem hiding this comment.
We need to handle FilePart and ThinkingPart as well
| @staticmethod | ||
| def _make_builtin_tool_call_id(provider_name: str | None, tool_call_id: str) -> str: | ||
| """Create a full builtin tool call ID from provider name and tool call ID.""" | ||
| return f'{BUILTIN_TOOL_CALL_ID_PREFIX}|{provider_name}|{tool_call_id}' |
There was a problem hiding this comment.
Can we dedupe this with the analogous logic in the event stream?
| @staticmethod | ||
| def _convert_tool_call(part: ToolCallPart) -> ToolCall: | ||
| """Convert a ToolCallPart to an AG-UI ToolCall.""" | ||
| args_str = part.args if isinstance(part.args, str) else json.dumps(part.args) |
There was a problem hiding this comment.
We should use part.args_as_json_str()
| @staticmethod | ||
| def _convert_builtin_tool_call(part: BuiltinToolCallPart) -> ToolCall: | ||
| """Convert a BuiltinToolCallPart to an AG-UI ToolCall.""" | ||
| args_str = part.args if isinstance(part.args, str) else json.dumps(part.args) |
| ) | ||
|
|
||
| @staticmethod | ||
| def _serialize_content(content: Any) -> str: |
There was a problem hiding this comment.
Can we use ToolReturnPart.model_response_str() and RetryPromptPart.model_response()?
DouweM
changed the title
AGUIAdapter.dump_messages to convert Pydantic AI messages to AG-UI messages
2 participants
Adds
dump_messagesfunction for AG-UI (introduced in https://github.com/pydantic/pydantic-ai/pull/3392/files).Assuming that PR goes out soon, I can rebase this branch after and resolve conflicts.
NOTE: I attempted to implement preservation of the AG_UI id's but the complexity and lack of
metadata(as discussed in the outdated comment below), made it more trouble than it was worth.Sounds like this is consistent with the behavior in the "Vercel AI" implementation, although I didn't add a
_id_generator()for testing purposes.