Integrate Vertex AI with Astra DB for search and chat
Vertex AI can integrate with Astra DB Serverless for two-way syncs between your Serverless (Vector) database and a Google Cloud Storage bucket. This integration allows you to build context-aware search and chat applications in Vertex AI with Astra DB as the vector database for the Gemini LLM.
This tutorial shows you how to integrate Astra DB Serverless with Vertex AI on Cloud Run or with a local Docker container.
Prerequisites
The steps in this guide assume the following:
-
You have an active Astra account.
-
You have created a Serverless (Vector) database.
-
You have created a collection in your database and loaded data into it.
-
-
You have a Google Cloud project with the following properties:
-
The Vertex AI API and Secret Manager API are enabled.
-
The Vertex AI Administrator and Vertex AI User roles are assigned to users or service accounts in your project.
-
-
You have the following software installed on your computer:
Create a secret for your database credentials
-
In the Astra Portal, go to Databases, and select your database.
-
Make sure the database is in Active status, and then, in the Database Details section, click Generate Token.
-
In the Application Token dialog, click
Copy, and then store the token securely. The token format isAstraCS:
followed by a unique token string.Application tokens created from Database Details have the Database Administrator role for the associated database.
-
In Database Details, copy your database’s API endpoint. The endpoint format is
https://ASTRA_DB_ID-ASTRA_DB_REGION.apps.astra.datastax.com
. -
In your Google Cloud project, create a secret for your database credentials.
-
Name your secret
DATASTAX_VERTEX_AI_TOKEN
. -
Add your database credentials, separated by semicolons, to the secret value:
TOKEN;API_ENDPOINT;COLLECTION
Replace the following:
-
TOKEN
: Your database API endpoint in the form ofAstraCS:
followed by a unique token string. -
API_ENDPOINT
: Your database application token in the form ofhttps://ASTRA_DB_ID-ASTRA_DB_REGION.apps.astra.datastax.com
. -
COLLECTION
: The name of the collection or table in your database.
-
-
-
Grant access to the secret.
You can grant access to your own account if you plan to register your Astra DB extension manually. Otherwise, you can grant access to a service account.
Create a Cloud Run service to sync a GCS bucket with Astra DB
Deploy a Docker container to Google Cloud Run to sync your Astra DB database with a GCS bucket.
-
In your GCP project, navigate to Google Cloud Run.
-
Click Create Service.
-
In the Container image URL field, select the image you want to use for the Cloud Run service.
-
To use the DataStax-hosted image from the Google Container Registry, enter
gcr.io/datastax-public/astra-search:latest
. -
To build your own image and push it to the container registry, see Build a local Docker container to sync a GCS bucket with Astra DB.
-
-
Select Require authentication to require authenticated access to the container.
-
Grant the Cloud Run service permissions to the following roles: Vertex AI Administrator, Secret Manager Secret Accessor, Storage Admin, and Discovery Engine Admin.
-
Click Create to deploy the Cloud Run service.
-
-
To sync your Astra DB database with Google Cloud Storage, call the
/sync
endpoint on the deployed Cloud Run service. ReplaceYOUR_CLOUD_RUN_URL
with the URL for your Cloud Run service.curl -H "Authorization: Bearer $(gcloud auth print-identity-token)" https://YOUR_CLOUD_RUN_URL-uc.a.run.app/sync
The default query parameters for the sync operation include the following:
Sync default parameterstable = params.get("table", "moviereviews") filter = params.get("filter", None) bucket = params.get("bucket", "astra_bucket") datastore = params.get("datastore", "astra_datastore")
To change which table or column is synced, add the
?table=
or?column=
query parameter to the URL. For example, this call will only sync records from theALTERNATE_ASTRA_TABLE
table:https://YOUR_CLOUD_RUN_URL-uc.a.run.app/sync?table=ALTERNATE_ASTRA_TABLE
To filter your Astra results, add the
?filter=
parameter to the URL. This call will only sync records with a rating value greater than 3:https://YOUR_CLOUD_RUN_URL-uc.a.run.app/sync?filter=rating>3
Deploy a Vertex AI search application
-
To deploy the Vertex AI application, navigate to Google Cloud Vertex AI App Builder.
-
Click Continue and Activate the API.
-
In the navigation menu, select Data Stores.
-
Click Create Data Store.
-
Select Cloud Storage as your data source.
-
Select the GCS bucket you want to use for the data store, and then click Continue.
-
Provide a name for your data store, and then click Create. The default values are acceptable for this tutorial.
-
-
To create a search application, navigate to Google Cloud Vertex AI App Builder, and then click Apps.
-
Click New App.
-
Select Search as the app type.
-
Select the Data Store you created previously, and then click Create.
-
To test your new search application, select Preview.
-
-
In the
Search here field, ask a question or search for a term. Your application should respond with data gleaned from your Astra DB database.
Build a local Docker container to sync a GCS bucket with Astra DB
Build a local Docker container that syncs a Google Cloud Storage bucket with your Astra DB database.
-
Clone the DataStax Vertex AI search extension repository:
git clone git@github.com:datastax/vertex-search-extension.git
-
Create and download a Service Account JSON key with the Vertex AI Administrator, Secret Manager Secret Accessor, Storage Admin, and Discovery Engine Admin permissions. For more information, see Creating and managing service account keys.
-
Put the key in the
astra-search-extension
folder so it is available within the container. -
Authenticate your Google Cloud account in the Google Cloud CLI:
gcloud auth login
-
Set your project ID:
gcloud config set project PROJECT_ID
-
Build the Docker image from the root of the repository. Replace
PROJECT_ID
with your Google Cloud project ID.docker build --platform linux/amd64 -t us-central1-docker.pkg.dev/[PROJECT_ID]/astra-api/astra-search astra-search-extension
-
Start your Docker container. Replace
PATH_TO_SERVICE_ACCOUNT
andGOOGLE_PROJECT_ID
with values from your GCP project.docker run -v PATH_TO_SERVICE_ACCOUNT.json:/secrets/service-account.json \ -e GOOGLE_APPLICATION_CREDENTIALS=/secrets/service-account.json \ --rm -p 8080:8080 -e GOOGLE_PROJECT_ID=GOOGLE_PROJECT_ID \ astra-search
-
To sync your Serverless (Vector) database to your GCS bucket, navigate to
http://localhost:8080/sync
. -
Optional: Push the Docker image to Google Artifact Registry for use in a Cloud Run application:
-
Create a new Artifact Hub repository:
gcloud artifacts repositories create astra-api \ --repository-format=docker \ --location=us-central1 \ --description="Vertex AI Containers for Astra DB" \ --async
-
Push the Docker image to the Artifact Hub. Replace
PROJECT_ID
with your GCP project ID.docker push us-central1-docker.pkg.dev/PROJECT_ID/astra-api/astra-search
-
Create a Cloud Run service to sync a GCS bucket with Astra DB
-
Remove the integration
To delete the Astra DB application, container, and Vertex AI token, do the following:
-
Locate your Astra DB application on your project’s Google Cloud Vertex AI App Builder page.
-
Click
More, and then click Delete. -
Locate your container image on your project’s Google Artifact Registry page.
-
Select the container, and then click
Delete. -
Locate and delete the
DATASTAX_VERTEX_AI_TOKEN
secret and any associated permissions.
The Astra DB integration is removed and Vertex AI can no longer access your Astra DB database.