fix(google): count Gemini thinking tokens in usage.output_tokens 🤖🤖🤖🤖 - #2622
fix(google): count Gemini thinking tokens in usage.output_tokens 🤖🤖🤖🤖#2622ebarkhordar wants to merge 2 commits into
Conversation
GoogleChatTranslator.from_google mapped output_tokens to candidates_token_count only, but google-genai defines total_token_count as prompt + candidates + tool_use_prompt + thoughts. When the model thinks, thoughts_token_count is generated (and billed) output that the mapping dropped, so input_tokens + output_tokens != total_tokens on the same Usage object. Fold thoughts_token_count into output_tokens and tool_use_prompt_token_count into input_tokens so the sum matches total_tokens, per Google's own definition. Add a regression test covering a thinking response.
There was a problem hiding this comment.
Code Review
This pull request updates the token usage calculation for Google Chat responses to include tool_use_prompt_token_count in input_tokens and thoughts_token_count in output_tokens, ensuring they sum up to total_tokens. A corresponding unit test was added to verify this logic. The reviewer suggested using getattr to access these new fields to prevent potential AttributeError issues on older versions of the google-genai SDK.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
Following up, no rush. Nothing has executed against this branch on your side: |
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Description
GoogleChatTranslator.from_googlebuildsUsagefrom Gemini'susage_metadatabut maps
output_tokenstocandidates_token_countalone:google-genai documents
total_token_countas the sum ofprompt_token_count,candidates_token_count,tool_use_prompt_token_countandthoughts_token_count.thoughts_token_countis the model's generated (and billed) reasoning output.When a Gemini model thinks, that count lands in
total_token_countbut in none ofinput_tokens/output_tokens, so within a singleUsageobjectinput_tokens + output_tokens != total_tokens.Invariant:
input_tokens + output_tokens == total_tokens, mapping each ofGoogle's four components to the correct side. The fix folds
thoughts_token_countinto
output_tokens(generated output) andtool_use_prompt_token_countintoinput_tokens(prompt-side), so the sum matchestotal_token_countbyconstruction.
Scope: this touches only the
acompletionpath intranslators/google_chat.py.The Interactions path in
google_response.pymaps a different SDK surface(
total_input_tokens/total_output_tokens) and is not changed here.Verification
Reproduced at HEAD in a clean
python:3.13-slimcontainer (google-genai only, noAPI key). A response with
prompt=10, candidates=5, thoughts=20, total=35maps toinput=10, output=5, total=35onmain(15 != 35), and toinput=10, output=25, total=35with the fix. The added regression testtest_from_google_thinking_tokens_counted_in_outputfails onmainand passes onthis branch; the existing 27 google translator tests still pass.
ruff check,ruff format --check, vermin (3.12 target) and basedpyright are clean on thechanged files.
Not verified by me: I did not make a live
gemini-3.5-flashcall (no key on thebox), so how often real responses carry a non-zero
thoughts_token_countis notmeasured here. The defect and the fix are exercised against the real google-genai
usage_metadatatype.tool_use_prompt_token_countis included for formulacompleteness; it is
0in giskard's current usage, which sends onlyfunction_declarations.Design note
Usagehas no reasoning field, so this is the minimal change that restores theinvariant. If you would rather surface reasoning tokens separately (a new
reasoning_tokensfield onUsage), that is a larger API change and I am happy tofollow that direction instead; this PR keeps to the smallest correct fix.
Related Issue
None.
Type of Change
Coding agents
Opened by an autonomous agent per AUTONOMOUS.md; the title carries the required
four-robot marker.
Checklist
CODE_OF_CONDUCT.mddocument.CONTRIBUTING.mdguide.uv.lock(not applicable:pyproject.tomlunchanged).