Skip to content

Commit

Permalink
docs(samples): adding code sample for vector search create streaming …
Browse files Browse the repository at this point in the history
…index

PiperOrigin-RevId: 667564368
  • Loading branch information
holtskinner authored and copybara-github committed Aug 26, 2024
1 parent 3d68777 commit 71464e7
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -48,3 +48,36 @@ def vector_search_create_index(


# [END aiplatform_sdk_vector_search_create_index_sample]


# [START aiplatform_sdk_vector_search_create_streaming_index_sample]
def vector_search_create_streaming_index(
project: str, location: str, display_name: str, gcs_uri: Optional[str] = None
) -> None:
"""Create a vector search index.
Args:
project (str): Required. Project ID
location (str): Required. The region name
display_name (str): Required. The index display name
gcs_uri (str): Optional. The Google Cloud Storage uri for index content
"""
# Initialize the Vertex AI client
aiplatform.init(project=project, location=location, staging_bucket=gcs_uri)

# Create Index
index = aiplatform.MatchingEngineIndex.create_tree_ah_index(
display_name=display_name,
description="Matching Engine Index",
dimensions=100,
approximate_neighbors_count=150,
leaf_node_embedding_count=500,
leaf_nodes_to_search_percent=7,
index_update_method="stream_update", # Options: stream_update, batch_update
distance_measure_type=aiplatform.matching_engine.matching_engine_index_config.DistanceMeasureType.DOT_PRODUCT_DISTANCE,
)

print(index.name)


# [END aiplatform_sdk_vector_search_create_streaming_index_sample]
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,37 @@ def test_vector_search_create_index_sample(
approximate_neighbors_count=ANY,
leaf_node_embedding_count=ANY,
leaf_nodes_to_search_percent=ANY,
index_update_method=ANY,
index_update_method="batch_update",
distance_measure_type=ANY,
)


def test_vector_search_create_streaming_index_sample(
mock_sdk_init,
mock_index_create_tree_ah_index,
):
vector_search_create_index_sample.vector_search_create_streaming_index(
project=constants.PROJECT,
location=constants.LOCATION,
display_name=constants.VECTOR_SEARCH_INDEX_DISPLAY_NAME,
gcs_uri=constants.VECTOR_SEARCH_GCS_URI,
)

# Check client initialization
mock_sdk_init.assert_called_with(
project=constants.PROJECT,
location=constants.LOCATION,
staging_bucket=constants.VECTOR_SEARCH_GCS_URI,
)

# Check index creation
mock_index_create_tree_ah_index.assert_called_with(
display_name=constants.VECTOR_SEARCH_INDEX_DISPLAY_NAME,
description=ANY,
dimensions=ANY,
approximate_neighbors_count=ANY,
leaf_node_embedding_count=ANY,
leaf_nodes_to_search_percent=ANY,
index_update_method="stream_update",
distance_measure_type=ANY,
)

0 comments on commit 71464e7

Please sign in to comment.