Handle complex agent interactions in `ag_ui` module

[lukekh]

Description

I've been trying to get some stuff to work with AG-UI and run into some annoying issues when trying to get more complex patterns to work with the ag_ui classes and methods exposed.

Here are some suggestions for things I think the module is currently lacking:

1. Parsing and holding the state: I can't see a clean way to use the current AGUIApp class or handle_ag_ui_request/run_ag_ui methods to:
   • convert an incoming request into data that pydantic_ai agents can natively handle without writing my own tedious methods which are extremely close to the methods in the ag_ui module
   • allow multiple agents to do things relating to an incoming request and have access to the events/messages that have been sent by other agents. (This potentially raises some design questions about the AGUIApp class having a single agent attribute.)
   • expose/handle the shared state data object (e.g. AgentState patterns) more cleanly
   • allow for some more interesting tool handling (I'm imagining being able to filter/alter the tools that are exposed)
2. Better error handling: My AGUIApp falls over due to errors all the time because I can't intercept the way the handle_ag_ui_request/run_ag_ui methods deal with an error
   • this might be a skill issue, but it's been a headache for me. I am having it fall over despite defining a retry_client a la retries
   • should there be a ag_ui method that can consume an arbitrary AsyncIterator[str] stream and interpret it as an agent output?
3. Methods for emitting events without tools: This seems to be a common issue for a lot of people using the AG-UI framework (take this for example) and not one I feel pydantic_ai has a clean solution for yet.
   • I think adding something that allows a really simple way to emit an event without it being a tool would be very useful
4. Support agents with structured outputs: this greatly impedes plugging agents into AGUIApps in creative ways.
   • There are ways around this and I've already litigated this issue here, but I think it should still be mentioned in this laundry list of ideas
5. Make it easier to return data in tool calls to front end client: It's a bit clunky getting your agent to sync up the results of tool calls with frontend actions at the moment.
   • this is another one that might just be a skill issue, but something I've found challenging presently is an agent rendering a large table of data to the frontend without that dataset living in my agents context window in subsequent requests. It may be that the docs just need some good examples (e.g. I've been pointed to this repo which I'm in the process of working through) for me to wrap my head around how this is supposed to work. But are we sure this is the pattern we want for simple frontend actions?
   • There are work arounds (e.g. AgentState patterns, etc), but I haven't found any patterns that are particularly nice to use yet.
   • Upon doing a little bit more reading, it seems most of what I'm asking for is answered in this bit of doco: https://ai.pydantic.dev/ag-ui/#events so perhaps this point is moot!
6. More explicit control of actions between events: run_ag_ui elegantly handles the basics, but perhaps it would be useful to have some with LifecycleEvent(...) as le: context management to begin/end event streams.

I hope I haven't overstepped here by making all of these suggestions and I want to conclude by saying that I've been so impressed with what this library has allowed me to do and I owe the main contributors a huge debt of gratitude for all the hard work they've put in!

References

1. Pydantic AI AG-UI Docs
2. AG-UI Protocol docs
3. ag_ui module
4. AgentState pattern
5. CrewAI custom copilotkit_emit_message method: Looking at the Issues and Discord of CopilotKit, this comes up a lot in user related issues
6. ag_ui module not working with structured output_type #3011 : A previous issue I raised about output_type causing problems
7. pydantic-ai x ag-ui demo that I was pointed to
   • And my e.g. Are we sure this is the pattern we want for simple frontend actions?

[DouweM]

I hope I haven't overstepped here by making all of these suggestions

@lukekh Definitely not, we can't make it better without feedback like this!

• Parsing and holding the state: I can't see a clean way to use the current AGUIApp class or handle_ag_ui_request/run_ag_ui methods to:

  • convert an incoming request into data that pydantic_ai agents can natively handle without writing my own tedious methods which are extremely close to the methods in the ag_ui module

Do you have an example of the type of processing you'd like to do? run_ag_ui takes a RunAgentInput object, that you could modify before passing it in. If you're thinking about server-side-stored message history specifically, I'd be open to adding message_history to the run_ag_ui and handle_ag_ui_request methods.

• allow multiple agents to do things relating to an incoming request and have access to the events/messages that have been sent by other agents. (This potentially raises some design questions about the AGUIApp class having a single agent attribute.)

An example would be useful. Note that the on_complete callback handle gives you access to the full AgentRunResult so you can access all_messages() etc.

• expose/handle the shared state data object (e.g. AgentState patterns) more cleanly

What does "more cleanly" mean to you? :) It's currently stored on deps, meaning you can access it from all functions registered to the agent that take RunContext.

• allow for some more interesting tool handling (I'm imagining being able to filter/alter the tools that are exposed)

That's a good idea, we can let you provide a function that takes the _AGUIFrontendToolset that holds all the frontend tools and can return a new toolset wrapping it, which can apply filters or other modifications.

• Better error handling: My AGUIApp falls over due to errors all the time because I can't intercept the way the handle_ag_ui_request/run_ag_ui methods deal with an error

  • this might be a skill issue, but it's been a headache for me. I am having it fall over despite defining a retry_client a la retries

Do you have an example of the type of error you see? They should all turn into RunErrorEvents sent to AG-UI at the least.

• should there be a ag_ui method that can consume an arbitrary AsyncIterator[str] stream and interpret it as an agent output?

Can you elaborate? What would the iterated strings hold? What would they be interpreted as?

• Methods for emitting events without tools: This seems to be a common issue for a lot of people using the AG-UI framework (take this for example) and not one I feel pydantic_ai has a clean solution for yet.

  • I think adding something that allows a really simple way to emit an event without it being a tool would be very useful

I plan to work on #2382 this month, but that's specifically about emitting custom events from places where RunContext is available (like tools) and the on_complete handler. Where else would you like to be able to emit events? I can see the value of something that wraps run_ag_ui (or some new method that yields event classes rather than str-encoded events) and lets you process/filter/add events before sending them on. Is that what you were referring to in the previous point?

• Support agents with structured outputs: this greatly impedes plugging agents into AGUIApps in creative ways.

  • There are ways around this and I've already litigated this issue here, but I think it should still be mentioned in this laundry list of ideas

In tool output mode, I can see why it would make sense to emit the ToolCallResultEvent. We could always do that or at least make it easy to do that from on_complete.

In the other output modes (native, prompted) the structured output comes in a text message from the model that's sent on to the frontend as a text message as well. In that scenario, I suppose you wouldn't want the text message event (we could skip it), and would want to send the structured output to the frontend with a custom event from on_complete. Would that work for you?

• Make it easier to return data in tool calls to front end client: It's a bit clunky getting your agent to sync up the results of tool calls with frontend actions at the moment.

  • this is another one that might just be a skill issue, but something I've found challenging presently is an agent rendering a large table of data to the frontend without that dataset living in my agents context window in subsequent requests. It may be that the docs just need some good examples (e.g. I've been pointed to this repo which I'm in the process of working through) for me to wrap my head around how this is supposed to work. But are we sure this is the pattern we want for simple frontend actions?
  • There are work arounds (e.g. AgentState patterns, etc), but I haven't found any patterns that are particularly nice to use yet.
  • Upon doing a little bit more reading, it seems most of what I'm asking for is answered in this bit of doco: ai.pydantic.dev/ag-ui#events so perhaps this point is moot!

Yeah I think this is covered by ToolReturn with events in the metadata. With #2382, you could also yield the custom events explicitly and then just return the text that should go to the model.

• More explicit control of actions between events: run_ag_ui elegantly handles the basics, but perhaps it would be useful to have some with LifecycleEvent(...) as le: context management to begin/end event streams.

Can you elaborate? When would you want to send your own lifecycle events?

[ElvisTheKing]

I hope you wouldn't mind if i share my thoughts, i was thinking on something in this area for couple of days.

To move from simple workflow which is demonstrated in every example, namely:

1. user asks questions using text prompt
2. agent runs some tools, streams some events using ag-ui
3. frontend happily renders events and changed state
4. goto 1

we need

Persistence:
In production(ish) solution we need to add persistence of messages and state, and don't fight either with pydantic-ai or ag-ui frontend implementation while doing this.

So production flow will look like this:

1. user ask question (ag ui sends all messages, new one already with generated id)

2. backend persists new message somewhere, processes new message and history (and ignores old ones, replacing them with its own from persistence layer, maybe altered using summarization or even with some smart memory layer) -

   This can be implemented just working with RunAgentInput instance, but there is some misalignment between pydantic and ag-ui, namely messages loose it's ids

3. backend runs agent, produces events, final response, sends to frontend.

4. Somewhere here we need to persist newly produced messages and there's currently no way to get new messages with ids generated other than to process response stream, decode it, and get generated message ids ( new ids are generated in ag-ui right before new messages are encoded, there's no access to ids in on_complete hook). Not sure how to approach this, maybe on_message hook added to

   pydantic-ai/pydantic_ai_slim/pydantic_ai/ag_ui.py

   Line 378 in efaf6be

   yield encoder.encode(event)

State changes:
I think state changes can already be handled, but require some manual work to correctly change state:

1. change state
2. produce JsonDiff and emit StateDeltaEvent
3. probably persist changed state

But it's up to developer how to approach this, not sure if it should be integrated to pydantic-ai.

Only issue i can see here is that StateDeltaEvent (or any other event) can be emitted only at the end of tool call, which is a problem if tool call is really long.
It's easy enough to work around with second event stream, which can be accessed via deps and then mixed in to main steam:

from asyncio import Queue

@dataclass
class CustomDeps(Generic[StateT]):
    state: StateT
    event_queue: Queue[BaseEvent] = field(default_factory=Queue)

There's some boilerplate, maybe this can be added to library

[DouweM]

4. Somewhere here we need to persist newly produced messages and there's currently no way to get new messages with ids generated other than to process response stream, decode it, and get generated message ids ( new ids are generated in ag-ui right before new messages are encoded, there's no access to ids in on_complete hook).

@ag14774 Can you help me understand what you need the AG-UI IDs for? I was expecting them to not be meaningful or need to be stable, just unique and matching between start/end events etc.

1. change state
2. produce JsonDiff and emit StateDeltaEvent
3. probably persist changed state

I was talking on Slack with @dorukgezici who has built a solution on this that we could consider incorporating in Pydantic AI. Doruk, would you mind filing an issue for that with context on what you built?

Only issue i can see here is that StateDeltaEvent (or any other event) can be emitted only at the end of tool call, which is a problem if tool call is really long.

Yep, custom events (including AG-UI events) during the tool call are going to be implemented in #2382!

[dorukgezici]

@DouweM shared it! Could you comment the problem you saw with deps, not sure I understood it fully. I can turn it into a PR afterwards #3078

[ElvisTheKing]

4. Somewhere here we need to persist newly produced messages and there's currently no way to get new messages with ids generated other than to process response stream, decode it, and get generated message ids ( new ids are generated in ag-ui right before new messages are encoded, there's no access to ids in on_complete hook).

@ag14774 Can you help me understand what you need the AG-UI IDs for? I was expecting them to not be meaningful or need to be stable, just unique and matching between start/end events etc.

This was probably meant for me, so I'll answer

Some kind of message id is needed to keep backend persistence layer and frontend message history in sync, in our workflow we generate ids for user messages on frontend and current ag-ui implementation generates ids for agent messages in _handle_model_request_event, which is ok, just need a way to access ids somewhere on backend after generation. At the moment we just use patched run_ag_ui (on event callback added)