feat: Populate per-modality prompt token count in embedding responses for gemini-embedding-2

google/genai/models.py

@@ -318,6 +318,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
 
 
@@ -1136,6 +1143,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

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

google/genai/types.py

@@ -8564,6 +8564,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):
@@ -8578,6 +8583,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