AG-UI - How to force a StateSnapshotEvent at the end of the agent execution?

[danielb-gyde]

Question

I'm defining a test state and a route inside our FastAPI backend via pydantic_ai.ag_ui.run_ag_ui

class TestState(BaseModel):
    userName: str
    counter: int


class TestStateSnapshot(BaseModel):
    testState: TestState = Field(
        default_factory=lambda: TestState(userName="Test User", counter=1),
        description="The current test state",
    )

@router.post("/chat")
async def run_agent(request: Request) -> Response:
    # [boilerplate]

    state = (
        TestStateSnapshot.model_validate(run_input.state)
        if run_input.state
        else TestStateSnapshot()
    )

    event_stream = run_ag_ui(
        main_agent,
        run_input,
        accept=accept,
        on_complete=on_request_complete,
        deps=StateDeps(state),
    )

    return StreamingResponse(event_stream, media_type=accept)

and I have a dummy-agent defining some tools and a sub-agent for agent orchestration:

main_agent = Agent(
    name="Routing Agent",
    model=model,
    system_prompt=MAIN_AGENT_PROMPT,
    deps_type=StateDeps[TestStateSnapshot],
)


@main_agent.instructions
def beg_llm_to_update_state(ctx: RunContext[StateDeps[TestStateSnapshot]]) -> str:
    return "always call update_state when you're done"


@main_agent.tool
async def roll_dice(ctx: RunContext[StateDeps[TestStateSnapshot]]) -> DiceAgentResponse:
    """Roll a six-sided die and return the random result."""
    result = 6
    ctx.deps.state.testState.counter = result
    return DiceAgentResponse(roll=result)


@main_agent.tool
async def update_state(
    ctx: RunContext[StateDeps[TestStateSnapshot]],
) -> StateSnapshotEvent:
    return StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=ctx.deps.state)


sub_agent = Agent(
    name="Pink Fluffy Unicorn agent",
    model=model, 
    deps_type=StateDeps[TestStateSnapshot],
)

@main_agent.tool
async def pink_fluffy_unicorns(ctx: RunContext[StateDeps[TestStateSnapshot]]):
'''Test agent that changes the count to 42 in one of it's tools'''
    response = await sub_agent.run(
        message_history=ctx.messages,
        deps=ctx.deps,
    )
    return response


@sub_agent.tool
async def get_pink_fluffy_unicorn_info(
    ctx: RunContext[StateDeps[TestStateSnapshot]], increment: int
):
    state = ctx.deps.state
    state.testState.counter = 42

In order to update the state and send it back to my frontend, main_agent needs to call the tool update_state at the end of it's execution. Here's the correct call stack:

 Routing Agent run
   chat gpt-4o
   running 2 tools
     running tool: roll_dice
     running tool: update_state
   chat gpt-4o

However, the update_state call is missing more often than not, even if I prompt the agent to do it. It seems to work more or less consistently if I specifically ask it in the chat. Here are some example screenshots:

Asking the agent to roll the dice, update_state is not called:

   running 1 tool
     running tool: roll_dice
   chat gpt-4o

Asking the agent to roll the dice and then update the state:

 Routing Agent run
   chat gpt-4o
   running 2 tools
     running tool: roll_dice
     running tool: update_state
   chat gpt-4o

Obviously not what you'd want for a user-facing application

It gets even flimsier when I ask it to call a sub-agent. In this example update_state is called before the tool call that changes the state to 42.

 Routing Agent run
   chat gpt-4o
   running 2 tools
     running tool: pink_fluffy_unicorns
       Pink Fluffy Unicorn agent run
     running tool: update_state
       running tool: pink_fluffy_unicorns
         Pink Fluffy Unicorn agent run
         running 2 tools
           running tool: pink_fluffy_unicorns
           running tool: update_state
         chat gpt-4o
         running 1 tool
           running tool: get_pink_fluffy_unicorn_info
         chat gpt-4o
   chat gpt-4o

My Question(s):

• Is updating the state via a tool call the only way to do it? Couldn't I send a custom event via run_ag_ui's on_complete callback parameter? Send custom events from tool function, for AG-UI and event_stream_handler #2382 (comment) <- this workaround seems to be the closest thing I've found so far
• CopilotKit Inserts a SystemMessage at the beginning of the message history. How does that interact with Agents (and sub-agents)?
• What's the best practice for this type of problem? (I don't really mind highly-opinionated takes on this)

If there's any input you could give me on this issue, I'd be stoked.

Additional Context

No response

[DouweM]

@danielb-gyde Indeed #2382 and passing the event bus to the on_complete handler would be the way to go. Since the handler will also want access to the deps (to get state), and tools will also want access to the event bus, we could pass in the entire RunContext, and handle the event_bus similar to how deps.state is auto-populated if it exists, with a new deps.event_bus.

As a workaround, I suggest following that custom event bus example. But if you happen to have the capacity to submit a PR to implement what I described above, that'd be amazing!

[ChuckJonas]

CopilotKit Inserts a SystemMessage at the beginning of the message history. How does that interact with Agents (and sub-agents)?

This was a major issue for us and we actually ended up parsing the incoming messages ourselves and filtering out everything except for the new user Request(s).

IMO, allowing the FE to control message history (especially system messages!) is a very bad design choice in AG-UI. It opens our system up to abuse and makes prompt injection much easier.

So instead we keep track of the message history on our side and append the latest messages from the AG-UI request (you also have to handle FE tools and state). However, this means we not able to use the built in run_ag_ui and instead we wrote a very simple agui_publisher that we can invoke to send AG-UI messages when different agent events occur.

This also allows us more flexibility in what we send and the ability to run multi-agent setups.

[DouweM]

@ChuckJonas I'd love to add support for those workarounds natively, if you'd like to share the code for inspiration (or submit a PR). A flag to ignore the system message, or some kind of history processor, would be great.

[ChuckJonas]

@DouweM ya, i'm trying to slowly work our specific implementation towards a general purpose one, at which point I'd love to review and see if there's something worth bringing into pydantic-ai.

Here's the list of use cases/concerns we've had to handle so far:

1. Parsing & Control over AG-UI request:

• Message History vs New Requests
• Removal of system messages
• Validation / Whitelisting of FE tools

2. Control over published AG UI events

• Control over which tool-calls & tool-responses get published as AG-UI messages
• Control over which assistant messages get published to AG-UI
• Control over State, Custom thinking messages, etc

3. Proper handling of DeferredTools (FE Tools)
4. Error handling

At a high level, our current implementation looks something like this:

#NOTE: modified code, incomplete / only a rough representation
async def agent_endpoint(
    request: Request,
) -> Response:
    try:
        body = await request.body()
        data = json.loads(body)
        run_input = RunAgentInput.model_validate(data)

        session_id = request.headers.get("x-session-id")
        if not session_id:
            raise ValueError("x-session-id header is required")
       
        ag_ui_messages = _messages_from_ag_ui(run_input.messages)

        thread_id = run_input.thread_id

        logger.info(f"Executing thread_id: {thread_id} run_id: {run_input.run_id}")

        # remove everything except for the new messages.  
        # Full message history will be pulled from persistent storage (via `thread id` by the agents as needed.
        new_messages = select_message_parts(
            get_new_messages(ag_ui_messages),
            exclude=[{"kind": "request", "parts": "system-prompt"}],
        )

        # Parse FE Actions tools from AG UI request
        fe_toolset: Optional[_AGUIFrontendToolset] = None
        if run_input.tools:
            fe_toolset = _AGUIFrontendToolset(run_input.tools)
            fe_toolset = fe_toolset.filtered(fe_tool_filter)  # type: ignore

        if len(new_messages) == 0:
            raise ValueError("No new inputs found input")

        logger.info(
            f"new_messages:\n{RootModel(new_messages).model_dump_json(indent=2)}"
        )

        publisher = AGUIPublisher(run_input)

        async def on_tool_call(tool_name: str, tool_args: Any):
            logger.info(f"Tool call: {tool_name} with args: {tool_args}")
            await publisher.send_state_snapshot(get_tool_call_progress_message(tool_name))

        async def before_final_response():
            await publisher.send_state_snapshot(get_stage_message("response_generation"))
            await publisher.start_response()

        async def on_final_stream_part(content: str):
            await publisher.send_response_part(content)

        async def after_final_response():
            await publisher.end_response()
            await publisher.send_state_snapshot({})

        async def on_deferred_tool_call(tool_call: ToolCallPart):
            await publisher.send_state_snapshot({})
            await publisher.send_tool_call(tool_call)

        async def work(pub: AGUIPublisher):
            
            # `run_agents` streams multiple agents with custom event callbacks using `iter` and `node.stream`.  
            #    - some agents connect to these callbacks, others are internal only
            await run_agents(
                thread_id,
                new_messages=new_messages,
                publisher=pub,
                fe_toolset=fe_toolset,
                on_tool_call=on_tool_call,
                before_final_response=before_final_response,
                on_final_stream_part=on_final_stream_part,
                after_final_response=after_final_response,
                on_deferred_tool_call=on_deferred_tool_call,
            )

        def get_error_message(e: Exception) -> str:
            return "Something went wrong.  Please try again or contact support if this persists."

        return StreamingResponse(
            stream_with(publisher, work, get_error_message),
            media_type="text/event-stream",
        )
    except (
        Exception
    ) as e:  # NOTE: Copilot kit isn't actually handling http error codes right now :(
        logger.exception("Error in radvice agent request")
        return JSONResponse({"error": str(e)}, status_code=400)

sorry, if this is too off-topic for the original thread, we can move this conversation elsewhere

[DouweM]

@ChuckJonas Thanks, that's helpful, if indeed a bit offtopic :) When you get a chance, I'd appreciate issues for the specific knobs we should add to our implementation

[danielb-gyde]

Thanks for the answers.

Here's what we ended up on for now:

We're calling update_state() at the end of every tool call of our main agent. The pattern looks something like this:

@main_agent.tool
async def sub_agent(ctx: RunContext[StateDeps[AgentStateSnapshot]]) -> StateSnapshotEvent:
    await sub_agent.run(message_history=ctx.messages, deps=ctx.deps)
    return await update_state(ctx)

async def update_state(
    ctx: RunContext[StateDeps[AgentStateSnapshot]],
) -> StateSnapshotEvent:
    return StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=ctx.deps.state)

this is IMHO still fairly clean and we don't have to rely on prompting anymore. We also don't have to modify any sub-agents if they want to modify the state.

The obvious downside is that we can't have live state updates while the sub-agent is running.

[DouweM]

The obvious downside is that we can't have live state updates while the sub-agent is running.

@danielb-gyde That'll come in #2382 which I expect to get to later this month!

[DouweM]

@ChuckJonas With the changes in #2923 (that are not just for Vercel AI, but also AG-UI), everything you asked for in #2906 (comment) should now be possible by using specific methods on AGUIAdapter (see https://ai.pydantic.dev/ui/overview/#advanced-usage) and/or subclassing AGUIAdapter and AGUIEventStream to add your own event processing logic. Let me know if this doesn't give you the knobs you're looking for yet.

That PR also lets the on_complete handler yield protocol-specific events like StateSnapshotEvent, but it does not yet have access to deps as currently with AG-UI state, the deps object you pass into the run would not be the same one that gets state set on it. Let me fix that.