State Management with AG-UI

[tschokokuki]

Dear agent-framework community,

my goal is to deterministically edit the state of an agent via a tool.

However, all I found in the demos was, how the agent state is edited by the predict_state via the predict_state_config.
As described in this demo: https://learn.microsoft.com/en-us/agent-framework/integrations/ag-ui/state-management?pivots=programming-language-python

However, the predict_state consists of non-deterministically generated tool arguments.
Thus, to my understanding predict state should be used for optimistic UI updates but not define final state.

Motivation:

• Lets say we want a get_weather tool to write the current weather into a global state shared with Backend and Frontend.
  The state cannot be predicted by the tool call arguments by the llm, but must be retrieved during the tool call.
• Furthermore for Human In The Loop (HITL) it would be required to use the Tool's returned results as a state, as the predict_state is emitted instantly as soon as available.

I found this pydantic_ai demo, where the tool emits a StateDeltaEvent, but the agent_framework tools would not let me return events.
Source: https://docs.copilotkit.ai/pydantic-ai/shared-state/predictive-state-updates

@agent.tool_plain
async def update_steps(steps: list[str]) -> StateSnapshotEvent:
    """Update the steps of the agent."""
    return StateSnapshotEvent(
        type=EventType.STATE_SNAPSHOT,
        snapshot={
            "observed_steps": steps
        }
    )

How would this be done using agent_framework?

Thanks for your help!

[moonbox3]

Hi @tschokokuki, thanks for bringing this up. It is a current gap. As a potential workaround, one could intercept the agent stream and inject StateSnapshotEvent when specific FunctionResultContent is available. Something like:

from ag_ui.core import StateSnapshotEvent

async def intercept_stream():
    async for update in agent.run_stream(messages):
        yield update
        for content in update.contents:
            if isinstance(content, FunctionResultContent) and content.name == "get_weather":
                # Emit state based on actual tool result
                yield StateSnapshotEvent(snapshot={"weather": content.result})

Otherwise we will need to add handling for tools to directly emit state events. Would you be interested in contributing to this? The relevant code is in agent_framework_ag_ui/_events.py, specifically _handle_function_result_content() which currently doesn't extract state from tool results.

[tschokokuki]

Hi @moonbox3,

Thank you very much for pointing me in the right direction.

I got the workaround running as you suggested.
I also began with extending the _handle_function_result_content().

From my side, I added one more requirement: I want to pass a custom message for the LLM rather than relying on a default one.

I wonder, what would be the best return format for the Tools in context of the agent_framework project.

Option 1: StateDeltaEvent

Pro:

• any Event could be emitted in _handle_function_result_content() without conversions

Cons:

• The StateDeltaEvent seems to create issues when serializing the agent`s thread state. => result: null
• result_message adds a new convention which must be picked up in _handle_function_result_content()

@ai_function()
def get_weather() -> StateDeltaEvent:
    weather = "sunny"
    return StateDeltaEvent(
        delta=[{"op": "replace", "path": "/weather", "value": weather}],
        result_message=f"Weather updated to: {weather}")

Option 2: Dictionary with a marker

Pro:

• The entire Dict is compatible with storing it as a result in the agent' thread state => result = "{ "agui_event": True, "event_type": ..... }"

Cons:

• Actually storing only the "result_message" as a result in thread state would be sufficient
  => This creates unwanted artefacts when converting the result from agent_framework_messages to ag_ui_messages
  => Thus we could extend the converting function to filter out events, which adds more complexity.

@ai_function()
def get_weather() -> Dict:
    weather = "sunny"
    return {
        "__agui_event__": True,
        "event_type": "StateDeltaEvent",
        "event_data": {
            "delta": [{"op": "replace", "path": "/weather", "value": weather}]
        },
        "result_message": f"Weather updated to: {weather}"
    }

In my opinion 1 would be preferable.
But I haven’t looked more deeply yet into when agent_framework loses the results, supposedly due to a serialization failure.

Before dedicating myself to the final solution, what would be the prefered option in your opinion?
Basically its about calculating the efforts for mitigating the contra arguments, which I cannot yet estimate at the moment.

Thanks a lot!