Streamlining VertexAI Integration: Hassle-Free API Calls

Source: Self generated using Dall-E 3

What is VertexAI?

It is a comprehensive AI platform for building, deploying, and managing machine learning models at scale. Besides your own trained models, VertexAI also provides a catalog of models to be used out of the box. With the amount of model options available on VertexAI, it is easily one of the platforms that any AI practitioner would want easy programmatic access to for development.

Challenge

Unfortunately, calling VertexAI API is not as straightforward as desired. Existing solutions for calling it can be overly complex and impractical for certain use cases. The OAuth2 authentication process, which is typically required, involves multiple intricate steps, such as redirecting users to an authorisation page, handling callback responses, and managing refresh tokens. This convoluted workflow becomes particularly cumbersome in server-to-server interactions where user interaction is not feasible or desirable.

Moreover, the need to periodically refresh access tokens adds an additional layer of complexity, requiring secure storage mechanisms and careful token management to avoid expiration issues. OAuth2 workflows often necessitate periodic user consent and manual approval, rendering them unsuitable for automated, large-scale deployments that demand rapid, seamless access to the VertexAI API.

This reliance on user interaction further highlights the impracticality of OAuth2 for many programmatic access scenarios, where minimizing human intervention is crucial for efficiency and reliability. These challenges underscore the limitations of OAuth2 as an authentication method for applications that require consistent, automated access to the VertexAI API, making it a sub-optimal choice in such contexts.

Solution

For server-to-server interactions, leveraging service account keys offers a more streamlined and efficient authentication mechanism. With service accounts, developers can generate key files that enable their applications to authenticate seamlessly, without the need for user interaction or manual approval processes.

This direct authentication method simplifies the implementation process significantly, reducing the overhead associated with token management strategies. Furthermore, it enhances security by eliminating the necessity to handle refresh tokens, which can be a potential vulnerability if not managed properly.

Service accounts are specifically designed to support automated, large-scale deployments, making them an ideal choice for backend systems that require consistent and reliable access to VertexAI and other GCP services. This streamlined approach not only improves operational efficiency but also ensures that applications can function smoothly and securely in a production environment without disruptions caused by complex authentication workflows.

Introduction

In this article, we’ll explore how to call the VertexAI API using Python’s request library and without the need for OAuth2 login. This approach is particularly useful for backend systems that require programmatic access to services like VertexAI, where triggering a login can be inconvenient compared to generating a token for authentication.

While Google Cloud Platform (GCP) documentation provides information on various ways to provide credentials, it can be challenging to navigate and time-consuming to parse through. This article aims to provide a quick and straightforward method to leverage Service Accounts to generate a Bearer token, which can then be used to call VertexAI using the requests package in Python.

Step 1: Creating a Service Account

The first step is to create a Service Account in GCP. Follow these steps:

1. Go to the IAM & Admin section in your GCP console.
2. Navigate to the Service Accounts tab.
3. Click on Create Service Account and provide a suitable name and description.
4. Grant the Service Account access to VertexAI by assigning the Vertex AI Service Agent role.

Initially, the Service Account won’t have a Key ID. To create one:

1. Click into your newly created Service Account.
2. Go to the Keys tab and click on Add Key.
3. Select Create new key and choose the JSON key type.
4. Download the JSON file, which contains information about your Key ID.

This JSON file is essential for the next step.

Step 2: Generating the Bearer Token

Using the JSON file obtained in the previous step, you can generate a Bearer token to pass as an authentication key for the VertexAI API. Here’s an example Python function that accomplishes this task:

import json
from google.auth.jwt import OnDemandCredentials

def get_oauth_token(sa_key_fpath: str) -> str:
    cred_file = json.load(open(sa_key_fpath))
    gcp_cred_object = OnDemandCredentials.from_service_account_info(cred_file)
    bearer_token = gcp_cred_object._get_jwt_for_audience(
        audience="https://aiplatform.googleapis.com/"
    )
    return bearer_token.decode()
bearer_token = get_oauth_token("service_account_key_file.json")

In this code snippet, we leverage Google’s Python SDK and the OnDemandCredentials object to create the Bearer token. Note that while we're using OnDemandCredentials here, the SDK provides other objects that serve similar purposes, so feel free to explore the codebase of the package.

Step 3: Calling the VertexAI API

With the Bearer token generated, you can now call the VertexAI API. The first step is to determine the appropriate endpoint by referring to the VertexAI API reference documentation.

Here’s an example Python code snippet that demonstrates how to call the VertexAI API using the requests library:

import requests

response = requests.post(
    url="https://asia-east1-aiplatform.googleapis.com/v1/projects/{project_name}/locations/{location}/publishers/google/models/text-bison@002:predict",
    headers={
        "Content-Type": "application/json",
        "authorization": f"Bearer {bearer_token}"
    },
    json={
        "instances": [
            {
                "role": "user",
                "content": "hello world!"
            }
        ],
        "parameters": {
            "max_tokens": 256,
            "temperature": 0.8,
            "frequency_penalty": 0,
            "presence_penalty": 0,
            "top_p": 0.95
        }
    }
)

In this example, we’re using the text-bison model in the asia-east1 region. Make sure to replace {project_name} and {location} with your appropriate values.

Conclusion

By following the steps outlined in this article, you can successfully call the VertexAI API using REST API without the need for OAuth2 login. This approach allows backend systems to programmatically access VertexAI services in a more convenient manner, eliminating the need for triggering a login process. With the Bearer token generated from the Service Account, you can make API requests to VertexAI using the requests package in Python or any other programming language that supports HTTP requests.