pydantic-ai-examples 0.1.2__py3-none-any.whl → 1.12.0__py3-none-any.whl

pydantic_ai_examples/ag_ui/__init__.py

@@ -0,0 +1,41 @@
+"""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
+
+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',
+)

pydantic_ai_examples/ag_ui/__main__.py

@@ -0,0 +1,9 @@
+"""Very simply CLI to run the AG-UI example.
+
+See https://ai.pydantic.dev/examples/ag-ui/ for more information.
+"""
+
+if __name__ == '__main__':
+    import uvicorn
+
+    uvicorn.run('pydantic_ai_examples.ag_ui:app', port=9000)

pydantic_ai_examples/ag_ui/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',
+]

pydantic_ai_examples/ag_ui/api/agentic_chat.py

@@ -0,0 +1,28 @@
+"""Agentic Chat feature."""
+
+from __future__ import annotations
+
+from datetime import datetime
+from zoneinfo import ZoneInfo
+
+from pydantic_ai import Agent
+from pydantic_ai.ui.ag_ui.app import AGUIApp
+
+agent = Agent('openai:gpt-5-mini')
+
+
+@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()
+
+
+app = AGUIApp(agent)

pydantic_ai_examples/ag_ui/api/agentic_generative_ui.py

@@ -0,0 +1,120 @@
+"""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
+from pydantic_ai.ui.ag_ui.app import AGUIApp
+
+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-5-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 = AGUIApp(agent)

pydantic_ai_examples/ag_ui/api/human_in_the_loop.py

@@ -0,0 +1,27 @@
+"""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
+from pydantic_ai.ui.ag_ui.app import AGUIApp
+
+agent = Agent(
+    'openai:gpt-5-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 = AGUIApp(agent)

pydantic_ai_examples/ag_ui/api/predictive_state_updates.py

@@ -0,0 +1,78 @@
+"""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.ui import StateDeps
+from pydantic_ai.ui.ag_ui.app import AGUIApp
+
+
+class DocumentState(BaseModel):
+    """State for the document being written."""
+
+    document: str = ''
+
+
+agent = Agent('openai:gpt-5-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 = AGUIApp(agent, deps=StateDeps(DocumentState()))

pydantic_ai_examples/ag_ui/api/shared_state.py

@@ -0,0 +1,139 @@
+"""Shared State feature."""
+
+from __future__ import annotations
+
+from enum import Enum
+from textwrap import dedent
+
+from pydantic import BaseModel, Field
+
+from ag_ui.core import EventType, StateSnapshotEvent
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.ui import StateDeps
+from pydantic_ai.ui.ag_ui.app import AGUIApp
+
+
+class SkillLevel(str, Enum):
+    """The level of skill required for the recipe."""
+
+    BEGINNER = 'Beginner'
+    INTERMEDIATE = 'Intermediate'
+    ADVANCED = 'Advanced'
+
+
+class SpecialPreferences(str, Enum):
+    """Special preferences for the recipe."""
+
+    HIGH_PROTEIN = 'High Protein'
+    LOW_CARB = 'Low Carb'
+    SPICY = 'Spicy'
+    BUDGET_FRIENDLY = 'Budget-Friendly'
+    ONE_POT_MEAL = 'One-Pot Meal'
+    VEGETARIAN = 'Vegetarian'
+    VEGAN = 'Vegan'
+
+
+class CookingTime(str, Enum):
+    """The cooking time of the recipe."""
+
+    FIVE_MIN = '5 min'
+    FIFTEEN_MIN = '15 min'
+    THIRTY_MIN = '30 min'
+    FORTY_FIVE_MIN = '45 min'
+    SIXTY_PLUS_MIN = '60+ min'
+
+
+class Ingredient(BaseModel):
+    """A class representing an ingredient in a recipe."""
+
+    icon: str = Field(
+        default='ingredient',
+        description="The icon emoji (not emoji code like '\x1f35e', but the actual emoji like 🥕) of the ingredient",
+    )
+    name: str
+    amount: str
+
+
+class Recipe(BaseModel):
+    """A class representing a recipe."""
+
+    skill_level: SkillLevel = Field(
+        default=SkillLevel.BEGINNER,
+        description='The skill level required for the recipe',
+    )
+    special_preferences: list[SpecialPreferences] = Field(
+        default_factory=list,
+        description='Any special preferences for the recipe',
+    )
+    cooking_time: CookingTime = Field(
+        default=CookingTime.FIVE_MIN, description='The cooking time of the recipe'
+    )
+    ingredients: list[Ingredient] = Field(
+        default_factory=list,
+        description='Ingredients for the recipe',
+    )
+    instructions: list[str] = Field(
+        default_factory=list, description='Instructions for the recipe'
+    )
+
+
+class RecipeSnapshot(BaseModel):
+    """A class representing the state of the recipe."""
+
+    recipe: Recipe = Field(
+        default_factory=Recipe, description='The current state of the recipe'
+    )
+
+
+agent = Agent('openai:gpt-5-mini', deps_type=StateDeps[RecipeSnapshot])
+
+
+@agent.tool_plain
+async def display_recipe(recipe: Recipe) -> StateSnapshotEvent:
+    """Display the recipe to the user.
+
+    Args:
+        recipe: The recipe to display.
+
+    Returns:
+        StateSnapshotEvent containing the recipe snapshot.
+    """
+    return StateSnapshotEvent(
+        type=EventType.STATE_SNAPSHOT,
+        snapshot={'recipe': recipe},
+    )
+
+
+@agent.instructions
+async def recipe_instructions(ctx: RunContext[StateDeps[RecipeSnapshot]]) -> str:
+    """Instructions for the recipe generation agent.
+
+    Args:
+        ctx: The run context containing recipe state information.
+
+    Returns:
+        Instructions string for the recipe generation agent.
+    """
+    return dedent(
+        f"""
+        You are a helpful assistant for creating recipes.
+
+        IMPORTANT:
+        - Create a complete recipe using the existing ingredients
+        - 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
+        detail or send it as a message to the user.
+
+        The current state of the recipe is:
+
+        {ctx.deps.state.recipe.model_dump_json(indent=2)}
+        """,
+    )
+
+
+app = AGUIApp(agent, deps=StateDeps(RecipeSnapshot()))

pydantic_ai_examples/ag_ui/api/tool_based_generative_ui.py

@@ -0,0 +1,12 @@
+"""Tool Based Generative UI feature.
+
+No special handling is required for this feature.
+"""
+
+from __future__ import annotations
+
+from pydantic_ai import Agent
+from pydantic_ai.ui.ag_ui.app import AGUIApp
+
+agent = Agent('openai:gpt-5-mini')
+app = AGUIApp(agent)

pydantic_ai_examples/bank_support.py

@@ -1,4 +1,4 @@
-"""Small but complete example of using PydanticAI to build a support agent for a bank.
+"""Small but complete example of using Pydantic AI to build a support agent for a bank.
 
 Run with:
 
@@ -7,7 +7,7 @@ Run with:
 
 from dataclasses import dataclass
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 
 from pydantic_ai import Agent, RunContext
 
@@ -26,8 +26,11 @@ class DatabaseConn:
 
     @classmethod
     async def customer_balance(cls, *, id: int, include_pending: bool) -> float:
-        if id == 123 and include_pending:
-            return 123.45
+        if id == 123:
+            if include_pending:
+                return 123.45
+            else:
+                return 100.00
         else:
             raise ValueError('Customer not found')
 
@@ -39,16 +42,19 @@ class SupportDependencies:
 
 
 class SupportOutput(BaseModel):
-    support_advice: str = Field(description='Advice returned to the customer')
-    block_card: bool = Field(description='Whether to block their card or not')
-    risk: int = Field(description='Risk level of query', ge=0, le=10)
+    support_advice: str
+    """Advice returned to the customer"""
+    block_card: bool
+    """Whether to block their card or not"""
+    risk: int
+    """Risk level of query"""
 
 
 support_agent = Agent(
-    'openai:gpt-4o',
+    'openai:gpt-5',
     deps_type=SupportDependencies,
     output_type=SupportOutput,
-    system_prompt=(
+    instructions=(
         'You are a support agent in our bank, give the '
         'customer support and judge the risk level of their query. '
         "Reply using the customer's name."
@@ -56,7 +62,7 @@ support_agent = Agent(
 )
 
 
-@support_agent.system_prompt
+@support_agent.instructions
 async def add_customer_name(ctx: RunContext[SupportDependencies]) -> str:
     customer_name = await ctx.deps.db.customer_name(id=ctx.deps.customer_id)
     return f"The customer's name is {customer_name!r}"

pydantic_ai_examples/chat_app.py

@@ -10,14 +10,14 @@ from __future__ import annotations as _annotations
 import asyncio
 import json
 import sqlite3
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Callable
 from concurrent.futures.thread import ThreadPoolExecutor
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from functools import partial
 from pathlib import Path
-from typing import Annotated, Any, Callable, Literal, TypeVar
+from typing import Annotated, Any, Literal, TypeVar
 
 import fastapi
 import logfire
@@ -25,21 +25,22 @@ from fastapi import Depends, Request
 from fastapi.responses import FileResponse, Response, StreamingResponse
 from typing_extensions import LiteralString, ParamSpec, TypedDict
 
-from pydantic_ai import Agent
-from pydantic_ai.exceptions import UnexpectedModelBehavior
-from pydantic_ai.messages import (
+from pydantic_ai import (
+    Agent,
     ModelMessage,
     ModelMessagesTypeAdapter,
     ModelRequest,
     ModelResponse,
     TextPart,
+    UnexpectedModelBehavior,
     UserPromptPart,
 )
 
 # 'if-token-present' means nothing will be sent (and the example will work) if you don't have logfire configured
 logfire.configure(send_to_logfire='if-token-present')
+logfire.instrument_pydantic_ai()
 
-agent = Agent('openai:gpt-4o', instrument=True)
+agent = Agent('openai:gpt-5')
 THIS_DIR = Path(__file__).parent
 
 
@@ -126,7 +127,7 @@ async def post_chat(
         messages = await database.get_messages()
         # run the agent with the user prompt and the chat history
         async with agent.run_stream(prompt, message_history=messages) as result:
-            async for text in result.stream(debounce_by=0.01):
+            async for text in result.stream_output(debounce_by=0.01):
                 # text here is a `str` and the frontend wants
                 # JSON encoded ModelResponse, so we create one
                 m = ModelResponse(parts=[TextPart(text)], timestamp=result.timestamp())