📜

AWS上でRAG環境を雑に構築してみた手順

に公開

本記事の前提

  • 「とりあえず動けばええわ」精神で環境を構築しているので、セキュリティ周りなど、すべてにおいて雑です。
  • 手元のメモと記憶を頼りに書いているので、ちょくちょく間違っているかもです。
    • 本記事を参考に環境を構築してみて、詰まったりしたらコメントいただけると嬉しいです。

モチベーション

  • 以前、ちょくちょく「それ、プロンプトで頑張ってチャットボットを作るのではなく、RAG環境を構築した方がいいんじゃね?」というお仕事の相談を受けることがあった。
  • RAG環境を構築したことがなかったのと、興味がわいてきたのでやってみたくなった。

構築手順

1. データソース用のS3バケット準備

  1. 任意の名前でS3バケットを作成する。
    • パブリックアクセスはブロックでOK。
  2. PDFなどのデータソース作成用ファイルを格納するディレクトリを作成する。(今回はdocsという名前で作成。)
  3. 2のディレクトリに適当なPDFファイルをアップロードする。

2. Bedrockで使用するロールを作成

  • 以下の信頼ポリシー設定JSONを使用し、任意の名前でロールを作成する。
信頼ポリシー
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "Service": "bedrock.amazonaws.com"
            },
            "Action": "sts:AssumeRole",
            "Condition": {
                "StringEquals": {
                    "aws:SourceAccount": "<アカウントID>"
                },
                "ArnLike": {
                    "AWS:SourceArn": "arn:aws:bedrock:ap-northeast-1:<アカウントID>:knowledge-base/*"
                }
            }
        }
    ]
}

3. Bedrockでナレッジベースを作成

言及していない項目は初期値のままでOK。

  1. Unstructured dataでナレッジベースを作成する。(今回は表形式ではないPDFドキュメントをソースに使うため。)
  2. IAM許可に作成済みのBedrockで使用するロールを指定する。
  3. 次へをクリックする。
  4. S3のURIにS3のデータソースファイルを指定する。
    • ファイルを直指定でも、フォルダ指定でも良い。
    • フォルダ指定の場合は、末尾に/を入れるのを忘れないこと。
  5. 次へをクリックする。
  6. 埋め込みモデルにTitan Text Embeddings V2を選択する。
  7. ベクトルデータベースに新しいベクトルストアをクイック作成、ベクトルストアにAmazon Aurora PostgreSQL Serverlessを選択する。
  8. 次へをクリックする。
  9. ナレッジベースを作成をクリックする。

あとは、ナレッジベースの作成が終わるまで待つ。

4. 作成済みのBedrock用に作成したロールの更新

以下の設定JSONを元に、作成済みのBedrock用に作成したロールへインラインポリシーを追加する。

インラインポリシー
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "BedrockModelAccess",
      "Effect": "Allow",
      "Action": [
        "bedrock:ListFoundationModels",
        "bedrock:ListCustomModels",
        "bedrock:InvokeModel"
      ],
      "Resource": [
        "arn:aws:bedrock:<リージョン>::foundation-model/amazon.titan-embed-text-v2:0",
        "arn:aws:bedrock:<リージョン>::foundation-model/cohere.embed-multilingual-v3"
      ]
    },

    {
      "Sid": "S3ListBucket",
      "Effect": "Allow",
      "Action": ["s3:ListBucket"],
      "Resource": ["arn:aws:s3:::<作成したS3のバケット>"]
    },
    {
      "Sid": "S3GetObject",
      "Effect": "Allow",
      "Action": ["s3:GetObject"],
      "Resource": ["arn:aws:s3:::<作成したS3のバケット>/docs/*"]
    },

    {
      "Sid": "RdsDescribeCluster",
      "Effect": "Allow",
      "Action": ["rds:DescribeDBClusters"],
      "Resource": ["arn:aws:rds:<リージョン>:<アカウントID>:cluster:<データソースのDBクラスター名>"]
    },
    {
      "Sid": "RdsDataApi",
      "Effect": "Allow",
      "Action": [
        "rds-data:ExecuteStatement",
        "rds-data:BatchExecuteStatement"
      ],
      "Resource": ["arn:aws:rds:<リージョン>:<アカウントID>:cluster:<データソースのDBクラスター名>"]
    },

    {
      "Sid": "RetrieveAndGenerate",
      "Effect": "Allow",
      "Action": ["bedrock:RetrieveAndGenerate", "bedrock:Retrieve"],
      "Resource": "arn:aws:bedrock:<リージョン>:<アカウントID>:knowledge-base/<ナレッジベースID>"
    }
  ]
}

5. 動作確認用Lambda作成

  1. 以下のPythonコードをコピペし、動作確認用のLambda関数を作成する。
動作確認用Lambdaコード
import json, os
import boto3

KB_ID = os.environ["KB_ID"]                      # 例: kb-xxxxxxxx
MODEL_ARN = os.environ["MODEL_ARN"]              # 例: arn:aws:bedrock:ap-northeast-1::foundation-model/amazon.titan-text-lite-v1
REGION = os.environ.get("AWS_REGION", "ap-northeast-1")

br = boto3.client("bedrock-agent-runtime", region_name=REGION)

def _resp(status, body):
    return {
        "statusCode": status,
        "headers": {
            "Content-Type": "application/json; charset=utf-8",
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Headers": "content-type",
            "Access-Control-Allow-Methods": "OPTIONS,POST,GET"
        },
        "body": json.dumps(body, ensure_ascii=False)
    }

def lambda_handler(event, context):
    # Function URL は HTTP イベント互換
    method = event.get("requestContext", {}).get("http", {}).get("method", "GET")
    if method == "OPTIONS":
        return _resp(204, {})

    if method == "GET":
        return _resp(200, {"ok": True, "message": "healthcheck"})

    # POST: { "question": "...", "sessionId": "optional" }
    try:
        body = json.loads(event.get("body") or "{}")
        question = body.get("question")
        session_id = body.get("sessionId")
        if not question:
            return _resp(400, {"error": "question is required"})

        cfg = {
            "type": "KNOWLEDGE_BASE",
            "knowledgeBaseConfiguration": {
                "knowledgeBaseId": KB_ID,
                "modelArn": MODEL_ARN
            }
        }
        if session_id:
            # 会話継続(任意)
            cfg["knowledgeBaseConfiguration"]["sessionId"] = session_id

        res = br.retrieve_and_generate(
            input={"text": question},
            retrieveAndGenerateConfiguration=cfg,
        )

        answer = res.get("output", {}).get("text", "")
        citations = []
        for c in res.get("citations", []):
            for ref in c.get("retrievedReferences", []):
                # 可能ならソースの URI/パスを返す
                uri = (
                    ref.get("location", {}).get("s3Location", {}).get("uri")
                    or ref.get("location", {}).get("webLocation", {}).get("url")
                    or ref.get("location", {}).get("type")
                )
                title = ref.get("metadata", {}).get("title") or ref.get("content", {}).get("text", "")[:80]
                citations.append({"uri": uri, "title": title})

        out = {
            "answer": answer,
            "citations": citations,
            "sessionId": res.get("sessionId")
        }
        return _resp(200, out)

    except Exception as e:
        return _resp(500, {"error": str(e)})

  1. Lambdaに割当たっているロールへインラインポリシーを追加する。
Lambdaインラインポリシー
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "AllowKBQuery",
      "Effect": "Allow",
      "Action": [
        "bedrock:Retrieve",
        "bedrock:RetrieveAndGenerate"
      ],
      "Resource": "arn:aws:bedrock:<リージョン>:<アカウントID>:knowledge-base/<ナレッジベースID>"
    }
  ]
}
  1. Lambdaへ環境変数を設定する。
    3-1. KB_IDへナレッジベースIDを設定する。
    3-2. MODEL_ARNへ arn:aws:bedrock:ap-northeast-1::foundation-model/anthropic.claude-3-5-sonnet-20240620-v1:0 を設定する。

6. 動作確認

以下のペイロードを使用して、Lambdaをテスト実行する。

ペイロード
{
  "version": "2.0",
  "routeKey": "$default",
  "rawPath": "/",
  "rawQueryString": "",
  "headers": { "content-type": "application/json" },
  "requestContext": {
    "http": {
      "method": "POST",
      "path": "/",
      "protocol": "HTTP/1.1",
      "sourceIp": "127.0.0.1",
      "userAgent": "LambdaConsoleTest"
    }
  },
  "isBase64Encoded": false,
  "body": "{\"question\":\"任意の聞きたい内容\"}"
}

やってみた所感

  • 簡単、本当に簡単。
  • データソースサイズに依存するのだろうけども、Lambda経由で質問して回答が来るまで約7秒だったので思ったよりはやい。
  • 今回触った範囲の料金周りよくわからん。
    • 設定時間アクセスが無ければ自動で停止状態になるRDS Auroraインスタンスって何。
    • データソース更新1回ごとにいくら掛かるの?データ量に依存する?
    • などなど
  • 今回触った範囲の権限周りよくわからん。

今後やりたいこと・課題

  • Terraformを使ってテンプレ化したい。
  • 権限周りの設定が雑すぎるので、調べなおして本運用できそうなレベルまで調整したい。
  • 運用費を正確に算出したい。
  • topkや温度など、回答の調整周りを試す。
  • S3へ大元のドキュメントをアップロードするたびに自動でデータソースを更新するようにしたい。

最後に

今後やりたいこと・課題で書いたことを調べたり試してみた結果を、続編として記事化するかも。

Discussion