Amazon Cognito のユーザーグループに応じて Bedrock Guardrails を動的に切り替える

はじめに

こんにちは、ほりぐちです。

みなさん、チャット AI は使っていますか？
生成 AI の代表的なユースケースのひとつであるチャット AI ですが、サービスを提供する上では無制限に回答させることが適切ではない場面が存在します。

そんなときにガードレールを設定することで、生成 AI の入出力に対して特定のトピックを拒否したり、コンテンツのフィルタリングなどの安全策を適用できます。

しかし、同じ生成 AI アプリケーションでも、利用者の権限や利用シーンに応じて適用するガードレールを変えたいケースがあります。例えば、一般ユーザーには金融に関する拒否トピックなどを設定しない基本的なガードレールを適用し、より制約の強い利用ポリシーが必要なユーザーには、個別の投資判断や価格予測を拒否する厳格なガードレールを適用する、といった使い分けです。

本記事では、Amazon Cognito のユーザーグループを利用して JWT の cognito:groups クレームを取得し、その内容に応じて AWS Lambda から Amazon Bedrock Guardrails を切り替える構成を試します。

ユーザーグループや認可情報に応じてガードレールを切り替えたいという方はぜひご覧ください。

今回の構成

※ 図中ではガードレールの選択を分岐として表現していますが、実装上は Lambda が Converse API の guardrailConfig に選択したガードレール ID とバージョンを指定します。

注意

今回はガードレール検証のために生成 AI に投資判断に関する質問をしています。
しかし、本ブログ記事は特定の金融商品の購入、またはその逆を推奨するものではありません。

2 種類のガードレールを作成する

それではまず Amazon Bedrock で金融に関する拒否トピックなどを設定しない基本的なガードレールと金融に関する質問やアドバイスを拒否する厳格なガードレールの 2 種類を作成します。
Amazon Bedrock のコンソール左のメニューから「ガードレール」を選択し、画面中央のボタンからガードレールの作成に入ります。

今回はこちらの拒否トピックを使って質問をブロックします。拒否トピックは自然言語に基づいて答えて欲しくないトピックを指定できるガードレールの機能です。
なお、日本語で拒否トピックを作成したい場合は Denied topics tier で右側の Standard を選ぶ必要があります。(参考)

今回は拒否トピックに設定する内容として金融商品に関する質問や投資判断に関するアドバイスを指定します。今回は以下のような拒否トピックを設定しました。

個別の金融商品、投資対象、または金融市場に関する売買判断、投資判断、取引戦略、価格予測、相場予測、資産配分、保有判断、売却判断、投資タイミングに関する助言または推奨。

対象には、株式、ETF、投資信託、債券、金・銀などの貴金属、暗号資産、FX・為替、商品先物、オプション、不動産投資、その他の金融資産を含む。

この時、自然言語で定義をするとともに、ブロックしてほしい具体的な質問の例を複数例示することで意図したトピックを検出しやすくなります。

作成したガードレールの ID とバージョンは後ほど利用するので控えておきます。

Cognito でユーザーとグループを作成する

次に、Cognito を利用して実際にリクエストを送るユーザーと、ユーザーが所属するグループを作成していきます。

まず Cognito のコンソールからユーザープールを作成します。
アプリケーションタイプはシングルページアプリケーション、サインイン識別子はメールアドレスを選択しました。

作成したユーザープールを選択し、「アプリケーションクライアントに関する情報」横の編集ボタンを押し、設定画面に入ります。

認証フローの項目で「ユーザー名とパスワード (ALLOW_USER_PASSWORD_AUTH)」にチェックを入れます。

次に、左側メニュー「ユーザー」からユーザーの作成に入ります。
今回は 2 種類のグループから Bedrock を利用するために 2 ユーザー作成します。

次に、グループを作成します。今回はわかりやすく、厳格なガードレールを設定する strict-guardrail グループと基本的なガードレールを設定する standard-guardrail グループを作成しました。

作成したグループに先ほどのユーザーを一人ずつ追加します。

このグループ名によって適用されるガードレールの強度が変わるように設定していきます。

Lambda 関数を作成する

先ほど作成したグループに応じて異なるガードレール ID へルーティングされるような Lambda 関数を作成します。

関数の内容自体はシンプルに、Bedrock の Converse API にメッセージを送って返答を受け取るだけの関数とします。

Lambda 関数はデフォルトのタイムアウトが 3 秒になっていますが、Bedrock を利用する際にはモデル応答が 3 秒以内で終わらないことを考えて 30 秒程度にしておきます。

import json
import os

import boto3

bedrock = boto3.client("bedrock-runtime")

STRICT_GROUP = "strict-guardrail"
STANDARD_GROUP = "standard-guardrail"

GUARDRAILS = {
    "strict": {
        "guardrailIdentifier": os.environ["STRICT_GUARDRAIL_ID"],
        "guardrailVersion": os.environ["STRICT_GUARDRAIL_VERSION"],
    },
    "standard": {
        "guardrailIdentifier": os.environ["STANDARD_GUARDRAIL_ID"],
        "guardrailVersion": os.environ["STANDARD_GUARDRAIL_VERSION"],
    },
}

def lambda_handler(event, context):
    message = json.loads(event["body"])["message"].strip()

    claims = event["requestContext"]["authorizer"]["jwt"]["claims"]
    raw_groups = claims.get("cognito:groups", "")

    if isinstance(raw_groups, list):
        groups = raw_groups
    else:
        groups = [
            group.strip().strip('"').strip("'")
            for group in str(raw_groups).strip("[]").split(",")
            if group.strip()
        ]

    policy = (
        "strict"
        if STRICT_GROUP in groups
        else "standard"
        if STANDARD_GROUP in groups
        else "strict"
    )

    result = bedrock.converse(
        modelId=os.environ["MODEL_ID"],
        messages=[{
            "role": "user",
            "content": [{"text": message}],
        }],
        guardrailConfig={
            **GUARDRAILS[policy],
            "trace": "enabled",
        },
        inferenceConfig={
            "maxTokens": 512,
            "temperature": 0.3,
        },
    )

    answer = "".join(
        block["text"]
        for block in result["output"]["message"]["content"]
        if "text" in block
    )

    return {
        "statusCode": 200,
        "headers": {
            "Content-Type": "application/json; charset=utf-8",
        },
        "body": json.dumps({
            "answer": answer,
            "policy": policy,

        }, ensure_ascii=False),
    }

また、環境変数として以下を設定します。

Lambda の実行ロールには bedrock:InvokeModel と bedrock:ApplyGuardrail を許可します。

さらに bedrock:GuardrailIdentifier 条件キーを利用し、strict / standard として作成したガードレールを指定する推論リクエストだけを許可します。これにより、コード上でガードレールの指定が漏れた場合も、ガードレールなしのモデル呼び出しを防止できます。

今回使用した Nova 2 Lite の US クロスリージョン推論プロファイルでは、推論プロファイル、推論先候補リージョンの Foundation Model、および Guardrail Profile を対象に含める必要があります。

API Gateway を作成する

次に、Lambda 関数にリクエストを送るための API Gateway を作成します。

API のタイプは HTTP API を選択します。

ルートは POST メソッドを /chat のパスに設定します。

今回は Cognito 経由でリクエストを送信するため、オーソライザとして JWT を選んで作成、アタッチします。

これで、Cognito で認証されたユーザーが API Gateway を通して Lambda 関数にアクセスし、Lambda がユーザーグループに応じたガードレールを指定して Bedrock を呼び出す構成が完成しました。

実際に試してみる

それでは、実際に異なるグループからリクエストを送った際に、アウトプットが異なるのかを試してみましょう。

まずは、Cognito の standard-guardrail グループに所属するユーザーで試してみます。なお、このユーザーには同名のガードレール（standard-guardrail）が適用されます。

今回は strict-guardrail として金融系の質問を拒否するガードレールを作成したので、差を見るために投資判断に関する質問をしてみます。（以下の操作は全て CloudShell から行います）

CloudShell
  ↓ ユーザー名・パスワード・Client ID を送信
Amazon Cognito の InitiateAuth API
  ↓ 認証成功
Access Token（JWT）を返す

まず、アクセストークンを取得します

export AWS_REGION="<リージョン>"
export CLIENT_ID="<Client ID>"
export USERNAME="<作成したユーザー名>"

read -s -p "User password: " PASSWORD
# 送信後、パスワードの入力を求められるのでユーザーのパスワードを入力します

# initiate-auth コマンドで必要情報を入力し、アクセストークンを取得します
export ACCESS_TOKEN=$(
  aws cognito-idp initiate-auth \
    --region "$AWS_REGION" \
    --auth-flow USER_PASSWORD_AUTH \
    --client-id "$CLIENT_ID" \
    --auth-parameters "USERNAME=$USERNAME,PASSWORD=$PASSWORD" \
    --query "AuthenticationResult.AccessToken" \
    --output text
)

アクセストークンが取得できているかを調べます。

echo "${ACCESS_TOKEN:0:30}..."

このコマンドでeyJ...のようにアクセストークンの序盤 30 文字が出力されれば成功です。

本記事では AWS CLI から簡単に検証できるよう、Client secret を生成しない App Client を使用しています。
Client secret を生成した App Client を利用する場合は、initiate-auth 実行時に SECRET_HASH の指定が必要です。

さて、それでは実際にリクエストを送ってみましょう。
リクエストに含まれている API_URL はAPI Gateway のコンソールから デプロイ > ステージと遷移し、「URL を呼び出す」に書かれている URL を利用します。

取得した URL はリクエストを送るときに使うので、API_URL として変数に格納しておきましょう。取得した呼び出し URL の末尾に、作成したルート /chat を付与して API_URL に格納します。

export API_URL="<取得した API URL>/chat"

今回は金融系の質問を拒否するガードレールを設定したため、実際に金融系の質問をしてみます。
特定の金融商品を買うべきか、という投資に関するアドバイスを求めてみましょう。

curl -sS -X POST "$API_URL" \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $ACCESS_TOKEN" \
  -d '{"message":"余剰資金があります。今週中に金ETFを買うべきですか？買うなら金額も決めてください。"}' \
  | jq -r '"policy: \(.policy)\n\n\(.answer)"'

こちらは基本的なガードレールなので、ブロックされずに返答が返ってくることが期待されます。
こちらを実際に送信すると、以下のような返答が返ってきました。

※以下は、ガードレール機能の動作確認を目的とした検証例です。記事内容は投資アドバイスではありません。

ほぼモザイクで恐縮ですが、基本的なガードレールでは投資判断を求める質問に対してモデルから具体的な回答が返されることを確認できました。

こちらのガードレールには投資関連のアドバイスをブロックする拒否トピックを設定していないので、想定通りの挙動です。

それでは、次に厳しい方のガードレール (strict-guardrail) を適用したグループで同じ質問をしてみます。

まず strict-guardrail グループに所属するユーザーのアクセストークンを取得します。
先ほどの基本的なガードレールのユーザーとは異なるユーザー名とパスワードを指定して、同じ手順でアクセストークンを再取得します。

アクセストークンを取得したら先ほどと同じリクエストを送ってみます。

curl -sS -X POST "$API_URL" \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $ACCESS_TOKEN" \
  -d '{"message":"余剰資金があります。今週中に金ETFを買うべきですか？買うなら金額も決めてください。"}' \
  | jq -r '"policy: \(.policy)\n\n\(.answer)"'

申し訳ありませんが、モデルはこの質問に回答できません。

こちらは想定通り、投資判断に関する質問に対して回答が拒否されることを確認できました。

同じ Lambda 関数でもログインしてきたユーザーの所属によって異なるガードレールにルーティングされていることが確認できました。

終わりに

おつかれさまでした！

今回は、Amazon Cognito のユーザーグループに基づいて、Amazon Bedrock Guardrails を動的に切り替える構成を試しました。ガードレールの選択は、API Gateway で検証された JWT の cognito:groups クレームをもとに Lambda 関数側で行います。そのため、クライアントは適用する Guardrail の種類を指定・変更できません。

クライアントから渡される値ではなく、認証済みユーザーのグループ情報に基づいて適用ポリシーを決定できる点が、この構成のポイントです。

リクエストごとのガードレール選択は動的に行われますが、グループとガードレールの対応付けは Lambda 関数のコードで静的に定義しています。そのため、どのガードレールを適用するかを生成 AI に判断させる必要がなく、想定外のポリシーが選択されることを防げます。

今回の構成では、以下のような使い分けを実装しました。

• standard-guardrail グループ：金融に関する拒否トピックを設定していない基本的なガードレールを適用
• strict-guardrail グループ：金融商品の売買判断や将来価格予測を拒否するガードレールを適用
• 未知のグループまたはグループ未所属ユーザー：安全側に倒して strict ポリシー を適用

今回は CloudShell からテストしましたが、フロントエンドを用意してチャット画面からメッセージを送るような場合にも適用可能です。

ガードレールを一律に適用するだけでなく、認証・認可情報と組み合わせて適用ポリシーを分けたい場合の参考になれば幸いです。

それでは