Skip to content

Python: count Gemini thinking tokens in the usage of Vertex AI completions - #14716

Open
alexcastgonz (pcbeingused333) wants to merge 1 commit into
microsoft:mainfrom
pcbeingused333:fix/vertex-thinking-completion-tokens
Open

alexcastgonz (pcbeingused333) wants to merge 1 commit into
microsoft:mainfrom
pcbeingused333:fix/vertex-thinking-completion-tokens

Conversation

@pcbeingused333

Copy link
Copy Markdown
Contributor

Motivation and Context

With a thinking Gemini model through the Vertex AI connector, the usage in the response metadata under-reports the completion tokens. Vertex AI reports thinking tokens in usage_metadata.thoughts_token_count, separately from candidates_token_count, and bills them as output tokens, but VertexAIChatCompletion and VertexAITextCompletion took completion_tokens from candidates_token_count alone. Telemetry and cost tracking that read this usage missed every thinking token. This is the Vertex counterpart of the Google AI fix in #14671.

Description

A new get_completion_token_count helper in vertex_ai/services/utils.py returns candidates_token_count + thoughts_token_count (either may be missing, and it is None only when both are). Both Vertex AI completion services use it in _get_metadata_from_response, which streaming also goes through.

Tests in test_vertex_ai_utils.py cover the helper (both counts, one zero, none, no usage metadata) and check that the chat and text services report 50 completion tokens for 10 candidate and 40 thinking tokens. test_vertex_ai_utils.py passes (17 tests; the chat service test fails without the fix) and ruff check and format are clean on the touched paths. A live Vertex AI call was not made.

Contribution Checklist

  • The code builds clean without any errors or warnings
  • The PR follows the SK Contribution Guidelines and the pre-submission formatting script raises no violations (ruff check and ruff format --check pass on the changed files; mypy not run)
  • All unit tests pass, and I have added new tests where possible
  • I didn't break anyone 😄

This change was prepared with the help of an AI assistant (Claude Code), and the author is responsible for the change.

🤖 Generated with Claude Code

Copilot AI balanced review requested due to automatic review settings November 3, 2026 11:54
@semantic-kernel-automation semantic-kernel-automation Bot added the python Pull requests for the Python Semantic Kernel label Nov 3, 2026

Copilot AI 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.

🟢 Approval recommended

No active review issues are recorded for the Vertex AI completion-token accounting change.

No open issues 🎉

What changed in this PR

This PR corrects under-reported completion usage for thinking Gemini models in the Python Vertex AI connector, extending the approach from the Google AI fix in #14671.

Changes:

  • Adds a shared helper that sums candidate and thinking token counts.
  • Uses the helper in chat and text response metadata, including streaming paths.
  • Adds tests for token-count combinations, absent metadata, and both services’ usage reporting.
File Description
python/​tests/​unit/​connectors/​ai/​google/​vertex_ai/​services/​test_vertex_ai_utils.py Adds helper and service-level token accounting tests.
python/​semantic_kernel/​connectors/​ai/​google/​vertex_ai/​services/​vertex_ai_text_completion.py Includes thinking tokens in text completion usage.
python/​semantic_kernel/​connectors/​ai/​google/​vertex_ai/​services/​vertex_ai_chat_completion.py Includes thinking tokens in chat completion usage.
python/​semantic_kernel/​connectors/​ai/​google/​vertex_ai/​services/​utils.py Adds the shared completion-token counting helper.

Review effort: Balanced

This branch was successfully deployed

1 active deployment
github-app-auth — bab750c8 Deployed Nov 3, 2026 by pcbeingused333 via add_label #29576
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

python Pull requests for the Python Semantic Kernel

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants