chore: Add more vector search samples

PiperOrigin-RevId: 688286858

google/cloud/aiplatform/matching_engine/matching_engine_index.py

@@ -339,9 +339,7 @@ def update_embeddings(
         Args:
             contents_delta_uri (str):
                 Required. Allows inserting, updating  or deleting the contents of the Matching Engine Index.
-                The string must be a valid Google Cloud Storage directory path. If this
-                field is set when calling IndexService.UpdateIndex, then no other
-                Index field can be  also updated as part of the same call.
+                The string must be a valid Google Cloud Storage directory path.
                 The expected structure and format of the files this URI points to is
                 described at
                 https://cloud.google.com/vertex-ai/docs/vector-search/setup/format-structure

samples/model-builder/conftest.py

@@ -1309,20 +1309,65 @@ def mock_index_endpoint():
     yield mock
 
 
+@pytest.fixture
+def mock_list_indexes(mock_index):
+    with patch.object(aiplatform.MatchingEngineIndex, "list") as mock:
+        mock.return_value = [mock_index, mock_index]
+        yield mock
+
+
 @pytest.fixture
 def mock_index_init(mock_index):
     with patch.object(aiplatform, "MatchingEngineIndex") as mock:
         mock.return_value = mock_index
         yield mock
 
 
+@pytest.fixture
+def mock_index_update_embeddings(mock_index):
+    with patch.object(mock_index, "update_embeddings") as mock:
+        mock.return_value = None
+        yield mock
+
+
+@pytest.fixture
+def mock_index_update_metadata(mock_index):
+    with patch.object(mock_index, "update_metadata") as mock:
+        mock.return_value = None
+        yield mock
+
+
 @pytest.fixture
 def mock_index_upsert_datapoints(mock_index):
     with patch.object(mock_index, "upsert_datapoints") as mock_upsert:
         mock_upsert.return_value = None
         yield mock_upsert
 
 
+@pytest.fixture
+def mock_index_remove_datapoints(mock_index):
+    with patch.object(mock_index, "remove_datapoints") as mock:
+        mock.return_value = None
+        yield mock
+
+
+@pytest.fixture
+def mock_index_delete(mock_index):
+    with patch.object(mock_index, "delete") as mock:
+        mock.return_value = None
+        yield mock
+
+
+@pytest.fixture
+def mock_list_index_endpoints(mock_index_endpoint):
+    with patch.object(aiplatform.MatchingEngineIndexEndpoint, "list") as mock:
+        mock.return_value = [
+            mock_index_endpoint,
+            mock_index_endpoint,
+        ]
+        yield mock
+
+
 @pytest.fixture
 def mock_index_endpoint_init(mock_index_endpoint):
     with patch.object(aiplatform, "MatchingEngineIndexEndpoint") as mock:
@@ -1360,3 +1405,24 @@ def mock_index_endpoint_deploy_index(mock_index_endpoint):
     with patch.object(mock_index_endpoint, "deploy_index") as mock_deploy_index:
         mock_deploy_index.return_value = mock_index_endpoint
         yield mock_deploy_index
+
+
+@pytest.fixture
+def mock_index_endpoint_undeploy_index(mock_index_endpoint):
+    with patch.object(mock_index_endpoint, "undeploy_index") as mock:
+        mock.return_value = mock_index_endpoint
+        yield mock
+
+
+@pytest.fixture
+def mock_index_endpoint_mutate_deployed_index(mock_index_endpoint):
+    with patch.object(mock_index_endpoint, "mutate_deployed_index") as mock:
+        mock.return_value = mock_index_endpoint
+        yield mock
+
+
+@pytest.fixture
+def mock_index_endpoint_delete(mock_index_endpoint):
+    with patch.object(mock_index_endpoint, "delete") as mock:
+        mock.return_value = None
+        yield mock

samples/model-builder/test_constants.py

@@ -378,6 +378,7 @@
 VECTOR_SERACH_INDEX_DATAPOINTS = [
     {"datapoint_id": "datapoint_id_1", "feature_vector": [0.1]}
 ]
+VECTOR_SEARCH_INDEX_DATAPOINT_IDS = ["datapoint_id_1", "datapoint_id_2"]
 VECTOR_SEARCH_INDEX_ENDPOINT = "456"
 VECTOR_SEARCH_DEPLOYED_INDEX_ID = "789"
 VECTOR_SERACH_INDEX_QUERIES = [[0.1]]
@@ -402,5 +403,7 @@
     ),
 ]
 VECTOR_SEARCH_INDEX_DISPLAY_NAME = "my-vector-search-index"
+VECTOR_SEARCH_INDEX_DESCRIPTION = "test description"
+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"

samples/model-builder/vector_search/vector_search_delete_index_endpoint_sample.py

@@ -0,0 +1,44 @@
+# Copyright 2024 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 google.cloud import aiplatform
+
+
+#  [START aiplatform_sdk_vector_search_delete_index_endpoint_sample]
+def vector_search_delete_index_endpoint(
+    project: str, location: str, index_endpoint_name: str, force: bool = False
+) -> None:
+    """Delete a vector search index endpoint.
+
+    Args:
+        project (str): Required. Project ID
+        location (str): Required. The region name
+        index_endpoint_name (str): Required. Index endpoint to run the query
+          against.
+        force (bool): Required. If true, undeploy any deployed indexes on this
+          endpoint before deletion.
+    """
+    # Initialize the Vertex AI client
+    aiplatform.init(project=project, location=location)
+
+    # Create the index endpoint instance from an existing endpoint
+    index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
+        index_endpoint_name=index_endpoint_name
+    )
+
+    # Delete the index endpoint
+    index_endpoint.delete(force=force)
+
+
+#  [END aiplatform_sdk_vector_search_delete_index_endpoint_sample]

samples/model-builder/vector_search/vector_search_delete_index_endpoint_sample_test.py

@@ -0,0 +1,43 @@
+# Copyright 2024 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_delete_index_endpoint_sample
+
+
+def test_vector_search_delete_index_endpoint_sample(
+    mock_sdk_init,
+    mock_index_endpoint_init,
+    mock_index_endpoint_delete,
+):
+    vector_search_delete_index_endpoint_sample.vector_search_delete_index_endpoint(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        index_endpoint_name=constants.VECTOR_SEARCH_INDEX_ENDPOINT,
+        force=True,
+    )
+
+    # 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 deletion
+    mock_index_endpoint_delete.assert_called_with(force=True)

samples/model-builder/vector_search/vector_search_delete_index_sample.py

@@ -0,0 +1,42 @@
+# Copyright 2024 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 google.cloud import aiplatform
+
+
+#  [START aiplatform_sdk_vector_search_delete_index_sample]
+def vector_search_delete_index(
+    project: str, location: str, index_name: str
+) -> None:
+    """Delete a vector search index.
+
+    Args:
+        project (str): Required. Project ID
+        location (str): Required. The region name
+        index_name (str): Required. The index to update. A fully-qualified index
+          resource name or a index ID.  Example:
+          "projects/123/locations/us-central1/indexes/my_index_id" or
+          "my_index_id".
+    """
+    # Initialize the Vertex AI client
+    aiplatform.init(project=project, location=location)
+
+    # Create the index instance from an existing index
+    index = aiplatform.MatchingEngineIndex(index_name=index_name)
+
+    # Delete the index
+    index.delete()
+
+
+#  [END aiplatform_sdk_vector_search_delete_index_sample]

samples/model-builder/vector_search/vector_search_delete_index_sample_test.py

@@ -0,0 +1,40 @@
+# Copyright 2024 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_delete_index_sample
+
+
+def test_vector_search_delete_index_sample(
+    mock_sdk_init,
+    mock_index_init,
+    mock_index_delete,
+):
+    vector_search_delete_index_sample.vector_search_delete_index(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        index_name=constants.VECTOR_SEARCH_INDEX,
+    )
+
+    # Check client initialization
+    mock_sdk_init.assert_called_with(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+    )
+
+    # Check index initialization with right index name
+    mock_index_init.assert_called_with(index_name=constants.VECTOR_SEARCH_INDEX)
+
+    # Check index deletion
+    mock_index_delete.assert_called()

samples/model-builder/vector_search/vector_search_deploy_index_sample.py

@@ -32,8 +32,10 @@ def vector_search_deploy_index(
           resource name or a index ID.  Example:
           "projects/123/locations/us-central1/indexes/my_index_id" or
           "my_index_id".
-        index_endpoint_name (str): Required. Index endpoint to deploy the index to.
-        deployed_index_id (str): Required. The user specified ID of the DeployedIndex.
+        index_endpoint_name (str): Required. Index endpoint to deploy the index
+          to.
+        deployed_index_id (str): Required. The user specified ID of the
+          DeployedIndex.
     """
     # Initialize the Vertex AI client
     aiplatform.init(project=project, location=location)
@@ -55,3 +57,55 @@ def vector_search_deploy_index(
 
 
 #  [END aiplatform_sdk_vector_search_deploy_index_sample]
+
+
+#  [START aiplatform_sdk_vector_search_deploy_autoscaling_index_sample]
+def vector_search_deploy_autoscaling_index(
+    project: str,
+    location: str,
+    index_name: str,
+    index_endpoint_name: str,
+    deployed_index_id: str,
+    min_replica_count: int,
+    max_replica_count: int,
+) -> None:
+    """Deploy a vector search index to a vector search index endpoint.
+
+    Args:
+        project (str): Required. Project ID
+        location (str): Required. The region name
+        index_name (str): Required. The index to update. A fully-qualified index
+          resource name or a index ID.  Example:
+          "projects/123/locations/us-central1/indexes/my_index_id" or
+          "my_index_id".
+        index_endpoint_name (str): Required. Index endpoint to deploy the index
+          to.
+        deployed_index_id (str): Required. The user specified ID of the
+          DeployedIndex.
+        min_replica_count (int): Required. The minimum number of replicas to
+          deploy.
+        max_replica_count (int): Required. The maximum number of replicas to
+          deploy.
+    """
+    # Initialize the Vertex AI client
+    aiplatform.init(project=project, location=location)
+
+    # Create the index instance from an existing index
+    index = aiplatform.MatchingEngineIndex(index_name=index_name)
+
+    # Create the index endpoint instance from an existing endpoint.
+    index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
+        index_endpoint_name=index_endpoint_name
+    )
+
+    # Deploy Index to Endpoint. Specifying min and max replica counts will
+    # enable autoscaling.
+    index_endpoint.deploy_index(
+        index=index,
+        deployed_index_id=deployed_index_id,
+        min_replica_count=min_replica_count,
+        max_replica_count=max_replica_count,
+    )
+
+
+#  [END aiplatform_sdk_vector_search_deploy_autoscaling_index_sample]