Integrate Vertex AI with Astra DB for search and conversation

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 Astra Portal, under Databases, navigate to your database.
Ensure the database is in Active status, and then select Generate Token. In the Application Token dialog, click content_paste Copy to copy the token (e.g. AstraCS:WSnyFUhRxsrg…). Store the token in a secure location before closing the dialog.

Your token is automatically assigned the Database Administrator role.
Copy your database’s API endpoint, located under Database Details > API Endpoint (e.g. 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.
  
  API_ENDPOINT: Your database application token.
  
  COLLECTION: The name of the collection or table in your database.
  Example secret value
  AstraCS:QCyMhoAQlqBtHoNoOsGPXxeU:09315c939c7d6b2a0b1d586356952bb3a7d8da38e9714a855306f26ac1a7e178;https://a71b76c8-5236-4c1e-adb2-9fdd6c85d60b-us-east-1.apps.astra.datastax.com;my_collection
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 from 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 are:
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://[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.
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 searchSearch 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

To start your Docker container, run the following command. 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, visit http://localhost:8080/sync.
Optionally, 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. See Cloud Run to continue with the deployment steps.

Remove the integration

To remove the Astra DB integration:

Locate your Astra DB application on your project’s Google Cloud Vertex AI App Builder page.
Click more_vert More, and then click Delete.

Your application is deleted.
Locate your container image on your project’s Google Artifact Registry page.
Select the container, and then click deleteDelete.

Your container is deleted.
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.