Generative AI on Google Cloud

51 / 80·notebook

Chapter 51

Agent Platform Vector Search Quickstart

NotebookPython 354 cells

In [ ]python · cell 1

python

# 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.

Agent Platform Vector Search Quickstart


Author(s)	Smitha Venkat, Kaz Sato, Zachary Layne

Prerequisites

This tutorial requires a Google Cloud project that is linked with a billing account. To create a new project, take a look at this document to create a project and setup a billing account for it. To get the permissions that you need to give a service account access to enable APIs and interact with Agent Platform resources, ask your administrator to grant you the Security Admin (roles/iam.securityAdmin) IAM role on your project. For more information about granting roles, see Manage access to projects, folders, and organizations.

Choose the runtime environment

This tutorial can be run on either Google Colab or Agent Platform Workbench.

To use Colab: Click this link to open the tutorial in Colab.
To use Workbench: If it is the first time to use Workbench in your Google Cloud project, open the Workbench console and click ENABLE button to enable Notebooks API. The click this link, and select an existing notebook or create a new notebook.

How much will this cost?

To go through this tutorial, it will cost roughly a few US dollars. The pricing of the Cloud services we will use in this tutorial are available in the following pages:

You can use the Pricing Calculator to generate a cost estimate based on your projected usage.

Warning: delete your objects after the tutorial

In case you are using your own Cloud project, please make sure to delete all the Indexes and Index Endpoints (and the Workbench instance if you use one) after finishing this tutorial. Otherwise the remaining assets would incur unexpected costs.

Setup

Before get started with the Agent Platform services, we need to setup the following.

Install Python SDK
Environment variables
Authentication (Colab only)
Enable APIs
Set IAM permissions

Install the Agent Platform SDK

Agent Platform APIs can be accessed with multiple ways including REST API and Python SDK. In this tutorial we will use the SDK.

In [ ]python · cell 9

python

%pip install --upgrade --quiet --user google-cloud-aiplatform

Restart current runtime

To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which will restart the current kernel.

For Colab users: you will see "Your session crushed" message when restarting the kernel, but it is expected. Continue to the following cells.

In [ ]python · cell 11

python

# Restart kernel after installs so that your environment can access the new packages
import IPython

app = IPython.Application.instance()
app.kernel.do_shutdown(True)

Output

{'status': 'ok', 'restart': True}

Environment variables

Sets environment variables. If asked, replace the following [your-project-id] with your project ID and run it.

In [ ]python · cell 14

python

# get project ID
PROJECT_ID = ! gcloud config get project
PROJECT_ID = PROJECT_ID[0]
LOCATION = "us-central1"
if PROJECT_ID == "(unset)":
    print("Please set the project ID manually below")

In [ ]python · cell 15

python

# define project information
if PROJECT_ID == "(unset)":
    PROJECT_ID = "[your-project-id]"  # @param {type:"string"}

# generate an unique id for this session
from datetime import datetime

UID = datetime.now().strftime("%m%d%H%M")

Authentication (Colab only)

If you are running this notebook on Colab, you will need to run the following cell authentication. This step is not required if you are using Agent Platform Workbench as it is pre-authenticated.

In [ ]python · cell 17

python

import sys

# if it's Colab runtime, authenticate the user with Google Cloud
if "google.colab" in sys.modules:
    from google.colab import auth

    auth.authenticate_user()

Output

WARNING: google.colab.auth.authenticate_user() is not supported in Colab Enterprise.

Set IAM permissions

Also, we need to add access permissions to the default service account for using those services.

Go to the IAM page in the Console
Look for the principal for default compute service account. It should look like: <project-number>-compute@developer.gserviceaccount.com
Click the edit button at right and click ADD ANOTHER ROLE to add Agent Platform User and Service Usage Admin roles to the account.

Enable APIs

Run the following to enable APIs for Compute Engine and Agent Platform with this Google Cloud project.

In [ ]python · cell 20

python

! gcloud services enable compute.googleapis.com aiplatform.googleapis.com --project "{PROJECT_ID}"

Output

Operation "operations/acat.p2-761793285222-effb80de-d42f-4a2f-9e81-0e997d5b8f19" finished successfully.

Prepare the sample data

In this tutorial, we will use TheLook dataset that has a products table with about 30,000 rows of synthetic product data for a fictious e-commerce clothing site.

From this table, we have prepared the product-embs.json file.

This file is in JSONL format and each row has id for the product id, name for the product name, and embedding for the embedding of the product name in 768 dimensions which was generated previously with Agent Platform Embeddings for Text.

The text embeddings represent the meaning of the clothing product names. In this tutorial, we will use Vector Search for completing a semantic search of the items. This sample code can be used as a basis for other simple recommendation system where you can quickly find "other products similar to this one".

To learn more about how to create the embeddings from the data on a BigQuery table and store them in a JSON file, see Getting Started with Text Embeddings + Agent Platform Vector Search.

Download the data

For building an index with Vector Search, download the JSONL file.

In [ ]python · cell 23

python

! gcloud storage cp "gs://github-repo/data/vs-quickstart/product-embs.json" .

Output

Copying gs://github-repo/data/vs-quickstart/product-embs.json...
\
Operation completed over 1 objects/79.3 MiB.

Build and Deploy a Vector Search Index

Create Index

Now it's ready to load the embeddings to Vector Search. Its APIs are available under the aiplatform package of the SDK.

In [ ]python · cell 25

python

# init the aiplatform package
from google.cloud import aiplatform

aiplatform.init(project=PROJECT_ID, location=LOCATION)

Create an MatchingEngineIndex with its create_tree_ah_index function (Matching Engine is the previous name of Vector Search).

In [ ]python · cell 27

python

# create Index
INDEX_NAME = f"vs-quickstart-index-{UID}"

my_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(
    display_name=INDEX_NAME,
    dimensions=768,
    approximate_neighbors_count=10,
    index_update_method="STREAM_UPDATE",
)

By calling the create_tree_ah_index function, it starts building an Index. This will take a few minutes if the dataset is small, otherwise about 60 minutes or more depending on the size of the dataset. You can check status of the index creation on the Vector Search Console > INDEXES tab.

The parameters for creating index

dimensions: Dimension size of each embedding. In this case, it is 768 as we are using the embeddings from the Text Embeddings API.
approximate_neighbors_count: how many similar items we want to retrieve in typical cases
index_update_method: STREAM_UPDATE allows real time update to the index.

See the document for more details on creating Index and the parameters.

Create Index Endpoint and deploy the Index

To use the Index, you need to create an Index Endpoint. It works as a server instance accepting query requests for your Index.

In [ ]python · cell 30

python

# create IndexEndpoint
my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(
    display_name=f"vs-quickstart-index-endpoint-{UID}", public_endpoint_enabled=True
)

This tutorial utilizes a Public Endpoint and does not support Virtual Private Cloud (VPC). Unless you have a specific requirement for VPC, we recommend using a Public Endpoint. Despite the term "public" in its name, it does not imply open access to the public internet. Rather, it functions like other endpoints in Agent Platform services, which are secured by default through IAM. Without explicit IAM permissions, as we have previously established, no one can access the endpoint.

With the Index Endpoint, deploy the Index by specifying an unique deployed index ID.

In [ ]python · cell 33

python

DEPLOYED_INDEX_ID = f"vs_quickstart_deployed_{UID}"

Note: There is a difference between INDEX_NAME and DEPLOYED_INDEX_ID. The first one may include dashes (-), while the second one only allows underscores (_).

In [ ]python · cell 35

python

# deploy the Index to the Index Endpoint
my_index_endpoint.deploy_index(index=my_index, deployed_index_id=DEPLOYED_INDEX_ID)

If it is the first time to deploy an Index to an Index Endpoint, it will take around 30 minutes to automatically build and initiate the backend for it. After the first deployment, it will finish in seconds. To see the status of the index deployment, open the Vector Search Console > INDEX ENDPOINTS tab and click the Index Endpoint.

Streaming update to the vector search index

At this time, the vector search index still empty. The following code will insert the product embeddings from the JSONL file by using upsert_datapoints method.

In [ ]python · cell 38

python

import json

# read product-embs.json and put them to a list
with open("product-embs.json") as f:
    datapoints = []
    for line in f:
        item = json.loads(line)
        datapoints.append(
            {"datapoint_id": str(item["id"]), "feature_vector": item["embedding"]}
        )

In [ ]python · cell 39

python

# insert datapoints to the index
for i in range(0, len(datapoints), 1000):
    my_index.upsert_datapoints(datapoints=datapoints[i : i + 1000])

The upsert_datapoints() method takes a list of dicts where each dict should have two properties: datapoint_id for ID of the item, and feature_vector for embeddings of the item. The method can take up to 1000 items at a time.

The upsert_datapoints() method adds the items in the list to the vector search index in near real-time. As long as you are adding a few items at a time, it should be reflected to the query results in a couple of seconds. So this is suitable for use cases such as product catalogs where the sellers want to add or update many items every day and want the search system to reflect it instantly.

See the document for more details on streaming update.

Run a Query with Vector Search

Finally it's ready to use Vector Search. To run a query, you need to get an embedding for a query item. For example, if you like to find products with similar names to "cloudveil women's excursion short", you need to get the embedding for it before running a query. In this tutorial, we will get the embedding from the product-embs.json file as follows:

In [ ]python · cell 42

python

# read product-embs.json and build dicts of product names and embeddings
with open("product-embs.json") as f:
    product_names = {}
    product_embs = {}
    for line in f:
        item = json.loads(line)
        product_names[str(item["id"])] = item["name"]
        product_embs[str(item["id"])] = item["embedding"]

With the product_embs dict, you can specify a product ID to get an embedding for it.

In [ ]python · cell 44

python

# get the embedding for ID 6523 "cloudveil women's excursion short"
# you can also try with other IDs such as 12711, 18090, 19536 and 11863
query_emb = product_embs["6523"]

Run a Query

Then, pass the embedding to find_neighbors function to find similar product names.

In [ ]python · cell 46

python

# run query
response = my_index_endpoint.find_neighbors(
    deployed_index_id=DEPLOYED_INDEX_ID, queries=[query_emb], num_neighbors=10
)

# show the results
for idx, neighbor in enumerate(response[0]):
    print(f"{neighbor.distance:.2f} {product_names[neighbor.id]}")

Output

1.00 cloudveil women's excursion short
0.82 quiksilver womens cruiser short
0.80 xcvi women's alisal short
0.80 cloudveil men's kahuna short
0.78 ibex women's gozo short
0.78 sanctuary clothing women's coquette short
0.78 sunner women's collins printed short
0.77 hurley lowrider cargo 2.5 short - women's
0.77 stitch's women's fox knee length short
0.77 sanctuary clothing women's passenger skirt

The find_neighbors function only takes milliseconds to fetch the similar items even when you have billions of items on the Index, thanks to the ScaNN algorithm. Vector Search also supports autoscaling which can automatically resize the number of nodes based on the demands of your workloads.

IMPORTANT: Cleaning Up

In case you are using your own Cloud project, not a temporary project on Qwiklab, please make sure to delete all the Indexes and Index Endpoints after finishing this tutorial. Otherwise the remaining objects would incur unexpected costs.

If you used Workbench, you may also need to delete the Notebooks from the console.

In [ ]python · cell 49

python

# wait for a confirmation
input("Press Enter to delete Index Endpoint and Index:")

# delete Index Endpoint
my_index_endpoint.undeploy_all()
my_index_endpoint.delete(force=True)

# delete Index
my_index.delete()

Utilities

It can take some time to create or deploy indexes, and in that time you might lose connection with the Colab runtime. If you lose connection, instead of creating or deploying your new index again, you can check the Vector Search Console and use the existing ones to continue.

Get an existing Index

To get an index object that already exists, replace the following [numeric-index-id] with the index ID and run the cell. You can check the ID on the Vector Search Console > INDEXES tab.

In [ ]python · cell 52

python

my_index_id = "[numeric-index-id]"  # @param {type:"string"}
my_index = aiplatform.MatchingEngineIndex(my_index_id)

Get an existing Index Endpoint

To get an index endpoint object that already exists, replace the following [numeric-index-endpoint-id] with the index endpoint ID and run the cell. You can check the ID on the Vector Search Console > INDEX ENDPOINTS tab.

In [ ]python · cell 54

python

my_index_endpoint_id = "[numeric-index-endpoint-id]"  # @param {type:"string"}
my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(my_index_endpoint_id)

← 返回课程

Generative AI on Google Cloud

51 / 80·notebook

Chapter 51

Agent Platform Vector Search Quickstart

NotebookPython 354 cells

In [ ]python · cell 1

python

# 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.

Agent Platform Vector Search Quickstart


Author(s)	Smitha Venkat, Kaz Sato, Zachary Layne

Prerequisites

Choose the runtime environment

This tutorial can be run on either Google Colab or Agent Platform Workbench.

To use Colab: Click this link to open the tutorial in Colab.
To use Workbench: If it is the first time to use Workbench in your Google Cloud project, open the Workbench console and click ENABLE button to enable Notebooks API. The click this link, and select an existing notebook or create a new notebook.

How much will this cost?

To go through this tutorial, it will cost roughly a few US dollars. The pricing of the Cloud services we will use in this tutorial are available in the following pages:

You can use the Pricing Calculator to generate a cost estimate based on your projected usage.

Warning: delete your objects after the tutorial

Setup

Before get started with the Agent Platform services, we need to setup the following.

Install Python SDK
Environment variables
Authentication (Colab only)
Enable APIs
Set IAM permissions

Install the Agent Platform SDK

Agent Platform APIs can be accessed with multiple ways including REST API and Python SDK. In this tutorial we will use the SDK.

In [ ]python · cell 9

python

%pip install --upgrade --quiet --user google-cloud-aiplatform

Restart current runtime

To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which will restart the current kernel.

For Colab users: you will see "Your session crushed" message when restarting the kernel, but it is expected. Continue to the following cells.

In [ ]python · cell 11

python

# Restart kernel after installs so that your environment can access the new packages
import IPython

app = IPython.Application.instance()
app.kernel.do_shutdown(True)

Output

{'status': 'ok', 'restart': True}

Environment variables

Sets environment variables. If asked, replace the following [your-project-id] with your project ID and run it.

In [ ]python · cell 14

python

# get project ID
PROJECT_ID = ! gcloud config get project
PROJECT_ID = PROJECT_ID[0]
LOCATION = "us-central1"
if PROJECT_ID == "(unset)":
    print("Please set the project ID manually below")

In [ ]python · cell 15

python

# define project information
if PROJECT_ID == "(unset)":
    PROJECT_ID = "[your-project-id]"  # @param {type:"string"}

# generate an unique id for this session
from datetime import datetime

UID = datetime.now().strftime("%m%d%H%M")

Authentication (Colab only)

If you are running this notebook on Colab, you will need to run the following cell authentication. This step is not required if you are using Agent Platform Workbench as it is pre-authenticated.

In [ ]python · cell 17

python

import sys

# if it's Colab runtime, authenticate the user with Google Cloud
if "google.colab" in sys.modules:
    from google.colab import auth

    auth.authenticate_user()

Output

WARNING: google.colab.auth.authenticate_user() is not supported in Colab Enterprise.

Set IAM permissions

Also, we need to add access permissions to the default service account for using those services.

Go to the IAM page in the Console
Look for the principal for default compute service account. It should look like: <project-number>-compute@developer.gserviceaccount.com
Click the edit button at right and click ADD ANOTHER ROLE to add Agent Platform User and Service Usage Admin roles to the account.

Enable APIs

Run the following to enable APIs for Compute Engine and Agent Platform with this Google Cloud project.

In [ ]python · cell 20

python

! gcloud services enable compute.googleapis.com aiplatform.googleapis.com --project "{PROJECT_ID}"

Output

Operation "operations/acat.p2-761793285222-effb80de-d42f-4a2f-9e81-0e997d5b8f19" finished successfully.

Prepare the sample data

In this tutorial, we will use TheLook dataset that has a products table with about 30,000 rows of synthetic product data for a fictious e-commerce clothing site.

From this table, we have prepared the product-embs.json file.

To learn more about how to create the embeddings from the data on a BigQuery table and store them in a JSON file, see Getting Started with Text Embeddings + Agent Platform Vector Search.

Download the data

For building an index with Vector Search, download the JSONL file.

In [ ]python · cell 23

python

! gcloud storage cp "gs://github-repo/data/vs-quickstart/product-embs.json" .

Output

Copying gs://github-repo/data/vs-quickstart/product-embs.json...
\
Operation completed over 1 objects/79.3 MiB.

Build and Deploy a Vector Search Index

Create Index

Now it's ready to load the embeddings to Vector Search. Its APIs are available under the aiplatform package of the SDK.

In [ ]python · cell 25

python

# init the aiplatform package
from google.cloud import aiplatform

aiplatform.init(project=PROJECT_ID, location=LOCATION)

Create an MatchingEngineIndex with its create_tree_ah_index function (Matching Engine is the previous name of Vector Search).

In [ ]python · cell 27

python

# create Index
INDEX_NAME = f"vs-quickstart-index-{UID}"

my_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(
    display_name=INDEX_NAME,
    dimensions=768,
    approximate_neighbors_count=10,
    index_update_method="STREAM_UPDATE",
)

The parameters for creating index

dimensions: Dimension size of each embedding. In this case, it is 768 as we are using the embeddings from the Text Embeddings API.
approximate_neighbors_count: how many similar items we want to retrieve in typical cases
index_update_method: STREAM_UPDATE allows real time update to the index.

See the document for more details on creating Index and the parameters.

Create Index Endpoint and deploy the Index

To use the Index, you need to create an Index Endpoint. It works as a server instance accepting query requests for your Index.

In [ ]python · cell 30

python

# create IndexEndpoint
my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(
    display_name=f"vs-quickstart-index-endpoint-{UID}", public_endpoint_enabled=True
)

With the Index Endpoint, deploy the Index by specifying an unique deployed index ID.

In [ ]python · cell 33

python

DEPLOYED_INDEX_ID = f"vs_quickstart_deployed_{UID}"

Note: There is a difference between INDEX_NAME and DEPLOYED_INDEX_ID. The first one may include dashes (-), while the second one only allows underscores (_).

In [ ]python · cell 35

python

# deploy the Index to the Index Endpoint
my_index_endpoint.deploy_index(index=my_index, deployed_index_id=DEPLOYED_INDEX_ID)

Streaming update to the vector search index

At this time, the vector search index still empty. The following code will insert the product embeddings from the JSONL file by using upsert_datapoints method.

In [ ]python · cell 38

python

import json

# read product-embs.json and put them to a list
with open("product-embs.json") as f:
    datapoints = []
    for line in f:
        item = json.loads(line)
        datapoints.append(
            {"datapoint_id": str(item["id"]), "feature_vector": item["embedding"]}
        )

In [ ]python · cell 39

python

# insert datapoints to the index
for i in range(0, len(datapoints), 1000):
    my_index.upsert_datapoints(datapoints=datapoints[i : i + 1000])

See the document for more details on streaming update.

Run a Query with Vector Search

In [ ]python · cell 42

python

# read product-embs.json and build dicts of product names and embeddings
with open("product-embs.json") as f:
    product_names = {}
    product_embs = {}
    for line in f:
        item = json.loads(line)
        product_names[str(item["id"])] = item["name"]
        product_embs[str(item["id"])] = item["embedding"]

With the product_embs dict, you can specify a product ID to get an embedding for it.

In [ ]python · cell 44

python

# get the embedding for ID 6523 "cloudveil women's excursion short"
# you can also try with other IDs such as 12711, 18090, 19536 and 11863
query_emb = product_embs["6523"]

Run a Query

Then, pass the embedding to find_neighbors function to find similar product names.

In [ ]python · cell 46

python

# run query
response = my_index_endpoint.find_neighbors(
    deployed_index_id=DEPLOYED_INDEX_ID, queries=[query_emb], num_neighbors=10
)

# show the results
for idx, neighbor in enumerate(response[0]):
    print(f"{neighbor.distance:.2f} {product_names[neighbor.id]}")

Output

1.00 cloudveil women's excursion short
0.82 quiksilver womens cruiser short
0.80 xcvi women's alisal short
0.80 cloudveil men's kahuna short
0.78 ibex women's gozo short
0.78 sanctuary clothing women's coquette short
0.78 sunner women's collins printed short
0.77 hurley lowrider cargo 2.5 short - women's
0.77 stitch's women's fox knee length short
0.77 sanctuary clothing women's passenger skirt

IMPORTANT: Cleaning Up

If you used Workbench, you may also need to delete the Notebooks from the console.

In [ ]python · cell 49

python

# wait for a confirmation
input("Press Enter to delete Index Endpoint and Index:")

# delete Index Endpoint
my_index_endpoint.undeploy_all()
my_index_endpoint.delete(force=True)

# delete Index
my_index.delete()

Utilities

Get an existing Index

To get an index object that already exists, replace the following [numeric-index-id] with the index ID and run the cell. You can check the ID on the Vector Search Console > INDEXES tab.

In [ ]python · cell 52

python

my_index_id = "[numeric-index-id]"  # @param {type:"string"}
my_index = aiplatform.MatchingEngineIndex(my_index_id)

Get an existing Index Endpoint

In [ ]python · cell 54

python

my_index_endpoint_id = "[numeric-index-endpoint-id]"  # @param {type:"string"}
my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(my_index_endpoint_id)