Fix AG-UI shared state example (#2272)

Author: stevenh

docs/ag-ui.md

@@ -150,15 +150,15 @@ app = agent.to_ag_ui(deps=StateDeps(DocumentState()))
 
 
 @agent.tool
-def update_state(ctx: RunContext[StateDeps[DocumentState]]) -> StateSnapshotEvent:
+async def update_state(ctx: RunContext[StateDeps[DocumentState]]) -> StateSnapshotEvent:
     return StateSnapshotEvent(
         type=EventType.STATE_SNAPSHOT,
         snapshot=ctx.deps.state,
     )
 
 
 @agent.tool_plain
-def custom_events() -> list[CustomEvent]:
+async def custom_events() -> list[CustomEvent]:
     return [
         CustomEvent(
             type=EventType.CUSTOM,

examples/pydantic_ai_examples/ag_ui/api/agentic_generative_ui.py

@@ -67,7 +67,7 @@ class JSONPatchOp(BaseModel):
 
 
 @agent.tool_plain
-def create_plan(steps: list[str]) -> StateSnapshotEvent:
+async def create_plan(steps: list[str]) -> StateSnapshotEvent:
     """Create a plan with multiple steps.
 
     Args:
@@ -86,7 +86,7 @@ def create_plan(steps: list[str]) -> StateSnapshotEvent:
 
 
 @agent.tool_plain
-def update_plan_step(
+async def update_plan_step(
     index: int, description: str | None = None, status: StepStatus | None = None
 ) -> StateDeltaEvent:
     """Update the plan with new steps or changes.

examples/pydantic_ai_examples/ag_ui/api/predictive_state_updates.py

@@ -23,7 +23,7 @@ class DocumentState(BaseModel):
 # Tools which return AG-UI events will be sent to the client as part of the
 # event stream, single events and iterables of events are supported.
 @agent.tool_plain
-def document_predict_state() -> list[CustomEvent]:
+async def document_predict_state() -> list[CustomEvent]:
     """Enable document state prediction.
 
     Returns:
@@ -45,7 +45,7 @@ def document_predict_state() -> list[CustomEvent]:
 
 
 @agent.instructions()
-def story_instructions(ctx: RunContext[StateDeps[DocumentState]]) -> str:
+async def story_instructions(ctx: RunContext[StateDeps[DocumentState]]) -> str:
     """Provide instructions for writing document if present.
 
     Args:

examples/pydantic_ai_examples/ag_ui/api/shared_state.py

@@ -88,7 +88,7 @@ class RecipeSnapshot(BaseModel):
 
 
 @agent.tool_plain
-def display_recipe(recipe: Recipe) -> StateSnapshotEvent:
+async def display_recipe(recipe: Recipe) -> StateSnapshotEvent:
     """Display the recipe to the user.
 
     Args:
@@ -104,7 +104,7 @@ def display_recipe(recipe: Recipe) -> StateSnapshotEvent:
 
 
 @agent.instructions
-def recipe_instructions(ctx: RunContext[StateDeps[RecipeSnapshot]]) -> str:
+async def recipe_instructions(ctx: RunContext[StateDeps[RecipeSnapshot]]) -> str:
     """Instructions for the recipe generation agent.
 
     Args:
@@ -122,6 +122,7 @@ def recipe_instructions(ctx: RunContext[StateDeps[RecipeSnapshot]]) -> str:
         - Append new ingredients to the existing ones
         - Use the `display_recipe` tool to present the recipe to the user
         - Do NOT repeat the recipe in the message, use the tool instead
+        - Do NOT run the `display_recipe` tool multiple times in a row
 
         Once you have created the updated recipe and displayed it to the user,
         summarise the changes in one sentence, don't describe the recipe in