fix: Add MatchingEngineIndexConfig enums to proto value converters

PiperOrigin-RevId: 700481419

Author: vertex-sdk-bot

google/cloud/aiplatform/matching_engine/matching_engine_index.py

@@ -443,7 +443,7 @@ def create_tree_ah_index(
                 contents_delta_uri="gs://my_bucket/embeddings",
                 dimensions=1,
                 approximate_neighbors_count=150,
-                distance_measure_type="SQUARED_L2_DISTANCE",
+                distance_measure_type=matching_engine_index_config.DistanceMeasureType.SQUARED_L2_DISTANCE,
                 leaf_node_embedding_count=100,
                 leaf_nodes_to_search_percent=50,
                 description="my description",
@@ -604,7 +604,7 @@ def create_brute_force_index(
                 contents_delta_uri="gs://my_bucket/embeddings",
                 dimensions=1,
                 approximate_neighbors_count=150,
-                distance_measure_type="SQUARED_L2_DISTANCE",
+                distance_measure_type=matching_engine_index_config.DistanceMeasureType.SQUARED_L2_DISTANCE,
                 description="my description",
                 labels={ "label_name": "label_value" },
             )

google/cloud/aiplatform/matching_engine/matching_engine_index_config.py

@@ -18,6 +18,7 @@
 import abc
 import enum
 from dataclasses import dataclass
+from google.protobuf.struct_pb2 import Value
 from typing import Any, Dict, Optional
 
 
@@ -34,6 +35,10 @@ class DistanceMeasureType(enum.Enum):
     # Cosine Distance. Defined as 1 - cosine similarity.
     COSINE_DISTANCE = "COSINE_DISTANCE"
 
+    def to_value(self) -> str:
+        """Returns the value of the distance measure type."""
+        return Value(string_value=self.name)
+
 
 class FeatureNormType(enum.Enum):
     """Type of normalization to be carried out on each vector."""
@@ -43,6 +48,10 @@ class FeatureNormType(enum.Enum):
     # No normalization type is specified.
     NONE = "NONE"
 
+    def to_value(self) -> str:
+        """Returns the value of the feature norm type."""
+        return Value(string_value=self.name)
+
 
 class AlgorithmConfig(abc.ABC):
     """Base class for configuration options for matching algorithm."""

tests/unit/aiplatform/test_matching_engine_index.py

@@ -32,6 +32,10 @@
     index_service_client,
 )
 
+from google.cloud.aiplatform.matching_engine import (
+    matching_engine_index_config,
+)
+
 from google.cloud.aiplatform.compat.types import (
     index as gca_index,
     encryption_spec as gca_encryption_spec,
@@ -52,8 +56,12 @@
     test_constants.MatchingEngineConstants._TEST_INDEX_DISPLAY_NAME
 )
 _TEST_CONTENTS_DELTA_URI = "gs://contents"
-_TEST_INDEX_DISTANCE_MEASURE_TYPE = "SQUARED_L2_DISTANCE"
-_TEST_INDEX_FEATURE_NORM_TYPE = "UNIT_L2_NORM"
+_TEST_INDEX_DISTANCE_MEASURE_TYPE = (
+    matching_engine_index_config.DistanceMeasureType.SQUARED_L2_DISTANCE
+)
+_TEST_INDEX_FEATURE_NORM_TYPE = (
+    matching_engine_index_config.FeatureNormType.UNIT_L2_NORM
+)
 
 _TEST_CONTENTS_DELTA_URI_UPDATE = "gs://contents_update"
 _TEST_IS_COMPLETE_OVERWRITE_UPDATE = True
@@ -510,8 +518,8 @@ def test_create_tree_ah_index_backward_compatibility(self, create_index_mock):
             contents_delta_uri=_TEST_CONTENTS_DELTA_URI,
             dimensions=_TEST_INDEX_CONFIG_DIMENSIONS,
             approximate_neighbors_count=_TEST_INDEX_APPROXIMATE_NEIGHBORS_COUNT,
-            distance_measure_type=_TEST_INDEX_DISTANCE_MEASURE_TYPE,
-            feature_norm_type=_TEST_INDEX_FEATURE_NORM_TYPE,
+            distance_measure_type=_TEST_INDEX_DISTANCE_MEASURE_TYPE.value,
+            feature_norm_type=_TEST_INDEX_FEATURE_NORM_TYPE.value,
             leaf_node_embedding_count=_TEST_LEAF_NODE_EMBEDDING_COUNT,
             leaf_nodes_to_search_percent=_TEST_LEAF_NODES_TO_SEARCH_PERCENT,
             description=_TEST_INDEX_DESCRIPTION,