Integrate Vertex AI with Astra DB for search and chat

query_builder 15 min

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:
- Google Cloud CLI
- Docker

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 is AstraCS: 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.
1. Name your secret DATASTAX_VERTEX_AI_TOKEN.
2. 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 of AstraCS: followed by a unique token string.
  
  API_ENDPOINT: Your database application token in the form of https://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.
1. In the Container image URL field, select the image you want to use for the Cloud Run service.
  1. To use the DataStax-hosted image from the Google Container Registry, enter gcr.io/datastax-public/astra-search:latest.
  2. 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.
2. Select Require authentication to require authenticated access to the container.
3. Grant the Cloud Run service permissions to the following roles: Vertex AI Administrator, Secret Manager Secret Accessor, Storage Admin, and Discovery Engine Admin.
4. 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. Replace YOUR_CLOUD_RUN_URL with the URL for your Cloud Run service.
```
curl -H "Authorization: Bearer $(gcloud auth print-identity-token)" https://CLOUD_RUN_URL-uc.a.run.app/sync
```
The default query parameters for the sync operation include the following:
Sync default parameters
```
table = 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 the ALTERNATE_ASTRA_TABLE table:
```
https://CLOUD_RUN_URL-uc.a.run.app/sync?table=ALTERNATE_ASTRA_TABLE
```
To filter your 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.
1. Click Continue and Activate the API.
2. In the navigation menu, select Data Stores.
3. Click Create Data Store.
4. Select Cloud Storage as your data source.
5. Select the GCS bucket you want to use for the data store, and then click Continue.
6. 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.
1. Click New App.
2. Select Search as the app type.
3. Select the Data Store you created previously, and then click Create.
4. 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 and GOOGLE_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:
1. 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
2. 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
3. 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.