chore(ag-ui): address PR feedback

* Split out coverage fixes to separate PRs
* Rename FastAGUI to AGUIApp
* Add thinking part delta test
* Remove example agent helper
* Use dedent to help create clean instructions
* Remove logging from examples
* Restore explicit re-raise variable
* Eliminate agent_to_ag_ui
* Ensure user toolsets are not overwritten incorrectly
* Use temporary vars to reduce repetitive indirection
* Raise ToolCallNotFoundError if tool result without matching call

Author: stevenh

.github/workflows/ci.yml

@@ -204,7 +204,6 @@ jobs:
           enable-cache: true
 
       - run: uv sync --package pydantic-ai-slim --only-dev
-      - run: rm coverage/.coverage.*-py3.9-* # Exclude 3.9 coverage as it gets the wrong line numbers, causing invalid failures.
       - run: uv run coverage combine coverage
 
       - run: uv run coverage html --show-contexts --title "PydanticAI coverage for ${{ github.sha }}"

docs/ag-ui.md

@@ -12,7 +12,7 @@ The team at [Rocket Science](https://www.rocketscience.gg/), contributed the
 protocol with PydanticAI agents.
 
 This also includes an [`Agent.to_ag_ui`][pydantic_ai.Agent.to_ag_ui] convenience
-method which simplifies the creation of [`FastAGUI`][pydantic_ai.ag_ui.FastAGUI]
+method which simplifies the creation of [`AGUIApp`][pydantic_ai.ag_ui.AGUIApp]
 for PydanticAI agents, which is built on top of [Starlette](https://www.starlette.io/),
 meaning it's fully compatible with any ASGI server.
 
@@ -88,7 +88,7 @@ A user request may require multiple round trips between client UI and PydanticAI
 server, depending on the tools and events needed.
 
 In addition to the [Adapter][pydantic_ai.ag_ui.Adapter] there is also
-[FastAGUI][pydantic_ai.ag_ui.FastAGUI] which is slim wrapper around
+[AGUIApp][pydantic_ai.ag_ui.AGUIApp] which is slim wrapper around
 [Starlette](https://www.starlette.io/) providing easy access to run a PydanticAI
 server with AG-UI support with any ASGI server.
 

examples/pydantic_ai_ag_ui_examples/api/agent.py

@@ -5,14 +5,23 @@
 from datetime import datetime
 from zoneinfo import ZoneInfo
 
-from pydantic_ai.ag_ui import FastAGUI
+from dotenv import load_dotenv
 
-from .agent import agent
+from pydantic_ai import Agent
+from pydantic_ai.ag_ui import AGUIApp
 
-app: FastAGUI = agent()
+# Ensure environment variables are loaded.
+load_dotenv()
 
+agent: Agent = Agent(
+    'openai:gpt-4o-mini',
+    output_type=str,
+)
 
-@app.adapter.agent.tool_plain
+app: AGUIApp = agent.to_ag_ui()
+
+
+@agent.tool_plain
 async def current_time(timezone: str = 'UTC') -> str:
     """Get the current time in ISO format.
 

examples/pydantic_ai_ag_ui_examples/api/agentic_chat.py

@@ -3,29 +3,40 @@
 from __future__ import annotations
 
 from enum import StrEnum
+from textwrap import dedent
 from typing import Any, Literal
 
 from ag_ui.core import EventType, StateDeltaEvent, StateSnapshotEvent
+from dotenv import load_dotenv
 from pydantic import BaseModel, Field
 
-from pydantic_ai.ag_ui import FastAGUI
-
-from .agent import agent
-
-app: FastAGUI = agent(
-    instructions="""When planning use tools only, without any other messages.
-IMPORTANT:
-- Use the `create_plan` tool to set the initial state of the steps
-- Use the `update_plan_step` tool to update the status of each step
-- Do NOT repeat the plan or summarise it in a message
-- Do NOT confirm the creation or updates in a message
-- Do NOT ask the user for additional information or next steps
-
-Only one plan can be active at a time, so do not call the `create_plan` tool
-again until all the steps in current plan are completed.
-"""
+from pydantic_ai import Agent
+from pydantic_ai.ag_ui import AGUIApp
+
+# Ensure environment variables are loaded.
+load_dotenv()
+
+agent: Agent = Agent(
+    'openai:gpt-4o-mini',
+    output_type=str,
+    instructions=dedent(
+        """
+        When planning use tools only, without any other messages.
+        IMPORTANT:
+        - Use the `create_plan` tool to set the initial state of the steps
+        - Use the `update_plan_step` tool to update the status of each step
+        - Do NOT repeat the plan or summarise it in a message
+        - Do NOT confirm the creation or updates in a message
+        - Do NOT ask the user for additional information or next steps
+
+        Only one plan can be active at a time, so do not call the `create_plan` tool
+        again until all the steps in current plan are completed.
+        """
+    ),
 )
 
+app: AGUIApp = agent.to_ag_ui()
+
 
 class StepStatus(StrEnum):
     """The status of a step in a plan."""
@@ -71,7 +82,7 @@ class JSONPatchOp(BaseModel):
     )
 
 
-@app.adapter.agent.tool_plain
+@agent.tool_plain
 def create_plan(steps: list[str]) -> StateSnapshotEvent:
     """Create a plan with multiple steps.
 
@@ -90,7 +101,7 @@ def create_plan(steps: list[str]) -> StateSnapshotEvent:
     )
 
 
-@app.adapter.agent.tool_plain
+@agent.tool_plain
 def update_plan_step(
     index: int, description: str | None = None, status: StepStatus | None = None
 ) -> StateDeltaEvent:

examples/pydantic_ai_ag_ui_examples/api/agentic_generative_ui.py

@@ -5,16 +5,29 @@
 
 from __future__ import annotations
 
-from pydantic_ai.ag_ui import FastAGUI
+from textwrap import dedent
 
-from .agent import agent
+from dotenv import load_dotenv
 
-app: FastAGUI = agent(
-    instructions="""When planning tasks use tools only, without any other messages.
-IMPORTANT:
-- Use the `generate_task_steps` tool to display the suggested steps to the user
-- Never repeat the plan, or send a message detailing steps
-- If accepted, confirm the creation of the plan and the number of selected (enabled) steps only
-- If not accepted, ask the user for more information, DO NOT use the `generate_task_steps` tool again
-"""
+from pydantic_ai import Agent
+from pydantic_ai.ag_ui import AGUIApp
+
+# Ensure environment variables are loaded.
+load_dotenv()
+
+agent: Agent = Agent(
+    'openai:gpt-4o-mini',
+    output_type=str,
+    instructions=dedent(
+        """
+        When planning tasks use tools only, without any other messages.
+        IMPORTANT:
+        - Use the `generate_task_steps` tool to display the suggested steps to the user
+        - Never repeat the plan, or send a message detailing steps
+        - If accepted, confirm the creation of the plan and the number of selected (enabled) steps only
+        - If not accepted, ask the user for more information, DO NOT use the `generate_task_steps` tool again
+        """
+    ),
 )
+
+app: AGUIApp = agent.to_ag_ui()

examples/pydantic_ai_ag_ui_examples/api/human_in_the_loop.py

@@ -2,42 +2,46 @@
 
 from __future__ import annotations
 
-import logging
 from typing import TYPE_CHECKING
 
 from ag_ui.core import CustomEvent, EventType
+from dotenv import load_dotenv
 from pydantic import BaseModel
 
-from pydantic_ai.ag_ui import FastAGUI, StateDeps
-
-from .agent import agent
+from pydantic_ai import Agent
+from pydantic_ai.ag_ui import AGUIApp, StateDeps
 
 if TYPE_CHECKING:  # pragma: no cover
     from pydantic_ai import RunContext
 
 
-_LOGGER: logging.Logger = logging.getLogger(__name__)
-
-
 class DocumentState(BaseModel):
     """State for the document being written."""
 
     document: str = ''
 
 
-app: FastAGUI = agent(deps=StateDeps(DocumentState()))
+# Ensure environment variables are loaded.
+load_dotenv()
+
+agent: Agent = Agent(
+    'openai:gpt-4o-mini',
+    output_type=str,
+    deps_type=StateDeps[DocumentState],
+)
+
+app: AGUIApp = agent.to_ag_ui(deps=StateDeps(DocumentState()))
 
 
 # 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.
-@app.adapter.agent.tool_plain
+@agent.tool_plain
 def document_predict_state() -> list[CustomEvent]:
     """Enable document state prediction.
 
     Returns:
         CustomEvent containing the event to enable state prediction.
     """
-    _LOGGER.info('enabling document state state prediction')
     return [
         CustomEvent(
             type=EventType.CUSTOM,
@@ -53,7 +57,7 @@ def document_predict_state() -> list[CustomEvent]:
     ]
 
 
-@app.adapter.agent.instructions()
+@agent.instructions()
 def story_instructions(ctx: RunContext[StateDeps[DocumentState]]) -> str:
     """Provide instructions for writing document if present.
 
@@ -63,8 +67,6 @@ def story_instructions(ctx: RunContext[StateDeps[DocumentState]]) -> str:
     Returns:
         Instructions string for the document writing agent.
     """
-    _LOGGER.info('story instructions document=%s', ctx.deps.state.document)
-
     return f"""You are a helpful assistant for writing documents.
 
 Before you start writing, you MUST call the `document_predict_state`

examples/pydantic_ai_ag_ui_examples/api/predictive_state_updates.py

@@ -2,23 +2,19 @@
 
 from __future__ import annotations
 
-import json
-import logging
 from enum import StrEnum
 from typing import TYPE_CHECKING
 
 from ag_ui.core import EventType, StateSnapshotEvent
+from dotenv import load_dotenv
 from pydantic import BaseModel, Field
 
-from pydantic_ai.ag_ui import FastAGUI, StateDeps
-
-from .agent import agent
+from pydantic_ai import Agent
+from pydantic_ai.ag_ui import AGUIApp, StateDeps
 
 if TYPE_CHECKING:  # pragma: no cover
     from pydantic_ai import RunContext
 
-_LOGGER: logging.Logger = logging.getLogger(__name__)
-
 
 class SkillLevel(StrEnum):
     """The level of skill required for the recipe."""
@@ -92,10 +88,19 @@ class RecipeSnapshot(BaseModel):
     )
 
 
-app: FastAGUI = agent(deps=StateDeps(RecipeSnapshot()))
+# Ensure environment variables are loaded.
+load_dotenv()
+
+agent: Agent = Agent(
+    'openai:gpt-4o-mini',
+    output_type=str,
+    deps_type=StateDeps[RecipeSnapshot],
+)
+
+app: AGUIApp = agent.to_ag_ui(deps=StateDeps(RecipeSnapshot()))
 
 
-@app.adapter.agent.tool_plain
+@agent.tool_plain
 def display_recipe(recipe: Recipe) -> StateSnapshotEvent:
     """Display the recipe to the user.
 
@@ -111,7 +116,7 @@ def display_recipe(recipe: Recipe) -> StateSnapshotEvent:
     )
 
 
-@app.adapter.agent.instructions
+@agent.instructions
 def recipe_instructions(ctx: RunContext[StateDeps[RecipeSnapshot]]) -> str:
     """Instructions for the recipe generation agent.
 
@@ -121,8 +126,6 @@ def recipe_instructions(ctx: RunContext[StateDeps[RecipeSnapshot]]) -> str:
     Returns:
         Instructions string for the recipe generation agent.
     """
-    _LOGGER.info('recipe instructions recipe=%s', ctx.deps.state.recipe)
-
     return f"""You are a helpful assistant for creating recipes.
 
 IMPORTANT:
@@ -135,10 +138,6 @@ def recipe_instructions(ctx: RunContext[StateDeps[RecipeSnapshot]]) -> str:
 summarise the changes in one sentence, don't describe the recipe in
 detail or send it as a message to the user.
 
-The structure of a recipe is as follows:
-
-{json.dumps(Recipe.model_json_schema(), indent=2)}
-
 The current state of the recipe is:
 
 {ctx.deps.state.recipe.model_dump_json(indent=2)}

examples/pydantic_ai_ag_ui_examples/api/shared_state.py

@@ -5,8 +5,17 @@
 
 from __future__ import annotations
 
-from pydantic_ai.ag_ui import FastAGUI
+from dotenv import load_dotenv
 
-from .agent import agent
+from pydantic_ai import Agent
+from pydantic_ai.ag_ui import AGUIApp
 
-app: FastAGUI = agent()
+# Ensure environment variables are loaded.
+load_dotenv()
+
+agent: Agent = Agent(
+    'openai:gpt-4o-mini',
+    output_type=str,
+)
+
+app: AGUIApp = agent.to_ag_ui()

examples/pydantic_ai_ag_ui_examples/api/tool_based_generative_ui.py

@@ -621,7 +621,7 @@ async def process_function_tools(  # noqa: C901
                 result_data = await toolset.call_tool(call, run_context)
             except exceptions.UnexpectedModelBehavior as e:
                 ctx.state.increment_retries(ctx.deps.max_result_retries, e)
-                raise  # pragma: no cover
+                raise e  # pragma: no cover
             except ToolRetryError as e:
                 ctx.state.increment_retries(ctx.deps.max_result_retries, e)
                 yield _messages.FunctionToolCallEvent(call)