fix stateful llama-index examples

typescript-sdk/apps/dojo/e2e/tests/llamaIndexTests/agenticGenUI.spec.ts

@@ -3,7 +3,7 @@ import { AgenticGenUIPage } from "../../pages/llamaIndexPages/AgenticUIGenPage";
 
 test.describe("Agent Generative UI Feature", () => {
   // Fails. Issue with integration or something.
-  test.fixme("[LlamaIndex] should interact with the chat to get a planner on prompt", async ({
+  test("[LlamaIndex] should interact with the chat to get a planner on prompt", async ({
     page,
   }) => {
     const genUIAgent = new AgenticGenUIPage(page);
@@ -37,7 +37,7 @@ test.describe("Agent Generative UI Feature", () => {
   });
 
   // Fails. Issue with integration or something.
-  test.fixme("[LlamaIndex] should interact with the chat using predefined prompts and perform steps", async ({
+  test("[LlamaIndex] should interact with the chat using predefined prompts and perform steps", async ({
     page,
   }) => {
     const genUIAgent = new AgenticGenUIPage(page);

typescript-sdk/apps/dojo/src/files.json

@@ -1263,7 +1263,7 @@
     },
     {
       "name": "agentic_generative_ui.py",
-      "content": "import asyncio\nimport copy\nimport jsonpatch\nfrom pydantic import BaseModel\n\nfrom llama_index.core.workflow import Context\nfrom llama_index.llms.openai import OpenAI\nfrom llama_index.protocols.ag_ui.router import get_ag_ui_workflow_router\nfrom llama_index.protocols.ag_ui.events import StateDeltaWorkflowEvent, StateSnapshotWorkflowEvent\n\nclass Step(BaseModel):\n    description: str\n\nclass Task(BaseModel):\n    steps: list[Step]\n\n# Genrative UI demo\nasync def run_task(\n    ctx: Context, task: Task,\n) -> str:\n    \"\"\"Execute any list of steps needed to complete a task. Useful for anything the user wants to do.\"\"\"\n    state = await ctx.get(\"state\", default={})\n    task = Task.model_validate(task)\n\n    state = {\n        \"steps\": [\n            {\n                \"description\": step.description,\n                \"status\": \"pending\"\n            }\n            for step in task.steps\n        ]\n    }\n\n    # Send initial state snapshot\n    ctx.write_event_to_stream(\n        StateSnapshotWorkflowEvent(\n            snapshot=state\n        )\n    )\n\n    # Sleep for 1 second\n    await asyncio.sleep(1.0)\n\n    # Create a copy to track changes for JSON patches\n    previous_state = copy.deepcopy(state)\n\n    # Update each step and send deltas\n    for i, step in enumerate(state[\"steps\"]):\n        step[\"status\"] = \"completed\"\n        \n        # Generate JSON patch from previous state to current state\n        patch = jsonpatch.make_patch(previous_state, state)\n        \n        # Send state delta event\n        ctx.write_event_to_stream(\n            StateDeltaWorkflowEvent(\n                delta=patch.patch\n            )\n        )\n        \n        # Update previous state for next iteration\n        previous_state = copy.deepcopy(state)\n        \n        # Sleep for 1 second\n        await asyncio.sleep(1.0)\n\n    # Optionally send a final snapshot to the client\n    ctx.write_event_to_stream(\n        StateSnapshotWorkflowEvent(\n            snapshot=state\n        )\n    )\n\n    return \"Done!\"\n\n\nagentic_generative_ui_router = get_ag_ui_workflow_router(\n    llm=OpenAI(model=\"gpt-4.1\"),\n    frontend_tools=[run_task],\n    initial_state={},\n    system_prompt=(\n        \"You are a helpful assistant that can help the user with their task. \"\n        \"If the user asks you to do any task, use the run_task tool to do it. \"\n        \"Use your best judgement to describe the steps.\"\n    )\n)\n",
+      "content": "import asyncio\nimport copy\nimport jsonpatch\nfrom pydantic import BaseModel\n\nfrom llama_index.core.workflow import Context\nfrom llama_index.llms.openai import OpenAI\nfrom llama_index.protocols.ag_ui.router import get_ag_ui_workflow_router\nfrom llama_index.protocols.ag_ui.events import StateDeltaWorkflowEvent, StateSnapshotWorkflowEvent\n\nclass Step(BaseModel):\n    description: str\n\nclass Task(BaseModel):\n    steps: list[Step]\n\n# Genrative UI demo\nasync def run_task(\n    ctx: Context, task: Task,\n) -> str:\n    \"\"\"Execute any list of steps needed to complete a task. Useful for anything the user wants to do.\"\"\"\n\n    async with ctx.store.edit_state() as global_state:\n        state = global_state.get(\"state\", {})\n        task = Task.model_validate(task)\n\n        state = {\n            \"steps\": [\n                {\n                    \"description\": step.description,\n                    \"status\": \"pending\"\n                }\n                for step in task.steps\n            ]\n        }\n\n        # Send initial state snapshot\n        ctx.write_event_to_stream(\n            StateSnapshotWorkflowEvent(\n                snapshot=state\n            )\n        )\n\n        # Sleep for 1 second\n        await asyncio.sleep(1.0)\n\n        # Create a copy to track changes for JSON patches\n        previous_state = copy.deepcopy(state)\n\n        # Update each step and send deltas\n        for i, step in enumerate(state[\"steps\"]):\n            step[\"status\"] = \"completed\"\n            \n            # Generate JSON patch from previous state to current state\n            patch = jsonpatch.make_patch(previous_state, state)\n            \n            # Send state delta event\n            ctx.write_event_to_stream(\n                StateDeltaWorkflowEvent(\n                    delta=patch.patch\n                )\n            )\n            \n            # Update previous state for next iteration\n            previous_state = copy.deepcopy(state)\n            \n            # Sleep for 1 second\n            await asyncio.sleep(1.0)\n\n        # Optionally send a final snapshot to the client\n        ctx.write_event_to_stream(\n            StateSnapshotWorkflowEvent(\n                snapshot=state\n            )\n        )\n\n        global_state[\"state\"] = state\n\n    return \"Task Done!\"\n\n\nagentic_generative_ui_router = get_ag_ui_workflow_router(\n    llm=OpenAI(model=\"gpt-4.1\"),\n    backend_tools=[run_task],\n    initial_state={},\n    system_prompt=(\n        \"You are a helpful assistant that can help the user with their task. \"\n        \"If the user asks you to do any task, use the run_task tool to do it. \"\n        \"Use your best judgement to describe the steps.\"\n    )\n)\n",
       "language": "python",
       "type": "file"
     }
@@ -1289,7 +1289,7 @@
     },
     {
       "name": "shared_state.py",
-      "content": "from typing import Literal, List\nfrom pydantic import BaseModel\n\nfrom llama_index.core.workflow import Context\nfrom llama_index.llms.openai import OpenAI\nfrom llama_index.protocols.ag_ui.events import StateSnapshotWorkflowEvent\nfrom llama_index.protocols.ag_ui.router import get_ag_ui_workflow_router\n\n\nclass Ingredient(BaseModel):\n    icon: str\n    name: str\n    amount: str\n\nclass Recipe(BaseModel):\n    skill_level: str\n    special_preferences: List[str]\n    cooking_time: str\n    ingredients: List[Ingredient]\n    instructions: List[str]\n\n\nasync def update_recipe(ctx: Context, recipe: Recipe) -> str:\n    \"\"\"Useful for recording a recipe to shared state.\"\"\"\n    recipe = Recipe.model_validate(recipe)\n\n    state = await ctx.get(\"state\")\n    if state is None:\n        state = {}\n\n    state[\"recipe\"] = recipe.model_dump()\n\n    ctx.write_event_to_stream(\n        StateSnapshotWorkflowEvent(\n            snapshot=state\n        )\n    )\n\n    await ctx.set(\"state\", state)\n\n    return \"Recipe updated!\"\n\n\nshared_state_router = get_ag_ui_workflow_router(\n    llm=OpenAI(model=\"gpt-4.1\"),\n    frontend_tools=[update_recipe],\n    initial_state={\n        \"recipe\": None,\n    }\n)\n\n\n",
+      "content": "from typing import Literal, List\nfrom pydantic import BaseModel\n\nfrom llama_index.core.workflow import Context\nfrom llama_index.llms.openai import OpenAI\nfrom llama_index.protocols.ag_ui.events import StateSnapshotWorkflowEvent\nfrom llama_index.protocols.ag_ui.router import get_ag_ui_workflow_router\n\n\nclass Ingredient(BaseModel):\n    icon: str\n    name: str\n    amount: str\n\nclass Recipe(BaseModel):\n    skill_level: str\n    special_preferences: List[str]\n    cooking_time: str\n    ingredients: List[Ingredient]\n    instructions: List[str]\n\n\nasync def update_recipe(ctx: Context, recipe: Recipe) -> str:\n    \"\"\"Useful for recording a recipe to shared state.\"\"\"\n    recipe = Recipe.model_validate(recipe)\n\n    async with ctx.store.edit_state() as global_state:\n        state = global_state.get(\"state\", {})\n        if state is None:\n            state = {}\n\n        state[\"recipe\"] = recipe.model_dump()\n\n        ctx.write_event_to_stream(\n            StateSnapshotWorkflowEvent(\n                snapshot=state\n            )\n        )\n\n        global_state[\"state\"] = state\n\n    return \"Recipe updated!\"\n\n\nshared_state_router = get_ag_ui_workflow_router(\n    llm=OpenAI(model=\"gpt-4.1\"),\n    frontend_tools=[update_recipe],\n    initial_state={\n        \"recipe\": None,\n    }\n)\n\n\n",
       "language": "python",
       "type": "file"
     }

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

@@ -9,9 +9,9 @@ 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-llms-openai>=0.5.0,<0.6.0",
+    "llama-index-protocols-ag-ui>=0.2.2,<0.3",
     "jsonpatch>=1.33",
     "uvicorn>=0.27.0",
     "fastapi>=0.100.0",
@@ -26,5 +26,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!"