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

pydantic_ai_examples/data_analyst.py

@@ -0,0 +1,107 @@
+from dataclasses import dataclass, field
+
+import datasets
+import duckdb
+import pandas as pd
+
+from pydantic_ai import Agent, ModelRetry, RunContext
+
+
+@dataclass
+class AnalystAgentDeps:
+    output: dict[str, pd.DataFrame] = field(default_factory=dict)
+
+    def store(self, value: pd.DataFrame) -> str:
+        """Store the output in deps and return the reference such as Out[1] to be used by the LLM."""
+        ref = f'Out[{len(self.output) + 1}]'
+        self.output[ref] = value
+        return ref
+
+    def get(self, ref: str) -> pd.DataFrame:
+        if ref not in self.output:
+            raise ModelRetry(
+                f'Error: {ref} is not a valid variable reference. Check the previous messages and try again.'
+            )
+        return self.output[ref]
+
+
+analyst_agent = Agent(
+    'openai:gpt-5',
+    deps_type=AnalystAgentDeps,
+    instructions='You are a data analyst and your job is to analyze the data according to the user request.',
+)
+
+
+@analyst_agent.tool
+def load_dataset(
+    ctx: RunContext[AnalystAgentDeps],
+    path: str,
+    split: str = 'train',
+) -> str:
+    """Load the `split` of dataset `dataset_name` from huggingface.
+
+    Args:
+        ctx: Pydantic AI agent RunContext
+        path: name of the dataset in the form of `<user_name>/<dataset_name>`
+        split: load the split of the dataset (default: "train")
+    """
+    # begin load data from hf
+    builder = datasets.load_dataset_builder(path)  # pyright: ignore[reportUnknownMemberType]
+    splits: dict[str, datasets.SplitInfo] = builder.info.splits or {}  # pyright: ignore[reportUnknownMemberType]
+    if split not in splits:
+        raise ModelRetry(
+            f'{split} is not valid for dataset {path}. Valid splits are {",".join(splits.keys())}'
+        )
+
+    builder.download_and_prepare()  # pyright: ignore[reportUnknownMemberType]
+    dataset = builder.as_dataset(split=split)
+    assert isinstance(dataset, datasets.Dataset)
+    dataframe = dataset.to_pandas()
+    assert isinstance(dataframe, pd.DataFrame)
+    # end load data from hf
+
+    # store the dataframe in the deps and get a ref like "Out[1]"
+    ref = ctx.deps.store(dataframe)
+    # construct a summary of the loaded dataset
+    output = [
+        f'Loaded the dataset as `{ref}`.',
+        f'Description: {dataset.info.description}'
+        if dataset.info.description
+        else None,
+        f'Features: {dataset.info.features!r}' if dataset.info.features else None,
+    ]
+    return '\n'.join(filter(None, output))
+
+
+@analyst_agent.tool
+def run_duckdb(ctx: RunContext[AnalystAgentDeps], dataset: str, sql: str) -> str:
+    """Run DuckDB SQL query on the DataFrame.
+
+    Note that the virtual table name used in DuckDB SQL must be `dataset`.
+
+    Args:
+        ctx: Pydantic AI agent RunContext
+        dataset: reference string to the DataFrame
+        sql: the query to be executed using DuckDB
+    """
+    data = ctx.deps.get(dataset)
+    result = duckdb.query_df(df=data, virtual_table_name='dataset', sql_query=sql)
+    # pass the result as ref (because DuckDB SQL can select many rows, creating another huge dataframe)
+    ref = ctx.deps.store(result.df())  # pyright: ignore[reportUnknownMemberType]
+    return f'Executed SQL, result is `{ref}`'
+
+
+@analyst_agent.tool
+def display(ctx: RunContext[AnalystAgentDeps], name: str) -> str:
+    """Display at most 5 rows of the dataframe."""
+    dataset = ctx.deps.get(name)
+    return dataset.head().to_string()  # pyright: ignore[reportUnknownMemberType]
+
+
+if __name__ == '__main__':
+    deps = AnalystAgentDeps()
+    result = analyst_agent.run_sync(
+        user_prompt='Count how many negative comments are there in the dataset `cornell-movie-review-data/rotten_tomatoes`',
+        deps=deps,
+    )
+    print(result.output)

pydantic_ai_examples/evals/agent.py

@@ -21,7 +21,7 @@ class TimeRangeDeps:
 
 
 time_range_agent = Agent[TimeRangeDeps, TimeRangeResponse](
-    'gpt-4o',
+    'gpt-5',
     output_type=TimeRangeResponse,  # type: ignore  # we can't yet annotate something as receiving a TypeForm
     deps_type=TimeRangeDeps,
     system_prompt="Convert the user's request into a structured time range.",

pydantic_ai_examples/evals/custom_evaluators.py

@@ -1,6 +1,11 @@
 from dataclasses import dataclass
 from datetime import timedelta
 
+from pydantic_ai_examples.evals.models import (
+    TimeRangeBuilderSuccess,
+    TimeRangeInputs,
+    TimeRangeResponse,
+)
 from pydantic_evals.evaluators import (
     Evaluator,
     EvaluatorContext,
@@ -8,12 +13,6 @@ from pydantic_evals.evaluators import (
 )
 from pydantic_evals.otel import SpanQuery
 
-from pydantic_ai_examples.evals.models import (
-    TimeRangeBuilderSuccess,
-    TimeRangeInputs,
-    TimeRangeResponse,
-)
-
 
 @dataclass
 class ValidateTimeRange(Evaluator[TimeRangeInputs, TimeRangeResponse]):

pydantic_ai_examples/evals/example_01_generate_dataset.py

@@ -2,11 +2,10 @@ import asyncio
 from pathlib import Path
 from types import NoneType
 
+from pydantic_ai_examples.evals.models import TimeRangeInputs, TimeRangeResponse
 from pydantic_evals import Dataset
 from pydantic_evals.generation import generate_dataset
 
-from pydantic_ai_examples.evals.models import TimeRangeInputs, TimeRangeResponse
-
 
 async def main():
     dataset = await generate_dataset(

pydantic_ai_examples/evals/example_02_add_custom_evaluators.py

@@ -1,8 +1,6 @@
 from pathlib import Path
 from types import NoneType
 
-from pydantic_evals import Dataset
-
 from pydantic_ai_examples.evals.custom_evaluators import (
     CUSTOM_EVALUATOR_TYPES,
     AgentCalledTool,
@@ -13,6 +11,7 @@ from pydantic_ai_examples.evals.models import (
     TimeRangeInputs,
     TimeRangeResponse,
 )
+from pydantic_evals import Dataset
 
 
 def main():

pydantic_ai_examples/evals/example_03_unit_testing.py

@@ -2,7 +2,6 @@ from pathlib import Path
 from types import NoneType
 
 import logfire
-from pydantic_evals import Dataset
 
 from pydantic_ai_examples.evals import infer_time_range
 from pydantic_ai_examples.evals.custom_evaluators import (
@@ -12,6 +11,7 @@ from pydantic_ai_examples.evals.models import (
     TimeRangeInputs,
     TimeRangeResponse,
 )
+from pydantic_evals import Dataset
 
 logfire.configure(
     send_to_logfire='if-token-present',
@@ -29,7 +29,9 @@ def evaluate_dataset():
     report = dataset.evaluate_sync(infer_time_range)
     print(report)
 
-    assertion_pass_rate = report.averages().assertions
+    averages = report.averages()
+    assert averages is not None
+    assertion_pass_rate = averages.assertions
     assert assertion_pass_rate is not None, 'There should be at least one assertion'
     assert assertion_pass_rate > 0.9, (
         f'The assertion pass rate was {assertion_pass_rate:.1%}; it should be above 90%.'

pydantic_ai_examples/evals/example_04_compare_models.py

@@ -2,7 +2,6 @@ from pathlib import Path
 from types import NoneType
 
 import logfire
-from pydantic_evals import Dataset
 
 from pydantic_ai_examples.evals import infer_time_range
 from pydantic_ai_examples.evals.agent import time_range_agent
@@ -13,6 +12,7 @@ from pydantic_ai_examples.evals.models import (
     TimeRangeInputs,
     TimeRangeResponse,
 )
+from pydantic_evals import Dataset
 
 logfire.configure(
     send_to_logfire='if-token-present',
@@ -28,8 +28,8 @@ def compare_models():
         dataset_path, custom_evaluator_types=CUSTOM_EVALUATOR_TYPES
     )
     with logfire.span('Comparing different models for time_range_agent'):
-        with time_range_agent.override(model='openai:gpt-4o'):
-            dataset.evaluate_sync(infer_time_range, name='openai:gpt-4o')
+        with time_range_agent.override(model='openai:gpt-5'):
+            dataset.evaluate_sync(infer_time_range, name='openai:gpt-5')
         with time_range_agent.override(model='openai:o1'):
             dataset.evaluate_sync(infer_time_range, name='openai:o1')
 

pydantic_ai_examples/flight_booking.py

@@ -11,12 +11,18 @@ import logfire
 from pydantic import BaseModel, Field
 from rich.prompt import Prompt
 
-from pydantic_ai import Agent, ModelRetry, RunContext
-from pydantic_ai.messages import ModelMessage
-from pydantic_ai.usage import Usage, UsageLimits
+from pydantic_ai import (
+    Agent,
+    ModelMessage,
+    ModelRetry,
+    RunContext,
+    RunUsage,
+    UsageLimits,
+)
 
 # '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()
 
 
 class FlightDetails(BaseModel):
@@ -43,19 +49,18 @@ class Deps:
 
 # This agent is responsible for controlling the flow of the conversation.
 search_agent = Agent[Deps, FlightDetails | NoFlightFound](
-    'openai:gpt-4o',
+    'openai:gpt-5',
     output_type=FlightDetails | NoFlightFound,  # type: ignore
     retries=4,
     system_prompt=(
         'Your job is to find the cheapest flight for the user on the given date. '
     ),
-    instrument=True,
 )
 
 
 # This agent is responsible for extracting flight details from web page text.
 extraction_agent = Agent(
-    'openai:gpt-4o',
+    'openai:gpt-5',
     output_type=list[FlightDetails],
     system_prompt='Extract all the flight details from the given text.',
 )
@@ -107,8 +112,8 @@ class Failed(BaseModel):
 
 # This agent is responsible for extracting the user's seat selection
 seat_preference_agent = Agent[None, SeatPreference | Failed](
-    'openai:gpt-4o',
-    output_type=SeatPreference | Failed,  # type: ignore
+    'openai:gpt-5',
+    output_type=SeatPreference | Failed,
     system_prompt=(
         "Extract the user's seat preference. "
         'Seats A and F are window seats. '
@@ -182,7 +187,7 @@ async def main():
         req_date=datetime.date(2025, 1, 10),
     )
     message_history: list[ModelMessage] | None = None
-    usage: Usage = Usage()
+    usage: RunUsage = RunUsage()
     # run the agent until a satisfactory flight is found
     while True:
         result = await search_agent.run(
@@ -213,7 +218,7 @@ async def main():
                 )
 
 
-async def find_seat(usage: Usage) -> SeatPreference:
+async def find_seat(usage: RunUsage) -> SeatPreference:
     message_history: list[ModelMessage] | None = None
     while True:
         answer = Prompt.ask('What seat would you like?')

pydantic_ai_examples/pydantic_model.py

@@ -1,4 +1,4 @@
-"""Simple example of using PydanticAI to construct a Pydantic model from a text input.
+"""Simple example of using Pydantic AI to construct a Pydantic model from a text input.
 
 Run with:
 
@@ -14,6 +14,7 @@ from pydantic_ai import Agent
 
 # '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()
 
 
 class MyModel(BaseModel):
@@ -21,9 +22,9 @@ class MyModel(BaseModel):
     country: str
 
 
-model = os.getenv('PYDANTIC_AI_MODEL', 'openai:gpt-4o')
+model = os.getenv('PYDANTIC_AI_MODEL', 'openai:gpt-5')
 print(f'Using model: {model}')
-agent = Agent(model, output_type=MyModel, instrument=True)
+agent = Agent(model, output_type=MyModel)
 
 if __name__ == '__main__':
     result = agent.run_sync('The windy city in the US of A.')

pydantic_ai_examples/question_graph.py

@@ -12,6 +12,8 @@ from pathlib import Path
 
 import logfire
 from groq import BaseModel
+
+from pydantic_ai import Agent, ModelMessage, format_as_xml
 from pydantic_graph import (
     BaseNode,
     End,
@@ -20,13 +22,11 @@ from pydantic_graph import (
 )
 from pydantic_graph.persistence.file import FileStatePersistence
 
-from pydantic_ai import Agent, format_as_xml
-from pydantic_ai.messages import ModelMessage
-
 # '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()
 
-ask_agent = Agent('openai:gpt-4o', output_type=str, instrument=True)
+ask_agent = Agent('openai:gpt-5', output_type=str)
 
 
 @dataclass
@@ -65,7 +65,7 @@ class EvaluationOutput(BaseModel, use_attribute_docstrings=True):
 
 
 evaluate_agent = Agent(
-    'openai:gpt-4o',
+    'openai:gpt-5',
     output_type=EvaluationOutput,
     system_prompt='Given a question and answer, evaluate if the answer is correct.',
 )

pydantic_ai_examples/rag.py

@@ -30,16 +30,17 @@ import asyncpg
 import httpx
 import logfire
 import pydantic_core
+from anyio import create_task_group
 from openai import AsyncOpenAI
 from pydantic import TypeAdapter
 from typing_extensions import AsyncGenerator
 
-from pydantic_ai import RunContext
-from pydantic_ai.agent import Agent
+from pydantic_ai import Agent, RunContext
 
 # '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_asyncpg()
+logfire.instrument_pydantic_ai()
 
 
 @dataclass
@@ -48,7 +49,7 @@ class Deps:
     pool: asyncpg.Pool
 
 
-agent = Agent('openai:gpt-4o', deps_type=Deps, instrument=True)
+agent = Agent('openai:gpt-5', deps_type=Deps)
 
 
 @agent.tool
@@ -114,7 +115,7 @@ async def build_search_db():
     async with httpx.AsyncClient() as client:
         response = await client.get(DOCS_JSON)
         response.raise_for_status()
-    sections = sessions_ta.validate_json(response.content)
+    sections = sections_ta.validate_json(response.content)
 
     openai = AsyncOpenAI()
     logfire.instrument_openai(openai)
@@ -126,9 +127,9 @@ async def build_search_db():
                     await conn.execute(DB_SCHEMA)
 
         sem = asyncio.Semaphore(10)
-        async with asyncio.TaskGroup() as tg:
+        async with create_task_group() as tg:
             for section in sections:
-                tg.create_task(insert_doc_section(sem, openai, pool, section))
+                tg.start_soon(insert_doc_section, sem, openai, pool, section)
 
 
 async def insert_doc_section(
@@ -182,7 +183,7 @@ class DocsSection:
         return '\n\n'.join((f'path: {self.path}', f'title: {self.title}', self.content))
 
 
-sessions_ta = TypeAdapter(list[DocsSection])
+sections_ta = TypeAdapter(list[DocsSection])
 
 
 # pyright: reportUnknownMemberType=false

pydantic_ai_examples/roulette_wheel.py

@@ -1,4 +1,4 @@
-"""Example demonstrating how to use PydanticAI to create a simple roulette game.
+"""Example demonstrating how to use Pydantic AI to create a simple roulette game.
 
 Run with:
     uv run -m pydantic_ai_examples.roulette_wheel
@@ -28,7 +28,6 @@ roulette_agent = Agent(
     system_prompt=(
         'Use the `roulette_wheel` function to determine if the customer has won based on the number they bet on.'
     ),
-    instrument=True,
 )
 
 

pydantic_ai_examples/slack_lead_qualifier/agent.py

@@ -0,0 +1,47 @@
+from textwrap import dedent
+from types import NoneType
+
+import logfire
+
+### [imports]
+from pydantic_ai import Agent, NativeOutput
+from pydantic_ai.common_tools.duckduckgo import duckduckgo_search_tool  ### [/imports]
+
+from .models import Analysis, Profile
+
+### [agent]
+agent = Agent(
+    'openai:gpt-5',
+    instructions=dedent(
+        """
+        When a new person joins our public Slack, please put together a brief snapshot so we can be most useful to them.
+
+        **What to include**
+
+        1. **Who they are:**  Any details about their professional role or projects (e.g. LinkedIn, GitHub, company bio).
+        2. **Where they work:**  Name of the organisation and its domain.
+        3. **How we can help:**  On a scale of 1–5, estimate how likely they are to benefit from **Pydantic Logfire**
+           (our paid observability tool) based on factors such as company size, product maturity, or AI usage.
+           *1 = probably not relevant, 5 = very strong fit.*
+
+        **Our products (for context only)**
+        • **Pydantic Validation** – Python data-validation (open source)
+        • **Pydantic AI** – Python agent framework (open source)
+        • **Pydantic Logfire** – Observability for traces, logs & metrics with first-class AI support (commercial)
+
+        **How to research**
+
+        • Use the provided DuckDuckGo search tool to research the person and the organization they work for, based on the email domain or what you find on e.g. LinkedIn and GitHub.
+        • If you can't find enough to form a reasonable view, return **None**.
+        """
+    ),
+    tools=[duckduckgo_search_tool()],
+    output_type=NativeOutput([Analysis, NoneType]),
+)  ### [/agent]
+
+
+### [analyze_profile]
+@logfire.instrument('Analyze profile')
+async def analyze_profile(profile: Profile) -> Analysis | None:
+    result = await agent.run(profile.as_prompt())
+    return result.output  ### [/analyze_profile]

pydantic_ai_examples/slack_lead_qualifier/app.py

@@ -0,0 +1,36 @@
+from typing import Any
+
+import logfire
+from fastapi import FastAPI, HTTPException, status
+from logfire.propagate import get_context
+
+from .models import Profile
+
+
+### [process_slack_member]
+def process_slack_member(profile: Profile):
+    from .modal import process_slack_member as _process_slack_member
+
+    _process_slack_member.spawn(
+        profile.model_dump(), logfire_ctx=get_context()
+    )  ### [/process_slack_member]
+
+
+### [app]
+app = FastAPI()
+logfire.instrument_fastapi(app, capture_headers=True)
+
+
+@app.post('/')
+async def process_webhook(payload: dict[str, Any]) -> dict[str, Any]:
+    if payload['type'] == 'url_verification':
+        return {'challenge': payload['challenge']}
+    elif (
+        payload['type'] == 'event_callback' and payload['event']['type'] == 'team_join'
+    ):
+        profile = Profile.model_validate(payload['event']['user']['profile'])
+
+        process_slack_member(profile)
+        return {'status': 'OK'}
+
+    raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY)  ### [/app]

pydantic_ai_examples/slack_lead_qualifier/functions.py

@@ -0,0 +1,85 @@
+import logfire
+
+### [imports]
+from .agent import analyze_profile
+from .models import Profile
+
+### [imports-daily_summary]
+from .slack import send_slack_message
+from .store import AnalysisStore  ### [/imports,/imports-daily_summary]
+
+### [constant-new_lead_channel]
+NEW_LEAD_CHANNEL = '#new-slack-leads'
+### [/constant-new_lead_channel]
+### [constant-daily_summary_channel]
+DAILY_SUMMARY_CHANNEL = '#daily-slack-leads-summary'
+### [/constant-daily_summary_channel]
+
+
+### [process_slack_member]
+@logfire.instrument('Process Slack member')
+async def process_slack_member(profile: Profile):
+    analysis = await analyze_profile(profile)
+    logfire.info('Analysis', analysis=analysis)
+
+    if analysis is None:
+        return
+
+    await AnalysisStore().add(analysis)
+
+    await send_slack_message(
+        NEW_LEAD_CHANNEL,
+        [
+            {
+                'type': 'header',
+                'text': {
+                    'type': 'plain_text',
+                    'text': f'New Slack member with score {analysis.relevance}/5',
+                },
+            },
+            {
+                'type': 'divider',
+            },
+            *analysis.as_slack_blocks(),
+        ],
+    )  ### [/process_slack_member]
+
+
+### [send_daily_summary]
+@logfire.instrument('Send daily summary')
+async def send_daily_summary():
+    analyses = await AnalysisStore().list()
+    logfire.info('Analyses', analyses=analyses)
+
+    if len(analyses) == 0:
+        return
+
+    sorted_analyses = sorted(analyses, key=lambda x: x.relevance, reverse=True)
+    top_analyses = sorted_analyses[:5]
+
+    blocks = [
+        {
+            'type': 'header',
+            'text': {
+                'type': 'plain_text',
+                'text': f'Top {len(top_analyses)} new Slack members from the last 24 hours',
+            },
+        },
+    ]
+
+    for analysis in top_analyses:
+        blocks.extend(
+            [
+                {
+                    'type': 'divider',
+                },
+                *analysis.as_slack_blocks(include_relevance=True),
+            ]
+        )
+
+    await send_slack_message(
+        DAILY_SUMMARY_CHANNEL,
+        blocks,
+    )
+
+    await AnalysisStore().clear()  ### [/send_daily_summary]

pydantic_ai_examples/slack_lead_qualifier/modal.py

@@ -0,0 +1,66 @@
+from typing import Any
+
+### [setup_modal]
+import modal
+
+image = modal.Image.debian_slim(python_version='3.13').pip_install(
+    'pydantic',
+    'pydantic_ai_slim[openai,duckduckgo]',
+    'logfire[httpx,fastapi]',
+    'fastapi[standard]',
+    'httpx',
+)
+app = modal.App(
+    name='slack-lead-qualifier',
+    image=image,
+    secrets=[
+        modal.Secret.from_name('logfire'),
+        modal.Secret.from_name('openai'),
+        modal.Secret.from_name('slack'),
+    ],
+)  ### [/setup_modal]
+
+
+### [setup_logfire]
+def setup_logfire():
+    import logfire
+
+    logfire.configure(service_name=app.name)
+    logfire.instrument_pydantic_ai()
+    logfire.instrument_httpx(capture_all=True)  ### [/setup_logfire]
+
+
+### [web_app]
+@app.function(min_containers=1)
+@modal.asgi_app()  # type: ignore
+def web_app():
+    setup_logfire()
+
+    from .app import app as _app
+
+    return _app  ### [/web_app]
+
+
+### [process_slack_member]
+@app.function()
+async def process_slack_member(profile_raw: dict[str, Any], logfire_ctx: Any):
+    setup_logfire()
+
+    from logfire.propagate import attach_context
+
+    from .functions import process_slack_member as _process_slack_member
+    from .models import Profile
+
+    with attach_context(logfire_ctx):
+        profile = Profile.model_validate(profile_raw)
+        await _process_slack_member(profile)  ### [/process_slack_member]
+
+
+### [send_daily_summary]
+@app.function(schedule=modal.Cron('0 8 * * *'))  # Every day at 8am UTC
+async def send_daily_summary():
+    setup_logfire()
+
+    from .functions import send_daily_summary as _send_daily_summary
+
+    await _send_daily_summary()  ### [/send_daily_summary]