Skip to content

fix(google): count Gemini thinking tokens in usage.output_tokens 🤖🤖🤖🤖 - #2622

Closed
ebarkhordar wants to merge 2 commits into
Giskard-AI:mainfrom
ebarkhordar:fix/gemini-thinking-tokens-output-count
Closed

fix(google): count Gemini thinking tokens in usage.output_tokens 🤖🤖🤖🤖#2622
ebarkhordar wants to merge 2 commits into
Giskard-AI:mainfrom
ebarkhordar:fix/gemini-thinking-tokens-output-count

Conversation

@ebarkhordar

Copy link
Copy Markdown

Description

GoogleChatTranslator.from_google builds Usage from Gemini's usage_metadata
but maps output_tokens to candidates_token_count alone:

usage = Usage(
    input_tokens=raw.usage_metadata.prompt_token_count or 0,
    output_tokens=raw.usage_metadata.candidates_token_count or 0,
    total_tokens=raw.usage_metadata.total_token_count or 0,
)

google-genai documents total_token_count as the sum of prompt_token_count,
candidates_token_count, tool_use_prompt_token_count and thoughts_token_count.
thoughts_token_count is the model's generated (and billed) reasoning output.
When a Gemini model thinks, that count lands in total_token_count but in none of
input_tokens/output_tokens, so within a single Usage object
input_tokens + output_tokens != total_tokens.

Invariant: input_tokens + output_tokens == total_tokens, mapping each of
Google's four components to the correct side. The fix folds thoughts_token_count
into output_tokens (generated output) and tool_use_prompt_token_count into
input_tokens (prompt-side), so the sum matches total_token_count by
construction.

Scope: this touches only the acompletion path in translators/google_chat.py.
The Interactions path in google_response.py maps 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-slim container (google-genai only, no
API key). A response with prompt=10, candidates=5, thoughts=20, total=35 maps to
input=10, output=5, total=35 on main (15 != 35), and to input=10, output=25, total=35 with the fix. The added regression test
test_from_google_thinking_tokens_counted_in_output fails on main and passes on
this branch; the existing 27 google translator tests still pass. ruff check,
ruff format --check, vermin (3.12 target) and basedpyright are clean on the
changed files.

Not verified by me: I did not make a live gemini-3.5-flash call (no key on the
box), so how often real responses carry a non-zero thoughts_token_count is not
measured here. The defect and the fix are exercised against the real google-genai
usage_metadata type. tool_use_prompt_token_count is included for formula
completeness; it is 0 in giskard's current usage, which sends only
function_declarations.

Design note

Usage has no reasoning field, so this is the minimal change that restores the
invariant. If you would rather surface reasoning tokens separately (a new
reasoning_tokens field on Usage), that is a larger API change and I am happy to
follow that direction instead; this PR keeps to the smallest correct fix.

Related Issue

None.

Type of Change

  • 🔧 Bug fix (non-breaking change which fixes an issue)

Coding agents

Opened by an autonomous agent per AUTONOMOUS.md; the title carries the required
four-robot marker.

Checklist

  • I've read the CODE_OF_CONDUCT.md document.
  • I've read the CONTRIBUTING.md guide.
  • I've written tests for all new methods and classes that I created.
  • I've written the docstring in NumPy format for all the methods and classes that I created or modified.
  • I've updated the uv.lock (not applicable: pyproject.toml unchanged).

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.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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.

Comment thread libs/giskard-llm/src/giskard/llm/translators/google_chat.py Outdated
@ebarkhordar

Copy link
Copy Markdown
Author

Following up, no rush. Nothing has executed against this branch on your side: authorize stops at External contributors require a maintainer to add the 'safe for build' label, and Main CI plus the security-analysis run are held pending approval of the fork workflow, so the only verification here is the local one in the description. Happy to rebase, it sits behind main now.

@davidberenstein1957
davidberenstein1957 self-requested a review July 29, 2026 14:31
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@davidberenstein1957

Copy link
Copy Markdown
Member

Opened a formatted follow-up so CI can pass: #2653

Same change as this PR, plus ruff format/lint fixes. Please use #2653 going forward — we can close this one once that is green.

@ebarkhordar

Copy link
Copy Markdown
Author

Thanks for the review and for taking the formatting on. I compared the functional diff on #2653 against this branch: google_chat.py and test_google_chat_return.py are identical, and lint and authorize are green there. Closing this one in favor of #2653.

@ebarkhordar
ebarkhordar deleted the fix/gemini-thinking-tokens-output-count branch July 29, 2026 15:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Development

Successfully merging this pull request may close these issues.

2 participants