Bring in the example from the Pydantic repo so you can see it locally. (#236)

* copy example from pydantic repo

* Make work without local pydantic repo

* Clean up feedback

Author: maxkorp

typescript-sdk/integrations/pydantic-ai/examples/.gitignore

@@ -0,0 +1 @@
+server.egg-info/

typescript-sdk/integrations/pydantic-ai/examples/README.md

@@ -0,0 +1,46 @@
+# Pydantic AI AG-UI Examples
+
+This directory contains example usage of the AG-UI adapter for Pydantic AI. It provides a FastAPI application that demonstrates how to use the Pydantic AI agent with the AG-UI protocol.
+
+## Features
+
+The examples include implementations for each of the AG-UI dojo features:
+- Agentic Chat
+- Human in the Loop
+- Agentic Generative UI
+- Tool Based Generative UI
+- Shared State
+- Predictive State Updates
+
+## Setup
+
+1. Install dependencies:
+   ```bash
+   uv sync
+   ```
+
+2. Run the development server:
+   ```bash
+   uv run dev
+   ```
+
+## Usage
+
+Once the server is running, launch the frontend dojo with:
+
+```bash
+cd ../../../
+pnpm install
+turbo run dev
+```
+
+and view it at http://localhost:3000.
+
+By default, the agents can be reached at:
+
+- `http://localhost:9000/agentic_chat` - Agentic Chat
+- `http://localhost:9000/agentic_generative_ui` - Agentic Generative UI
+- `http://localhost:9000/human_in_the_loop` - Human in the Loop
+- `http://localhost:9000/predictive_state_updates` - Predictive State Updates
+- `http://localhost:9000/shared_state` - Shared State
+- `http://localhost:9000/tool_based_generative_ui` - Tool Based Generative UI

typescript-sdk/integrations/pydantic-ai/examples/pyproject.toml

@@ -0,0 +1,20 @@
+tool.uv.package = true
+
+[project]
+name = "server"
+version = "0.1.0"
+description = "Example usage of the AG-UI adapter for Pydantic AI"
+license = "MIT"
+
+readme = "README.md"
+requires-python = ">=3.9"
+dependencies = [
+    "fastapi>=0.104.0",
+    "uvicorn[standard]>=0.24.0",
+    "pydantic-ai-slim[openai,ag-ui]>=0.1.0",
+]
+authors = []
+
+[project.scripts]
+dev = "server:main"
+

typescript-sdk/integrations/pydantic-ai/examples/server/__init__.py

@@ -0,0 +1,53 @@
+"""Example usage of the AG-UI adapter for Pydantic AI.
+
+This provides a FastAPI application that demonstrates how to use the
+Pydantic AI agent with the AG-UI protocol. It includes examples for
+each of the AG-UI dojo features:
+- Agentic Chat
+- Human in the Loop
+- Agentic Generative UI
+- Tool Based Generative UI
+- Shared State
+- Predictive State Updates
+"""
+
+from __future__ import annotations
+
+from fastapi import FastAPI
+import uvicorn
+
+
+from .api import (
+    agentic_chat_app,
+    agentic_generative_ui_app,
+    human_in_the_loop_app,
+    predictive_state_updates_app,
+    shared_state_app,
+    tool_based_generative_ui_app,
+)
+
+app = FastAPI(title='Pydantic AI AG-UI server')
+app.mount('/agentic_chat', agentic_chat_app, 'Agentic Chat')
+app.mount('/agentic_generative_ui', agentic_generative_ui_app, 'Agentic Generative UI')
+app.mount('/human_in_the_loop', human_in_the_loop_app, 'Human in the Loop')
+app.mount(
+    '/predictive_state_updates',
+    predictive_state_updates_app,
+    'Predictive State Updates',
+)
+app.mount('/shared_state', shared_state_app, 'Shared State')
+app.mount(
+    '/tool_based_generative_ui',
+    tool_based_generative_ui_app,
+    'Tool Based Generative UI',
+)
+
+
+def main():
+    """Main function to start the FastAPI server."""
+    uvicorn.run(app, host="0.0.0.0", port=9000)
+
+if __name__ == "__main__":
+    main()
+
+__all__ = ["main"]

typescript-sdk/integrations/pydantic-ai/examples/server/api/__init__.py

@@ -0,0 +1,19 @@
+"""Example API for a AG-UI compatible Pydantic AI Agent UI."""
+
+from __future__ import annotations
+
+from .agentic_chat import app as agentic_chat_app
+from .agentic_generative_ui import app as agentic_generative_ui_app
+from .human_in_the_loop import app as human_in_the_loop_app
+from .predictive_state_updates import app as predictive_state_updates_app
+from .shared_state import app as shared_state_app
+from .tool_based_generative_ui import app as tool_based_generative_ui_app
+
+__all__ = [
+    'agentic_chat_app',
+    'agentic_generative_ui_app',
+    'human_in_the_loop_app',
+    'predictive_state_updates_app',
+    'shared_state_app',
+    'tool_based_generative_ui_app',
+]

typescript-sdk/integrations/pydantic-ai/examples/server/api/agentic_chat.py

@@ -0,0 +1,25 @@
+"""Agentic Chat feature."""
+
+from __future__ import annotations
+
+from datetime import datetime
+from zoneinfo import ZoneInfo
+
+from pydantic_ai import Agent
+
+agent = Agent('openai:gpt-4o-mini')
+app = agent.to_ag_ui()
+
+
+@agent.tool_plain
+async def current_time(timezone: str = 'UTC') -> str:
+    """Get the current time in ISO format.
+
+    Args:
+        timezone: The timezone to use.
+
+    Returns:
+        The current time in ISO format string.
+    """
+    tz: ZoneInfo = ZoneInfo(timezone)
+    return datetime.now(tz=tz).isoformat()

typescript-sdk/integrations/pydantic-ai/examples/server/api/agentic_generative_ui.py

@@ -0,0 +1,119 @@
+"""Agentic Generative UI feature."""
+
+from __future__ import annotations
+
+from textwrap import dedent
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+from ag_ui.core import EventType, StateDeltaEvent, StateSnapshotEvent
+from pydantic_ai import Agent
+
+StepStatus = Literal['pending', 'completed']
+
+
+class Step(BaseModel):
+    """Represents a step in a plan."""
+
+    description: str = Field(description='The description of the step')
+    status: StepStatus = Field(
+        default='pending',
+        description='The status of the step (e.g., pending, completed)',
+    )
+
+
+class Plan(BaseModel):
+    """Represents a plan with multiple steps."""
+
+    steps: list[Step] = Field(default_factory=list, description='The steps in the plan')
+
+
+class JSONPatchOp(BaseModel):
+    """A class representing a JSON Patch operation (RFC 6902)."""
+
+    op: Literal['add', 'remove', 'replace', 'move', 'copy', 'test'] = Field(
+        description='The operation to perform: add, remove, replace, move, copy, or test',
+    )
+    path: str = Field(description='JSON Pointer (RFC 6901) to the target location')
+    value: Any = Field(
+        default=None,
+        description='The value to apply (for add, replace operations)',
+    )
+    from_: str | None = Field(
+        default=None,
+        alias='from',
+        description='Source path (for move, copy operations)',
+    )
+
+
+agent = Agent(
+    'openai:gpt-4o-mini',
+    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.
+        """
+    ),
+)
+
+
+@agent.tool_plain
+async def create_plan(steps: list[str]) -> StateSnapshotEvent:
+    """Create a plan with multiple steps.
+
+    Args:
+        steps: List of step descriptions to create the plan.
+
+    Returns:
+        StateSnapshotEvent containing the initial state of the steps.
+    """
+    plan: Plan = Plan(
+        steps=[Step(description=step) for step in steps],
+    )
+    return StateSnapshotEvent(
+        type=EventType.STATE_SNAPSHOT,
+        snapshot=plan.model_dump(),
+    )
+
+
+@agent.tool_plain
+async def update_plan_step(
+    index: int, description: str | None = None, status: StepStatus | None = None
+) -> StateDeltaEvent:
+    """Update the plan with new steps or changes.
+
+    Args:
+        index: The index of the step to update.
+        description: The new description for the step.
+        status: The new status for the step.
+
+    Returns:
+        StateDeltaEvent containing the changes made to the plan.
+    """
+    changes: list[JSONPatchOp] = []
+    if description is not None:
+        changes.append(
+            JSONPatchOp(
+                op='replace', path=f'/steps/{index}/description', value=description
+            )
+        )
+    if status is not None:
+        changes.append(
+            JSONPatchOp(op='replace', path=f'/steps/{index}/status', value=status)
+        )
+    return StateDeltaEvent(
+        type=EventType.STATE_DELTA,
+        delta=changes,
+    )
+
+
+app = agent.to_ag_ui()

typescript-sdk/integrations/pydantic-ai/examples/server/api/human_in_the_loop.py

@@ -0,0 +1,26 @@
+"""Human in the Loop Feature.
+
+No special handling is required for this feature.
+"""
+
+from __future__ import annotations
+
+from textwrap import dedent
+
+from pydantic_ai import Agent
+
+agent = Agent(
+    'openai:gpt-4o-mini',
+    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 = agent.to_ag_ui()

typescript-sdk/integrations/pydantic-ai/examples/server/api/predictive_state_updates.py

@@ -0,0 +1,77 @@
+"""Predictive State feature."""
+
+from __future__ import annotations
+
+from textwrap import dedent
+
+from pydantic import BaseModel
+
+from ag_ui.core import CustomEvent, EventType
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.ag_ui import StateDeps
+
+
+class DocumentState(BaseModel):
+    """State for the document being written."""
+
+    document: str = ''
+
+
+agent = Agent('openai:gpt-4o-mini', deps_type=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.
+@agent.tool_plain
+async def document_predict_state() -> list[CustomEvent]:
+    """Enable document state prediction.
+
+    Returns:
+        CustomEvent containing the event to enable state prediction.
+    """
+    return [
+        CustomEvent(
+            type=EventType.CUSTOM,
+            name='PredictState',
+            value=[
+                {
+                    'state_key': 'document',
+                    'tool': 'write_document',
+                    'tool_argument': 'document',
+                },
+            ],
+        ),
+    ]
+
+
+@agent.instructions()
+async def story_instructions(ctx: RunContext[StateDeps[DocumentState]]) -> str:
+    """Provide instructions for writing document if present.
+
+    Args:
+        ctx: The run context containing document state information.
+
+    Returns:
+        Instructions string for the document writing agent.
+    """
+    return dedent(
+        f"""You are a helpful assistant for writing documents.
+
+        Before you start writing, you MUST call the `document_predict_state`
+        tool to enable state prediction.
+
+        To present the document to the user for review, you MUST use the
+        `write_document` tool.
+
+        When you have written the document, DO NOT repeat it as a message.
+        If accepted briefly summarize the changes you made, 2 sentences
+        max, otherwise ask the user to clarify what they want to change.
+
+        This is the current document:
+
+        {ctx.deps.state.document}
+        """
+    )
+
+
+app = agent.to_ag_ui(deps=StateDeps(DocumentState()))