Google Gemini Enterprise Agent Platform (Vertex AI) - Bring Your Own Key (BYOK)

​

Requirements

• In Omni:
  • An Omni Enterprise plan. Contact your Omni Account Executive or Omni support to learn more.
  • An AWS-hosted Omni instance
  • Organization Admin permissions
• In Google Cloud Platform, permissions that allow you to:
  • Create and configure workload identity federation providers
  • Configure trust policies for external AWS identities
  • Grant IAM permissions for Vertex AI API access
  • Enable Vertex AI API in your GCP project
• The gcloud CLI

​

Setting up GCP

1

Enable Vertex AI API

Ensure the Vertex AI API is enabled in your GCP project. You can enable it through the Google Cloud Console or using the gcloud CLI:

gcloud services enable aiplatform.googleapis.com --project=<GCP_PROJECT_ID>

Replace <GCP_PROJECT_ID> with the ID of your GCP project.

2

Create a workload identity federation provider

Create a workload identity pool and provider to allow Omni’s AWS ECS tasks to authenticate to your GCP project.

1

Retrieve Omni's AWS identity information

1. In Omni, navigate to AI Hub > General and click the LLM settings tab.
2. In the Providers section, click the Your API keys tab, then click Gemini Enterprise Agent Platform to expand it.
3. The fields required to configure the provider will display, including a field with Omni’s AWS identity information. It will contain values similar to the following:
   Federated Identity example

   arn:aws:iam::123456789012:role/Prod-Prod
   
   Attribute condition prefix example

   arn:aws:sts::123456789012:assumed-role/Prod-Prod
   
   These values identify Omni to Google Cloud. Note that:
   • It contains Omni’s AWS account ID. This is the 123456789012 segment. You’ll need this when creating the provider in the next step.
   • Both values are prefixes that match all AI-calling services. Omni uses multiple ECS task roles for AI-calling services that all share this prefix. Your attribute condition must use a starts-with comparison to match all of them. An exact-match condition will never match.

2

Create a workload identity pool and AWS provider

In your Google Cloud project, create a workload identity pool and an AWS provider that trusts the identity you copied in the previous step.

• Console
• gcloud CLI

1. In the Google Cloud Console, navigate to IAM & Admin > Workload Identity Federation.
2. Click Create pool.
3. Enter a name, such as omni, and continue.
4. Add a provider to the pool:
   • Provider type - AWS
   • Provider name - Enter a name, such as Omni
   • AWS account ID - The account ID from the value you copied in step 1 of this section
5. Configure the provider attribute mappings. AWS presents Omni’s identity in a different format than the value shown in Omni, so the mappings reconstruct the role ARN before comparing it.
   1. Click Edit mapping.
   2. For google.subject, paste the following into the corresponding AWS field:
      assertion.arn.extract('assumed-role/{role}/')
      
   3. For attribute.aws_role_arn, paste the following into the corresponding AWS field:
      'arn:aws:iam::' + assertion.account + ':role/' + assertion.arn.extract('assumed-role/{role}/')
      
   The mappings should look like the following when finished:

   

6. Configure the provider attribute conditions:
   1. Click Add condition.
   2. In the text field, add the attribute condition. Replace OMNI_ATTRIBUTE_CONDITION_PREFIX with the Attribute condition prefix value you copied from Omni in step 1 of this section:
      attribute.aws_role_arn.startsWith('OMNI_ATTRIBUTE_CONDITION_PREFIX')
      
      For example:
      attribute.aws_role_arn.startsWith('arn:aws:sts::123456789012:assumed-role/Prod-Prod')
      
      The attribute conditions should look like the following when finished:

7. Click Save.

Using the gcloud CLI, run the following to create a workload identity pool:

gcloud iam workload-identity-pools create omni \
  --project=GCP_PROJECT_ID \
  --location=global \
  --display-name="Omni"

Replace GCP_PROJECT_ID with the ID of your GCP project.Then, create the provider:

gcloud iam workload-identity-pools providers create-aws omni-aws \
  --project=GCP_PROJECT_ID \
  --location=global \
  --workload-identity-pool=omni \
  --account-id=OMNI_AWS_ACCOUNT_ID \
  --attribute-mapping="google.subject=assertion.arn.extract('assumed-role/{role}/'),attribute.aws_role_arn='arn:aws:iam::'+assertion.account+':role/'+assertion.arn.extract('assumed-role/{role}/')" \
  --attribute-condition="attribute.aws_role_arn.startsWith('OMNI_ATTRIBUTE_CONDITION_PREFIX')"

Replace the following values:

Placeholder	Replace with
GCP_PROJECT_ID	The ID of your GCP project
OMNI_AWS_ACCOUNT_ID	The account ID from the value you copied in step 1 of this section
OMNI_ATTRIBUTE_CONDITION_PREFIX	The entire Attribute condition prefix value you copied in step 1 of this section

3

Grant the workload identity pool access to Vertex AI

Give the workload identity pool permission to call Vertex AI.Choose one of the following options:

Option 1 - Grant direct access

Grant the pool the Vertex AI User role directly on your project.

Grant the pool the Vertex AI User role (roles/aiplatform.user) directly on your project, replacing GCP_PROJECT_ID with the ID of your GCP project:

Grant the pool Vertex AI access

PROJECT_NUMBER=$(gcloud projects describe GCP_PROJECT_ID --format='value(projectNumber)')

gcloud projects add-iam-policy-binding GCP_PROJECT_ID \
  --role="roles/aiplatform.user" \
  --member="principalSet://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/omni/*"

When you configure Omni, leave the Service account impersonation URL field blank.

Option 2 - Impersonation through a service account

Have the pool impersonate a named service account that holds the Vertex AI User role.

If your organization requires Vertex AI access to be held by a named service account (for example, to satisfy least-privilege or audit policies), grant the pool permission to impersonate a service account that holds roles/aiplatform.user, instead of granting the pool directly.Run the following, replacing GCP_PROJECT_ID with the ID of your GCP project:

Grant access through a service account

# Create a service account (or use an existing one)
gcloud iam service-accounts create omni-vertex-sa \
  --display-name="Omni Vertex AI Service Account" \
  --project=GCP_PROJECT_ID

# Grant it the Vertex AI User role
gcloud projects add-iam-policy-binding GCP_PROJECT_ID \
  --member="serviceAccount:omni-vertex-sa@GCP_PROJECT_ID.iam.gserviceaccount.com" \
  --role="roles/aiplatform.user"

# Allow the pool to impersonate the service account
PROJECT_NUMBER=$(gcloud projects describe GCP_PROJECT_ID --format='value(projectNumber)')

gcloud iam service-accounts add-iam-policy-binding \
  omni-vertex-sa@GCP_PROJECT_ID.iam.gserviceaccount.com \
  --role="roles/iam.workloadIdentityUser" \
  --member="principalSet://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/omni/*" \
  --project=GCP_PROJECT_ID

When you configure Omni, set the Service account impersonation URL to the following, replacing GCP_PROJECT_ID with the ID of your GCP project:

https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/omni-vertex-sa@GCP_PROJECT_ID.iam.gserviceaccount.com:generateAccessToken

4

Get the workload identity provider audience

Retrieve the audience value for your workload identity provider. This will be used in the Omni configuration:

gcloud iam workload-identity-pools providers describe omni-aws \
  --project=GCP_PROJECT_ID \
  --location=global \
  --workload-identity-pool=omni \
  --format="value(name)"

The output will be in the format:

//iam.googleapis.com/projects/GCP_PROJECT_ID/locations/global/workloadIdentityPools/omni/providers/omni-aws

5

Enable the Claude models and confirm quota

Before Omni can generate responses, the Claude models you plan to use must be enabled for your GCP project and your region must have quota for them.

1. In the Vertex AI Model Garden, enable each Claude model you plan to use with Omni and accept its terms. Anthropic models aren’t enabled by default.
2. Confirm your project has non-zero quota for those models in the region you’ll configure in Omni. You can review and request quota changes under IAM & Admin > Quotas in the Google Cloud console.

See Google’s documentation on using Claude models on Vertex AI for more information.

​

Enabling Gemini Enterprise Agent Platform as a provider

​

Disabling Gemini Enterprise Agent Platform as a provider

​

Troubleshooting

• Authentication errors (403): Verify that your workload identity pool uses the role prefix shown in Omni’s AI settings and that the attribute condition uses a starts-with comparison.
• Permission denied errors: Ensure the service account has the roles/aiplatform.user role and that the workload identity pool has permission to impersonate the service account (roles/iam.workloadIdentityUser).
• Model not found errors: Model availability varies by GCP region. Verify that the model IDs you’ve configured are available in your selected region, or enable cross-region access in your GCP project.
• Invalid audience: Double-check that the WIF Provider Audience matches the full path from the command output in the Get the workload identity provider audience step.
• Rate limit or quota errors: If you receive an error like Your application has exceeded its allocated rate limit or quota for a specific API, verify that the models you’re using have a positive quota balance.

​

Related

• AI in Omni
• AI settings
• Bring Your Own Key providers