Repository navigation
Python: fix(gemini): include thinking and tool-use prompt tokens in usage totals - #9989
lalith0407 wants to merge 2 commits into
Conversation
…otals GeminiChatClient mapped candidates_token_count to output_token_count and prompt_token_count to input_token_count. Gemini reports thinking tokens (thoughts_token_count) and tool-use prompt tokens (tool_use_prompt_token_count) outside those two fields, so on thinking models output_token_count was smaller than reasoning_output_token_count and input + output did not add up to total_token_count. Fold both into the totals, matching the other connectors and the GenAI semconv (reasoning.output_tokens SHOULD be included in output_tokens).
|
lalith0407 please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
1 similar comment
|
lalith0407 please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
There was a problem hiding this comment.
🟡 Changes recommended
The input-token condition still omits tool-use tokens when prompt_token_count is None, leaving input_token_count absent and token totals inconsistent.
1 open issue
What changed in this PR
This PR addresses #9912 by correcting Gemini token accounting in the Python connector so usage reporting includes thinking and tool-use prompt tokens.
Changes:
- Include thinking tokens in output totals and tool-use prompt tokens in input totals.
- Preserve output usage for responses cut off while thinking.
- Add four regression tests covering non-streaming, streaming, and thinking-only responses, and correct an existing mock total.
| File | Description |
|---|---|
python/packages/gemini/tests/test_gemini_client.py |
Extend usage mocks and test inclusive token totals and cut-off responses. |
python/packages/gemini/agent_framework_gemini/_chat_client.py |
Fold thinking and tool-use prompt counts into output and input usage. |
Review effort: Balanced
| details["input_token_count"] = v | ||
| if (v := usage.candidates_token_count) is not None: | ||
| details["output_token_count"] = v | ||
| details["input_token_count"] = v + (tool_use_prompt_tokens or 0) |
…ount; fix test typing

Motivation & Context
Gemini reports thinking tokens (
thoughts_token_count) and tool-use prompt tokens (tool_use_prompt_token_count) outsidecandidates_token_countandprompt_token_count.GeminiChatClientcopied only those two intooutput_token_count/input_token_count. On thinking models,gen_ai.usage.output_tokenstherefore left out the (billed) thinking tokens and came out smaller thanreasoning_output_token_count(80 vs 900 in the repro), and input + output didn't add up tototal_token_count. Every other connector reports full totals, and the GenAI semconv says reasoning tokens SHOULD be included inoutput_tokens.HUMAN: I noticed this in traces from Gemini 2.5 thinking models, where
output_tokenscame out lower than the reasoning token count, which shouldn't be possible. Looking at the response, the thinking tokens sit inthoughts_token_countand the client never added them in. I checked the google-genaitotal_token_countdefinition and the GenAI semconv, and both count thinking as output, so I made the totals add up and added tests for normal, streaming and cut-off responses.Description & Review Guide
_parse_usagesetsoutput_token_count = candidates_token_count + thoughts_token_countandinput_token_count = prompt_token_count + tool_use_prompt_token_count(missing values count as 0). If only thinking tokens are reported (aMAX_TOKENScut-off mid-thought),output_token_countis still set.reasoning_output_token_count,cache_read_input_token_countandtotal_token_countare unchanged.output_token_count - reasoning_output_token_count. The package isbeta.tool_use_prompt_token_countinto input in this PR (I think it belongs here, since the SDK defines total as the sum of all four, but I can split it out). Streaming uses the same_parse_usageon the final chunk, so it is covered by the same change (with a test).Tests (
packages/gemini): 276 passed, 10 skipped. The 4 new tests fail on main and pass with the fix. One existing test's mock total was updated from 28 to 34 so the mock matches real Gemini (20 + 8 + 6 thoughts). ruff check/format is clean, and pyright reports 0 errors onagent_framework_gemini.Related Issue
Fixes #9912
AI Assistance
AI assistance details: AI helped with the investigation, the offline repro script, the implementation, the tests and drafting this description. I reviewed and re-ran everything.
Contribution Checklist