Bedrock、Lambda、Kendra、S3を使用したRAGをSAMで実装してみた

こんにちは、つくぼし(tsukuboshi0755)です！

Retrieval-Augmented Generation(以下RAG)が実際にどのような挙動になるのか、パッと確認してみたいという要望を持つ方はいらっしゃるのではないでしょうか。

今回は、Bedrock、Lambda、Kendra及びS3を使用したRAGについて、誰でも手軽にデプロイ/削除できるようにSAMを使って実装してみたいと思います！

前提条件

今回は下記のソフトウェアの使用を前提としています。

足りない方は個別で導入/設定をご実施ください。

構成

今回の構成図は以下になります。

RAGの検証手順としては以下です。

1. 上記の構成図をSAMでデプロイ
2. S3バケットに特定のPDFファイルをアップロード
3. Kendraでデータ同期を実施
4. Lambdaでテスト実行

なお今回の構成では、Lambdaにプロンプトを直接渡しています。

実際のシステムでは、API Gateway等を用いて、プロンプトを入力するインターフェース(チャットアプリやウェブフォーム等)とLambdaとの連携が必要になるためご注意ください。

テンプレート

全体のコードは、以下のGitHubリポジトリに格納しています。

今回はKendra及びS3を作成し、新たにBedrock及びKendraにアクセスするためのLambdaを追加する形になっています。

なおKendra及びS3の設定については、下記のブログで紹介したものを流用しているので、別途ご参照ください。

今回は新規のLambdaに関する設定についてのみ紹介します。

Lambdaロール

template.yaml

  RagFunctionRole:
    Type: 'AWS::IAM::Role'
    Properties:
      RoleName: !Sub '${SysName}-${Env}-rag-function-role'
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Sid: ''
            Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: 'sts:AssumeRole'
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'

  RagFunctionPolicy:
    Type: 'AWS::IAM::ManagedPolicy'
    Properties:
      ManagedPolicyName: !Sub '${SysName}-${Env}-rag-function-policy'
      Roles:
        - !Ref RagFunctionRole
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Resource: !Sub 'arn:aws:bedrock:${BedrockRegion}::foundation-model/${BedrockModelId}'
            Action: 'bedrock:InvokeModel'
          - Effect: Allow
            Resource: !Sub
              - 'arn:aws:kendra:${AWS::Region}:${AWS::AccountId}:index/${IndexId}'
              -  IndexId: !GetAtt KendraIndex.Id
            Action: 'kendra:Query'

Lambda用のIAMロールを作成します。

ここでは、BedrockのAnthropic Claudeモデルのみ、及び作成するKendraインデックスのみに対してアクセス可能な権限をアタッチします。

Lambda関数

template.yaml

  RagFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: function/
      Environment:
        Variables:
          KENDRA_INDEX_ID: !GetAtt KendraIndex.Id
          BEDROCK_MODEL_ID: !Ref BedrockModelId
      FunctionName: !Sub '${SysName}-${Env}-rag-function'
      Handler: index.lambda_handler
      Role: !GetAtt RagFunctionRole.Arn
      Runtime: python3.11
      Timeout: 900
      Layers:
        - !Ref Boto3Layer

Lambda関数を作成します。

Environmentでは、作成したKendraインデックスのID、及び指定したBedrockモデルのIDを環境変数で指定します。

この設定により、Kendraインデックス及びBedrockモデルに対するLambdaのアクセス先を、動的に変更できるようになります。

Lambdaコード

function/index.py

import json
import logging
import os
from typing import Any, Dict

import boto3

logger = logging.getLogger()

kendra_client = boto3.client("kendra")
bedrock_runtime_client = boto3.client(
    service_name="bedrock-runtime", region_name=os.getenv("BEDROCK_REGION")
)


# Kendra から検索結果を取得
def get_retrieval_result(query_text: Any | None, index_id: str | None) -> str:
    response = kendra_client.query(
        QueryText=query_text,
        IndexId=index_id,
        AttributeFilter={
            "EqualsTo": {
                "Key": "_language_code",
                "Value": {"StringValue": "ja"},
            },
        },
    )

    # Kendra の応答から最初の5つの結果を抽出
    results = response["ResultItems"][:5] if response["ResultItems"] else []

    # 結果からドキュメントの抜粋部分のテキストを抽出
    for i in range(len(results)):
        results[i] = (
            results[i].get("DocumentExcerpt", {}).get("Text", "").replace("\\n", " ")
        )

    print("Received results:" + json.dumps(results, ensure_ascii=False))

    return json.dumps(results, ensure_ascii=False)


# Lambda のハンドラー関数
def lambda_handler(event: Dict[Any, Any], context: Any) -> Any:
    user_prompt = event.get("user_prompt")
    index_id = os.getenv("KENDRA_INDEX_ID")

    prompt = f"""\n\nHuman:
    [参考]情報をもとに[質問]に適切に答えてください。
    [質問]
    {user_prompt}
    [参考]
    {get_retrieval_result(user_prompt,index_id)}
    Assistant:
    """

    # 各種パラメーターの指定
    modelId = os.getenv("BEDROCK_MODEL_ID")
    accept = "application/json"
    contentType = "application/json"

    body = json.dumps(
        {
            "prompt": prompt,
            "max_tokens_to_sample": 600,
        }
    )

    response = bedrock_runtime_client.invoke_model(
        modelId=modelId, accept=accept, contentType=contentType, body=body
    )

    response_body = json.loads(response.get("body").read())

    print("Received response_body:" + json.dumps(response_body, ensure_ascii=False))

    return response_body.get("completion")

入力されたプロンプトに対して、BedrockとKendraにアクセスし、レスポンスを返すLambdaコードを作成します。

コードの内容については、以下のブログを参照しています。

(なおSAMデプロイ後に使用しやすいように、一部の変数やログ設定を変更しています)

Lambdaレイヤー

template.yaml

  Boto3Layer:
    Type: AWS::Serverless::LayerVersion
    Properties:
      LayerName: !Sub '${SysName}-${Env}-boto3'
      ContentUri: layer/
      CompatibleRuntimes:
        - python3.11
    Metadata:
      BuildMethod: python3.11

2023/11時点で、Python3.11のLambdaに内蔵されるboto3のバージョン(1.27.1)は、Bedrockにアクセス可能なboto3のバージョン(1.28.57以降)に対応していません。

そのため別途requirements.txtを以下のように記載し、Lambda関数に対してカスタムレイヤーを付与する事でBedrockにアクセスできるようにします。

layer/requirements.txt

boto3 >= 1.28.57

Lambdaログ(任意)

template.yaml

  RagFunctionLog:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub
        - '/aws/Lambda/${FunctionName}'
        -  FunctionName: !Ref RagFunction
      RetentionInDays: !Ref LogRetentionDays

Lambda用のCloudWatch Logsロググループを作成します。

必要に応じて、RetentionInDaysでLambdaログの保持期間を指定します。

RAGの動作確認

上記のSAMコードをデプロイする事で、S3にアップロードされたPDFファイルを参照し、RAGが正しい回答を返答できるか確認します。

Bedrockで使用可能なモデルの確認

まず前提として、SAMをデプロイするリージョンで、BedrockのClaudeにアクセス可能か確認してください。

もしアクセスしたいモデルのステータスがAccess grantedになっていなければ、別途モデルアクセスの申請を実施しておきましょう。

なお今回のテンプレートでは、デフォルト値としてオレゴンus-west-2のBedrockを使用します。

SAMアプリのデプロイ

次にリポジトリをクローンし、クローンしたディレクトリに移動します。

# リポジトリをクローン
git clone https://github.com/tsukuboshi/rag-with-bedrock

# ディレクトリを移動
cd rag-with-bedrock

ディレクトリ内で、SAMを用いてRAGシステムをデプロイします。

# SAMアプリをビルド
sam build

# SAMアプリをデプロイ
sam deploy

なおデプロイ時に、--parameter-overridesオプションで、以下のパラメータを指定する事も可能です。

※2023/11時点では、東京リージョンにおいてanthropic.claude-v2は使用不可のため注意

S3バケットへのデータ投入

次にCloudShell上で以下のコマンドを実施し、PDFファイルを作成したS3バケットにアップロードします。

# S3 バケット名を設定
BUCKET_NAME=<バケット名>

# S3プレフィックスを設定
BUCKET_PREFIX=awsdoc

# AWS の公式ドキュメントの PDF ファイルをダウンロード
mkdir ${BUCKET_PREFIX}
cd ${BUCKET_PREFIX}
wget https://docs.aws.amazon.com/ja_jp/amazondynamodb/latest/developerguide/dynamodb-dg.pdf -O DynamoDB.pdf
wget https://docs.aws.amazon.com/ja_jp/lambda/latest/dg/lambda-dg.pdf -O Lambda.pdf
wget https://docs.aws.amazon.com/ja_jp/vpc/latest/userguide/vpc-ug.pdf -O VPC.pdf
wget https://docs.aws.amazon.com/ja_jp/kendra/latest/dg/kendra-dg.pdf -O Kendra.pdf
wget https://docs.aws.amazon.com/ja_jp/Route53/latest/DeveloperGuide/route53-dg.pdf -O Route53.pdf
cd ..

# S3 バケットに PDF ファイルをアップロード
aws s3 cp ${BUCKET_PREFIX} s3://${BUCKET_NAME}/${BUCKET_PREFIX}/ --recursive

なお今回の確認で使用する文書データの例として、以下の記事で利用しているPDFファイルを流用しています。

ファイルのアップロードが成功すると、以下の通り作成したS3バケットのawsdocプレフィックス配下に5つのPDFファイルが作成されます。

Kendraデータソースの同期

次に作成したKendraのデータソースに移動し、Sync nowボタンを押すと、S3にアップロードしたPDFファイルがKendraに同期されます。

データの同期が成功すると、Sync Historyタブに以下のようなCompleted履歴が表示されます。

Lambdaのテスト実行

最後に作成したLambda関数をクリックし、コードタブにてConfigure test eventをクリックし、テストイベントを作成します。

今回イベントJSONには、キーにuser_promptを指定し、値にRAGに対するユーザーからの任意の質問を入力してください。

例えば以下のような形で指定します。

{
  "user_prompt": "Lambda 関数で使用できるメモリの最大値を教えてください。"
}

イベントJSONを変更したら、呼び出すボタンをクリックしてLambda関数を実行しましょう。

実行後の結果で、以下のようにS3にアップロードしたPDFファイルが参照され、想定されるレスポンスが返る事が確認できればOKです！

Lambda 関数で使用できるメモリの最大値は 10,240 MB(10GB)です。
Lambda 関数で使用できるメモリの量は、関数の実行時に割り当てられます。
このメモリ容量は関数の処理能力と直接関係しており、更に大きなメモリを必要とする関数ではこの値を上限の 10GB まで設定できます。

最後に

今回は、Bedrock、Lambda、Kendra及びS3を使用したRAGについて、SAMを使って実装してみました！

RAGを検証してみたいけど手頃に使用できる環境がない...という方がいらっしゃれば、ぜひ使ってみてください。

以上、つくぼし(tsukuboshi0755)でした！