Python: count Gemini thinking tokens in the usage of Google AI completions

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),
             ),
         }
 

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),
             ),
         }
 

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)

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)