Productionise genAI directly in dbt

Photo by Steve Johnson on Unsplash

Google created the connection between BigQuery and Vertex AI enables to use generative AI (genAI) directly in BigQuery. This article explains how to productionise such a solution using dbt and infrastructure as code.

The method is based on another medium post, that explains the steps needed to query a LLM in BigQuery by clicking in the GCP console. Most enterprise data pipelines require or highly benefits from infrastructure as code (IaC) and scheduling, instead of clicking in an interface. Luckily, at Nordnet, we figured out how to do schedule the job and take care of the infrastructure with code.

At Nordnet, we have built an end-to-end data pipeline, based on the approach linked above, and will solve:

• How to schedule the code
• How to set up the infrastructure as code

This article talks about all steps needed to productionise Vertex AI LLMs in dbt using BigQuery and set up this with infrastructure as code using Terraform. The advantages:

• Able to use genAI by simply writing SQL
• Short time from idea to production
• Able to iterate and test quickly
• Reproducible infrastructure

Initially, the tech setup at Nordnet is discussed and the architecture visualised. Afterwards, each step of the solution is explained in detail.

While this article builds on basic knowledge of dbt, Terraform, BigQuery, Vertex AI, and GCP, please note it may be challenging if you are not yet familiar with these concepts.

Tech setup

Nordnet is a FinTech bank in Scandinavia and leverages a modern tech stack on GCP, similar to other GCP based companies.

Data: The data is ingested into BigQuery.

Tech infrastructure: Luckily, Nordnet has a modern tech stack:

• GCP
• Terraform for Infrastructure as Code
• dbt for scheduling and maintaining the data warehouse in BigQuery

genAI: Model Garden in Vertex AI, including Gemini Pro.

What architecture can we create to connect all pieces and productionise a solution? At Nordnet, we implemented the architecture as shown in the diagram below.

Architecture drawing

Detailed guide

In this section a step by step guide is given to setup the connection between BigQuery and Vertex AI using dbt and Terraform. This is shown in five steps: setup of the connection, granting permissions, setup of BigQuery Model, SQL call to Gemini and running dbt.

Step 1: Setup of connection

Firstly, we need to create a connection between Vertex AI and BigQuery

resource "google_bigquery_connection" "connection" {
  project = my_project_id
  connection_id = "external-connection-vertex-ai"
  location = "europe-north1"
  Description = "An external connection to Vertex AI"
  cloud_resource {
  }
}

After applying the terraform code, you should be able to see this in BigQuery.

What if the external connection is not shown in the GCP console? Typically, as user, you can be missing the role bigquery.connectionUser which is needed to view connections.

Step 2: Granting permissions

The external connection also creates a Service Account that is used to access Vertex AI. This service account needs to have the role Vertex AI User.

In Terraform, we can grant the service account the Vertex AI User role as following.

resource "google_project_iam_member" "external_connection_vertex_ai_iam" {
  project = my_project_id
  role = "roles/aiplatform.user"
  member = "serviceAccount:${google_bigquery_connection.connection.cloud_resource[0].service_account_id}"
  }

Step 3: Setup of a BigQuery model

Now that we have a connection to Vertex AI setup in BigQuery, we need to create a bigQuery model that uses this connection to connect to a Vertex AI model.

The code for a BigQuery model in BigQuery is quite simple.

CREATE OR REPLACE MODEL `my_project_name.{{dataset}}.{{name}}`
REMOTE WITH CONNECTION `projects/my_project_name/locations/europe-north1/connections/external-connection-vertex-ai` OPTIONS (ENDPOINT = '{{endpoint}}');

However, how do we codify this in dbt or Terraform to have Infrastructure as Code? In Terraform, there is currently not a module available.

dbt usually does not allow to run CREATE AND REPLACE MODEL statements. At moment of writing, the dbt_ml package does not work with models with an external connection. The solution is to utilise a lovely dbt feature: pre_hook.

For simplicity, I created a macro that takes as input the location of the model name and the connection created in step 1.

{% macro createGenAIModel(dataset, name, endpoint) -%}
  CREATE OR REPLACE MODEL `my_project_name.{{dataset}}.{{name}}`
  REMOTE WITH CONNECTION `projects/my_project_name/locations/europe-north1/connections/external-connection-vertex-ai` OPTIONS (ENDPOINT = '{{endpoint}}');
{%- endmacro %}

In dbt, we can create a model and call this macro as pre_hook in the config header.

{{
  config(
  materialized = "table",
  pre_hook = createGenAIModel('my_dataset_name','gemini_pro','gemini-1.5-pro-preview-0514')
  )
}}

The config has a pre_hook where we send in the location of the BigQuery model to be created and the name of the Vertex AI model to be used.

The model will be created in my_dataset_name.gemini_pro and will connect to Gemini in Vertex AI. In the model garden of Vertex AI, there are many different models to choose from and by changing this parameter you will simply select a different model.

Step 4: SQL call to Gemini

Now that we have created a connection and a BigQuery Model, the next step is to actually write a prompt to send to Gemini in dbt.

{{
  config(
  materialized = "table",
  pre_hook = createGenAIModel('my_dataset_name','gemini_pro','gemini-1.5-pro-preview-0514')
  )
}}

SELECT ml_generate_text_llm_result, prompt
 FROM ML.GENERATE_TEXT(
  MODEL `my_project_name.my_dataset_name.gemini_pro`,
  (SELECT prompt FROM my_prompt_data_set),
  STRUCT(
    0.2 AS temperature, 100 AS max_output_tokens, 0.5 AS top_p, 40 AS top_k, TRUE AS flatten_json_output
    )
  )

The config is as explained in step 3.

The SQL part uses the ML.GENERATE_TEXT function to send a prompt into the LLM and receive a result. The input dataset needs to contain a column called prompt.

Step 5: Running dbt

In an usual dbt configuration, there is a separate service account to run dbt jobs. This service accounts needs to following roles:

• Connection User to be able to see the connection to Vertex AI
• Vertex AI User to be able to send queries to Vertex AI.

The below Terraform code grants the permissions.

resource "google_project_iam_member" "dbt_service_account_connection_user" {
  project = my_project_id
  for_each = toset([
  "roles/bigquery.connectionUser",
  "roles/aiplatform.user"
  ])
  role = each.value
  member = "serviceAccount:my_dbt_service_account@my_project_id.iam.gserviceaccount.com"
}

Conclusion

This blogs showed how in a common GCP tech stack, you can utilise dbt to schedule genAI models and setup the required infrastructure as code in five steps:

1. Setup of connection between BigQuery and Vertex AI
2. Granting needed permissions
3. Setup of a BigQuery Model
4. SQL call to Gemini
5. Running dbt

Now you are able to get started on productionising Large Language Models using dbt and infrastructure as code and quickly iterating over solutions.