Skip to content

Commit 0b626f8

Browse files
dmontagusamuelcolvin
dmontagu
and
samuelcolvin
authored
Add timezones (#25)
Co-authored-by: Samuel Colvin <[email protected]>
1 parent 7f9e9df commit 0b626f8

File tree

8 files changed

+95
-66
lines changed

8 files changed

+95
-66
lines changed

‎pydantic_ai/messages.py‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
import json
44
from dataclasses import dataclass, field
5-
from datetime import datetime
5+
from datetime import datetime, timezone
66
from typing import Annotated, Any, Literal, Union
77

88
import pydantic
@@ -11,6 +11,10 @@
1111
from . import _pydantic
1212

1313

14+
def _now_utc() -> datetime:
15+
return datetime.now(tz=timezone.utc)
16+
17+
1418
@dataclass
1519
class SystemPrompt:
1620
content: str
@@ -20,7 +24,7 @@ class SystemPrompt:
2024
@dataclass
2125
class UserPrompt:
2226
content: str
23-
timestamp: datetime = field(default_factory=datetime.now)
27+
timestamp: datetime = field(default_factory=_now_utc)
2428
role: Literal['user'] = 'user'
2529

2630

@@ -32,7 +36,7 @@ class ToolReturn:
3236
tool_name: str
3337
content: str | dict[str, Any]
3438
tool_id: str | None = None
35-
timestamp: datetime = field(default_factory=datetime.now)
39+
timestamp: datetime = field(default_factory=_now_utc)
3640
role: Literal['tool-return'] = 'tool-return'
3741

3842
def model_response_str(self) -> str:
@@ -54,7 +58,7 @@ class RetryPrompt:
5458
content: list[pydantic_core.ErrorDetails] | str
5559
tool_name: str | None = None
5660
tool_id: str | None = None
57-
timestamp: datetime = field(default_factory=datetime.now)
61+
timestamp: datetime = field(default_factory=_now_utc)
5862
role: Literal['retry-prompt'] = 'retry-prompt'
5963

6064
def model_response(self) -> str:
@@ -68,7 +72,7 @@ def model_response(self) -> str:
6872
@dataclass
6973
class LLMResponse:
7074
content: str
71-
timestamp: datetime = field(default_factory=datetime.now)
75+
timestamp: datetime = field(default_factory=_now_utc)
7276
role: Literal['llm-response'] = 'llm-response'
7377

7478

@@ -102,7 +106,7 @@ def from_object(cls, tool_name: str, args_object: dict[str, Any]) -> ToolCall:
102106
@dataclass
103107
class LLMToolCalls:
104108
calls: list[ToolCall]
105-
timestamp: datetime = field(default_factory=datetime.now)
109+
timestamp: datetime = field(default_factory=_now_utc)
106110
role: Literal['llm-tool-calls'] = 'llm-tool-calls'
107111

108112

‎pydantic_ai/models/openai.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from collections.abc import Mapping, Sequence
44
from dataclasses import dataclass, field
5-
from datetime import datetime
5+
from datetime import datetime, timezone
66
from typing import Literal
77

88
from httpx import AsyncClient as AsyncHTTPClient
@@ -91,7 +91,7 @@ async def request(self, messages: list[Message]) -> tuple[LLMMessage, shared.Cos
9191
@staticmethod
9292
def process_response(response: chat.ChatCompletion) -> LLMMessage:
9393
choice = response.choices[0]
94-
timestamp = datetime.fromtimestamp(response.created)
94+
timestamp = datetime.fromtimestamp(response.created, tz=timezone.utc)
9595
if choice.message.tool_calls is not None:
9696
return LLMToolCalls(
9797
[ToolCall.from_json(c.function.name, c.function.arguments, c.id) for c in choice.message.tool_calls],

‎tests/models/test_gemini.py‎

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import json
44
from collections.abc import Callable
55
from dataclasses import dataclass
6+
from datetime import timezone
67

78
import httpx
89
import pytest
@@ -361,8 +362,8 @@ async def test_request_simple_success(get_gemini_client: GetGeminiClient):
361362
assert result.response == 'Hello world'
362363
assert result.message_history == snapshot(
363364
[
364-
UserPrompt(content='Hello', timestamp=IsNow()),
365-
LLMResponse(content='Hello world', timestamp=IsNow()),
365+
UserPrompt(content='Hello', timestamp=IsNow(tz=timezone.utc)),
366+
LLMResponse(content='Hello world', timestamp=IsNow(tz=timezone.utc)),
366367
]
367368
)
368369
assert result.cost == snapshot(Cost(request_tokens=1, response_tokens=2, total_tokens=3))
@@ -382,15 +383,15 @@ async def test_request_structured_response(get_gemini_client: GetGeminiClient):
382383
assert result.response == [1, 2, 123]
383384
assert result.message_history == snapshot(
384385
[
385-
UserPrompt(content='Hello', timestamp=IsNow()),
386+
UserPrompt(content='Hello', timestamp=IsNow(tz=timezone.utc)),
386387
LLMToolCalls(
387388
calls=[
388389
ToolCall(
389390
tool_name='final_result',
390391
args=ArgsObject(args_object={'response': [1, 2, 123]}),
391392
)
392393
],
393-
timestamp=IsNow(),
394+
timestamp=IsNow(tz=timezone.utc),
394395
),
395396
]
396397
)
@@ -426,28 +427,30 @@ async def get_location(loc_name: str) -> str:
426427
assert result.message_history == snapshot(
427428
[
428429
SystemPrompt(content='this is the system prompt'),
429-
UserPrompt(content='Hello', timestamp=IsNow()),
430+
UserPrompt(content='Hello', timestamp=IsNow(tz=timezone.utc)),
430431
LLMToolCalls(
431432
calls=[
432433
ToolCall(
433434
tool_name='get_location',
434435
args=ArgsObject(args_object={'loc_name': 'San Fransisco'}),
435436
)
436437
],
437-
timestamp=IsNow(),
438+
timestamp=IsNow(tz=timezone.utc),
439+
),
440+
RetryPrompt(
441+
tool_name='get_location', content='Wrong location, please try again', timestamp=IsNow(tz=timezone.utc)
438442
),
439-
RetryPrompt(tool_name='get_location', content='Wrong location, please try again', timestamp=IsNow()),
440443
LLMToolCalls(
441444
calls=[
442445
ToolCall(
443446
tool_name='get_location',
444447
args=ArgsObject(args_object={'loc_name': 'London'}),
445448
)
446449
],
447-
timestamp=IsNow(),
450+
timestamp=IsNow(tz=timezone.utc),
448451
),
449-
ToolReturn(tool_name='get_location', content='{"lat": 51, "lng": 0}', timestamp=IsNow()),
450-
LLMResponse(content='final response', timestamp=IsNow()),
452+
ToolReturn(tool_name='get_location', content='{"lat": 51, "lng": 0}', timestamp=IsNow(tz=timezone.utc)),
453+
LLMResponse(content='final response', timestamp=IsNow(tz=timezone.utc)),
451454
]
452455
)
453456
assert result.cost == snapshot(Cost(request_tokens=3, response_tokens=6, total_tokens=9))

‎tests/models/test_model_function.py‎

Lines changed: 31 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
import json
2+
import re
23
from dataclasses import asdict
4+
from datetime import timezone
35

46
import pydantic_core
57
import pytest
8+
from dirty_equals import IsStr
69
from inline_snapshot import snapshot
710
from pydantic import BaseModel
811

@@ -38,12 +41,12 @@ def test_simple():
3841
[
3942
UserPrompt(
4043
content='Hello',
41-
timestamp=IsNow(),
44+
timestamp=IsNow(tz=timezone.utc),
4245
role='user',
4346
),
4447
LLMResponse(
4548
content="content='Hello' role='user' message_count=1",
46-
timestamp=IsNow(),
49+
timestamp=IsNow(tz=timezone.utc),
4750
role='llm-response',
4851
),
4952
]
@@ -55,22 +58,22 @@ def test_simple():
5558
[
5659
UserPrompt(
5760
content='Hello',
58-
timestamp=IsNow(),
61+
timestamp=IsNow(tz=timezone.utc),
5962
role='user',
6063
),
6164
LLMResponse(
6265
content="content='Hello' role='user' message_count=1",
63-
timestamp=IsNow(),
66+
timestamp=IsNow(tz=timezone.utc),
6467
role='llm-response',
6568
),
6669
UserPrompt(
6770
content='World',
68-
timestamp=IsNow(),
71+
timestamp=IsNow(tz=timezone.utc),
6972
role='user',
7073
),
7174
LLMResponse(
7275
content="content='World' role='user' message_count=3",
73-
timestamp=IsNow(),
76+
timestamp=IsNow(tz=timezone.utc),
7477
role='llm-response',
7578
),
7679
]
@@ -128,16 +131,19 @@ def test_weather():
128131
[
129132
UserPrompt(
130133
content='London',
131-
timestamp=IsNow(),
134+
timestamp=IsNow(tz=timezone.utc),
132135
role='user',
133136
),
134137
LLMToolCalls(
135138
calls=[ToolCall.from_json('get_location', '{"location_description": "London"}')],
136-
timestamp=IsNow(),
139+
timestamp=IsNow(tz=timezone.utc),
137140
role='llm-tool-calls',
138141
),
139142
ToolReturn(
140-
tool_name='get_location', content='{"lat": 51, "lng": 0}', timestamp=IsNow(), role='tool-return'
143+
tool_name='get_location',
144+
content='{"lat": 51, "lng": 0}',
145+
timestamp=IsNow(tz=timezone.utc),
146+
role='tool-return',
141147
),
142148
LLMToolCalls(
143149
calls=[
@@ -146,18 +152,18 @@ def test_weather():
146152
'{"lat": 51, "lng": 0}',
147153
)
148154
],
149-
timestamp=IsNow(),
155+
timestamp=IsNow(tz=timezone.utc),
150156
role='llm-tool-calls',
151157
),
152158
ToolReturn(
153159
tool_name='get_weather',
154160
content='Raining',
155-
timestamp=IsNow(),
161+
timestamp=IsNow(tz=timezone.utc),
156162
role='tool-return',
157163
),
158164
LLMResponse(
159165
content='Raining in London',
160-
timestamp=IsNow(),
166+
timestamp=IsNow(tz=timezone.utc),
161167
role='llm-response',
162168
),
163169
]
@@ -198,12 +204,14 @@ def get_var_args(ctx: CallContext[int], *args: int):
198204
def test_var_args():
199205
result = var_args_agent.run_sync('{"function": "get_var_args", "arguments": {"args": [1, 2, 3]}}')
200206
response_data = json.loads(result.response)
207+
# Can't parse ISO timestamps with trailing 'Z' in older versions of python:
208+
response_data['timestamp'] = re.sub('Z$', '+00:00', response_data['timestamp'])
201209
assert response_data == snapshot(
202210
{
203211
'tool_name': 'get_var_args',
204212
'content': '{"args": [1, 2, 3]}',
205213
'tool_id': None,
206-
'timestamp': IsNow(iso_string=True),
214+
'timestamp': IsStr() & IsNow(iso_string=True, tz=timezone.utc),
207215
'role': 'tool-return',
208216
}
209217
)
@@ -317,7 +325,7 @@ def test_call_all():
317325
assert result.message_history == snapshot(
318326
[
319327
SystemPrompt(content='foobar'),
320-
UserPrompt(content='Hello', timestamp=IsNow()),
328+
UserPrompt(content='Hello', timestamp=IsNow(tz=timezone.utc)),
321329
LLMToolCalls(
322330
calls=[
323331
ToolCall.from_object('foo', {'x': 0}),
@@ -326,14 +334,16 @@ def test_call_all():
326334
ToolCall.from_object('qux', {'x': 0}),
327335
ToolCall.from_object('quz', {'x': 'a'}),
328336
],
329-
timestamp=IsNow(),
337+
timestamp=IsNow(tz=timezone.utc),
338+
),
339+
ToolReturn(tool_name='foo', content='1', timestamp=IsNow(tz=timezone.utc)),
340+
ToolReturn(tool_name='bar', content='2', timestamp=IsNow(tz=timezone.utc)),
341+
ToolReturn(tool_name='baz', content='3', timestamp=IsNow(tz=timezone.utc)),
342+
ToolReturn(tool_name='qux', content='4', timestamp=IsNow(tz=timezone.utc)),
343+
ToolReturn(tool_name='quz', content='a', timestamp=IsNow(tz=timezone.utc)),
344+
LLMResponse(
345+
content='{"foo":"1","bar":"2","baz":"3","qux":"4","quz":"a"}', timestamp=IsNow(tz=timezone.utc)
330346
),
331-
ToolReturn(tool_name='foo', content='1', timestamp=IsNow()),
332-
ToolReturn(tool_name='bar', content='2', timestamp=IsNow()),
333-
ToolReturn(tool_name='baz', content='3', timestamp=IsNow()),
334-
ToolReturn(tool_name='qux', content='4', timestamp=IsNow()),
335-
ToolReturn(tool_name='quz', content='a', timestamp=IsNow()),
336-
LLMResponse(content='{"foo":"1","bar":"2","baz":"3","qux":"4","quz":"a"}', timestamp=IsNow()),
337347
]
338348
)
339349

‎tests/models/test_model_test.py‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
from __future__ import annotations as _annotations
44

5+
from datetime import timezone
56
from typing import Annotated, Any, Literal
67

78
import pytest
@@ -83,15 +84,15 @@ async def my_ret(x: int) -> str:
8384
assert result.response == snapshot('{"my_ret":"1"}')
8485
assert result.message_history == snapshot(
8586
[
86-
UserPrompt(content='Hello', timestamp=IsNow()),
87+
UserPrompt(content='Hello', timestamp=IsNow(tz=timezone.utc)),
8788
LLMToolCalls(
8889
calls=[ToolCall.from_object('my_ret', {'x': 0})],
89-
timestamp=IsNow(),
90+
timestamp=IsNow(tz=timezone.utc),
9091
),
91-
RetryPrompt(tool_name='my_ret', content='First call failed', timestamp=IsNow()),
92-
LLMToolCalls(calls=[ToolCall.from_object('my_ret', {'x': 0})], timestamp=IsNow()),
93-
ToolReturn(tool_name='my_ret', content='1', timestamp=IsNow()),
94-
LLMResponse(content='{"my_ret":"1"}', timestamp=IsNow()),
92+
RetryPrompt(tool_name='my_ret', content='First call failed', timestamp=IsNow(tz=timezone.utc)),
93+
LLMToolCalls(calls=[ToolCall.from_object('my_ret', {'x': 0})], timestamp=IsNow(tz=timezone.utc)),
94+
ToolReturn(tool_name='my_ret', content='1', timestamp=IsNow(tz=timezone.utc)),
95+
LLMResponse(content='{"my_ret":"1"}', timestamp=IsNow(tz=timezone.utc)),
9596
]
9697
)
9798

‎tests/models/test_openai.py‎

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,7 @@ async def test_request_structured_response():
115115
assert result.response == [1, 2, 123]
116116
assert result.message_history == snapshot(
117117
[
118-
UserPrompt(content='Hello', timestamp=IsNow()),
118+
UserPrompt(content='Hello', timestamp=IsNow(tz=datetime.timezone.utc)),
119119
LLMToolCalls(
120120
calls=[
121121
ToolCall(
@@ -124,7 +124,7 @@ async def test_request_structured_response():
124124
tool_id='123',
125125
)
126126
],
127-
timestamp=datetime.datetime(2024, 1, 1),
127+
timestamp=datetime.datetime(2024, 1, 1, tzinfo=datetime.timezone.utc),
128128
),
129129
]
130130
)
@@ -188,7 +188,7 @@ async def get_location(loc_name: str) -> str:
188188
assert result.message_history == snapshot(
189189
[
190190
SystemPrompt(content='this is the system prompt'),
191-
UserPrompt(content='Hello', timestamp=IsNow()),
191+
UserPrompt(content='Hello', timestamp=IsNow(tz=datetime.timezone.utc)),
192192
LLMToolCalls(
193193
calls=[
194194
ToolCall(
@@ -197,10 +197,13 @@ async def get_location(loc_name: str) -> str:
197197
tool_id='1',
198198
)
199199
],
200-
timestamp=datetime.datetime(2024, 1, 1, 0, 0),
200+
timestamp=datetime.datetime(2024, 1, 1, 0, 0, tzinfo=datetime.timezone.utc),
201201
),
202202
RetryPrompt(
203-
tool_name='get_location', content='Wrong location, please try again', tool_id='1', timestamp=IsNow()
203+
tool_name='get_location',
204+
content='Wrong location, please try again',
205+
tool_id='1',
206+
timestamp=IsNow(tz=datetime.timezone.utc),
204207
),
205208
LLMToolCalls(
206209
calls=[
@@ -210,15 +213,17 @@ async def get_location(loc_name: str) -> str:
210213
tool_id='2',
211214
)
212215
],
213-
timestamp=datetime.datetime(2024, 1, 1, 0, 0),
216+
timestamp=datetime.datetime(2024, 1, 1, 0, 0, tzinfo=datetime.timezone.utc),
214217
),
215218
ToolReturn(
216219
tool_name='get_location',
217220
content='{"lat": 51, "lng": 0}',
218221
tool_id='2',
219-
timestamp=IsNow(),
222+
timestamp=IsNow(tz=datetime.timezone.utc),
223+
),
224+
LLMResponse(
225+
content='final response', timestamp=datetime.datetime(2024, 1, 1, 0, 0, tzinfo=datetime.timezone.utc)
220226
),
221-
LLMResponse(content='final response', timestamp=datetime.datetime(2024, 1, 1, 0, 0)),
222227
]
223228
)
224229
assert result.cost == snapshot(

0 commit comments

Comments
 (0)