Skip to content

Commit 09e0821

Browse files
authored
simplify weather example (#2129)
1 parent 2e12934 commit 09e0821

File tree

2 files changed

+31
-88
lines changed

2 files changed

+31
-88
lines changed

examples/pydantic_ai_examples/weather_agent.py

Lines changed: 30 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -12,16 +12,14 @@
1212
from __future__ import annotations as _annotations
1313

1414
import asyncio
15-
import os
16-
import urllib.parse
1715
from dataclasses import dataclass
1816
from typing import Any
1917

2018
import logfire
21-
from devtools import debug
2219
from httpx import AsyncClient
20+
from pydantic import BaseModel
2321

24-
from pydantic_ai import Agent, ModelRetry, RunContext
22+
from pydantic_ai import Agent, RunContext
2523

2624
# 'if-token-present' means nothing will be sent (and the example will work) if you don't have logfire configured
2725
logfire.configure(send_to_logfire='if-token-present')
@@ -31,51 +29,38 @@
3129
@dataclass
3230
class Deps:
3331
client: AsyncClient
34-
weather_api_key: str | None
35-
geo_api_key: str | None
3632

3733

3834
weather_agent = Agent(
39-
'openai:gpt-4o',
35+
'openai:gpt-4.1-mini',
4036
# 'Be concise, reply with one sentence.' is enough for some models (like openai) to use
4137
# the below tools appropriately, but others like anthropic and gemini require a bit more direction.
42-
instructions=(
43-
'Be concise, reply with one sentence.'
44-
'Use the `get_lat_lng` tool to get the latitude and longitude of the locations, '
45-
'then use the `get_weather` tool to get the weather.'
46-
),
38+
instructions='Be concise, reply with one sentence.',
4739
deps_type=Deps,
4840
retries=2,
4941
)
5042

5143

44+
class LatLng(BaseModel):
45+
lat: float
46+
lng: float
47+
48+
5249
@weather_agent.tool
53-
async def get_lat_lng(
54-
ctx: RunContext[Deps], location_description: str
55-
) -> dict[str, float]:
50+
async def get_lat_lng(ctx: RunContext[Deps], location_description: str) -> LatLng:
5651
"""Get the latitude and longitude of a location.
5752
5853
Args:
5954
ctx: The context.
6055
location_description: A description of a location.
6156
"""
62-
if ctx.deps.geo_api_key is None:
63-
# if no API key is provided, return a dummy response (London)
64-
return {'lat': 51.1, 'lng': -0.1}
65-
66-
params = {'access_token': ctx.deps.geo_api_key}
67-
loc = urllib.parse.quote(location_description)
57+
# NOTE: the response here will be random, and is not related to the location description.
6858
r = await ctx.deps.client.get(
69-
f'https://api.mapbox.com/geocoding/v5/mapbox.places/{loc}.json', params=params
59+
'https://demo-endpoints.pydantic.workers.dev/latlng',
60+
params={'location': location_description},
7061
)
7162
r.raise_for_status()
72-
data = r.json()
73-
74-
if features := data['features']:
75-
lat, lng = features[0]['center']
76-
return {'lat': lat, 'lng': lng}
77-
else:
78-
raise ModelRetry('Could not find the location')
63+
return LatLng.model_validate_json(r.content)
7964

8065

8166
@weather_agent.tool
@@ -87,70 +72,32 @@ async def get_weather(ctx: RunContext[Deps], lat: float, lng: float) -> dict[str
8772
lat: Latitude of the location.
8873
lng: Longitude of the location.
8974
"""
90-
if ctx.deps.weather_api_key is None:
91-
# if no API key is provided, return a dummy response
92-
return {'temperature': '21 °C', 'description': 'Sunny'}
93-
94-
params = {
95-
'apikey': ctx.deps.weather_api_key,
96-
'location': f'{lat},{lng}',
97-
'units': 'metric',
98-
}
99-
with logfire.span('calling weather API', params=params) as span:
100-
r = await ctx.deps.client.get(
101-
'https://api.tomorrow.io/v4/weather/realtime', params=params
102-
)
103-
r.raise_for_status()
104-
data = r.json()
105-
span.set_attribute('response', data)
106-
107-
values = data['data']['values']
108-
# https://docs.tomorrow.io/reference/data-layers-weather-codes
109-
code_lookup = {
110-
1000: 'Clear, Sunny',
111-
1100: 'Mostly Clear',
112-
1101: 'Partly Cloudy',
113-
1102: 'Mostly Cloudy',
114-
1001: 'Cloudy',
115-
2000: 'Fog',
116-
2100: 'Light Fog',
117-
4000: 'Drizzle',
118-
4001: 'Rain',
119-
4200: 'Light Rain',
120-
4201: 'Heavy Rain',
121-
5000: 'Snow',
122-
5001: 'Flurries',
123-
5100: 'Light Snow',
124-
5101: 'Heavy Snow',
125-
6000: 'Freezing Drizzle',
126-
6001: 'Freezing Rain',
127-
6200: 'Light Freezing Rain',
128-
6201: 'Heavy Freezing Rain',
129-
7000: 'Ice Pellets',
130-
7101: 'Heavy Ice Pellets',
131-
7102: 'Light Ice Pellets',
132-
8000: 'Thunderstorm',
133-
}
75+
# NOTE: the responses here will be random, and are not related to the lat and lng.
76+
temp_response, descr_response = await asyncio.gather(
77+
ctx.deps.client.get(
78+
'https://demo-endpoints.pydantic.workers.dev/number',
79+
params={'min': 10, 'max': 30},
80+
),
81+
ctx.deps.client.get(
82+
'https://demo-endpoints.pydantic.workers.dev/weather',
83+
params={'lat': lat, 'lng': lng},
84+
),
85+
)
86+
temp_response.raise_for_status()
87+
descr_response.raise_for_status()
13488
return {
135-
'temperature': f'{values["temperatureApparent"]:0.0f}°C',
136-
'description': code_lookup.get(values['weatherCode'], 'Unknown'),
89+
'temperature': f'{temp_response.text} °C',
90+
'description': descr_response.text,
13791
}
13892

13993

14094
async def main():
14195
async with AsyncClient() as client:
14296
logfire.instrument_httpx(client, capture_all=True)
143-
# create a free API key at https://www.tomorrow.io/weather-api/
144-
weather_api_key = os.getenv('WEATHER_API_KEY')
145-
# create a free API key at https://www.mapbox.com/
146-
geo_api_key = os.getenv('GEO_API_KEY')
147-
deps = Deps(
148-
client=client, weather_api_key=weather_api_key, geo_api_key=geo_api_key
149-
)
97+
deps = Deps(client=client)
15098
result = await weather_agent.run(
15199
'What is the weather like in London and in Wiltshire?', deps=deps
152100
)
153-
debug(result)
154101
print('Response:', result.output)
155102

156103

examples/pydantic_ai_examples/weather_agent_gradio.py

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
from __future__ import annotations as _annotations
22

33
import json
4-
import os
54

65
from httpx import AsyncClient
76

@@ -18,10 +17,7 @@
1817
TOOL_TO_DISPLAY_NAME = {'get_lat_lng': 'Geocoding API', 'get_weather': 'Weather API'}
1918

2019
client = AsyncClient()
21-
weather_api_key = os.getenv('WEATHER_API_KEY')
22-
# create a free API key at https://geocode.maps.co/
23-
geo_api_key = os.getenv('GEO_API_KEY')
24-
deps = Deps(client=client, weather_api_key=weather_api_key, geo_api_key=geo_api_key)
20+
deps = Deps(client=client)
2521

2622

2723
async def stream_from_agent(prompt: str, chatbot: list[dict], past_messages: list):

0 commit comments

Comments
 (0)