はじめに
Bedrock をキャッチアップしていますが、今回は Agents for Amazon Bedrock を触ってみました。指定したモデルに対して自然言語による指示で特定のアクションを実行させます。やることとしては、全リージョンの EC2 インスタンスに対して怪しいインスタンスがないかチェックしてもらいます。アクションは以下を設定します。
- 全リージョンのインスタンス総数、実行中インスタンス数を調べる
- 全リージョンで Owner タグのついていないインスタンスを調べる
- 全リージョンでインバウンドが解放された (0.0.0.0/0 が許可された) インスタンスを調べる
構成
以下のような構成です。今回も CDK で作ります (リポジトリ)
Streamlit を使う
前回に引き続き Streamlit を Docker コンテナで動かします。今回は認証機能をつけたりチャット形式の UI にしたりと、そこそこ作ってみました。パスワード等の機密情報が必要なので、前回の VPC エンドポイントに加えて AWS Secrets Manager のエンドポイントも必要です。
アクショングループに登録する Lambda 関数とスキーマ定義
エージェントを構築する前にアクショングループに登録する Lambda 関数を作ります。
複数のアクショングループを用意し、Lambda 関数を 1 対 1 で紐づけてそれぞれエージェントに設定します。呼び出し時はこれらを自然言語の指示で出し分ける形です。
Lambda 関数はコンテナでデプロイします。なお今回のアクションはいずれも全リージョンに対して特定の条件で照会をかけるような内容なので、関数に引数はありません。
.
├── get_instances_count
│ ├── Dockerfile
│ ├── agent.py
│ ├── requirements.txt
│ └── schema.yaml
├── get_instances_with_open_permission
│ ├── Dockerfile
│ ├── agent.py
│ ├── requirements.txt
│ └── schema.yaml
└── get_instances_without_owner
├── Dockerfile
├── agent.py
├── requirements.txt
└── schema.yaml
Dockerfile
Dockerfile はいずれも同じ内容になります。Lambda 用のイメージに資材をコピーして実行コマンドを設定しているだけです。build/push するローカルマシンからの通信がプロキシを通るので、HTTP_PROXY を外から渡しています。流用の際はご注意ください。
FROM public.ecr.aws/lambda/python:3.12
ARG HTTP_PROXY
COPY requirements.txt ${LAMBDA_TASK_ROOT}
COPY agent.py ${LAMBDA_TASK_ROOT}
RUN pip install -r requirements.txt --proxy ${HTTP_PROXY}
CMD [ "agent.handler" ]
requirements.txt
どの関数も、今回インストールするのは boto3 だけです。Bedrock まわりはアップデートの頻度が高く急に新しい API が生えたりするので、なるべく最新を入れましょう。
boto3==1.34.124
全リージョンのインスタンス総数、実行中インスタンス数を調べる関数
コード
要点は処理結果を指定された形式のレスポンスに埋め込むところです。アクショングループの定義方法は 2 つあり、どちらで定義したかによってイベントの形式が変わります。
今回は API スキーマで定義しているので、Lambda 関数の内容もそのようにしています。
import json
import boto3
client = boto3.client("ec2")
regions = client.describe_regions()["Regions"]
def handler(event, context):
body = run_with_regions()
response = {
"messageVersion": "1.0",
"response": {
"actionGroup": event["actionGroup"],
"apiPath": event["apiPath"],
"httpMethod": event["httpMethod"],
"httpStatusCode": 200,
"responseBody": {"application/json": {"body": body}}, # <- これ
},
}
return response
def run_with_regions():
global results
results = []
for region in regions:
get_instances_count(region["RegionName"])
return json.dumps(obj=results, ensure_ascii=False)
def get_instances_count(region_name: str):
try:
client = boto3.client("ec2", region_name=region_name)
instances = client.describe_instances()
instance_count = 0
instance_running = 0
for reservation in instances["Reservations"]:
instance_count += len(reservation["Instances"])
for instance in reservation["Instances"]:
if instance["State"]["Name"] == "running":
instance_running += 1
result = {
"region": region_name,
"instance_count": instance_count,
"instance_running": instance_running,
}
results.append(result)
except Exception as e:
print("Error: {}: {}".format(region_name, e))
API スキーマ
前述の API スキーマ定義です。以下の要点に注意が必要です。
- Bedrock に関数の内容を伝えるために
descriptionをなるべく詳細に書く parametersを実際の引数に対応させてきちんと書く (今回は関数に引数がないので省略しています)responsesを実際の戻り値に対応させてきちんと書く
openapi: 3.0.0
info:
title: Lambda
version: 1.0.0
paths:
/count:
get:
summary: Check EC2 instances count
description: リージョンごとにインスタンスの総数、実行中のインスタンスの数を集計してJSONで返します。
operationId: get_instances_count
responses:
"200":
description: インスタンス数の確認が成功しました。
content:
application/json:
schema:
type: array
items:
type: object
properties:
region:
type: string
description: リージョン名
instance_count:
type: integer
description: インスタンスの総数
instance_running:
type: integer
description: 実行中のインスタンスの数
全リージョンで Owner タグのついていないインスタンスを調べる関数
コード
...
def run_with_regions():
global results
results = []
for region in regions:
get_instances_without_owner(region["RegionName"])
return json.dumps(obj=results, ensure_ascii=False)
def get_instances_without_owner(region_name):
try:
client = boto3.client("ec2", region_name=region_name)
instances = client.describe_instances()
for reservation in instances["Reservations"]:
for instance in reservation["Instances"]:
if get_instance_tag_value("Owner", instance) == "":
result = {
"region": region_name,
"instance_id": instance["InstanceId"],
"instance_name": get_instance_tag_value("Name", instance),
}
results.append(result)
except Exception as e:
print("Error: {}: {}".format(region_name, e))
def get_instance_tag_value(key, instance) -> str:
return next(
(
tag["Value"]
for tag in instance.get("Tags")
if tag["Key"].lower() == key.lower()
),
"",
)
API スキーマ
openapi: 3.0.0
info:
title: Lambda
version: 1.0.0
paths:
/check-without-owner:
get:
summary: Check EC2 instances without owner tag
description: リージョンごとにOwnerタグが付与されているか確認し、リージョン名、インスタンスID、インスタンス名を収集してJSONで返します。
operationId: get_instances_without_owner
responses:
"200":
description: Ownerタグが付与されていないインスタンスの存在確認が成功しました。
content:
application/json:
schema:
type: array
items:
type: object
properties:
region:
type: string
description: リージョン名
instance_id:
type: string
description: インスタンスID
instance_name:
type: string
description: インスタンス名
全リージョンでインバウンドが解放された (0.0.0.0/0 が許可された) インスタンスを調べる関数
コード
...
def run_with_regions():
global results
results = []
for region in regions:
get_instances_with_open_permission(region["RegionName"])
return json.dumps(obj=results, ensure_ascii=False)
def get_instances_with_open_permission(region_name: str):
try:
client = boto3.client("ec2", region_name=region_name)
security_groups = client.describe_security_groups(
Filters=[{"Name": "ip-permission.cidr", "Values": ["0.0.0.0/0"]}]
)
if not security_groups["SecurityGroups"]:
return None
instances = client.describe_instances(
Filters=[
{
"Name": "instance.group-id",
"Values": [
security_group["GroupId"]
for security_group in security_groups["SecurityGroups"]
],
}
]
)
for reservation in instances["Reservations"]:
for instance in reservation["Instances"]:
permissions = []
for security_group in instance["SecurityGroups"]:
security_group_details = client.describe_security_groups(
GroupIds=[security_group["GroupId"]]
)
for security_group_detail in security_group_details[
"SecurityGroups"
]:
for permission in security_group_detail["IpPermissions"]:
for ip_range in permission.get("IpRanges"):
if ip_range.get("CidrIp") == "0.0.0.0/0":
permission_detail = {
"protocol": permission.get("IpProtocol"),
"from_port": permission.get("FromPort"),
"to_port": permission.get("ToPort"),
"allow_from": security_group["GroupName"],
}
permissions.append(permission_detail)
if permissions:
result = {
"region": region_name,
"instance_id": instance["InstanceId"],
"instance_name": get_instance_tag_value("Name", instance),
"permissions": permissions,
}
results.append(result)
except Exception as e:
print("Error: {}: {}".format(region_name, e))
...
API スキーマ
openapi: 3.0.0
info:
title: Lambda
version: 1.0.0
paths:
/check-open-permission:
get:
summary: Check EC2 instances with open permission
description: リージョンごとにインバウンド通信で0.0.0.0/0が許可されているインスタンスがあるかを確認し、リージョン名、インスタンスID、インスタンス名、許可されているプロトコル、開始ポート、終了ポート、どのセキュリティグループからの許可かを示すセキュリティグループ名を収集してJSONで返します。
operationId: get_instances_with_open_permission
responses:
"200":
description: インバウンド通信が解放されたインスタンスの存在確認が成功しました。
content:
application/json:
schema:
type: array
items:
type: object
properties:
region:
type: string
description: リージョン名
instance_id:
type: string
description: インスタンスID
instance_name:
type: string
description: インスタンス名
permissions:
type: array
items:
type: object
properties:
protocol:
type: string
description: プロトコル
from_port:
type: string
description: 開始ポート
to_port:
type: string
description: 終了ポート
allow_from:
type: string
description: どのセキュリティグループからの許可かを示すセキュリティグループ名
CDK コード
lib/image/agent/src 配下にさきほどの Lambda 関数の構成を腹持ちしています。
. ├── LICENSE ├── README.md ├── bin │ └── main.ts ├── cdk-bedrock.code-workspace ├── cdk.context.json ├── cdk.example.json ├── cdk.json ├── docs │ └── diagram.png ├── jest.config.js ├── lib │ ├── constructs │ │ ├── bedrock.ts │ │ ├── ecs.ts │ │ └── function.ts │ ├── image │ │ ├── agent │ │ │ ├── README.md │ │ │ ├── pyproject.toml │ │ │ ├── requirements-dev.lock │ │ │ ├── requirements.lock │ │ │ └── src │ │ │ └── agent │ │ │ ├── get_instances_count │ │ │ │ ├── Dockerfile │ │ │ │ ├── agent.py │ │ │ │ ├── requirements.txt │ │ │ │ └── schema.yaml │ │ │ ├── get_instances_with_open_permission │ │ │ │ ├── Dockerfile │ │ │ │ ├── agent.py │ │ │ │ ├── requirements.txt │ │ │ │ └── schema.yaml │ │ │ └── get_instances_without_owner │ │ │ ├── Dockerfile │ │ │ ├── agent.py │ │ │ ├── requirements.txt │ │ │ └── schema.yaml │ │ └── app │ │ ├── README.md │ │ ├── pyproject.toml │ │ ├── requirements-dev.lock │ │ ├── requirements.lock │ │ └── src │ │ └── app │ │ ├── Dockerfile │ │ ├── app.py │ │ └── requirements.txt │ └── stack.ts ├── package-lock.json ├── package.json ├── test │ └── cdk-bedrock.test.ts └── tsconfig.json
前準備
前回と同じですが、cdk.json に必要な情報を入れます。パブリックホストゾーンと ECR リポジトリは既存のものを使う前提なので、ここで設定します。
// 例
{
...
"context": {
...
"owner": "user",
"serviceName": "my-service",
"hostZoneName": "example.com",
"allowedIps": ["0.0.0.0/0"],
"httpProxy": "http://my-proxy.com:port",
"repository": "user/reponame"
}
}
なお現在 CDK のベストプラクティスは、可変パラメータも TypeScript で定義して型の恩恵にあずかろうというのが主流です。今回そうしていないのには以下の理由があります。
- 機密情報とまではいかないがなるべく ignore しておきたい情報をそこそこ含んでいる
- dev/prod などでコンテキストを切り替える想定がない
このため cdk.json は ignore し、ガイドとして cdk.example.json を配置しています。
Lambda
コード。Lambda コンテナイメージのビルドから関数エイリアスの発行までを行います。
IAM ロール
Logs に記録する権限のほかに EC2 に対して Describe 系の API を実行する権限が必要です。
const functionRole = new cdk.aws_iam.Role(this, "Role", {
roleName: `${props.serviceName}-function-role`,
assumedBy: new cdk.aws_iam.ServicePrincipal("lambda.amazonaws.com"),
inlinePolicies: {
FunctionPolicy: new cdk.aws_iam.PolicyDocument({
statements: [
new cdk.aws_iam.PolicyStatement({
effect: cdk.aws_iam.Effect.ALLOW,
actions: ["logs:CreateLogGroup", "logs:CreateLogStream", "logs:PutLogEvents"],
resources: ["arn:aws:logs:*:*:*"],
}),
new cdk.aws_iam.PolicyStatement({
effect: cdk.aws_iam.Effect.ALLOW,
actions: ["ec2:Describe*"],
resources: ["*"],
}),
],
}),
},
});
クラスの定義
関数が複数あるので、クラスを作ってコードの繰り返しを削減してみました。ここで定義した Description はのちほどプロンプトで使います。
export interface FunctionConfig {
BaseName: String;
FunctionId: string;
FunctionName: string;
Path: string;
SchemaFilePath: string;
Description: string;
Alias?: cdk.aws_lambda.Alias;
}
class funcionConfig {
getConfig = (prefix: string, name: string, description: string): FunctionConfig => {
const kebabToPascalCase = (name: string): string => {
return name
.split("-")
.map((word) => word.charAt(0).toUpperCase() + word.slice(1).toLowerCase())
.join("");
};
const kebabToSnakeCase = (name: string): string => {
return name.split("-").join("_");
};
const path = `lib/image/agent/src/agent/${kebabToSnakeCase(name)}`;
return {
BaseName: name,
FunctionId: kebabToPascalCase(name),
FunctionName: `${prefix}-${name}`,
Path: path,
Description: description,
SchemaFilePath: `${path}/schema.yaml`,
};
};
public Config(prefix: string): FunctionConfig[] {
return [
this.getConfig(prefix, "get-instances-count", "インスタンス数の取得"),
this.getConfig(prefix, "get-instances-without-owner", "Ownerタグのないインスタンスの取得"),
this.getConfig(prefix, "get-instances-with-open-permission", "0.0.0.0/0が許可されたインスタンスの取得"),
];
}
}
関数のデプロイ
上記クラスを使い、map で回すことでコンテナのビルドから関数エイリアスの発行までを行います。timeout は 5 分にしています。今回のケースではデフォルト値だと確実にタイムアウトします。また architecture はビルド環境に合ったものを選択する必要があります。
const cfg = new funcionConfig();
this.functionConfig = [];
cfg.Config(props.serviceName).map((obj) => {
const fn = new cdk.aws_lambda.DockerImageFunction(this, `${obj.FunctionId}Function`, {
functionName: obj.FunctionName,
description: obj.FunctionName,
code: cdk.aws_lambda.DockerImageCode.fromImageAsset(obj.Path, {
buildArgs: {
HTTP_PROXY: props.httpProxy,
},
}),
architecture: cdk.aws_lambda.Architecture.ARM_64,
role: functionRole,
logRetention: cdk.aws_logs.RetentionDays.THREE_DAYS,
currentVersionOptions: {
removalPolicy: cdk.RemovalPolicy.RETAIN,
},
timeout: cdk.Duration.minutes(5),
});
obj.Alias = new cdk.aws_lambda.Alias(this, `${obj.FunctionId}Alias`, {
aliasName: "live",
version: fn.currentVersion,
});
this.functionConfig.push(obj);
});
Agents for Amazon Bedrock
コード。やはり @cdklabs/generative-ai-cdk-constructs が有用です。
Instruction
いちばん大事なのは自然言語での指示です。今回は以下のようにしてみました。改善の余地はあると思います (関数名で伝えるのはアリなのか? など)
あなたはAWSに精通したソリューションアーキテクトです。いくつかの関数を使い分け、ユーザーの要求に日本語で回答してください。
ただし、コードやデータ構造については日本語ではなくそのまま回答してください。タスク1:
もし、例えば「インスタンスの数を教えてください」というように、インスタンスの数について聞かれたら、kawashima-get-instances-countというLambda関数を実行してください。タスク2:
もし、例えば「Ownerタグの付与されていないインスタンスの情報を教えてください」というように、Ownerタグが付与されていないインスタンスの有無について聞かれたら、kawashima-get-instances-without-ownerというLambda関数を実行してください。タスク3:
もし、例えば「インバウンド通信で0.0.0.0/0が許可されているインスタンスの情報を教えてください」というように、インバウンド通信が解放されたインスタンスの有無について聞かれたら、kawashima-get-instances-with-open-permissionというLambda関数を実行してください。タスク4:
タスク1、タスク2、タスク3以外の場合は、Lambda関数を実行せずに一般的な質問への回答をしてください。わからない質問には「その質問には回答できません」と回答してください。出力形式:
出力形式の指示が特にない場合は、項目を網羅して整形し、リスト形式で結果を返してください。
例えば「マークダウンの表で回答してください。」や「結果のJSONをそのまま返してください。」など、ユーザーから出力形式が指定された場合に限って、指示に合わせた回答をしてください。特記事項:
関数の実行結果に関する質問をユーザーから追加で受けた場合は、再度同様の関数を実行することなく、前回の結果を参照し、文脈にあった回答をしてください。
IAM ロール
ロールも書きました。エージェントがモデルにアクセスする権限が必要です。マネコンだと裏で勝手に作られるかと思います。
const agentRole = new cdk.aws_iam.Role(this, "Role", {
roleName: `${props.serviceName}-agent-role`,
assumedBy: new cdk.aws_iam.ServicePrincipal("bedrock.amazonaws.com"),
inlinePolicies: {
AgentPolicy: new cdk.aws_iam.PolicyDocument({
statements: [
new cdk.aws_iam.PolicyStatement({
effect: cdk.aws_iam.Effect.ALLOW,
actions: ["bedrock:InvokeModel"],
resources: [`arn:aws:bedrock:${stack.region}::foundation-model/*`],
}),
],
}),
},
});
Agent
母体となるエージェントです。エイリアスという概念があり、バージョン管理できます。ただし Lambda と違うのは、エイリアスを作成しないとバージョンが払い出されない点です。aliasName を設定することでエイリアスを作成し、shouldPrepareAgent: true でそのエイリアスを使用可能な状態にしています。
this.agent = new bedrock.Agent(this, "Agent", {
name: `${props.serviceName}-agent`,
aliasName: "v1",
existingRole: agentRole,
foundationModel: bedrock.BedrockFoundationModel.ANTHROPIC_CLAUDE_HAIKU_V1_0,
instruction: instruction,
shouldPrepareAgent: true,
enableUserInput: true,
idleSessionTTL: cdk.Duration.minutes(15),
});
ドキュメントに記載がありました。
Creating an agent alias will not prepare the agent, so if you create an alias with addAlias or by providing an aliasName when creating the agent then you should set shouldPrepareAgent to true.
エージェントのエイリアスを作成しても、エージェントの準備は行われません。したがって、addAlias でエイリアスを作成するか、エージェントの作成時に aliasName を指定する場合は、shouldPrepareAgent を true に設定する必要があります。
function.ts から props で渡した functionConfig を再度 map で回し、以下を実行します。
- 関数ごとにアクショングループを定義し、Lambda 関数のエイリアスを指定する
- バケット経由ではなくインラインで API スキーマ定義を渡す
- エージェントにアクショングループを登録する
- Lambda 関数のリソースベースポリシーでエージェントからの関数呼び出しを許可する (これ重要)
props.functionConfig.map((obj) => {
const actionGroup = new bedrock.AgentActionGroup(this, `${obj.FunctionId}ActionGroup`, {
actionGroupName: `${obj.FunctionName}-actiongroup`,
description: `${obj.FunctionName}-actiongroup`,
actionGroupExecutor: {
lambda: obj.Alias,
},
actionGroupState: "ENABLED",
apiSchema: bedrock.ApiSchema.fromAsset(obj.SchemaFilePath),
skipResourceInUseCheckOnDelete: false,
});
this.agent.addActionGroup(actionGroup)
obj.Alias?.addPermission(`${obj.FunctionId}Permission`, {
principal: new cdk.aws_iam.ServicePrincipal("bedrock.amazonaws.com"),
action: "lambda:InvokeFunction",
sourceArn: this.agent.agentArn,
});
});
ECS
コード。前回とほぼほぼ同じなので、追加部分を抜粋します。
VPC
VPC には追加で AWS Secrets Manager のエンドポイントを生やす必要があります。
...
vpc.addInterfaceEndpoint("SecretsManagerEndpoint", {
securityGroups: [endpointSecurityGroup],
service: cdk.aws_ec2.InterfaceVpcEndpointAwsService.SECRETS_MANAGER,
});
...
Secrets
Steramlit に認証画面を実装したので、以下 2 つのシークレットを定義します。
// ユーザー用
const userSecret = new cdk.aws_secretsmanager.Secret(this, "UserSecret", {
secretName: `${props.serviceName}-user-secret`,
description: `${props.serviceName}-user-secret`,
generateSecretString: {
generateStringKey: "password",
excludePunctuation: true,
passwordLength: 32,
secretStringTemplate: JSON.stringify({ username: "admin" }),
},
})
// 署名付き Cookie 用
const cookieSecret = new cdk.aws_secretsmanager.Secret(this, "CookieSecret", {
secretName: `${props.serviceName}-cookie-secret`,
description: `${props.serviceName}-cookie-secret`,
generateSecretString: {
generateStringKey: "key",
excludePunctuation: true,
passwordLength: 256,
secretStringTemplate: JSON.stringify({ name: "id" }),
},
});
コンテナの設定
環境変数とシークレットを設定します。シークレットは環境変数を経由して渡されます。
taskDefinition.addContainer("Container", {
containerName: contanerName,
image: cdk.aws_ecs.ContainerImage.fromEcrRepository(repository, tag),
logging: cdk.aws_ecs.LogDrivers.awsLogs({
logGroup: logGroup,
streamPrefix: "logs",
}),
environment: {
TARGET_REGION: stack.region,
AGENT_ID: props.agent.agentId,
AGENT_ALIAS_ID: props.agent.aliasId!,
ACTION_LABELS: JSON.stringify(props.functionConfig.map((item) => item.Description)),
},
secrets: {
USERNAME: cdk.aws_ecs.Secret.fromSecretsManager(userSecret, "username"),
PASSWORD: cdk.aws_ecs.Secret.fromSecretsManager(userSecret, "password"),
COOKIE_NAME: cdk.aws_ecs.Secret.fromSecretsManager(cookieSecret, "name"),
COOKIE_KEY: cdk.aws_ecs.Secret.fromSecretsManager(cookieSecret, "key"),
},
portMappings: [
{
containerPort: containerPort,
protocol: cdk.aws_ecs.Protocol.TCP,
},
],
});
関数のところで作った functionConfig から Description を抜き出して JSON で渡しています。これをアプリ側でセレクトボックスで選択させ、そのまま指定のアクションを呼ぶためのプロンプトにすることがねらいです。
environment: {
...
ACTION_LABELS: JSON.stringify(props.functionConfig.map((item) => item.Description)),
},
アプリケーション
コード。少し触れましたが、前回よりは Streamlit のコードを拡張しています。 LangChain を使わずに実装してみました。
import json
import os
import uuid
import boto3
import streamlit as st
import streamlit_authenticator as sa
def main():
if "session_id" not in st.session_state:
st.session_state.session_id = str(uuid.uuid4())
if "messages" not in st.session_state:
st.session_state.messages = []
if "selected_preset" not in st.session_state:
st.session_state.selected_preset = ""
if "use_preset" not in st.session_state:
st.session_state.use_preset = False
auth = sa.Authenticate(
credentials={
"usernames": {
"admin": {
"name": os.environ["USERNAME"],
"password": os.environ["PASSWORD"],
}
}
},
cookie_name=os.environ["COOKIE_NAME"],
cookie_key=os.environ["COOKIE_KEY"],
cookie_expiry_days=1,
)
auth.login()
if st.session_state["authentication_status"] is True:
st.title("AWSのことなんでもこたえるマン")
set_sidebar(auth)
set_messages()
if prompt := get_user_prompt():
handle_prompt(prompt)
elif st.session_state["authentication_status"] is False:
st.error("Error: Logion failed.")
def set_sidebar(auth):
actions = json.loads(os.environ["ACTION_LABELS"])
with st.sidebar:
st.session_state.selected_preset = st.selectbox("リソース調査", actions)
if st.button("依頼"):
st.session_state.use_preset = True
st.divider()
auth.logout("Logout", "sidebar")
def set_messages():
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.write(message["content"])
def get_user_prompt():
prompt = ""
if st.session_state.use_preset:
prompt = st.session_state.selected_preset
if user_input := st.chat_input("なんでもきいてください"):
prompt = user_input
return prompt
def handle_prompt(prompt):
with st.chat_message("Human"):
st.markdown(prompt)
st.session_state.messages.append({"role": "Human", "content": prompt})
with st.chat_message("Assistant"):
with st.spinner("回答を準備中..."):
response = invoke_agent(prompt)
result = ""
if stream := response.get("completion"):
for event in stream:
if chunk := event.get("chunk"):
if bytes := chunk.get("bytes"):
result += bytes.decode("utf-8")
st.markdown(result)
st.session_state.messages.append(
{"role": "Assistant", "content": result}
)
def invoke_agent(prompt):
try:
client = boto3.client(
service_name="bedrock-agent-runtime",
region_name=os.environ["TARGET_REGION"],
)
response = client.invoke_agent(
inputText=prompt,
agentId=os.environ["AGENT_ID"],
agentAliasId=os.environ["AGENT_ALIAS_ID"],
sessionId=st.session_state.session_id,
enableTrace=False,
endSession=False,
)
except Exception as e:
print("Error: {}".format(e))
raise
return response
if __name__ == "__main__":
main()
関数の呼び出し
サイドバーに設置したセレクトボックスで、関数の説明 ACTION_LABELS をそのままプロンプトとして設定できるようにしています。エージェントで定義したモデルが設定済みのインストラクションを参照して推論し、ACTION_LABELS から設定したプロンプトであればその内容に沿って Lambda 関数を出し分けます。そしてその結果も自然言語で返します。関係ない質問をした場合は関数を実行せず文脈に沿って回答します。
会話履歴の保持
invoke_agent は sessionId を設定することで DynamoDB などを使わずに会話履歴の保持が可能です。今回は uuid で払い出し、st.session_state に保持しておいてそれを invoke_agent に渡すようにしました。
動かしてみる
デプロイします。
cdk synth cdk deploy
無事に画面が表示されました。権限のあるユーザーで AWS Secrets Manager にアクセスしてパスワードを取得し、ログインします。
ログインできました。
名前を覚えているか確認してみましょう。会話履歴も保持されていますね。
ではサイドバーのセレクトボックスから「インスタンス数の取得」を実行してみましょう。回答が要約で返ってきました。結果も正しいです。
関数を再度実行させずに、出力形式の変更を指示してみましょう。これはねらい通りになりました。インストラクションの大切さがわかりますね。
別のアクションを実行させます。会話履歴を引きずっているからか、マークダウンのテーブルで返してきました (自分のインスタンス以外もあるのでマスクしています)
最後のアクションを実行させてみます。マスク範囲が大きめで視認性が悪いですが、結果は正しいです。
少し踏み込んだ内容を聞いてみましょう。「前回」と指示していますが、総括を求められたと判断したようです。文面はまともな回答です。
おわりに
今回試したのは Agents for Amazon Bedrock のさわりの部分だけですが、それでも可能性を感じる結果でした。
個人的には API スキーマと Lambda 関数を共通化して単一のアクショングループに設定しても指示次第で単一の関数の中から機能を呼び分けられるのか? という点が気になったので、近いうちに試してみたいと思います。
