From 2ba0154a41654dc2e999c6242633e2b2eedaa5e7 Mon Sep 17 00:00:00 2001
From: Ayush Agrawal <ayushagra@google.com>
Date: Sat, 11 Jul 2026 13:42:44 -0700
Subject: [PATCH] feat: Populate per-modality prompt token count in embedding
 responses for gemini-embedding-2

PiperOrigin-RevId: 946297224
---
 google/genai/models.py                          | 9 +++++++++
 google/genai/tests/models/test_embed_content.py | 6 ++++++
 google/genai/types.py                           | 9 +++++++++
 3 files changed, 24 insertions(+)

diff --git a/google/genai/models.py b/google/genai/models.py
index 69151333a..7a35a8b6c 100644
--- a/google/genai/models.py
+++ b/google/genai/models.py
@@ -319,6 +319,13 @@ def _ContentEmbeddingStatistics_from_vertex(
   if getv(from_object, ['token_count']) is not None:
     setv(to_object, ['token_count'], getv(from_object, ['token_count']))
 
+  if getv(from_object, ['tokensDetails']) is not None:
+    setv(
+        to_object,
+        ['tokens_details'],
+        [item for item in getv(from_object, ['tokensDetails'])],
+    )
+
   return to_object
 
 
@@ -1137,6 +1144,8 @@ def _EmbedContentResponse_from_vertex(
       stats = {}
       if usage_metadata and usage_metadata.get('promptTokenCount'):
         stats['token_count'] = usage_metadata['promptTokenCount']
+      if usage_metadata and usage_metadata.get('promptTokensDetails'):
+        stats['tokensDetails'] = usage_metadata['promptTokensDetails']
       if truncated:
         stats['truncated'] = truncated
       embedding['statistics'] = stats
diff --git a/google/genai/tests/models/test_embed_content.py b/google/genai/tests/models/test_embed_content.py
index 02cc81f71..8f565efef 100644
--- a/google/genai/tests/models/test_embed_content.py
+++ b/google/genai/tests/models/test_embed_content.py
@@ -247,6 +247,12 @@ def test_gemini_embedding_2_content_combination(client):
   assert response.embeddings is not None
   assert len(response.embeddings) == 1
   assert len(response.embeddings[0].values) == 100
+  if client._api_client.vertexai:
+    statistics = response.embeddings[0].statistics
+    assert statistics is not None
+    assert statistics.token_count is not None
+    assert statistics.tokens_details is not None
+    assert len(statistics.tokens_details) > 0
 
 
 @pytest.mark.asyncio
diff --git a/google/genai/types.py b/google/genai/types.py
index 611c82d0a..2db20ec64 100644
--- a/google/genai/types.py
+++ b/google/genai/types.py
@@ -8670,6 +8670,11 @@ class ContentEmbeddingStatistics(_common.BaseModel):
       description="""Gemini Enterprise Agent Platform only. Number of tokens of the input text.
       """,
   )
+  tokens_details: Optional[list[ModalityTokenCount]] = Field(
+      default=None,
+      description="""Gemini Enterprise Agent Platform only. List of modalities and their token count for the input content.
+      """,
+  )
 
 
 class ContentEmbeddingStatisticsDict(TypedDict, total=False):
@@ -8684,6 +8689,10 @@ class ContentEmbeddingStatisticsDict(TypedDict, total=False):
   """Gemini Enterprise Agent Platform only. Number of tokens of the input text.
       """
 
+  tokens_details: Optional[list[ModalityTokenCountDict]]
+  """Gemini Enterprise Agent Platform only. List of modalities and their token count for the input content.
+      """
+
 
 ContentEmbeddingStatisticsOrDict = Union[
     ContentEmbeddingStatistics, ContentEmbeddingStatisticsDict