Skip to content

Commit 8f63ba3

Browse files
⚡️ Speed up method Usage.opentelemetry_attributes by 85% (#2198)
Co-authored-by: codeflash-ai[bot] <148906541+codeflash-ai[bot]@users.noreply.github.com>
1 parent e6396cc commit 8f63ba3

File tree

2 files changed

+22
-7
lines changed

2 files changed

+22
-7
lines changed

pydantic_ai_slim/pydantic_ai/usage.py

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -57,13 +57,19 @@ def __add__(self, other: Usage) -> Usage:
5757

5858
def opentelemetry_attributes(self) -> dict[str, int]:
5959
"""Get the token limits as OpenTelemetry attributes."""
60-
result = {
61-
'gen_ai.usage.input_tokens': self.request_tokens,
62-
'gen_ai.usage.output_tokens': self.response_tokens,
63-
}
64-
for key, value in (self.details or {}).items():
65-
result[f'gen_ai.usage.details.{key}'] = value # pragma: no cover
66-
return {k: v for k, v in result.items() if v}
60+
result: dict[str, int] = {}
61+
if self.request_tokens:
62+
result['gen_ai.usage.input_tokens'] = self.request_tokens
63+
if self.response_tokens:
64+
result['gen_ai.usage.output_tokens'] = self.response_tokens
65+
details = self.details
66+
if details:
67+
prefix = 'gen_ai.usage.details.'
68+
for key, value in details.items():
69+
# Skipping check for value since spec implies all detail values are relevant
70+
if value:
71+
result[prefix + key] = value
72+
return result
6773

6874
def has_values(self) -> bool:
6975
"""Whether any values are set and non-zero."""

tests/test_usage_limits.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -176,6 +176,15 @@ async def delegate_to_other_agent2(ctx: RunContext[None], sentence: str) -> int:
176176
# confirm the usage from result2 is the sum of the usage from result1
177177
assert result2.usage() == functools.reduce(operator.add, run_1_usages)
178178

179+
result1_usage = result1.usage()
180+
result1_usage.details = {'custom1': 10, 'custom2': 20, 'custom3': 0}
181+
assert result1_usage.opentelemetry_attributes() == {
182+
'gen_ai.usage.input_tokens': 103,
183+
'gen_ai.usage.output_tokens': 13,
184+
'gen_ai.usage.details.custom1': 10,
185+
'gen_ai.usage.details.custom2': 20,
186+
}
187+
179188

180189
async def test_multi_agent_usage_sync():
181190
"""As in `test_multi_agent_usage_async`, with a sync tool."""

0 commit comments

Comments
 (0)