Skip to content

Commit d8265a0

Browse files
DouweMDouweM
authored
Include all usage fields in OTel attributes (#3221)
1 parent 6adf303 commit d8265a0

File tree

2 files changed

+57
-13
lines changed

2 files changed

+57
-13
lines changed

‎pydantic_ai_slim/pydantic_ai/usage.py‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,18 @@ def opentelemetry_attributes(self) -> dict[str, int]:
7272
result['gen_ai.usage.input_tokens'] = self.input_tokens
7373
if self.output_tokens:
7474
result['gen_ai.usage.output_tokens'] = self.output_tokens
75-
details = self.details
75+
76+
details = self.details.copy()
77+
if self.cache_write_tokens:
78+
details['cache_write_tokens'] = self.cache_write_tokens
79+
if self.cache_read_tokens:
80+
details['cache_read_tokens'] = self.cache_read_tokens
81+
if self.input_audio_tokens:
82+
details['input_audio_tokens'] = self.input_audio_tokens
83+
if self.cache_audio_read_tokens:
84+
details['cache_audio_read_tokens'] = self.cache_audio_read_tokens
85+
if self.output_audio_tokens:
86+
details['output_audio_tokens'] = self.output_audio_tokens
7687
if details:
7788
prefix = 'gen_ai.usage.details.'
7889
for key, value in details.items():

‎tests/models/test_instrumented.py‎

Lines changed: 45 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,16 @@ async def request(
8787
TextPart('text2'),
8888
{}, # test unexpected parts # type: ignore
8989
],
90-
usage=RequestUsage(input_tokens=100, output_tokens=200),
90+
usage=RequestUsage(
91+
input_tokens=100,
92+
output_tokens=200,
93+
cache_write_tokens=10,
94+
cache_read_tokens=20,
95+
input_audio_tokens=10,
96+
cache_audio_read_tokens=5,
97+
output_audio_tokens=30,
98+
details={'reasoning_tokens': 30},
99+
),
91100
model_name='gpt-4o-2024-11-20',
92101
provider_details=dict(finish_reason='stop', foo='bar'),
93102
provider_response_id='response_id',
@@ -190,9 +199,15 @@ async def test_instrumented_model(capfire: CaptureLogfire):
190199
'logfire.span_type': 'span',
191200
'gen_ai.response.model': 'gpt-4o-2024-11-20',
192201
'gen_ai.response.id': 'response_id',
202+
'gen_ai.usage.details.reasoning_tokens': 30,
203+
'gen_ai.usage.details.cache_write_tokens': 10,
204+
'gen_ai.usage.details.cache_read_tokens': 20,
205+
'gen_ai.usage.details.input_audio_tokens': 10,
206+
'gen_ai.usage.details.cache_audio_read_tokens': 5,
207+
'gen_ai.usage.details.output_audio_tokens': 30,
193208
'gen_ai.usage.input_tokens': 100,
194209
'gen_ai.usage.output_tokens': 200,
195-
'operation.cost': 0.00225,
210+
'operation.cost': 0.00188125,
196211
},
197212
},
198213
]
@@ -724,11 +739,17 @@ async def test_instrumented_model_attributes_mode(capfire: CaptureLogfire, instr
724739
'event.name': 'gen_ai.choice',
725740
},
726741
],
742+
'gen_ai.usage.details.reasoning_tokens': 30,
743+
'gen_ai.usage.details.cache_write_tokens': 10,
744+
'gen_ai.usage.details.cache_read_tokens': 20,
745+
'gen_ai.usage.details.input_audio_tokens': 10,
746+
'gen_ai.usage.details.cache_audio_read_tokens': 5,
747+
'gen_ai.usage.details.output_audio_tokens': 30,
727748
'logfire.json_schema': {
728749
'type': 'object',
729750
'properties': {'events': {'type': 'array'}, 'model_request_parameters': {'type': 'object'}},
730751
},
731-
'operation.cost': 0.00225,
752+
'operation.cost': 0.00188125,
732753
'gen_ai.response.id': 'response_id',
733754
},
734755
},
@@ -821,6 +842,12 @@ async def test_instrumented_model_attributes_mode(capfire: CaptureLogfire, instr
821842
'gen_ai.system_instructions': [{'type': 'text', 'content': 'instructions'}],
822843
'gen_ai.usage.input_tokens': 100,
823844
'gen_ai.usage.output_tokens': 200,
845+
'gen_ai.usage.details.reasoning_tokens': 30,
846+
'gen_ai.usage.details.cache_write_tokens': 10,
847+
'gen_ai.usage.details.cache_read_tokens': 20,
848+
'gen_ai.usage.details.input_audio_tokens': 10,
849+
'gen_ai.usage.details.cache_audio_read_tokens': 5,
850+
'gen_ai.usage.details.output_audio_tokens': 30,
824851
'logfire.json_schema': {
825852
'type': 'object',
826853
'properties': {
@@ -830,7 +857,7 @@ async def test_instrumented_model_attributes_mode(capfire: CaptureLogfire, instr
830857
'model_request_parameters': {'type': 'object'},
831858
},
832859
},
833-
'operation.cost': 0.00225,
860+
'operation.cost': 0.00188125,
834861
'gen_ai.response.id': 'response_id',
835862
},
836863
},
@@ -908,14 +935,14 @@ async def test_instrumented_model_attributes_mode(capfire: CaptureLogfire, instr
908935
'start_time_unix_nano': IsInt(),
909936
'time_unix_nano': IsInt(),
910937
'count': 1,
911-
'sum': 0.00025,
938+
'sum': 0.00018125,
912939
'scale': 20,
913940
'zero_count': 0,
914-
'positive': {'offset': -12547035, 'bucket_counts': [1]},
941+
'positive': {'offset': -13033519, 'bucket_counts': [1]},
915942
'negative': {'offset': 0, 'bucket_counts': [0]},
916943
'flags': 0,
917-
'min': 0.00025,
918-
'max': 0.00025,
944+
'min': 0.00018125,
945+
'max': 0.00018125,
919946
'exemplars': [],
920947
},
921948
{
@@ -929,14 +956,14 @@ async def test_instrumented_model_attributes_mode(capfire: CaptureLogfire, instr
929956
'start_time_unix_nano': IsInt(),
930957
'time_unix_nano': IsInt(),
931958
'count': 1,
932-
'sum': 0.002,
959+
'sum': 0.0017,
933960
'scale': 20,
934961
'zero_count': 0,
935-
'positive': {'offset': -9401307, 'bucket_counts': [1]},
962+
'positive': {'offset': -9647161, 'bucket_counts': [1]},
936963
'negative': {'offset': 0, 'bucket_counts': [0]},
937964
'flags': 0,
938-
'min': 0.002,
939-
'max': 0.002,
965+
'min': 0.0017,
966+
'max': 0.0017,
940967
'exemplars': [],
941968
},
942969
],
@@ -1490,6 +1517,12 @@ async def test_response_cost_error(capfire: CaptureLogfire, monkeypatch: pytest.
14901517
},
14911518
'gen_ai.usage.input_tokens': 100,
14921519
'gen_ai.usage.output_tokens': 200,
1520+
'gen_ai.usage.details.reasoning_tokens': 30,
1521+
'gen_ai.usage.details.cache_write_tokens': 10,
1522+
'gen_ai.usage.details.cache_read_tokens': 20,
1523+
'gen_ai.usage.details.input_audio_tokens': 10,
1524+
'gen_ai.usage.details.cache_audio_read_tokens': 5,
1525+
'gen_ai.usage.details.output_audio_tokens': 30,
14931526
'gen_ai.response.model': 'gpt-4o-2024-11-20',
14941527
'gen_ai.response.id': 'response_id',
14951528
},

0 commit comments

Comments
 (0)