feat: Add support for self-signed JWT for queries on private endpoints

PiperOrigin-RevId: 689941402

google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py

@@ -1387,6 +1387,7 @@ def find_neighbors(
         return_full_datapoint: bool = False,
         numeric_filter: Optional[List[NumericNamespace]] = None,
         embedding_ids: Optional[List[str]] = None,
+        signed_jwt: Optional[str] = None,
     ) -> List[List[MatchNeighbor]]:
         """Retrieves nearest neighbors for the given embedding queries on the
         specified deployed index which is deployed to either public or private
@@ -1456,6 +1457,9 @@ def find_neighbors(
                `embedding_ids` to lookup embedding values from dataset, if embedding
                with `embedding_ids` exists in the dataset, do nearest neighbor search.
 
+            signed_jwt (str):
+               Optional. A signed JWT for accessing the private endpoint.
+
         Returns:
             List[List[MatchNeighbor]] - A list of nearest neighbors for each query.
         """
@@ -1471,6 +1475,7 @@ def find_neighbors(
                 approx_num_neighbors=approx_num_neighbors,
                 fraction_leaf_nodes_to_search_override=fraction_leaf_nodes_to_search_override,
                 numeric_filter=numeric_filter,
+                signed_jwt=signed_jwt,
             )
 
         # Create the FindNeighbors request
@@ -1570,6 +1575,7 @@ def read_index_datapoints(
         *,
         deployed_index_id: str,
         ids: List[str] = [],
+        signed_jwt: Optional[str] = None,
     ) -> List[gca_index_v1beta1.IndexDatapoint]:
         """Reads the datapoints/vectors of the given IDs on the specified
         deployed index which is deployed to public or private endpoint.
@@ -1587,6 +1593,8 @@ def read_index_datapoints(
                 Required. The ID of the DeployedIndex to match the queries against.
             ids (List[str]):
                 Required. IDs of the datapoints to be searched for.
+            signed_jwt (str):
+               Optional. A signed JWT for accessing the private endpoint.
         Returns:
             List[gca_index_v1beta1.IndexDatapoint] - A list of datapoints/vectors of the given IDs.
         """
@@ -1595,6 +1603,7 @@ def read_index_datapoints(
             embeddings = self._batch_get_embeddings(
                 deployed_index_id=deployed_index_id,
                 ids=ids,
+                signed_jwt=signed_jwt,
             )
 
             response = []
@@ -1641,6 +1650,7 @@ def _batch_get_embeddings(
         *,
         deployed_index_id: str,
         ids: List[str] = [],
+        signed_jwt: Optional[str] = None,
     ) -> List[match_service_pb2.Embedding]:
         """
         Reads the datapoints/vectors of the given IDs on the specified index
@@ -1651,6 +1661,8 @@ def _batch_get_embeddings(
                 Required. The ID of the DeployedIndex to match the queries against.
             ids (List[str]):
                 Required. IDs of the datapoints to be searched for.
+            signed_jwt:
+               Optional. A signed JWT for accessing the private endpoint.
         Returns:
             List[match_service_pb2.Embedding] - A list of datapoints/vectors of the given IDs.
         """
@@ -1665,7 +1677,10 @@ def _batch_get_embeddings(
 
         for id in ids:
             batch_request.id.append(id)
-        response = stub.BatchGetEmbeddings(batch_request)
+        metadata = None
+        if signed_jwt:
+            metadata = (("authorization", f"Bearer: {signed_jwt}"),)
+        response = stub.BatchGetEmbeddings(batch_request, metadata=metadata)
 
         return response.embeddings
 
@@ -1680,6 +1695,7 @@ def match(
         fraction_leaf_nodes_to_search_override: Optional[float] = None,
         low_level_batch_size: int = 0,
         numeric_filter: Optional[List[NumericNamespace]] = None,
+        signed_jwt: Optional[str] = None,
     ) -> List[List[MatchNeighbor]]:
         """Retrieves nearest neighbors for the given embedding queries on the
         specified deployed index for private endpoint only.
@@ -1729,6 +1745,8 @@ def match(
                 results. For example:
                 [NumericNamespace(name="cost", value_int=5, op="GREATER")]
                 will match datapoints that its cost is greater than 5.
+            signed_jwt (str):
+               Optional. A signed JWT for accessing the private endpoint.
 
         Returns:
             List[List[MatchNeighbor]] - A list of nearest neighbors for each query.
@@ -1809,7 +1827,10 @@ def match(
         batch_request.requests.append(batch_request_for_index)
 
         # Perform the request
-        response = stub.BatchMatch(batch_request)
+        metadata = None
+        if signed_jwt:
+            metadata = (("authorization", f"Bearer: {signed_jwt}"),)
+        response = stub.BatchMatch(batch_request, metadata=metadata)
 
         # Wrap the results in MatchNeighbor objects and return
         match_neighbors_response = []

samples/model-builder/conftest.py

@@ -1382,6 +1382,20 @@ def mock_index_endpoint_find_neighbors(mock_index_endpoint):
         yield mock_find_neighbors
 
 
+@pytest.fixture
+def mock_index_endpoint_match(mock_index_endpoint):
+    with patch.object(mock_index_endpoint, "match") as mock:
+        mock.return_value = None
+        yield mock
+
+
+@pytest.fixture
+def mock_index_endpoint_read_index_datapoints(mock_index_endpoint):
+    with patch.object(mock_index_endpoint, "read_index_datapoints") as mock:
+        mock.return_value = None
+        yield mock
+
+
 @pytest.fixture
 def mock_index_create_tree_ah_index(mock_index):
     with patch.object(

samples/model-builder/test_constants.py

@@ -407,3 +407,4 @@
 VECTOR_SEARCH_INDEX_LABELS = {"my_key": "my_value"}
 VECTOR_SEARCH_GCS_URI = "gs://fake-dir"
 VECTOR_SEARCH_INDEX_ENDPOINT_DISPLAY_NAME = "my-vector-search-index-endpoint"
+VECTOR_SEARCH_PRIVATE_ENDPOINT_SIGNED_JWT = "fake-signed-jwt"

samples/model-builder/vector_search/vector_search_find_neighbors_sample.py

@@ -84,3 +84,53 @@ def vector_search_find_neighbors(
     print(hybrid_resp)
 
 #  [END aiplatform_sdk_vector_search_find_neighbors_sample]
+
+
+#  [START aiplatform_sdk_vector_search_find_neighbors_jwt_sample]
+def vector_search_find_neighbors_jwt(
+    project: str,
+    location: str,
+    index_endpoint_name: str,
+    deployed_index_id: str,
+    queries: List[List[float]],
+    num_neighbors: int,
+    signed_jwt: str,
+) -> List[List[aiplatform.matching_engine.matching_engine_index_endpoint.MatchNeighbor]]:
+    """Query the vector search index.
+
+    Args:
+        project (str): Required. Project ID
+        location (str): Required. The region name
+        index_endpoint_name (str): Required. Index endpoint to run the query
+        against.
+        deployed_index_id (str): Required. The ID of the DeployedIndex to run
+        the queries against.
+        queries (List[List[float]]): Required. A list of queries. Each query is
+        a list of floats, representing a single embedding.
+        num_neighbors (int): Required. The number of neighbors to return.
+        signed_jwt (str): Required. The signed JWT token for the private
+        endpoint. The endpoint must be configured to accept tokens from JWT's
+        issuer and encoded audience.
+
+    Returns:
+        List[List[aiplatform.matching_engine.matching_engine_index_endpoint.MatchNeighbor]] - A list of nearest neighbors for each query.
+    """
+    # Initialize the Vertex AI client
+    aiplatform.init(project=project, location=location)
+
+    # Create the index endpoint instance from an existing endpoint.
+    my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
+        index_endpoint_name=index_endpoint_name
+    )
+
+    # Query the index endpoint for the nearest neighbors.
+    resp = my_index_endpoint.find_neighbors(
+        deployed_index_id=deployed_index_id,
+        queries=queries,
+        num_neighbors=num_neighbors,
+        signed_jwt=signed_jwt,
+    )
+    return resp
+
+#  [END aiplatform_sdk_vector_search_find_neighbors_jwt_sample]
+

samples/model-builder/vector_search/vector_search_find_neighbors_sample_test.py

@@ -55,3 +55,34 @@ def test_vector_search_find_neighbors_sample(
         ],
         any_order=False,
     )
+
+
+def test_vector_search_find_neighbors_jwt_sample(
+    mock_sdk_init, mock_index_endpoint_init, mock_index_endpoint_find_neighbors
+):
+    vector_search_find_neighbors_sample.vector_search_find_neighbors_jwt(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        index_endpoint_name=constants.VECTOR_SEARCH_INDEX_ENDPOINT,
+        deployed_index_id=constants.VECTOR_SEARCH_DEPLOYED_INDEX_ID,
+        queries=constants.VECTOR_SERACH_INDEX_QUERIES,
+        num_neighbors=10,
+        signed_jwt=constants.VECTOR_SEARCH_PRIVATE_ENDPOINT_SIGNED_JWT,
+    )
+
+    # Check client initialization
+    mock_sdk_init.assert_called_with(
+        project=constants.PROJECT, location=constants.LOCATION
+    )
+
+    # Check index endpoint initialization with right index endpoint name
+    mock_index_endpoint_init.assert_called_with(
+        index_endpoint_name=constants.VECTOR_SEARCH_INDEX_ENDPOINT)
+
+    # Check index_endpoint.find_neighbors is called with right params.
+    mock_index_endpoint_find_neighbors.assert_called_with(
+        deployed_index_id=constants.VECTOR_SEARCH_DEPLOYED_INDEX_ID,
+        queries=constants.VECTOR_SERACH_INDEX_QUERIES,
+        num_neighbors=10,
+        signed_jwt=constants.VECTOR_SEARCH_PRIVATE_ENDPOINT_SIGNED_JWT,
+    )

samples/model-builder/vector_search/vector_search_match_sample.py

@@ -0,0 +1,67 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List
+
+from google.cloud import aiplatform
+
+
+#  [START aiplatform_sdk_vector_search_match_jwt_sample]
+def vector_search_match_jwt(
+    project: str,
+    location: str,
+    index_endpoint_name: str,
+    deployed_index_id: str,
+    queries: List[List[float]],
+    num_neighbors: int,
+    signed_jwt: str,
+) -> List[List[aiplatform.matching_engine.matching_engine_index_endpoint.MatchNeighbor]]:
+    """Query the vector search index.
+
+    Args:
+        project (str): Required. Project ID
+        location (str): Required. The region name
+        index_endpoint_name (str): Required. Index endpoint to run the query
+        against. The endpoint must be a private endpoint.
+        deployed_index_id (str): Required. The ID of the DeployedIndex to run
+        the queries against.
+        queries (List[List[float]]): Required. A list of queries. Each query is
+        a list of floats, representing a single embedding.
+        num_neighbors (int): Required. The number of neighbors to return.
+        signed_jwt (str): Required. The signed JWT token for the private
+        endpoint. The endpoint must be configured to accept tokens from JWT's
+        issuer and encoded audience.
+
+    Returns:
+        List[List[aiplatform.matching_engine.matching_engine_index_endpoint.MatchNeighbor]] - A list of nearest neighbors for each query.
+    """
+    # Initialize the Vertex AI client
+    aiplatform.init(project=project, location=location)
+
+    # Create the index endpoint instance from an existing endpoint.
+    my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
+        index_endpoint_name=index_endpoint_name
+    )
+
+    # Query the index endpoint for matches.
+    resp = my_index_endpoint.match(
+        deployed_index_id=deployed_index_id,
+        queries=queries,
+        num_neighbors=num_neighbors,
+        signed_jwt=signed_jwt,
+    )
+    return resp
+
+#  [END aiplatform_sdk_vector_search_match_jwt_sample]
+

samples/model-builder/vector_search/vector_search_match_sample_test.py

@@ -0,0 +1,47 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import test_constants as constants
+from vector_search import vector_search_match_sample
+
+
+def test_vector_search_match_jwt_sample(
+    mock_sdk_init, mock_index_endpoint_init, mock_index_endpoint_match
+):
+    vector_search_match_sample.vector_search_match_jwt(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        index_endpoint_name=constants.VECTOR_SEARCH_INDEX_ENDPOINT,
+        deployed_index_id=constants.VECTOR_SEARCH_DEPLOYED_INDEX_ID,
+        queries=constants.VECTOR_SERACH_INDEX_QUERIES,
+        num_neighbors=10,
+        signed_jwt=constants.VECTOR_SEARCH_PRIVATE_ENDPOINT_SIGNED_JWT,
+    )
+
+    # Check client initialization
+    mock_sdk_init.assert_called_with(
+        project=constants.PROJECT, location=constants.LOCATION
+    )
+
+    # Check index endpoint initialization with right index endpoint name
+    mock_index_endpoint_init.assert_called_with(
+        index_endpoint_name=constants.VECTOR_SEARCH_INDEX_ENDPOINT)
+
+    # Check index_endpoint.match is called with right params.
+    mock_index_endpoint_match.assert_called_with(
+        deployed_index_id=constants.VECTOR_SEARCH_DEPLOYED_INDEX_ID,
+        queries=constants.VECTOR_SERACH_INDEX_QUERIES,
+        num_neighbors=10,
+        signed_jwt=constants.VECTOR_SEARCH_PRIVATE_ENDPOINT_SIGNED_JWT,
+    )