Integrate with Vertex AI

This page details how to set up the integration between AlloyDB for PostgreSQL and Vertex AI, letting you apply queries to cloud-stored large language models (LLMs) to your data.

These instructions are specific to using AlloyDB, and not AlloyDB Omni. To instead integrate a local installation of AlloyDB Omni using Vertex AI, see Install AlloyDB Omni with AlloyDB AI.

For more information about using ML models with AlloyDB, see Build generative AI applications using AlloyDB AI.

For more information about Vertex AI, see Introduction to Vertex AI.

Enable database integration with Vertex AI

To enable database integration with Vertex AI, grant the AlloyDB service agent Identity and Access Management (IAM) permissions to access Vertex AI.

  1. Add Vertex AI permissions to the AlloyDB service agent for the project where the AlloyDB database's cluster is located:

    Console

    1. Get the project number of the project that has AlloyDB clusters or instances by following the instructions available in Identifying projects.
    2. In the Google Cloud console, go to the IAM page.

      Go to IAM

    3. Select the project that has Vertex AI endpoints.
    4. Enable the Include Google-provided role grants checkbox.
    5. Click Add.
    6. In the New principals field, enter the following:

      service-PROJECT_NUMBER@gcp-sa-alloydb.iam.gserviceaccount.com

      Replace PROJECT_NUMBER with the project number.

    7. In the Role field, enter Vertex AI User.
    8. Click Save.

    gcloud

    gcloud projects add-iam-policy-binding PROJECT_ID \
  --member="serviceAccount:service-PROJECT_NUMBER@gcp-sa-alloydb.iam.gserviceaccount.com" \
  --role="roles/aiplatform.user"
    • PROJECT_ID: The ID of the project that has the Vertex AI endpoint.
    • PROJECT_NUMBER: The project number of the project that has AlloyDB clusters or instances.
  2. Install the google_ml_integration extension in the database that contains the data that you want to run predictions on:
    1. Connect a psql client to the cluster's primary instance, as described in Connect a psql client to an instance.
    2. At the psql command prompt, connect to the database: 
      \c DB_NAME

      Replace DB_NAME with the name of the database that you want to install the extension onto.

    3. Install the extension: 
      CREATE EXTENSION IF NOT EXISTS google_ml_integration CASCADE;

What's next