Skip to content

fix(python): stop double-counting cached tokens in Usage.tally() - #125

Draft
rootkiller6788 wants to merge 1 commit into
google:mainfrom
rootkiller6788:fix-usage-tally-double-count
Draft

fix(python): stop double-counting cached tokens in Usage.tally()#125
rootkiller6788 wants to merge 1 commit into
google:mainfrom
rootkiller6788:fix-usage-tally-double-count

Conversation

@rootkiller6788

Copy link
Copy Markdown

Summary

Usage.tally() double-counts cached tokens when aggregating usage across messages/turns.

The method currently does:

self.total_tokens += subusage.total_tokens
...
self.total_tokens += subusage.cached_token_count

total_tokens is documented as "Total number of tokens used in the request (prompt + generation)". Cached tokens are a subset of the prompt tokens, so they are already included in total_tokens. Adding subusage.cached_token_count a second time inflates the aggregated total.

Repro

from sec_gemini.models.usage import Usage

total = Usage()
total.tally(Usage(
    prompt_tokens=100,
    generated_tokens=50,
    total_tokens=150,
    cached_token_count=30,
))
assert total.total_tokens == 180  # bug: expected 150

Fix

Remove the redundant self.total_tokens += subusage.cached_token_count line so total_tokens is only accumulated once, and add tests/test_usage.py covering single and multiple usage accumulation (including a cached-token case that would fail before this fix).

Verification

  • New unit tests pass locally.
  • No change to the separate Usage.cost() pricing logic.

Usage.tally() added subusage.total_tokens and then added
subusage.cached_token_count again on top. cached_token_count is already
included in total_tokens (cached tokens are part of the prompt), so the
aggregate total was inflated by the cached token count on every tally.

Remove the redundant addition and add unit tests covering single and
multiple usage accumulation.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant