This commit is contained in:
Timothy Jaeryang Baek
2026-07-26 18:54:07 -04:00
parent 5efe0951d5
commit df94268e89
2 changed files with 23 additions and 7 deletions
+12 -3
View File
@@ -55,8 +55,6 @@ USAGE_TOKEN_KEYS = {
'input_tokens',
'output_tokens',
'total_tokens',
'prompt_tokens',
'completion_tokens',
}
USAGE_COST_KEYS = {
@@ -105,7 +103,7 @@ def merge_usage(current: dict | None, incoming: dict | None) -> dict:
"""
Merge usage payloads from multiple model calls into one cumulative usage dict.
Token fields are additive; non-numeric metadata keeps the latest provider value.
Canonical token fields are additive; provider aliases keep the latest value.
"""
current_usage = normalize_usage(current or {}) if current else {}
incoming_usage = normalize_usage(incoming or {}) if incoming else {}
@@ -133,6 +131,17 @@ def merge_usage(current: dict | None, incoming: dict | None) -> dict:
incoming_usage.get(key) if isinstance(incoming_usage.get(key), dict) else {},
)
result['prompt_tokens'] = (
incoming_usage.get('prompt_tokens')
or incoming_usage.get('input_tokens')
or current_usage.get('prompt_tokens', 0)
)
result['completion_tokens'] = (
incoming_usage.get('completion_tokens')
or incoming_usage.get('output_tokens')
or current_usage.get('completion_tokens', 0)
)
return result