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