Skip to content

Commit 05cfb68

Browse files
committed
adding usage to default context fields
1 parent 359c6d2 commit 05cfb68

File tree

2 files changed

+22
-0
lines changed

2 files changed

+22
-0
lines changed

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_run_context.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ def serialize_run_context(cls, ctx: RunContext[Any]) -> dict[str, Any]:
5151
'max_retries': ctx.max_retries,
5252
'run_step': ctx.run_step,
5353
'partial_output': ctx.partial_output,
54+
'usage': ctx.usage,
5455
}
5556

5657
@classmethod

tests/test_temporal.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2191,6 +2191,27 @@ def test_temporal_run_context_preserves_run_id():
21912191
assert reconstructed.run_id == 'run-123'
21922192

21932193

2194+
def test_temporal_run_context_serializes_usage():
2195+
ctx = RunContext(
2196+
deps=None,
2197+
model=TestModel(),
2198+
usage=RunUsage(
2199+
requests=2,
2200+
tool_calls=1,
2201+
input_tokens=123,
2202+
output_tokens=456,
2203+
details={'foo': 1},
2204+
),
2205+
run_id='run-123',
2206+
)
2207+
2208+
serialized = TemporalRunContext.serialize_run_context(ctx)
2209+
assert serialized['usage'] == ctx.usage
2210+
2211+
reconstructed = TemporalRunContext.deserialize_run_context(serialized, deps=None)
2212+
assert reconstructed.usage == ctx.usage
2213+
2214+
21942215
fastmcp_agent = Agent(
21952216
model,
21962217
name='fastmcp_agent',

0 commit comments

Comments
 (0)