fix stateful examples

Author: logan-markewich

typescript-sdk/integrations/llamaindex/server-py/pyproject.toml

@@ -9,9 +9,8 @@ description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.9, <3.14"
 dependencies = [
-    "llama-index-core>=0.12.41,<0.13",
-    "llama-index-agent-openai>=0.4.9,<0.5",
-    "llama-index-protocols-ag-ui>=0.1.2",
+    "llama-index-core>=0.14.0,<0.15",
+    "llama-index-protocols-ag-ui>=0.2.0,<0.3",
     "jsonpatch>=1.33",
     "uvicorn>=0.27.0",
     "fastapi>=0.100.0",
@@ -26,5 +25,8 @@ include = ["server/"]
 [tool.hatch.build.targets.wheel]
 include = ["server/"]
 
+[tool.hatch.metadata]
+allow-direct-references = true
+
 [project.scripts]
 dev = "server:main"

typescript-sdk/integrations/llamaindex/server-py/server/routers/agentic_generative_ui.py

@@ -19,65 +19,69 @@ async def run_task(
     ctx: Context, task: Task,
 ) -> str:
     """Execute any list of steps needed to complete a task. Useful for anything the user wants to do."""
-    state = await ctx.get("state", default={})
-    task = Task.model_validate(task)
 
-    state = {
-        "steps": [
-            {
-                "description": step.description,
-                "status": "pending"
-            }
-            for step in task.steps
-        ]
-    }
+    async with ctx.store.edit_state() as global_state:
+        state = global_state.get("state", {})
+        task = Task.model_validate(task)
 
-    # Send initial state snapshot
-    ctx.write_event_to_stream(
-        StateSnapshotWorkflowEvent(
-            snapshot=state
-        )
-    )
-
-    # Sleep for 1 second
-    await asyncio.sleep(1.0)
-
-    # Create a copy to track changes for JSON patches
-    previous_state = copy.deepcopy(state)
+        state = {
+            "steps": [
+                {
+                    "description": step.description,
+                    "status": "pending"
+                }
+                for step in task.steps
+            ]
+        }
 
-    # Update each step and send deltas
-    for i, step in enumerate(state["steps"]):
-        step["status"] = "completed"
-        
-        # Generate JSON patch from previous state to current state
-        patch = jsonpatch.make_patch(previous_state, state)
-        
-        # Send state delta event
+        # Send initial state snapshot
         ctx.write_event_to_stream(
-            StateDeltaWorkflowEvent(
-                delta=patch.patch
+            StateSnapshotWorkflowEvent(
+                snapshot=state
             )
         )
-        
-        # Update previous state for next iteration
-        previous_state = copy.deepcopy(state)
-        
+
         # Sleep for 1 second
         await asyncio.sleep(1.0)
 
-    # Optionally send a final snapshot to the client
-    ctx.write_event_to_stream(
-        StateSnapshotWorkflowEvent(
-            snapshot=state
+        # Create a copy to track changes for JSON patches
+        previous_state = copy.deepcopy(state)
+
+        # Update each step and send deltas
+        for i, step in enumerate(state["steps"]):
+            step["status"] = "completed"
+            
+            # Generate JSON patch from previous state to current state
+            patch = jsonpatch.make_patch(previous_state, state)
+            
+            # Send state delta event
+            ctx.write_event_to_stream(
+                StateDeltaWorkflowEvent(
+                    delta=patch.patch
+                )
+            )
+            
+            # Update previous state for next iteration
+            previous_state = copy.deepcopy(state)
+            
+            # Sleep for 1 second
+            await asyncio.sleep(1.0)
+
+        # Optionally send a final snapshot to the client
+        ctx.write_event_to_stream(
+            StateSnapshotWorkflowEvent(
+                snapshot=state
+            )
         )
-    )
 
-    return "Done!"
+        global_state["state"] = state
+
+    return "Task Done!"
 
 
 agentic_generative_ui_router = get_ag_ui_workflow_router(
     llm=OpenAI(model="gpt-4.1"),
-    frontend_tools=[run_task],
+    backend_tools=[run_task],
     initial_state={},
     system_prompt=(
         "You are a helpful assistant that can help the user with their task. "

typescript-sdk/integrations/llamaindex/server-py/server/routers/shared_state.py

@@ -24,19 +24,20 @@ async def update_recipe(ctx: Context, recipe: Recipe) -> str:
     """Useful for recording a recipe to shared state."""
     recipe = Recipe.model_validate(recipe)
 
-    state = await ctx.get("state")
-    if state is None:
-        state = {}
+    async with ctx.store.edit_state() as global_state:
+        state = global_state.get("state", {})
+        if state is None:
+            state = {}
 
-    state["recipe"] = recipe.model_dump()
+        state["recipe"] = recipe.model_dump()
 
-    ctx.write_event_to_stream(
-        StateSnapshotWorkflowEvent(
-            snapshot=state
+        ctx.write_event_to_stream(
+            StateSnapshotWorkflowEvent(
+                snapshot=state
+            )
         )
-    )
 
-    await ctx.set("state", state)
+        global_state["state"] = state
 
     return "Recipe updated!"