diff --git a/google/genai/models.py b/google/genai/models.py index 9ff421480..84854361c 100644 --- a/google/genai/models.py +++ b/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 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 35269bcdf..8e62eb078 100644 --- a/google/genai/types.py +++ b/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