add pydantic demos, improve weather demo

Author: samuelcolvin

logfire-hello-world/main.py

@@ -1,5 +1,5 @@
 import logfire
 
-logfire.configure(environment='hello-world')
+logfire.configure(service_name='hello-world')
 
 logfire.info('hello {place}', place='world')

pai-pydantic/retry.py

@@ -0,0 +1,31 @@
+from datetime import date
+
+import logfire
+from pydantic import BaseModel, field_validator
+from pydantic_ai import Agent
+
+logfire.configure(service_name='pai-pydantic-retry')
+logfire.instrument_pydantic_ai()
+
+
+class Person(BaseModel):
+    """Definition of an historic person"""
+
+    name: str
+    dob: date
+    city: str
+
+    @field_validator('dob')
+    def validate_dob(cls, v: date) -> date:
+        if v >= date(1900, 1, 1):
+            raise ValueError('The person must be born in the 19th century')
+        return v
+
+
+agent = Agent(
+    'openai:gpt-4o',
+    output_type=Person,
+    instructions='Extract information about the person',
+)
+result = agent.run_sync("Samuel lived in London and was born on Jan 28th '87")
+print(repr(result.output))

pai-pydantic/simple.py

@@ -0,0 +1,23 @@
+from datetime import date
+
+import logfire
+from pydantic import BaseModel
+from pydantic_ai import Agent
+
+logfire.configure(service_name='pai-pydantic-simple')
+logfire.instrument_pydantic_ai()
+
+
+class Person(BaseModel):
+    name: str
+    dob: date
+    city: str
+
+
+agent = Agent(
+    'openai:gpt-4o',
+    output_type=Person,
+    instructions='Extract information about the person',
+)
+result = agent.run_sync("Samuel lived in London and was born on Jan 28th '87")
+print(repr(result.output))

pai-weather/main.py

@@ -1,65 +1,53 @@
 from __future__ import annotations as _annotations
 
 import asyncio
-import os
-import urllib.parse
 from dataclasses import dataclass
+from random import randint
 from typing import Any
 
 import logfire
 from httpx import AsyncClient
-from pydantic_ai import Agent, ModelRetry, RunContext
+from pydantic import BaseModel
+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()
+logfire.configure(service_name='pai-weather')
 logfire.instrument_pydantic_ai()
-logfire.instrument_httpx(capture_all=True)
 
 
 @dataclass
 class Deps:
     client: AsyncClient
-    weather_api_key: str | None
-    geo_api_key: str | None
 
 
 weather_agent = Agent(
-    'openai:gpt-4o',
+    'openai:gpt-4.1-mini',
     # 'Be concise, reply with one sentence.' is enough for some models (like openai) to use
     # the below tools appropriately, but others like anthropic and gemini require a bit more direction.
-    instructions=(
-        'Be concise, reply with one sentence.'
-        '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.'
-    ),
+    instructions='Be concise, reply with one sentence.',
     deps_type=Deps,
     retries=2,
 )
 
 
+class LatLng(BaseModel):
+    lat: float
+    lng: float
+
+
 @weather_agent.tool
-async def get_lat_lng(ctx: RunContext[Deps], location_description: str) -> dict[str, float]:
+async def get_lat_lng(ctx: RunContext[Deps], location_description: str) -> LatLng:
     """Get the latitude and longitude of a location.
 
     Args:
         ctx: The context.
         location_description: A description of a location.
     """
-    if ctx.deps.geo_api_key is None:
-        # if no API key is provided, return a dummy response (London)
-        return {'lat': 51.1, 'lng': -0.1}
-
-    params = {'access_token': ctx.deps.geo_api_key}
-    loc = urllib.parse.quote(location_description)
-    r = await ctx.deps.client.get(f'https://api.mapbox.com/geocoding/v5/mapbox.places/{loc}.json', params=params)
+    r = await ctx.deps.client.get(
+        'https://demo-endpoints.pydantic.workers.dev/latlng',
+        params={'location': location_description, 'sleep': randint(200, 1200)},
+    )
     r.raise_for_status()
-    data = r.json()
-
-    if features := data['features']:
-        lat, lng = features[0]['center']
-        return {'lat': lat, 'lng': lng}
-    else:
-        raise ModelRetry('Could not find the location')
+    return LatLng.model_validate_json(r.content)
 
 
 @weather_agent.tool
@@ -71,62 +59,25 @@ async def get_weather(ctx: RunContext[Deps], lat: float, lng: float) -> dict[str
         lat: Latitude of the location.
         lng: Longitude of the location.
     """
-    if ctx.deps.weather_api_key is None:
-        # if no API key is provided, return a dummy response
-        return {'temperature': '21 °C', 'description': 'Sunny'}
-
-    params = {
-        'apikey': ctx.deps.weather_api_key,
-        'location': f'{lat},{lng}',
-        'units': 'metric',
-    }
-    with logfire.span('calling weather API', params=params) as span:
-        r = await ctx.deps.client.get('https://api.tomorrow.io/v4/weather/realtime', params=params)
-        r.raise_for_status()
-        data = r.json()
-        span.set_attribute('response', data)
-
-    values = data['data']['values']
-    # https://docs.tomorrow.io/reference/data-layers-weather-codes
-    code_lookup = {
-        1000: 'Clear, Sunny',
-        1100: 'Mostly Clear',
-        1101: 'Partly Cloudy',
-        1102: 'Mostly Cloudy',
-        1001: 'Cloudy',
-        2000: 'Fog',
-        2100: 'Light Fog',
-        4000: 'Drizzle',
-        4001: 'Rain',
-        4200: 'Light Rain',
-        4201: 'Heavy Rain',
-        5000: 'Snow',
-        5001: 'Flurries',
-        5100: 'Light Snow',
-        5101: 'Heavy Snow',
-        6000: 'Freezing Drizzle',
-        6001: 'Freezing Rain',
-        6200: 'Light Freezing Rain',
-        6201: 'Heavy Freezing Rain',
-        7000: 'Ice Pellets',
-        7101: 'Heavy Ice Pellets',
-        7102: 'Light Ice Pellets',
-        8000: 'Thunderstorm',
-    }
-    return {
-        'temperature': f'{values["temperatureApparent"]:0.0f}°C',
-        'description': code_lookup.get(values['weatherCode'], 'Unknown'),
-    }
+    temp_response, descr_response = await asyncio.gather(
+        ctx.deps.client.get(
+            'https://demo-endpoints.pydantic.workers.dev/number',
+            params={'min': 10, 'max': 30, 'sleep': randint(200, 1200)},
+        ),
+        ctx.deps.client.get(
+            'https://demo-endpoints.pydantic.workers.dev/weather',
+            params={'lat': lat, 'lng': lng, 'sleep': randint(200, 1200)},
+        ),
+    )
+    temp_response.raise_for_status()
+    descr_response.raise_for_status()
+    return {'temperature': f'{temp_response.text} °C', 'description': descr_response.text}
 
 
 async def main():
     async with AsyncClient() as client:
         logfire.instrument_httpx(client, capture_all=True)
-        # create a free API key at https://www.tomorrow.io/weather-api/
-        weather_api_key = os.getenv('WEATHER_API_KEY')
-        # create a free API key at https://www.mapbox.com/
-        geo_api_key = os.getenv('GEO_API_KEY')
-        deps = Deps(client=client, weather_api_key=weather_api_key, geo_api_key=geo_api_key)
+        deps = Deps(client=client)
         result = await weather_agent.run('What is the weather like in London and in Wiltshire?', deps=deps)
         print('Response:', result.output)