chore: Include type for vector_search_upsert_datapoints sample

PiperOrigin-RevId: 697694053

Author: vertex-sdk-bot

samples/model-builder/test_constants.py

@@ -380,8 +380,9 @@
 
 # Vector Search
 VECTOR_SEARCH_INDEX = "123"
-VECTOR_SERACH_INDEX_DATAPOINTS = [
-    {"datapoint_id": "datapoint_id_1", "feature_vector": [0.1]}
+VECTOR_SEARCH_INDEX_DATAPOINTS = [
+    aiplatform.compat.types.index_v1beta1.IndexDatapoint(datapoint_id="datapoint_id_1", feature_vector=[0.1, 0.2]),
+    aiplatform.compat.types.index_v1beta1.IndexDatapoint(datapoint_id="datapoint_id_2", feature_vector=[0.3, 0.4]),
 ]
 VECTOR_SEARCH_INDEX_DATAPOINT_IDS = ["datapoint_id_1", "datapoint_id_2"]
 VECTOR_SEARCH_INDEX_ENDPOINT = "456"

samples/model-builder/vector_search/vector_search_stream_update_sample.py

@@ -18,10 +18,13 @@
 
 
 #  [START aiplatform_sdk_vector_search_stream_update_sample]
-def stream_update_vector_search_index(
-    project: str, location: str, index_name: str, datapoints: Sequence[dict]
+def vector_search_upsert_datapoints(
+    project: str,
+    location: str,
+    index_name: str,
+    datapoints: Sequence[aiplatform.compat.types.index_v1beta1.IndexDatapoint],
 ) -> None:
-    """Stream update an existing vector search index
+    """Upsert datapoints to the index.
 
     Args:
       project (str): Required. The Project ID
@@ -30,8 +33,10 @@ def stream_update_vector_search_index(
         resource name or a index ID.  Example:
         "projects/123/locations/us-central1/indexes/my_index_id" or
         "my_index_id".
-      datapoints: Sequence[dict]: Required. The datapoints to be updated. The dict
-        element should be of the IndexDatapoint type.
+      datapoints: Sequence[IndexDatapoint]: Required. The datapoints to be
+        updated. For example:
+        [IndexDatapoint(datapoint_id="1", feature_vector=[1.0, 2.0, 3.0]),
+        IndexDatapoint(datapoint_id="2", feature_vector=[4.0, 5.0, 6.0])]
     """
     # Initialize the Vertex AI client
     aiplatform.init(project=project, location=location)

samples/model-builder/vector_search/vector_search_stream_update_sample_test.py

@@ -19,11 +19,11 @@
 def test_vector_search_stream_update_sample(
     mock_sdk_init, mock_index_init, mock_index_upsert_datapoints
 ):
-    vector_search_stream_update_sample.stream_update_vector_search_index(
+    vector_search_stream_update_sample.vector_search_upsert_datapoints(
         project=constants.PROJECT,
         location=constants.LOCATION,
         index_name=constants.VECTOR_SEARCH_INDEX,
-        datapoints=constants.VECTOR_SERACH_INDEX_DATAPOINTS,
+        datapoints=constants.VECTOR_SEARCH_INDEX_DATAPOINTS,
     )
 
     # Check client initialization
@@ -36,5 +36,5 @@ def test_vector_search_stream_update_sample(
 
     # Check index.upsert_datapoints is called with right params.
     mock_index_upsert_datapoints.assert_called_with(
-        datapoints=constants.VECTOR_SERACH_INDEX_DATAPOINTS
+        datapoints=constants.VECTOR_SEARCH_INDEX_DATAPOINTS
     )