How To Manage an Amazon Bedrock Agent Using Terraform

How To Manage an Amazon Bedrock Agent Using Terraform

Learn how to use the new Agents for Amazon Bedrock Terraform resources to manage an agent using a forex rate assistant as an example.

·

12 min read

Introduction

In the previous blog post Building a Basic Forex Rate Assistant Using Agents for Amazon Bedrock, I demonstrated how to create a Bedrock agent in the AWS Management Console and outlined some ideas on improving the solution. Before further experimentation, it makes sense to automate the deployment of the solution to enable quicker updates as we go through trail and error in fine-tuning an agent.

In this blog post, we will automate the deployment of the basic forex rate assistant in Terraform using the resources that were recently released in v5.47.0 of the Terraform AWS Provider. Let's start by looking at the AWS resources in the AWS Management Console.

Taking inventory of the required resources

By examining the agent we previously built, we see that it is comprised of the following AWS resources:

  1. The agent itself

  2. The agent resource role which is an IAM service role that provides the agent with access to other AWS services and resources

    The agent and its resource role

  3. The action group that defines API actions that the agent can perform

    The action group

  4. The Lambda function associated with the action group, which itself requires an execution role and a resource policy that allows the agent to invoke the function

    The Lambda execution role and the resource policy

With the list of resources we need to provision, we can begin creating the Terraform configuration starting with the resources that the agent depends on.

Defining resources for the IAM and Lambda dependencies

For the agent resource role, the documentation already provides the trust policy and the permissions required. It also specifies that the prefix AmazonBedrockExecutionRoleForAgents_ must be used for the role name.

The permission requires the foundational model's ARN, so we need at least the model's ID, which in our case is anthropic.claude-3-haiku-20240307-v1:0 for Claude 3 Haiku. For consistency, we will use the aws_bedrock_foundational_model data source to look up its ARN. Thus we can define the Terraform configuration for the agent resource role as follows using the aws_iam_role resource and the aws_iam_role_policy resource:

# Use data sources to get common information about the environment
data "aws_caller_identity" "this" {}
data "aws_partition" "this" {}
data "aws_region" "this" {}
locals {
  account_id = data.aws_caller_identity.this.account_id
  partition  = data.aws_partition.this.partition
  region     = data.aws_region.this.name
}

data "aws_bedrock_foundation_model" "this" {
  model_id = "anthropic.claude-3-haiku-20240307-v1:0"
}
# Agent resource role
resource "aws_iam_role" "bedrock_agent_forex_asst" {
  name = "AmazonBedrockExecutionRoleForAgents_ForexAssistant"
  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Action = "sts:AssumeRole"
        Effect = "Allow"
        Principal = {
          Service = "bedrock.amazonaws.com"
        }
        Condition = {
          StringEquals = {
            "aws:SourceAccount" = local.account_id
          }
          ArnLike = {
            "aws:SourceArn" = "arn:${local.partition}:bedrock:${local.region}:${local.account_id}:agent/*"
          }
        }
      }
    ]
  })
}

resource "aws_iam_role_policy" "bedrock_agent_forex_asst" {
  name = "AmazonBedrockAgentBedrockFoundationModelPolicy_ForexAssistant"
  role = aws_iam_role.bedrock_agent_forex_asst.name
  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Action   = "bedrock:InvokeModel"
        Effect   = "Allow"
        Resource = data.aws_bedrock_foundation_model.this.model_arn
      }
    ]
  })
}

Next, we will define the Lambda execution role which just needs the basic permissions to write logs to CloudWatch that the AWS-managed IAM policy AWSLambdaBasicExecutionRole provides. The Terraform configuration for this IAM role can be defined as follows:

data "aws_iam_policy" "lambda_basic_execution" {
  name = "AWSLambdaBasicExecutionRole"
}

# Action group Lambda execution role
resource "aws_iam_role" "lambda_forex_api" {
  name = "FunctionExecutionRoleForLambda_ForexAPI"
  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Action = "sts:AssumeRole"
        Effect = "Allow"
        Principal = {
          Service = "lambda.amazonaws.com"
        }
        Condition = {
          StringEquals = {
            "aws:SourceAccount" = "${local.account_id}"
          }
        }
      }
    ]
  })
  managed_policy_arns = [data.aws_iam_policy.lambda_basic_execution.arn]
}

We will then define the Terraform configuration for the Lambda function and its resource policy. Here is the source code for the Forex API Lambda function from the previous blog post for reference:

import json
import urllib.parse # urllib is available in Lambda runtime w/o needing a layer
import urllib.request

def lambda_handler(event, context):
    agent = event['agent']
    actionGroup = event['actionGroup']
    apiPath = event['apiPath']
    httpMethod =  event['httpMethod']
    parameters = event.get('parameters', [])
    requestBody = event.get('requestBody', {})

    # Read and process input parameters
    code = None
    for parameter in parameters:
        if (parameter["name"] == "code"):
            # Just in case, convert to lowercase as expected by the API
            code = parameter["value"].lower()

    # Execute your business logic here. For more information, refer to: https://docs.aws.amazon.com/bedrock/latest/userguide/agents-lambda.html
    apiPathWithParam = apiPath
    # Replace URI path parameters
    if code is not None:
        apiPathWithParam = apiPathWithParam.replace("{code}", urllib.parse.quote(code))

    # TODO: Use a environment variable or Parameter Store to set the URL
    url = "https://cdn.jsdelivr.net/npm/@fawazahmed0/currency-api@latest/v1{apiPathWithParam}.min.json".format(apiPathWithParam = apiPathWithParam)

    # Call the currency exchange rates API based on the provided path and wrap the response
    apiResponse = urllib.request.urlopen(
        urllib.request.Request(
            url=url,
            headers={"Accept": "application/json"},
            method="GET"
        )
    )
    responseBody =  {
        "application/json": {
            "body": apiResponse.read()
        }
    }

    action_response = {
        'actionGroup': actionGroup,
        'apiPath': apiPath,
        'httpMethod': httpMethod,
        'httpStatusCode': 200,
        'responseBody': responseBody

    }

    api_response = {'response': action_response, 'messageVersion': event['messageVersion']}
    print("Response: {}".format(api_response))

    return api_response

We will save this source code into a file called index.py in the lambda/forex_api directory in the same directory as the Terraform configuration, which will be packaged as a zip file using the archive_file data source to pass as an argument to the aws_lambda_function resource.

Here is the Terraform configuration for the Lambda function based on my battle-tested templates:

# Action group Lambda function
data "archive_file" "forex_api_zip" {
  type             = "zip"
  source_file      = "${path.module}/lambda/forex_api/index.py"
  output_path      = "${path.module}/tmp/forex_api.zip"
  output_file_mode = "0666"
}

resource "aws_lambda_function" "forex_api" {
  function_name = "ForexAPI"
  role          = aws_iam_role.lambda_forex_api.arn
  description   = "A Lambda function for the forex API action group"
  filename      = data.archive_file.forex_api_zip.output_path
  handler       = "index.lambda_handler"
  runtime       = "python3.12"
  # source_code_hash is required to detect changes to Lambda code/zip
  source_code_hash = data.archive_file.forex_api_zip.output_base64sha256
}

Lastly, we will set the Lambda resource policy using the aws_lambda_permission resource according to the specifications in the documentation:

resource "aws_lambda_permission" "forex_api" {
  action         = "lambda:invokeFunction"
  function_name  = aws_lambda_function.forex_api.function_name
  principal      = "bedrock.amazonaws.com"
  source_account = local.account_id
  source_arn     = "arn:aws:bedrock:${local.region}:${local.account_id}:agent/*"
}

Defining the agent and action group resources

With the dependencies out of the way, we can now define the Terraform resource for the agent with the new aws_bedrockagent_agent resource, which is rather straightforward:

resource "aws_bedrockagent_agent" "forex_asst" {
  agent_name              = "ForexAssistant"
  agent_resource_role_arn = aws_iam_role.bedrock_agent_forex_asst.arn
  description             = "An assisant that provides forex rate information."
  foundation_model        = data.aws_bedrock_foundation_model.this.model_id
  instruction             = "You are an assistant that looks up today's currency exchange rates. A user may ask you what the currency exchange rate is for one currency to another. They may provide either the currency name or the three-letter currency code. If they give you a name, you may first need to first look up the currency code by its name."
}

The action group can be defined in the agent using the aws_bedrockagent_action_group resource. We will need the OpenAPI schema YAML file from the previous blog post, which is included below for reference:

openapi: 3.0.0
info:
  title: Currency API
  description: Provides information about different currencies.
  version: 1.0.0
servers:
  - url: https://cdn.jsdelivr.net/npm/@fawazahmed0/currency-api@latest/v1
paths:
  /currencies:
    get:
      description: |
        List all available currencies
      responses:
        "200":
          description: Successful response
          content:
            application/json:
              schema:
                type: object
                description: |
                  A map where the key refers to the lowercase three-letter currency code and the value to the currency name in English.
                additionalProperties:
                  type: string
  /currencies/{code}:
    get:
      description: |
        List the exchange rates of all available currencies with the currency specified by the given currency code in the URL path parameter as the base currency
      parameters:
        - in: path
          name: code
          required: true
          description: The lowercase three-letter code of the base currency for which to fetch exchange rates
          schema:
            type: string
      responses:
        "200":
          description: Successful response
          content:
            application/json:
              schema:
                type: object
                description: |
                  A map where the key refers to the three-letter currency code of the target currency and the value to the exchange rate to the target currency.
                additionalProperties:
                  type: number
                  format: float

We will save the file as schema.yaml in the lambda/forex_api directory for the Lambda function, since they somewhat go together. Since we are providing the OpenAPI schema in-line, the Terraform resource can be defined as follows:

resource "aws_bedrockagent_agent_action_group" "forex_api" {
  action_group_name          = "ForexAPI"
  agent_id                   = aws_bedrockagent_agent.forex_asst.id
  agent_version              = "DRAFT"
  description                = "The currency exchange rates API"
  skip_resource_in_use_check = true
  action_group_executor {
    lambda = aws_lambda_function.forex_api.arn
  }
  api_schema {
    payload = file("${path.module}/lambda/forex_api/schema.yaml")
  }
}

Testing the configuration

Now that the full Terraform configuration is developed, we can apply it and make sure that it is working correctly. For me it took less than a minute to complete - here is the output for reference:

aws_iam_role.bedrock_agent_forex_asst: Creating...
aws_iam_role.lambda_forex_api: Creating...
aws_iam_role.bedrock_agent_forex_asst: Creation complete after 0s [id=AmazonBedrockExecutionRoleForAgents_ForexAssistant]
aws_iam_role_policy.bedrock_agent_forex_asst: Creating...
aws_bedrockagent_agent.forex_asst: Creating...
aws_iam_role.lambda_forex_api: Creation complete after 1s [id=FunctionExecutionRoleForLambda_ForexAPI]
aws_lambda_function.forex_api: Creating...
aws_iam_role_policy.bedrock_agent_forex_asst: Creation complete after 1s [id=AmazonBedrockExecutionRoleForAgents_ForexAssistant:AmazonBedrockAgentBedrockFoundationModelPolicy_ForexAssistant]
aws_bedrockagent_agent.forex_asst: Creation complete after 4s [id=LTR1P1OJUC]
aws_lambda_function.forex_api: Still creating... [10s elapsed]
aws_lambda_function.forex_api: Creation complete after 14s [id=ForexAPI]
aws_lambda_permission.forex_api: Creating...
aws_bedrockagent_agent_action_group.forex_api: Creating...
aws_lambda_permission.forex_api: Creation complete after 0s [id=terraform-20240430193700768300000002]
aws_bedrockagent_agent_action_group.forex_api: Creation complete after 0s [id=W1PDUUCT8P,LTR1P1OJUC,DRAFT]

Apply complete! Resources: 7 added, 0 changed, 0 destroyed.

In the Bedrock console, we can see that the agent ForexAssistant is ready for testing. Using the test chat interface, I asked:

What is the exchange rate from US Dollar to Canadian Dollar?

However, I got the following unexpected answer:

I apologize, but I am unable to look up the current exchange rate between US Dollar and Canadian Dollar. There seems to be an issue with the function call format that I am unable to resolve. I cannot provide the exchange rate information you requested.

Looking at the trace, it seems that the agent was not given the tool list and it tried to make up random functions to call, leading to errors:

Trace showing the model's attempt to call an unknown function

On closer look, it seems that this is because there are pending changes in the agent which is requires preparation as indicated in the Bedrock console:

Agent needed to be prepared

This tells me that Terraform is not performing the preparation. In any case, once I click Prepare and ask the same question again in a new session, the agent responds with the currency exchange rate I asked for:

The exchange rate from US Dollar (USD) to Canadian Dollar (CAD) is 1 USD = 1.36660199 CAD.

This is also confirmed in the trace which I will not show for brevity. Now we are one step away from an end-to-end IaC solution for the forex rate assistant, so let's try to address the issue.

Workaround for agent preparation using a null resource

💡
2024-05-23: As of Terraform AWS Provider v5.49.0, the aws_bedrockagent_agent resource has a prepare_agent argument (true by default) that controls whether the agent is prepared after the agent is created or updated. The Terraform configuration in the GitHub repository has been updated to account for this enhancement. However, the null resource is still required for action groups since aws_bedrockagent_action_group still does not prepare the agent.

Looking at the Terraform AWS Provider documentation, I couldn't find any resource that supports preparation. As well, the aws_bedrockagent_agent resource and the aws_bedrockagent_action_group resource don't seem to have any argument that controls the preparation behavior. To be fair, the action is implemented as a separate API action called PrepareAgent in the Agents for Bedrock API, which does not directly fit into the resource concept in Terraform.

While I opened an issue in the hashicorp/terraform-provider-aws GitHub repository, one quick workaround I can think of is to use a null resource with the local-exec provisioner to run the equivalent AWS CLI command for the PrepareAgent API, which is the aws bedrock-agent prepare-agent command.

Our objective is to trigger this null resource to be rerun (technically replaced) every time there are changes to the agent, which also extends to the action group. It is inefficient to simply prepare every time you apply the Terraform configuration, and if anything it is just one more moving part that can break. With that in mind, I devised the following resource that serve the purpose well.

resource "null_resource" "forex_asst_prepare" {
  triggers = {
    forex_asst_state = sha256(jsonencode(aws_bedrockagent_agent.forex_asst))
    forex_api_state  = sha256(jsonencode(aws_bedrockagent_agent_action_group.forex_api))
  }
  provisioner "local-exec" {
    command = "aws bedrock-agent prepare-agent --agent-id ${aws_bedrockagent_agent.forex_asst.id}"
  }
  depends_on = [
    aws_bedrockagent_agent.forex_asst,
    aws_bedrockagent_agent_action_group.forex_api
  ]
}

As you can see, I am using the triggers argument in the null resource to control when the resource should be replaced. We target the two main sources of change, which is the agent and the action group. Since trigger requires a string, a good candidate is to use the two resource's state somehow, as long as it doesn't contain any attributes that change every time Terraform is run. To keep the string short, we simply derive the SHA256 checksum from the resource state JSON as the triggers. The local-exec provisioner simply calls the AWS CLI command with the agent ID from aws_bedrockagent_agent.forex_asst.

With this change, we will run terraform destroy and then terraform apply to ensure full validity of the re-test. After Terraform completes successfully, we first check the agent in the Bedrock console to ensure that the Prepare button is no longer shown. As well, we ask our question to hopefully receive an expected result, which we did:

Prepare button no visible and agent responds correctly

So there you have it, a functional Terraform configuration to deploy a basic forex rate assistant implemented using Agents for Amazon Bedrock!

For reference, I've dressed up the Terraform solution with variables and such, and checked in the final artifacts to the 1_basic directory in this repository. Feel free to check it out and use it as the basis for your Bedrock experimentation.

Current limitations (it's brand new after all)

It is not unexpected that we encounter some issues with brand new features, such as what we encountered in this blog post with the Agents for Amazon Bedrock resources. I myself dove a bit deeper and found a few more issues which I reported. I encourage you to report any issues that you see as you work more with the Terraform resources.

Meanwhile, there are still a couple resources related to Knowledge bases for Amazon Bedrock still under development. I plan to integrate knowledge bases to our forex rate assistant, so I will eagerly wait for the Terraform resources to be ready for my next step in my Bedrock journey.

Summary

In this blog post, we developed the Terraform configuration for the basic forex rate assistant that we created interactively in the blog post Building a Basic Forex Rate Assistant Using Agents for Amazon Bedrock. While we encountered some issues, we were able to work around it as the community continues to build out the features in the Terraform AWS Provider. For now, I will pivot to enhancing the forex rate agent to add new capabilities and to address some of its known shortcomings.

If you like this blog post, please be sure to check out other helpful articles on AWS, Terraform, and other DevOps topics in the Avangards Blog.