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 94da9ac225..e2cb9efc2b 100644 --- a/google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py +++ b/google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py @@ -1865,7 +1865,7 @@ def find_neighbors( dimensions=query.sparse_embedding_dimensions, ), ) - if query.rrf_ranking_alpha: + if query.rrf_ranking_alpha is not None: find_neighbors_query.rrf = ( gca_match_service_v1beta1.FindNeighborsRequest.Query.RRF( alpha=query.rrf_ranking_alpha, @@ -2200,7 +2200,7 @@ def match( match_service_pb2.MatchRequest.RRF( alpha=query.rrf_ranking_alpha, ) - if query_is_hybrid and query.rrf_ranking_alpha + if query_is_hybrid and query.rrf_ranking_alpha is not None else None ), ) diff --git a/tests/unit/aiplatform/test_matching_engine_index_endpoint.py b/tests/unit/aiplatform/test_matching_engine_index_endpoint.py index bb8d3642a8..15267f9d0c 100644 --- a/tests/unit/aiplatform/test_matching_engine_index_endpoint.py +++ b/tests/unit/aiplatform/test_matching_engine_index_endpoint.py @@ -258,6 +258,11 @@ sparse_embedding_dimensions=[1, 2, 3], sparse_embedding_values=[0.1, 0.2, 0.3], ), + HybridQuery( + sparse_embedding_dimensions=[1, 2, 3], + sparse_embedding_values=[0.1, 0.2, 0.3], + rrf_ranking_alpha=0.0, # Test edge case where alpha is 0 + ), ] _TEST_NUM_NEIGHBOURS = 1 _TEST_FILTER = [ @@ -1539,7 +1544,7 @@ def test_private_service_access_hybrid_search_match_queries( match_service_pb2.MatchRequest.RRF( alpha=_TEST_HYBRID_QUERIES[i].rrf_ranking_alpha, ) - if _TEST_HYBRID_QUERIES[i].rrf_ranking_alpha + if _TEST_HYBRID_QUERIES[i].rrf_ranking_alpha is not None else None ), ) @@ -2093,6 +2098,27 @@ def test_index_public_endpoint_find_neighbors_queries( approximate_neighbor_count=_TEST_APPROX_NUM_NEIGHBORS, fraction_leaf_nodes_to_search_override=_TEST_FRACTION_LEAF_NODES_TO_SEARCH_OVERRIDE, ), + gca_match_service_v1beta1.FindNeighborsRequest.Query( + neighbor_count=_TEST_NUM_NEIGHBOURS, + datapoint=gca_index_v1beta1.IndexDatapoint( + restricts=[ + gca_index_v1beta1.IndexDatapoint.Restriction( + namespace="class", + allow_list=["token_1"], + deny_list=["token_2"], + ) + ], + sparse_embedding=gca_index_v1beta1.IndexDatapoint.SparseEmbedding( + values=[0.1, 0.2, 0.3], dimensions=[1, 2, 3] + ), + ), + rrf=gca_match_service_v1beta1.FindNeighborsRequest.Query.RRF( + alpha=0.0, + ), + per_crowding_attribute_neighbor_count=_TEST_PER_CROWDING_ATTRIBUTE_NUM_NEIGHBOURS, + approximate_neighbor_count=_TEST_APPROX_NUM_NEIGHBORS, + fraction_leaf_nodes_to_search_override=_TEST_FRACTION_LEAF_NODES_TO_SEARCH_OVERRIDE, + ), ], return_full_datapoint=_TEST_RETURN_FULL_DATAPOINT, )