Installing and Setting-up Toolbox for your Gen AI & Agentic Applications on AlloyDB

1. Overview

Gen AI Toolbox for Databases is an open source server from Google that makes it easier to build Gen AI tools for interacting with databases. It enables you to develop tools easier, faster, and more securely by handling the complexities such as connection pooling, authentication, and more. It helps you build Gen AI tools that let your agents access data in your database. Toolbox provides:

Simplified development: Integrate tools to your agent in less than 10 lines of code, reuse tools between multiple agents or frameworks, and deploy new versions of tools more easily.

Better performance: Best practices such as connection pooling, authentication, and more.

Enhanced security: Integrated auth for more secure access to your data.

End-to-end observability: Out of the box metrics and tracing with built-in support for OpenTelemetry.

Toolbox sits between your application's orchestration framework and your database, providing a control plane that is used to modify, distribute, or invoke tools. It simplifies the management of your tools by providing you with a centralized location to store and update tools, allowing you to share tools between agents and applications and update those tools without necessarily redeploying your application.

What you'll build

As part of this lab, you will build an application that uses a tool to perform a simple database (AlloyDB) query that can be invoked from your agent or the generative AI application. For this you will

1. Install Toolbox
2. Set up the tool (that is designed to do a task in AlloyDB) on the Toolbox server
3. Deploy Toolbox on Cloud Run
4. Test the tool with its deployed Cloud Run endpoint
5. Build the Cloud run Function to invoke the Toolbox

Requirements

• A browser, such as Chrome or Firefox
• A Google Cloud project with billing enabled (steps in the next section).

2. Before you begin

Create a project

1. In the Google Cloud Console, on the project selector page, select or create a Google Cloud project.
2. Make sure that billing is enabled for your Cloud project. Learn how to check if billing is enabled on a project.
3. You'll use Cloud Shell, a command-line environment running in Google Cloud. Click Activate Cloud Shell at the top of the Google Cloud console.

4. Once connected to Cloud Shell, check whether you're already authenticated and whether the project is set to the correct project ID using the following command:

gcloud auth list

5. Run the following command in Cloud Shell to confirm that the gcloud command knows about your project.

gcloud config list project

6. If your project is not set, use the following command to set it:

gcloud config set project <YOUR_PROJECT_ID>

7. Enable the required APIs by running the following commands one by one in your Cloud Shell Terminal:

There is also a single command to run the below, but if you are a trial account user, you might encounter quota issues trying to enable these in bulk. That is why the commands are singled out one per line.

gcloud services enable alloydb.googleapis.com
gcloud services enable compute.googleapis.com 
gcloud services enable cloudresourcemanager.googleapis.com 
gcloud services enable servicenetworking.googleapis.com 
gcloud services enable run.googleapis.com 
gcloud services enable cloudbuild.googleapis.com 
gcloud services enable cloudfunctions.googleapis.com 
gcloud services enable aiplatform.googleapis.com

The alternative to the gcloud command is through the console by searching for each product or using this link.

If any API is missed, you can always enable it during the course of the implementation.

Refer documentation for gcloud commands and usage.

3. Database setup

In this lab we'll use AlloyDB as the database to hold the retail data. It uses clusters to hold all of the resources, such as databases and logs. Each cluster has a primary instance that provides an access point to the data. Tables will hold the actual data.

Let's create an AlloyDB cluster, instance and table where the ecommerce dataset will be loaded.

Create a cluster and instance

1. Navigate the AlloyDB page in the Cloud Console.

An easy way to find most pages in Cloud Console is to search for them using the search bar of the console.

2. Select CREATE CLUSTER from that page:

3. You'll see a screen like the one below. Create a cluster and instance with the following values (Make sure the values match in case you are cloning the application code from the repo):

• cluster id: "vector-cluster"
• password: "alloydb"
• PostgreSQL 15 compatible
• Region: "us-central1"
• Networking: "default"

4. When you select the default network, you'll see a screen like the one below. Select SET UP CONNECTION.
   
5. From there, select "Use an automatically allocated IP range" and Continue. After reviewing the information, select CREATE CONNECTION.
6. Once your network is set up, you can continue to create your cluster. Click on CREATE CLUSTER to complete setting up of the cluster as shown below:

Make sure to change the instance id to "

vector-instance"

.

Note that the Cluster creation will take around 10 minutes. Once it is successful, you should see a screen that shows the overview of your cluster you just created.

4. Data ingestion

Now it's time to add a table with the data about the store. Navigate to AlloyDB, select the primary cluster and then AlloyDB Studio:

You may need to wait for your instance to finish being created. Once it is completed, sign into AlloyDB using the credentials you created during cluster creation. Use the following data for authenticating to PostgreSQL:

• Username : "postgres"
• Database : "postgres"
• Password : "alloydb"

Once you have authenticated successfully into AlloyDB Studio, SQL commands can be entered in the Editor. You can add multiple Editor windows using the plus to the right of the last window.

You can enter commands for AlloyDB in editor windows, using the Run, Format, and Clear options as necessary.

Enable Extensions

For building this app, we will use the extensions pgvector and google_ml_integration. The pgvector extension allows you to store and search vector embeddings. The google_ml_integration extension provides functions you use to access Vertex AI prediction endpoints to get predictions in SQL. Enable these extensions by running the following DDLs:

CREATE EXTENSION IF NOT EXISTS google_ml_integration CASCADE;
CREATE EXTENSION IF NOT EXISTS vector;

If you would like to check the extensions that have been enabled on your database, run this SQL command:

select extname, extversion from pg_extension;

Create a table

Create a table using the DDL statement below:

CREATE TABLE toys ( id VARCHAR(25), name VARCHAR(25), description VARCHAR(20000), quantity INT, price FLOAT, image_url VARCHAR(200), text_embeddings vector(768)) ;

On successful execution of the above command, you should be able to view the table in the database.

Ingest data

For this lab, we have test data of about 72 records in this SQL file. It contains the id, name, description, quantity, price, image_url fields. The other fields will be filled in later in the lab.

Copy the lines/insert statements from there and then paste those lines in a blank editor tab and select RUN.

To see the table contents, expand the Explorer section until you can see the table named apparels. Select the tricolon (⋮) to see the option to Query the table. A SELECT statement will open in a new Editor tab.

Grant Permission

Run the below statement to grant execute rights on the embedding function to the user postgres:

GRANT EXECUTE ON FUNCTION embedding TO postgres;

Grant Vertex AI User ROLE to the AlloyDB service account

Go to Cloud Shell terminal and give the following command:

PROJECT_ID=$(gcloud config get-value project)

gcloud projects add-iam-policy-binding $PROJECT_ID \
  --member="serviceAccount:service-$(gcloud projects describe $PROJECT_ID --format="value(projectNumber)")@gcp-sa-alloydb.iam.gserviceaccount.com" \
--role="roles/aiplatform.user"

5. Create embeddings for the context

It's much easier for computers to process numbers than to process text. An embedding system converts text into a series of floating point numbers, called vector embeddings, that should represent the text, no matter how it's worded, what language it uses, etc.

For instance, a seaside location might be called "on the water", "beachfront", "walk from your room to the ocean", "sur la mer", "на берегу океана" etc. These terms all look different, but their semantic meaning or in machine learning terminology, their embeddings should be very close to each other.

Now that the data and context are ready, we will run the SQL to add the embeddings of the product description to the table in the field embedding. There are a variety of embedding models you can use. We're using text-embedding-005 from Vertex AI. Be sure to use the same embedding model throughout the project!

Note: If you are using an old Google Cloud Project, you might need to continue to use older versions of the text-embedding model like textembedding-gecko.

Return to the AlloyDB Studio tab and type the following DML:

UPDATE toys set text_embeddings = embedding( 'text-embedding-005', description);

Look at the toys table again to see some embeddings. Be sure to rerun the SELECT statement to see the changes.

SELECT id, name, description, price, quantity, image_url, text_embeddings FROM toys;

This should return the embeddings vector, that looks like an array of floats, for the toy description as shown below:

Note: Newly created Google Cloud Projects under the free tier might face quota issues when it comes to the number of embedding requests allowed per second to the Embedding models. We suggest that you use a filter query for the ID and then selectively choose 1-5 records and so on, while generating the embedding.

6. Perform Vector search

Now that the table, data, and embeddings are all ready, let's perform the real time vector search for the user search text.

Suppose the user asks:

"I want a white plush teddy bear toy with a floral pattern."

You can find matches for this by running the query below:

select * from toys
ORDER BY text_embeddings <=> CAST(embedding('text-embedding-005', 'I want a white plush teddy bear toy with a floral pattern') as vector(768))
LIMIT 5;

Let's look at this query in detail:

In this query,

1. The user's search text is: "I want a white plush teddy bear toy with a floral pattern."
2. We are converting it to embeddings in the embedding() method using the model: text-embedding-005. This step should look familiar after the last step, where we applied the embedding function to all of the items in the table.
3. "<=>" represents the use of the COSINE SIMILARITY distance method. You can find all the similarity measures available in the documentation of pgvector.
4. We are converting the embedding method's result to vector type to make it compatible with the vectors stored in the database.
5. LIMIT 5 represents that we want to extract 5 nearest neighbors for the search text.

Result looks like this:

As you can observe in your results, the matches are pretty close to the search text. Try changing the text to see how the results change.

7. Preparing AlloyDB for the Toolbox Interaction

In preparation for setting up Toolbox, let's enable public IP connectivity in our AlloyDB instance so the new tool can access the database.

1. Go to your AlloyDB instance, click EDIT and land on the Edit primary instance page.
2. Go to the Public IP connectivity section, check the Enable public IP checkbox and enter the IP address of your Cloud Shell machine.
3. To get your IP of your Cloud Shell machine, go to Cloud Shell Terminal and enter ifconfig. From the result identify the eth0 inet address and replace the last 2 digits with 0.0 with a mask size ‘/16'. For Example it would look like "XX.XX.0.0/16" where XX are numbers.
4. Paste this IP in the Authorized external networks "Networks" textbox of the edit instance page.

5. Click UPDATE INSTANCE once done.

It will take a few minutes to complete.

8. Toolbox Installation

1. You can create a project folder for storing the tool details. In this case since we are working on toy store data, let's create a folder named "toystore" and navigate into it. Navigate to Cloud Shell Terminal and make sure your project is selected and shown in the terminal's prompt. Run the below command from your Cloud Shell Terminal:

mkdir toystore

cd toystore

2. Run the below command to download and install toolbox in your new folder:

# see releases page for other versions
export VERSION=0.1.0
curl -O https://storage.googleapis.com/genai-toolbox/v$VERSION/linux/amd64/toolbox
chmod +x toolbox

3. Toggle to Cloud Shell Editor. Expand the newly created folder "toystore" and create a new file called tools.yaml. Copy the content below. Replace YOUR_PROJECT_ID and check if all the other connection details are accurate.

sources:
    alloydb-toys:
        kind: "alloydb-postgres"
        project: "YOUR_PROJECT_ID"
        region: "us-central1"
        cluster: "vector-cluster"
        instance: "vector-instance"
        database: "postgres"
        user: "postgres"
        password: "alloydb"

tools:
  get-toy-price:
    kind: postgres-sql
    source: alloydb-toys
    description: Get the price of a toy based on a description.
    parameters:
      - name: description
        type: string
        description: A description of the toy to search for.
    statement: |
      SELECT price FROM toys
      ORDER BY text_embeddings <=> CAST(embedding('text-embedding-005', $1) AS vector(768))
      LIMIT 1;

In this tool, we are just finding the closest match to the user's search text (custom toy description) and returning its price. You can also modify it to find the average price of the top 5 closest matching toys:

select avg(price) from ( SELECT price FROM toys ORDER BY text_embeddings <=> CAST(embedding(‘text-embedding-005', $1) AS vector(768)) LIMIT 5 ) as price;

You are all set with the tool definition!

For more details on configuring your tools.yaml, refer to this documentation.

4. Toggle to Cloud Shell Terminal and enter the following command to start the toolbox server with your tools configuration:

./toolbox --tools_file "tools.yaml"

5. Now if you open the server in a web preview mode on the cloud, you should be able to see the Toolbox server up and running with your new tool named get-toy-price.

9. Cloud Run Deployment of the Toolbox

Let's deploy it to Cloud Run so you can take this tool to actual use.

1. Follow the instructions in this page one by one until you reach the gcloud run deploy toolbox command that is in the 3rd point under the section "Deploy to Cloud Run". You need the first option and not the second one that is for when you are using a VPC network method.
2. Once deployed successfully, you will receive a Cloud Run deployed endpoint of your Toolbox server. Test it with a CURL command.

You are all set to use your newly deployed tool in your agentic application!!!

10. Connect your App with Toolbox

In this part, we will build a small application to test your tool to interact with the application's needs and retrieve response.

1. Navigate to Google Colab and open a new notebook.
2. Run the following in your notebook

pip install toolbox-langchain
from toolbox_langchain import ToolboxClient

# Replace with your Toolbox service's URL
toolbox = ToolboxClient("https://toolbox-*****-uc.a.run.app")
tool = toolbox.load_tool("get-toy-price")

# Invoke the tool with a search text to pass as the parameter
result = tool.invoke({"description": "white plush toy"})

# Print result
print(result)

3. You should result like this:

This is the tool being invoked explicitly in a Python application that uses the toolkit toolbox-langchain.

4. If you want to use this tool and bind it to an agent within a LangGraph integrated application, you could do that easily with the langgraph toolkit.
5. Refer to the code snippets for this.

11. Take it to Cloud!!!

Let's wrap this Python code snippet in a Cloud Run Functions to make it serverless!

1. Copy the source from the code repo folder for getting this to Cloud Functions.
2. Go to Cloud Run Functions console and click CREATE FUNCTION.
3. Keep it unauthenticated for the demo application and select the Python 3.11 runtime in the next page.
4. Copy the main.py and requirements.txt files from the source repo shared in step 1 and paste in respective files.
5. Deploy the function and you have got yourself a REST endpoint for the price prediction tool to be accessed in the toystore web application.
6. Your endpoint should be looking like this:

https://us-central1-*****.cloudfunctions.net/toolbox-toys

2. You can test it directly on the Cloud Functions console by navigating to the TESTING tab and entering the following as request input:

{

               "`search`"`:` "`White plush toy`"

}

3. Click TEST THE FUNCTION or run in Cloud Shell Terminal whatever you choose to use. You should see the result on the right under the title "Output":

12. Congratulations

Congratulations! You have successfully created a robust and truly modular tool that can interact across databases, platforms and Generative AI orchestration frameworks to help create your agentic application.