From d9c6eb199b6ccc1fae417463e1b374574f2ae2f8 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Fri, 19 Dec 2025 10:38:58 -0800 Subject: [PATCH] fix: return embedding metadata if available PiperOrigin-RevId: 846788453 --- .../matching_engine/matching_engine_index_endpoint.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py b/google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py index b4737dd488..48743613b3 100644 --- a/google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py +++ b/google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py @@ -16,7 +16,7 @@ # from dataclasses import dataclass, field -from typing import Dict, List, Optional, Sequence, Tuple, Union +from typing import Dict, List, Optional, Sequence, Tuple, Union, Any, Mapping from google.auth import credentials as auth_credentials from google.cloud.aiplatform import base @@ -208,6 +208,8 @@ class MatchNeighbor: For example, values [1,2,3] with dimensions [4,5,6] means value 1 is of the 4th dimension, value 2 is of the 4th dimension, and value 3 is of the 6th dimension. + embedding_metadata (Mapping[str, Any]): + Optional. The corresponding embedding metadata of the matching datapoint. """ @@ -220,6 +222,7 @@ class MatchNeighbor: numeric_restricts: Optional[List[NumericNamespace]] = None sparse_embedding_values: Optional[List[float]] = None sparse_embedding_dimensions: Optional[List[int]] = None + embedding_metadata: Optional[Mapping[str, Any]] = None def from_index_datapoint( self, index_datapoint: gca_index_v1beta1.IndexDatapoint @@ -276,6 +279,9 @@ def from_index_datapoint( self.sparse_embedding_dimensions = ( index_datapoint.sparse_embedding.dimensions ) + # retrieve embedding metadata + if index_datapoint.embedding_metadata is not None: + self.embedding_metadata = index_datapoint.embedding_metadata return self def from_embedding(self, embedding: match_service_pb2.Embedding) -> "MatchNeighbor":