improve how state snapshotting works in AG-UI (#19934)

improve how state snapshotting works

Author: logan-markewich

llama-index-integrations/protocols/llama-index-protocols-ag-ui/llama_index/protocols/ag_ui/agent.py

@@ -12,6 +12,7 @@
 from llama_index.core.workflow.events import Event, StartEvent, StopEvent
 from llama_index.protocols.ag_ui.events import (
     MessagesSnapshotWorkflowEvent,
+    StateSnapshotWorkflowEvent,
     TextMessageChunkWorkflowEvent,
     ToolCallChunkWorkflowEvent,
 )
@@ -134,6 +135,7 @@ async def chat(
 
             # Save state to context for tools to use
             await ctx.store.set("state", state)
+            ctx.write_event_to_stream(StateSnapshotWorkflowEvent(snapshot=state))
 
             if state:
                 for msg in chat_history[::-1]:
@@ -250,6 +252,13 @@ async def handle_tool_call(
                 kwargs[tool.ctx_param_name] = ctx
 
             tool_output = await tool.acall(**kwargs)
+
+            # Update the state snapshot
+            current_state = await ctx.store.get("state", default={})
+            ctx.write_event_to_stream(
+                StateSnapshotWorkflowEvent(snapshot=current_state)
+            )
+
             return ToolCallResultEvent(
                 tool_call_id=ev.tool_call_id,
                 tool_name=ev.tool_name,

llama-index-integrations/protocols/llama-index-protocols-ag-ui/llama_index/protocols/ag_ui/router.py

@@ -75,10 +75,6 @@ async def stream_response():
                 # Finish the run
                 _ = await handler
 
-                # Update the state
-                state = await handler.ctx.store.get("state", default={})
-                yield workflow_event_to_sse(StateSnapshotWorkflowEvent(snapshot=state))
-
                 yield workflow_event_to_sse(
                     RunFinishedWorkflowEvent(
                         timestamp=timestamp(),

llama-index-integrations/protocols/llama-index-protocols-ag-ui/pyproject.toml

@@ -28,7 +28,7 @@ dev = [
 
 [project]
 name = "llama-index-protocols-ag-ui"
-version = "0.2.1"
+version = "0.2.2"
 description = "llama-index protocols AG-UI integration"
 authors = [
     {name = "Logan Markewich", email = "[email protected]"},