Improve examples (#18)

Co-authored-by: hyperlint-ai[bot] <154288675+hyperlint-ai[bot]@users.noreply.github.com>

Makefile

@@ -1,5 +1,4 @@
 .DEFAULT_GOAL := all
-sources = pydantic_ai tests
 
 .PHONY: .uv  # Check that uv is installed
 .uv:
@@ -16,13 +15,13 @@ install: .uv .pre-commit
 
 .PHONY: format  # Format the code
 format:
-	uv run ruff format $(sources)
-	uv run ruff check --fix --fix-only $(sources)
+	uv run ruff format
+	uv run ruff check --fix --fix-only
 
 .PHONY: lint  # Lint the code
 lint:
-	uv run ruff format --check $(sources)
-	uv run ruff check $(sources)
+	uv run ruff format --check
+	uv run ruff check
 
 .PHONY: typecheck-pyright
 typecheck-pyright:

examples/README.md

@@ -0,0 +1,73 @@
+# Pydantic AI Examples
+
+Examples of how to use Pydantic AI and what it can do.
+
+## Usage
+
+To run the examples, run:
+
+```bash
+uv run -m examples.<example_module_name>
+```
+
+## Examples
+
+### `pydantic_model.py`
+
+(Demonstrates: custom `result_type`)
+
+Simple example of using Pydantic AI to construct a Pydantic model from a text input.
+
+```bash
+uv run --extra examples -m examples.pydantic_model
+```
+
+This examples uses `openai:gpt-4o` by default but it works well with other models, e.g. you can run it
+with Gemini using:
+
+```bash
+PYDANTIC_AI_MODEL=gemini-1.5-pro uv run --extra examples -m examples.pydantic_model
+```
+
+(or `PYDANTIC_AI_MODEL=gemini-1.5-flash...`)
+
+### `sql_gen.py`
+
+(Demonstrates: custom `result_type`, dynamic system prompt, result validation, agent deps)
+
+Example demonstrating how to use Pydantic AI to generate SQL queries based on user input.
+
+```bash
+uv run --extra examples -m examples.sql_gen
+```
+
+or to use a custom prompt:
+
+```bash
+uv run --extra examples -m examples.sql_gen "find me whatever"
+```
+
+This model uses `gemini-1.5-flash` by default since Gemini is good at single shot queries.
+
+### `weather.py`
+
+(Demonstrates: retrievers, multiple retrievers, agent deps)
+
+Example of Pydantic AI with multiple tools which the LLM needs to call in turn to answer a question.
+
+In this case the idea is a "weather" agent — the user can ask for the weather in multiple cities,
+the agent will use the `get_lat_lng` tool to get the latitude and longitude of the locations, then use
+the `get_weather` tool to get the weather.
+
+To run this example properly, you'll need two extra API keys:
+* A weather API key from [tomorrow.io](https://www.tomorrow.io/weather-api/) set via `WEATHER_API_KEY`
+* A geocoding API key from [geocode.maps.co](https://geocode.maps.co/) set via `GEO_API_KEY`
+
+**(Note if either key is missing, the code will fall back to dummy data.)**
+
+```bash
+uv run --extra examples -m examples.weather
+```
+
+This example uses `openai:gpt-4o` by default. Gemini seems to be unable to handle the multiple tool
+calls.

examples/pydantic_model.py

@@ -0,0 +1,31 @@
+"""Simple example of using Pydantic AI to construct a Pydantic model from a text input.
+
+Run with:
+
+    uv run --extra examples -m examples.pydantic_model
+"""
+
+import os
+from typing import cast
+
+import logfire
+from pydantic import BaseModel
+
+from pydantic_ai import Agent
+from pydantic_ai.agent import KnownModelName
+
+# 'if-token-present' means nothing will be sent (and the example wil work) if you don't have logfire set up
+logfire.configure(send_to_logfire='if-token-present')
+
+
+class MyModel(BaseModel):
+    city: str
+    country: str
+
+
+model = cast(KnownModelName, os.getenv('PYDANTIC_AI_MODEL', 'openai:gpt-4o'))
+agent = Agent(model, result_type=MyModel, deps=None)
+
+if __name__ == '__main__':
+    result = agent.run_sync('The windy city in the US of A.')
+    print(result.response)

examples/sql_gen.py

@@ -1,45 +1,80 @@
+"""Example demonstrating how to use Pydantic AI to generate SQL queries based on user input.
+
+Run with:
+
+    uv run --extra examples -m examples.sql_gen
+"""
+
+import asyncio
+import os
+import sys
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
 from dataclasses import dataclass
+from datetime import date
+from typing import Annotated, Any, cast
 
-from pydantic_ai import Agent
+import asyncpg
+import logfire
+from annotated_types import MinLen
 from devtools import debug
 
-system_prompt = """\
-Given the following PostgreSQL table of records, your job is to write a SQL query that suits the user's request.
+from pydantic_ai import Agent, CallContext, ModelRetry
+from pydantic_ai.agent import KnownModelName
+
+# 'if-token-present' means nothing will be sent (and the example wil work) if you don't have logfire set up
+logfire.configure()
 
-CREATE TABLE records AS (
-    start_timestamp timestamp with time zone,
-    created_at timestamp with time zone,
+DB_SCHEMA = """
+CREATE TABLE IF NOT EXISTS records (
+    created_at timestamptz,
+    start_timestamp timestamptz,
+    end_timestamp timestamptz,
     trace_id text,
     span_id text,
     parent_span_id text,
-    kind span_kind,
-    end_timestamp timestamp with time zone,
-    level smallint,
+    level log_level,
     span_name text,
     message text,
     attributes_json_schema text,
     attributes jsonb,
     tags text[],
-    otel_links jsonb,
-    otel_events jsonb,
     is_exception boolean,
-    otel_status_code status_code,
     otel_status_message text,
-    otel_scope_name text,
-    otel_scope_version text,
-    otel_scope_attributes jsonb,
-    service_namespace text,
-    service_name text,
-    service_version text,
-    service_instance_id text,
-    process_pid integer
+    service_name text
 );
+"""
+
+
+@dataclass
+class Response:
+    sql_query: Annotated[str, MinLen(1)]
+
+
+@dataclass
+class Deps:
+    conn: asyncpg.Connection
+
 
-today's date = 2024-10-09
+model = cast(KnownModelName, os.getenv('PYDANTIC_AI_MODEL', 'gemini-1.5-flash'))
+agent: Agent[Deps, Response] = Agent(model, result_type=Response)
+
+
+@agent.system_prompt
+async def system_prompt() -> str:
+    return f"""\
+Given the following PostgreSQL table of records, your job is to write a SQL query that suits the user's request.
+
+{DB_SCHEMA}
+
+today's date = {date.today()}
 
 Example
     request: show me records where foobar is false
-    response: SELECT * FROM records WHERE attributes->>'foobar' = false'
+    response: SELECT * FROM records WHERE attributes->>'foobar' = false
+Example
+    request: show me records where attributes include the key "foobar"
+    response: SELECT * FROM records WHERE attributes ? 'foobar'
 Example
     request: show me records from yesterday
     response: SELECT * FROM records WHERE start_timestamp::date > CURRENT_TIMESTAMP - INTERVAL '1 day'
@@ -49,15 +84,59 @@
 """
 
 
-@dataclass
-class Response:
-    sql_query: str
+@agent.result_validator
+async def validate_result(ctx: CallContext[Deps], result: Response) -> Response:
+    result.sql_query = result.sql_query.replace('\\', '')
+    lower_query = result.sql_query.lower()
+    if not lower_query.startswith('select'):
+        raise ModelRetry('Please a SELECT query')
 
+    try:
+        await ctx.deps.conn.execute(f'EXPLAIN {result.sql_query}')
+    except asyncpg.exceptions.PostgresError as e:
+        raise ModelRetry(f'Invalid query: {e}') from e
+    else:
+        return result
 
-agent = Agent('gemini-1.5-flash', result_type=Response, system_prompt=system_prompt, deps=None)
 
+async def main():
+    if len(sys.argv) == 1:
+        prompt = 'show me logs from yesterday, with level "error"'
+    else:
+        prompt = sys.argv[1]
 
-if __name__ == '__main__':
-    with debug.timer('SQL Generation'):
-        result = agent.run_sync('show me logs from yesterday, with level "error"')
+    async with database_connect('postgresql://postgres@localhost', 'pydantic_ai_sql_gen') as conn:
+        deps = Deps(conn)
+        result = await agent.run(prompt, deps=deps)
     debug(result.response.sql_query)
+
+
+# pyright: reportUnknownMemberType=false
+# pyright: reportUnknownVariableType=false
+@asynccontextmanager
+async def database_connect(server_dsn: str, database: str) -> AsyncGenerator[Any, None]:
+    with logfire.span('check and create DB'):
+        conn = await asyncpg.connect(server_dsn)
+        try:
+            db_exists = await conn.fetchval('SELECT 1 FROM pg_database WHERE datname = $1', database)
+            if not db_exists:
+                await conn.execute(f'CREATE DATABASE {database}')
+        finally:
+            await conn.close()
+
+    conn = await asyncpg.connect(f'{server_dsn}/{database}')
+    try:
+        with logfire.span('create schema'):
+            async with conn.transaction():
+                if not db_exists:
+                    await conn.execute(
+                        "CREATE TYPE log_level AS ENUM ('debug', 'info', 'warning', 'error', 'critical')"
+                    )
+                await conn.execute(DB_SCHEMA)
+        yield conn
+    finally:
+        await conn.close()
+
+
+if __name__ == '__main__':
+    asyncio.run(main())