From 3d04c6e783e26739a2c26f59ad2bcff9fae39101 Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Tue, 27 Oct 2026 22:53:15 -0400 Subject: [PATCH] Python: count Gemini thinking tokens in the usage of Google AI completions Co-Authored-By: Claude Sonnet 5.5 --- .../services/google_ai_chat_completion.py | 3 +- .../services/google_ai_text_completion.py | 3 +- .../ai/google/google_ai/services/utils.py | 21 ++++++++ .../services/test_google_ai_utils.py | 52 ++++++++++++++++++- 4 files changed, 76 insertions(+), 3 deletions(-) diff --git a/python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_chat_completion.py b/python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_chat_completion.py index a3bd3468a130..f210c19885b8 100644 --- a/python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_chat_completion.py +++ b/python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_chat_completion.py @@ -22,6 +22,7 @@ format_assistant_message, format_tool_message, format_user_message, + get_completion_token_count, update_settings_from_function_choice_configuration, ) from semantic_kernel.connectors.ai.google.shared_utils import ( @@ -416,7 +417,7 @@ def _get_metadata_from_response(self, response: GenerateContentResponse) -> dict "prompt_feedback": response.prompt_feedback, "usage": CompletionUsage( prompt_tokens=response.usage_metadata.prompt_token_count if response.usage_metadata else None, - completion_tokens=response.usage_metadata.candidates_token_count if response.usage_metadata else None, + completion_tokens=get_completion_token_count(response.usage_metadata), ), } diff --git a/python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_text_completion.py b/python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_text_completion.py index 37a79cd5a767..452e27363f41 100644 --- a/python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_text_completion.py +++ b/python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_text_completion.py @@ -14,6 +14,7 @@ ) from semantic_kernel.connectors.ai.google.google_ai.google_ai_settings import GoogleAISettings from semantic_kernel.connectors.ai.google.google_ai.services.google_ai_base import GoogleAIBase +from semantic_kernel.connectors.ai.google.google_ai.services.utils import get_completion_token_count from semantic_kernel.connectors.ai.text_completion_client_base import TextCompletionClientBase from semantic_kernel.contents import TextContent from semantic_kernel.contents.streaming_text_content import StreamingTextContent @@ -249,7 +250,7 @@ def _get_metadata_from_response(self, response: GenerateContentResponse) -> dict "prompt_feedback": response.prompt_feedback, "usage": CompletionUsage( prompt_tokens=response.usage_metadata.prompt_token_count if response.usage_metadata else None, - completion_tokens=response.usage_metadata.candidates_token_count if response.usage_metadata else None, + completion_tokens=get_completion_token_count(response.usage_metadata), ), } diff --git a/python/semantic_kernel/connectors/ai/google/google_ai/services/utils.py b/python/semantic_kernel/connectors/ai/google/google_ai/services/utils.py index 02633e3815ab..949ead94c10d 100644 --- a/python/semantic_kernel/connectors/ai/google/google_ai/services/utils.py +++ b/python/semantic_kernel/connectors/ai/google/google_ai/services/utils.py @@ -214,3 +214,24 @@ def _create_image_part(image_content: ImageContent) -> Part: raise ServiceInvalidRequestError( "ImageContent without data_uri in User message while formatting chat history for Google AI" ) + + +def get_completion_token_count(usage_metadata: Any) -> int | None: + """Get the number of tokens the model generated, including its thinking tokens. + + Gemini reports the thinking tokens in ``thoughts_token_count``, apart from ``candidates_token_count``, + but bills them as output tokens, so both count toward the completion tokens. + + Args: + usage_metadata: The usage metadata of a response, if it has one. + + Returns: + The completion token count, or None if the response reports neither count. + """ + if not usage_metadata: + return None + candidates_tokens = usage_metadata.candidates_token_count + thoughts_tokens = usage_metadata.thoughts_token_count + if candidates_tokens is None and thoughts_tokens is None: + return None + return (candidates_tokens or 0) + (thoughts_tokens or 0) diff --git a/python/tests/unit/connectors/ai/google/google_ai/services/test_google_ai_utils.py b/python/tests/unit/connectors/ai/google/google_ai/services/test_google_ai_utils.py index 733276128d52..a8306dd7bad8 100644 --- a/python/tests/unit/connectors/ai/google/google_ai/services/test_google_ai_utils.py +++ b/python/tests/unit/connectors/ai/google/google_ai/services/test_google_ai_utils.py @@ -1,12 +1,13 @@ # Copyright (c) Microsoft. All rights reserved. import pytest -from google.genai.types import FinishReason, Part +from google.genai.types import FinishReason, GenerateContentResponse, GenerateContentResponseUsageMetadata, Part from semantic_kernel.connectors.ai.google.google_ai.services.utils import ( finish_reason_from_google_ai_to_semantic_kernel, format_assistant_message, format_user_message, + get_completion_token_count, kernel_function_metadata_to_google_ai_function_call_format, ) from semantic_kernel.contents.chat_message_content import ChatMessageContent @@ -251,3 +252,52 @@ def test_google_ai_function_call_format_all_params_excluded_gives_no_parameters( ) result = kernel_function_metadata_to_google_ai_function_call_format(metadata) assert result["parameters"] is None + + +@pytest.mark.parametrize( + ("candidates_tokens", "thoughts_tokens", "expected"), + [ + (10, 40, 50), + (10, None, 10), + (None, 40, 40), + (0, 0, 0), + (None, None, None), + ], +) +def test_get_completion_token_count_includes_thinking_tokens(candidates_tokens, thoughts_tokens, expected) -> None: + usage_metadata = GenerateContentResponseUsageMetadata( + prompt_token_count=5, candidates_token_count=candidates_tokens, thoughts_token_count=thoughts_tokens + ) + assert get_completion_token_count(usage_metadata) == expected + + +def test_get_completion_token_count_without_usage_metadata() -> None: + assert get_completion_token_count(None) is None + + +def test_chat_completion_usage_counts_thinking_tokens(google_ai_unit_test_env) -> None: + from semantic_kernel.connectors.ai.google.google_ai.services.google_ai_chat_completion import ( + GoogleAIChatCompletion, + ) + + response = GenerateContentResponse( + usage_metadata=GenerateContentResponseUsageMetadata( + prompt_token_count=5, candidates_token_count=10, thoughts_token_count=40, total_token_count=55 + ) + ) + usage = GoogleAIChatCompletion()._get_metadata_from_response(response)["usage"] + assert (usage.prompt_tokens, usage.completion_tokens) == (5, 50) + + +def test_text_completion_usage_counts_thinking_tokens(google_ai_unit_test_env) -> None: + from semantic_kernel.connectors.ai.google.google_ai.services.google_ai_text_completion import ( + GoogleAITextCompletion, + ) + + response = GenerateContentResponse( + usage_metadata=GenerateContentResponseUsageMetadata( + prompt_token_count=5, candidates_token_count=10, thoughts_token_count=40, total_token_count=55 + ) + ) + usage = GoogleAITextCompletion()._get_metadata_from_response(response)["usage"] + assert (usage.prompt_tokens, usage.completion_tokens) == (5, 50)