Create Agents for Bedrock using CDK and check that Lambda functions are called based on input

2024-08-27 aws llm

Agents for Amazon Bedrock is a feature that builds generative AI agents using various foundation models of Bedrock to perform multi-step processes. It can call Lambda functions as needed and connect with services like OpenSearch Serverless to perform RAG.

The Lambda functions called by Agents need to return responses in the following format. Also, it is necessary to set permissions for Bedrock for invocation in the Lambda’s Resource-based policy, not in the agentResourceRole.

import * as lambda from "aws-cdk-lib/aws-lambda"

const actionFunction = new lambda.Function(this, "ActionFunction", {
  runtime: lambda.Runtime.NODEJS_20_X,
  handler: "index.handler",
  code: lambda.Code.fromInline(`
    exports.handler = async (event) => {
      const parameters = event['parameters'] || []
      const names = parameters.filter((param) => param.name === 'name')
      const name = (names.length > 0) ? names[0].value : 'Unknown'

      response_body = {
        'TEXT': {
          'body': 'おはこんハロチャオー！' + name + 'さん！'
        }
      }

      const function_response = {
        'actionGroup': event['actionGroup'],
        'function': event['function'],
        'functionResponse': {
          'responseBody': response_body
        }
      }
      
      return {
        'messageVersion': '1.0', 
        'response': function_response,
        'sessionAttributes': event['sessionAttributes'],
        'promptSessionAttributes': event['promptSessionAttributes']
      }
    }
  `),
  description: "Lambda function for Bedrock Agent action",
})

actionFunction.addPermission('BedrockInvoke', {
  principal: new iam.ServicePrincipal('bedrock.amazonaws.com'),
  action: 'lambda:InvokeFunction',
  sourceArn: `arn:aws:bedrock:${this.region}:${this.account}:agent/*`,
})

While the function’s API can be set up using OpenAPI schema, there’s also a simplified method as follows. The Agent decides which Actions to call by reading the description.

return {
  actionGroupName: "TestActionGroup",

  actionGroupExecutor: {
    lambda: actionFunction.functionArn,
  },
  functionSchema: {
    functions: [
      {
        name: "TestActionGroupFunction1",
        description: "Return a greeting message",
        parameters: {
          name: {
            type: "string",
            description: "Name of the user",
            required: true,
          },
        },
      },
    ],
  }
}

This is passed to actionGroups to create an Agent.

import * as iam from "aws-cdk-lib/aws-iam"
import * as bedrock from "aws-cdk-lib/aws-bedrock"

const agentResourceRole = new iam.Role(this, "BedrockAgentResourceRole", {
  assumedBy: new iam.ServicePrincipal("bedrock.amazonaws.com"),
  inlinePolicies: {
    BedrockAgentPolicy: new iam.PolicyDocument({
      statements: [
        new iam.PolicyStatement({
          actions: [
            "bedrock:InvokeModel",
            "bedrock:Retrieve",
            "bedrock:RetrieveGenerate",
            "bedrock:AssociateThirdPartyKnowledgeBase",
          ],
          resources: ["*"],
        }),
      ],
    }),
  },
})

new bedrock.CfnAgent(this, "TestBedrockAgent", {
  agentName: "TestAgent",
  instruction:
    "You are a good chatter person. You respond to trivial topics in appropriate responses that are neither too long nor too short.",
  foundationModel: "anthropic.claude-v2",
  idleSessionTtlInSeconds: 1800, // A user interaction remains active for the amount of time specified
  actionGroups: [this.makeHelloAction()],
  // PrepareAgentRequest caught an error. Agent roleArn cannot be null.
  agentResourceRoleArn: agentResourceRole.roleArn,
  autoPrepare: true,
})

We can confirm that values are passed to Lambda, and responses are given based on the output.