CDK で Agents for Bedrock を作成し入力に基づいて Lambda 関数が呼び出されることを確認する

2024-08-27 aws llm

Agents for Amazon Bedrock は Bedrock の各種基盤モデルを用いて、複数のステップからなる処理を行う生成AIのエージェントを構築する機能。必要に応じて Lambda を呼び出したり、OpenSearch Serverless などと接続して RAG を行ったりできる。

Agents から呼び出す Lambda 関数は次のようなフォーマットでレスポンスを返す必要がある。また、Bedrock から呼び出せるような権限を agentResourceRole ではなく Lambda の Resource based policy に設定する必要があった。

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/*`,
})

関数の API は OpenAPI schema で設定することもできるが、次のような簡易版の方法もある。Agent は description を読んでどの Actions を呼び出すか決定する。

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,
          },
        },
      },
    ],
  }
}

これを actionGroups に渡して 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,
})

Lambdaに値が渡り、出力に基づいた返答がされることを確認できる。