Import, manage, and run extensions

Extensions are connections to external APIs. They can process real-time data and connect to external APIs to perform real-world actions. Vertex AI provides an extension service that can import, manage, and run extensions.

This extension service can be linked to an application that processes user queries and communicates with an LLM. The extension service can provide the application with a set of extensions and the operations that the extensions can perform. Whenever the application receives a user query, it provides this information to the LLM. If the model determines that it should delegate query processing to an extension, then it provides a requested extension call and the associated parameter values. The application relays this request to the extension service, which then executes the extension. Finally, the extension service submits the result of this operation to the application, which then relays it back to the model.

This document shows you the key functionality of the Vertex AI extension service:

To learn how to import and run an extension provided by Google, see the following:

Create and import extensions

This document assumes that you already have a running API service that can back an extension. To create an extension, you must define its interface with an external API in an API specification file. You must upload this specification file to a Cloud Storage bucket or convert it into a string. You must then define an extension manifest, include the specification file, and send a registration request to the extension service.

Create an API specification file

An extension can be created by anyone through files that define and describe the extensions's API endpoints. The API endpoints can be public or private and hosted on any cloud or on-premises.

An API specification file describes the interface of a API service. You must provide an API specification file in YAML format that is compatible with OpenAPI 3.0. This specification file must define the following:

  • A server object. This object must define an API server URL. The Vertex AI extension service does not support multiple servers.

    servers:
      - url: API_SERVICE_URL
    
  • A paths object. This object must describe the various operations provided by the API service and the input parameters that correspond to each operation. Each operation must have a unique identifier and a response.

    paths:
      ...
        get:
          operationId: API_SERVICE_OPERATION_ID
          ...
          parameters:
            - name: API_SERVICE_INPUT_VAR
              ...
          responses:
          ...
    
  • A components object. This object is optional. You can use the components object to define reusable objects. For example, you can use the components object to provide a definition of the object schemas that are defined in the paths object. You can also use the components object to describe the output parameters of the API service.

    components:
      schemas:
        Result:
          ...
          properties:
            API_SERVICE_OUTPUT_VAR:
            ...
    

To learn more about OpenAPI, see OpenAPI Specification.

The following example is an API specification file for an API service that says "hello" in the requested language:

  openapi: "3.0.0"
  info:
    version: 1.0.0
    title: Hello Extension
    description: Learn to build Vertex AI extensions
  servers:
    - url: [API_SERVICE_URL]
  paths:
    /hello:
      get:
        operationId: "say_hello"
        description: "Say hello in prompted language."
        parameters:
          - name: "apiServicePrompt"
            in: query
            description: "Language"
            required: true
            schema:
              type: string
        responses:
          '200':
            description: Successful operation.
            content:
              application/json:
                schema:
                  $ref: "#/components/schemas/Result"
  components:
    schemas:
      Result:
        description: "Hello in the requested language."
        properties:
          apiServiceOutput:
            type: string

Upload the specification file

You can either upload the specification file to a Cloud Storage bucket or convert it into a string.

If you upload the specification file to a Cloud Storage bucket, grant the Vertex AI Extension Service Agent service account (service-PROJECT_NUMBER@gcp-sa-vertex-ex.iam.gserviceaccount.com) the Storage Object Viewer role. To learn how to list the buckets in your project, see Listing buckets. To learn how to copy an object to a Cloud Storage bucket, see Copy, rename, and move objects.

Define an extension import request

After creating an API specification file, you can define an extension import request in a JSON file. An extension import request must contain a reference to your API specification file (apiSpec) and the authentication configuration (authConfig). To connect the extension to a large language model (LLM) to see how the extension works, include the optional toolUseExamples parameter. If you want to only run the extension, don't include the toolUseExamples parameter.

An extension import request has the following format:

{
  "displayName":  "DISPLAY_NAME_HUMAN",
  "description": "DESCRIPTION_HUMAN",
  "manifest": {
    "name": "EXTENSION_NAME_LLM",
    "description": "DESCRIPTION_LLM",
    "apiSpec": { ... },
    "authConfig": { ... },
  }
  "toolUseExamples": [ ... ],
}
  • DISPLAY_NAME_HUMAN: The name of the extension that's displayed to users.
  • DESCRIPTION_HUMAN: The description of the extension that is displayed to users.
  • EXTENSION_NAME_LLM: The name of the extension that is used by the LLM for reasoning.
  • DESCRIPTION_LLM: The description of the extension that is used by the LLM for reasoning. You should provide a meaningful and informative description.

Reference to your API specification file

Your extension import request must contain a reference to your API specification file. You can provide the specification file in two ways:

  • Use openApiGcsUri to pass in the Cloud Storage URI of the YAML file.

    "apiSpec": {
      "openApiGcsUri": "gs://BUCKET_NAME/SPECIFICATION_FILE_NAME.yaml"
    },
    
    • BUCKET_NAME: The name of the Cloud Storage bucket that stores the specification file.
    • SPECIFICATION_FILE_NAME: The name of the API specification file.
  • Use openApiYaml to pass in the YAML file as a string.

Authentication configuration

Extensions can be public, available for any user to use, or private, only available to authorized users within one or more organizations.

An extension import request must contain an authentication configuration. You can choose between the following authentication methods:

  • NO_AUTH: No authentication
  • API_KEY_AUTH: API key authentication
  • HTTP_BASIC_AUTH: HTTP basic authentication
  • OAUTH: OAuth authentication
  • OIDC_AUTH: OIDC authentication

To learn more about authentication configurations, see Specify an authentication configuration.

Examples that demonstrate how the extension works

For best results, an extension import request should contain examples that demonstrate how the extension works. Use the toolUseExamples parameter to provide these examples.

The following code shows the format of toolUseExamples for a single example, with a single input parameter and a single output parameter. In this example, both the request and the response parameters are of string type.

"toolUseExamples": [
  {
      "extensionOperation": {
        "operationId": "API_SERVICE_OPERATION_ID",
      },
      "displayName": "EXAMPLE_DISPLAY_NAME",
      "query": "EXAMPLE_QUERY",
      "requestParams": {
        "fields": [
          {
            "key": "API_SERVICE_INPUT_VAR",
            "value": {
              "string_value": "EXAMPLE_INPUT",
            }
          }
        ]
      },
      "responseParams": {
        "fields": [
          {
            "key": "API_SERVICE_OUTPUT_VAR",
            "value": {
              "string_value": "EXAMPLE_OUTPUT",
            },
          }
        ],
      },
      "responseSummary": "EXAMPLE_SUMMARY"
    }
],
  • query: An example of a query that can take advantage of this extension. Use EXAMPLE_QUERY to provide the query text.
  • extensionOperation: An extension operation that is suitable for answering the query. Use API_SERVICE_OPERATION_ID to provide the ID of an extension operation defined in the API specification file.
  • displayName: A display name for the example. Use EXAMPLE_DISPLAY_NAME to provide a brief description.
  • requestParams: The request parameters that are necessary for the extensionOperation and example values, in key-value format. Use API_SERVICE_INPUT_VAR to provide an input parameter that is defined in the API specification file and corresponds with API_SERVICE_OPERATION_ID. Use EXAMPLE_INPUT to provide an example of an input value that corresponds with EXAMPLE_QUERY.
  • responseParams: The response parameters of the extensionOperation and example values in key-value format. Use API_SERVICE_OUTPUT_VAR to provide an output parameter that is defined in the API specification file and corresponds with the API service. Use EXAMPLE_OUTPUT to provide an example of an output value that corresponds with EXAMPLE_INPUT.
  • responseSummary: An example of a summary that the application might provide in response to the query. Use EXAMPLE_SUMMARY to provide the summary text.

The following is an example of toolUseExamples for an API service that says "hello" in the requested language:

"toolUseExamples": [
  {
      "extensionOperation": {
        "operationId": "say_hello",
      },
      "displayName": "Say hello in the requested language",
      "query": "Say hello in French",
      "requestParams": {
        "fields": [
          {
            "key": "apiServicePrompt",
            "value": {
              "string_value": "French",
            }
          }
        ]
      },
      "responseParams": {
        "fields": [
          {
            "key": "apiServiceOutput",
            "value": {
              "string_value": "bonjour",
            },
          }
        ],
      },
      "responseSummary": "Bonjour"
    }
],

Specify an authentication configuration

You must specify an authentication configuration when you define an extension import request.

If your extension does not require authentication, set the authType variable to NO_AUTH:

"authConfig": {
  "authType": "NO_AUTH"
}

If your extension requires authentication, then you must set the authentication type in the authType variable and supply an authentication configuration. You can choose between following authentication methods:

API key authentication

To support API key authentication, Vertex AI integrates with SecretManager for secret storage and access. The Vertex AI Extensions platform does not store the secret data directly. You have the responsibility to manage the lifecycle of your SecretManager resource.

Specify authConfig as follows:

"authConfig": {
  "authType": "API_KEY_AUTH",
  "apiKeyConfig": {
    "name": "API_KEY_CONFIG_NAME",
    "apiKeySecret": "API_KEY_SECRET",
    "httpElementLocation": "HTTP_ELEMENT_LOCATION",
  },
}
  • API_KEY_CONFIG_NAME: The name of the API key. For example, in the API request https://example.com/act?api_key=<API KEY>, API_KEY_CONFIG_NAME corresponds with api_key.
  • API_KEY_SECRET: SecretManager secret version resource that stores the key. This parameter has the following format: projects/PROJECT_ID/secrets/SECRET_ID/versions/VERSION.
  • HTTP_ELEMENT_LOCATION: The location of the API key in the HTTP request. Possible values are:

    • HTTP_IN_QUERY
    • HTTP_IN_HEADER
    • HTTP_IN_PATH
    • HTTP_IN_BODY
    • HTTP_IN_COOKIE

    To learn more, see Describing parameters.

HTTP basic authentication

To support HTTP basic authentication, Vertex AI integrates with SecretManager for secret storage and access. The Vertex AI Extensions platform does not store the secret data directly. You must manage the lifecycle of your SecretManager resource on your own.

Specify authConfig as follows:

"authConfig": {
  "authType": "HTTP_BASIC_AUTH",
  "httpBasicAuthConfig": {
    "credentialSecret": "CREDENTIAL_SECRET"
  },
}
  • CREDENTIAL_SECRET: SecretManager secret version resource that stores the base64-encoded credential. This parameter has the following format: projects/PROJECT_ID/secrets/SECRET_ID/versions/VERSION.

OAuth authentication

Vertex AI supports two methods of OAuth authentication: access token and service account.

Access token

Specify authConfig as follows:

"authConfig": {
  "authType": "OAUTH",
  "oauthConfig": {}
}

Leave the oauthConfig field blank when you import the extension. If you choose to run a registered extension, you must provide an access token in the oauthConfig field of the execution request. To learn more, see Run the extension.

Service account

Specify authConfig as follows:

"authConfig": {
  "authType": "OAUTH",
  "oauthConfig": {"service_account": "SERVICE_ACCOUNT_NAME"}
}
  • SERVICE_ACCOUNT_NAME: Vertex AI uses this service account to generate access tokens.

Perform the following steps to allow Vertex AI Extension Service Agent to get access tokens from SERVICE_ACCOUNT_NAME.

  1. Go to the IAM page.

    Go to IAM

  2. Select the Service Accounts tab.

  3. Click your service account. The value of SERVICE_ACCOUNT_NAME in authConfig must correspond with the name of your service account.

  4. Click the Permissions tab.

  5. Click Grant Access.

  6. In the Add principals section, in the New principals field, enter service-PROJECT_NUMBER@gcp-sa-vertex-ex.iam.gserviceaccount.com. This principal corresponds with the Vertex AI Extension Service Agent service account.

  7. In the Assign roles section, find and select the Service Account Token Creator role. This role includes the iam.serviceAccounts.getAccessToken permission.

  8. Click the Save button.

OIDC authentication

Vertex AI supports two methods of OIDC authentication: ID token and service account.

ID token

Specify authConfig as follows:

"authConfig": {
  "authType": "OIDC_AUTH",
  "oidcConfig": {}
}

Leave the oidcConfig field blank when you import the extension. If you choose to run a registered extension, you must provide an ID token in the oidcConfig field of the execution request. To learn more, see Run the extension.

Service account

Specify authConfig as follows:

"authConfig": {
  "authType": "OIDC_AUTH",
  "oidcConfig": {"service_account": "SERVICE_ACCOUNT_NAME"}
}
  • SERVICE_ACCOUNT_NAME: Vertex AI uses this service account to generate OpenID Connect (OIDC) tokens. Vertex AI sets the audience for the token to API_SERVICE_URL, as defined in the API specification file.

Perform the following steps to allow Vertex AI Extension Service Agent to get access tokens from SERVICE_ACCOUNT_NAME.

  1. Go to the IAM page.

    Go to IAM

  2. Select the Service Accounts tab.

  3. Click your service account. The value of SERVICE_ACCOUNT_NAME in authConfig must correspond with the name of your service account.

  4. Click the Permissions tab.

  5. Click Grant Access.

  6. In the Add principals section, in the New principals field, enter service-PROJECT_NUMBER@gcp-sa-vertex-ex.iam.gserviceaccount.com. This principal corresponds with the Vertex AI Extension Service Agent service account.

  7. In the Assign roles section, find and select the Service Account Token Creator role. This role includes the iam.serviceAccounts.getOpenIdToken permission.

  8. Click the Save button.

Import the extension with Vertex AI

After defining an extension import request, you can import the extension with Vertex AI.

  1. Set the following shell variables:

    ENDPOINT="LOCATION-aiplatform.googleapis.com"
    URL="https://${ENDPOINT}/v1beta1/projects/PROJECT_ID/locations/LOCATION"
    
    • PROJECT_ID: Your project.
    • LOCATION: A region of your choice. If you are not sure, choose us-central1.
  2. Run the following curl command to submit the import request:

    curl -X POST \
      -H "Authorization: Bearer $(gcloud auth print-access-token)" \
      -H "Content-Type: application/json; charset=utf-8" \
      -d @IMPORT_REQUEST.json "${URL}/extensions:import"
    

    The response has the following format:

    {
      "name": "projects/[PROJECT_NUMBER]/locations/[LOCATION]/extensions/[EXTENSION_ID]/operations/[IMPORT_OPERATION_ID]",
      "metadata": {
        "@type": "type.googleapis.com/google.cloud.aiplatform.v1beta1.ImportExtensionOperationMetadata",
        "genericMetadata": {
          "createTime": "[CREATE_TIME]",
          "updateTime": "[UPDATE_TIME]"
        }
      }
    }
    
  3. Set shell variables based on the output of the import request:

    EXTENSION_ID=EXTENSION_ID
    IMPORT_OPERATION_ID=IMPORT_OPERATION_ID
    
  4. To check the status of your import, run the following curl command:

    curl -X GET \
      -H "Authorization: Bearer $(gcloud auth print-access-token)" \
      -H "Content-Type: application/json; charset=utf-8" \
    "${URL}/operations/${IMPORT_OPERATION_ID}"
    

Manage extensions

To list all registered extensions, run the following curl command:

curl -X GET \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "Content-Type: application/json; charset=utf-8" \
"${URL}/extensions"

To get an extension, run the following curl command:

curl -X GET \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "Content-Type: application/json; charset=utf-8" \
"${URL}/extensions/${EXTENSION_ID}"

You can update the extension's displayName, description or toolUseExamples. If you specify toolUseExamples when you update an extension, then the update replaces the examples. For example, if you have examples a and b, then update the extension with example c, then the updated extension contains only example c.To update an extension description, run the following curl command:

curl -X PATCH \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json" \
${URL}/extensions/${EXTENSION_ID}?update_mask="description" \
-d '{
  "description": "A nice tool.",
}'

To delete an extension, run the following curl command:

curl \
 -X DELETE \
 -H "Authorization: Bearer $(gcloud auth print-access-token)" \
 -H "Content-Type: application/json" \
${URL}/extensions/${EXTENSION_ID}

Run an extension

If your extension uses OAuth authentication and an access token, see Run an extension with OAuth authentication and an access token.

If your extension uses OIDC authentication and an ID token, see Run an extension with OIDC authentication and an ID token.

Otherwise, you can run it using the following steps:

  1. Create a file named execute-extension.json with the following contents:

    {
      "operation_id": "API_SERVICE_OPERATION_ID",
      "operation_params": {
        "API_SERVICE_INPUT_VAR": "API_SERVICE_INPUT_VALUE"
      }
    }
    
    • API_SERVICE_OPERATION_ID: The ID of the API service operation you want to run. API service operations are defined in the API specification file.
    • API_SERVICE_INPUT_VAR: An input variable that corresponds with API_SERVICE_OPERATION_ID and is defined in the API specification file.
    • API_SERVICE_INPUT_VALUE: An input value for the extension.
  2. Run the following curl command:

    curl \
      -H "Authorization: Bearer $(gcloud auth print-access-token)" \
      -H "Content-Type: application/json; charset=utf-8" -d @execute-extension.json \
    "${URL}/extensions/${EXTENSION_ID}:execute"
    

    The response has the following format:

    {
      "output": {
        "content": "{\"API_SERVICE_OUTPUT_VAR\": \"API_SERVICE_OUTPUT_VALUE\"}"
      }
    }
    
    • API_SERVICE_OUTPUT_VAR: An output parameter that is defined in the API specification file and corresponds with the API service.
    • API_SERVICE_OUTPUT_VALUE: A string value that is a serialization of the response object. If your API specification file defines a JSON response schema, you must parse this output string into JSON on your own.

Run an extension with OAuth authentication and an access token

If your extension uses OAuth authentication and an access token, you can run it using the following steps:

  1. Create a file named execute-extension.json with the following contents:

    {
      "operation_id": "API_SERVICE_OPERATION_ID",
      "operation_params": {...},
      "runtime_auth_config": {
        "authType": "OAUTH",
        "oauth_config": {"access_token": "'$(gcloud auth print-access-token)'"}
      }
    }
    
    • API_SERVICE_OPERATION_ID: The ID of the API service operation you want to run. API service operations are defined in the API specification file.
  2. Run the following curl command:

    curl \
      -H "Authorization: Bearer $(gcloud auth print-access-token)" \
      -H "Content-Type: application/json; charset=utf-8" -d @execute-extension.json \
    "${URL}/extensions/${EXTENSION_ID}:execute"
    

Run an extension with OIDC authentication and an ID token

If your extension uses OIDC authentication and an ID token, you can run it using the following steps:

  1. Create a file named execute-extension.json with the following contents:

    {
      "operation_id": "API_SERVICE_OPERATION_ID",
      "operation_params": {...},
      "runtime_auth_config": {
        "authType": "OIDC_AUTH",
        "oidc_config": {"id_token": "$(gcloud auth print-identity-token)"}
      }
    }
    
    • API_SERVICE_OPERATION_ID: The ID of the API service operation you want to run. API service operations are defined in the API specification file.
  2. Run the following curl command:

    curl \
      -H "Authorization: Bearer $(gcloud auth print-access-token)" \
      -H "Content-Type: application/json; charset=utf-8" -d @execute-extension.json \
    "${URL}/extensions/${EXTENSION_ID}:execute"
    

What's next