Skip to content

⚡️ Speed up function _map_usage by 172% #2197

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open

⚡️ Speed up function _map_usage by 172% #2197

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a903c73
⚡️ Speed up function `_map_usage` by 172%
codeflash-ai[bot] Jul 12, 2025
2645ee5
Update pydantic_ai_slim/pydantic_ai/models/anthropic.py
misrasaurabh1 Jul 14, 2025
329b2ac
fix pyright
misrasaurabh1 Jul 14, 2025
215f995
return a new empty usage object
misrasaurabh1 Jul 15, 2025
44ceeb1
Merge branch 'main' into codeflash/optimize-_map_usage-md0d6dgh
misrasaurabh1 Jul 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 18 additions & 8 deletions pydantic_ai_slim/pydantic_ai/models/anthropic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -425,12 +425,19 @@ def _map_tool_definition(f: ToolDefinition) -> BetaToolParam:


def _map_usage(message: BetaMessage | BetaRawMessageStreamEvent) -> usage.Usage:
if isinstance(message, BetaMessage):
response_usage = message.usage
elif isinstance(message, BetaRawMessageStartEvent):
response_usage = message.message.usage
elif isinstance(message, BetaRawMessageDeltaEvent):
response_usage = message.usage
"""Maps Anthropic API message object to pydantic-ai Usage object, extracting integer-type usage statistics.

Handles BetaMessage, BetaRawMessageStartEvent, BetaRawMessageDeltaEvent.
Returns empty Usage if type doesn't contain usage info.
"""
msg_type = type(message)

if msg_type is BetaMessage:
response_usage = cast(BetaMessage, message).usage
elif msg_type is BetaRawMessageStartEvent:
response_usage = cast(BetaRawMessageStartEvent, message).message.usage
elif msg_type is BetaRawMessageDeltaEvent:
response_usage = cast(BetaRawMessageDeltaEvent, message).usage
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are the performance improvement metrics in the description still accurate after the latest changes to this PR?

I think this change reduces readability (and consistency with other code) significantly so I'd only want to do it if the performance impact (on real examples, not just the test_usage test) are significant. And if it is significant, should we do this everywhere we currently use isinstance (and don't care about subclasses)? Otherwise it seems odd to do it only in this specific place.

else:
# No usage information provided in:
# - RawMessageStopEvent
Expand All @@ -454,10 +461,13 @@ def _map_usage(message: BetaMessage | BetaRawMessageStreamEvent) -> usage.Usage:
+ details.get('cache_read_input_tokens', 0)
)

output_tokens = response_usage.output_tokens
total_tokens = request_tokens + output_tokens

return usage.Usage(
request_tokens=request_tokens or None,
response_tokens=response_usage.output_tokens,
total_tokens=request_tokens + response_usage.output_tokens,
response_tokens=output_tokens,
total_tokens=total_tokens,
details=details or None,
)

Expand Down