chore: Add vector search sample for PSC automation deployment

PiperOrigin-RevId: 698867449

Author: vertex-sdk-bot

samples/model-builder/test_constants.py

@@ -416,3 +416,7 @@
 VECTOR_SEARCH_PRIVATE_ENDPOINT_SIGNED_JWT = "fake-signed-jwt"
 VECTOR_SEARCH_VPC_NETWORK = "vpc-network"
 VECTOR_SEARCH_PSC_PROJECT_ALLOWLIST = ["test-project", "test-project-2"]
+VECTOR_SEARCH_PSC_AUTOMATION_CONFIGS = [
+    ("test-project", "network1"),
+    ("test-project2", "network2"),
+]

samples/model-builder/vector_search/vector_search_deploy_index_sample.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Sequence, Tuple
+
 from google.cloud import aiplatform
 
 
@@ -109,3 +111,53 @@ def vector_search_deploy_autoscaling_index(
 
 
 #  [END aiplatform_sdk_vector_search_deploy_autoscaling_index_sample]
+
+
+#  [START aiplatform_sdk_vector_search_deploy_psc_automation_sample]
+def vector_search_deploy_psc_automation_index(
+    project: str,
+    location: str,
+    index_name: str,
+    index_endpoint_name: str,
+    deployed_index_id: str,
+    psc_automation_configs: Sequence[Tuple[str, str]],
+) -> None:
+    """Deploy a vector search index to an index endpoint using PSC automation.
+
+    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.
+        psc_automation_config (Sequence[Tuple[str, str]]): Required. A list of
+          (project_id, network) pairs where PSC endpoints will be setup for the
+          deployed index. Example:
+          [("123", "{projects/123/global/networks/my-network1"),
+          ("123", "{projects/123/global/networks/my-network2")]
+    """
+    # 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 with PSC automation enabled.
+    index_endpoint.deploy_index(
+        index=index,
+        deployed_index_id=deployed_index_id,
+        psc_automation_configs=psc_automation_configs,
+    )
+
+
+#  [END aiplatform_sdk_vector_search_deploy_psc_automation_sample]

samples/model-builder/vector_search/vector_search_deploy_index_sample_test.py

@@ -86,3 +86,39 @@ def test_vector_search_deploy_autoscaling_index_sample(
         min_replica_count=constants.MIN_REPLICA_COUNT,
         max_replica_count=constants.MAX_REPLICA_COUNT,
     )
+
+
+def test_vector_search_deploy_psc_automation_index_sample(
+    mock_sdk_init,
+    mock_index_init,
+    mock_index_endpoint_init,
+    mock_index_endpoint_deploy_index,
+    mock_index,
+):
+    vector_search_deploy_index_sample.vector_search_deploy_psc_automation_index(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        index_name=constants.VECTOR_SEARCH_INDEX,
+        index_endpoint_name=constants.VECTOR_SEARCH_INDEX_ENDPOINT,
+        deployed_index_id=constants.VECTOR_SEARCH_DEPLOYED_INDEX_ID,
+        psc_automation_configs=constants.VECTOR_SEARCH_PSC_AUTOMATION_CONFIGS,
+    )
+
+    # 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 endpoint initialization with right index endpoint name
+    mock_index_endpoint_init.assert_called_with(
+        index_endpoint_name=constants.VECTOR_SEARCH_INDEX_ENDPOINT
+    )
+    # Check index deployment
+    mock_index_endpoint_deploy_index.assert_called_with(
+        index=mock_index,
+        deployed_index_id=constants.VECTOR_SEARCH_DEPLOYED_INDEX_ID,
+        psc_automation_configs=constants.VECTOR_SEARCH_PSC_AUTOMATION_CONFIGS,
+    )