見出し画像

AWS Bedrock と Aurora でRAGを作る


本記事の概要

前置き

AWS の Bedrock には Knowledge Bases という機能がある。
同機能を使えば手軽に RAG の仕組みを構築できる。

Bedrock 上で管理されている LLM や Embedding モデルを用いつつ、データストア(ベクトルDB)は別で準備する形となる。
そこら辺は下記でまとめてある。

上記資料に記載の事情から、ベクトルDBは OpenSearch でも Pinecone でもなく Aurora Serverless V2 PostgreSQL を用いることとなった。

本資料の記載範囲

実際に、Bedrockその他の機能を用いてRAG機能のベースを構築するところまでのアレヤコレヤをまとめる。
構築に際してはTerraformを用いており、Terraformのコードもここに記載する。


構成

すごく簡単に、構成図

  1. 専用S3バケット(データソース)を準備。ここにナレッジの素材となるファイルをアップロードする

  2. Knowledge Base は、同期処理を行うと同データソース内のデータを取得する

  3. そのデータをチャンクに分解し、Embeddingモデルを用いてベクトルデータ化する

  4. このベクトルデータをAurora Serverless V2 PostgreSQL内に作成した専用テーブルに格納する

  5. アプリケーションからの問い合わせ文字列に対して、ベクトルDB検索を行い、結果を返す

といった流れになる。DB接続時にはSecrets Manager に保存されている username, password を用いる。

参考にした記事

こちらを参考にさせてもらった。

公式も参照

構築

Terraform ベースで構築内容を説明

※前提条件

  • Aurora Serverless V2 PostgreSQL と S3 バケットは作成済みとする。

  • 開発作業の都合上、PostgreSQL はパブリック環境からの接続が可能(パブリックサブネットに配置)となっている。
    プライベートサブネット配備の場合にBedrockからPostgreSQLに接続するために追加でなにか設定が必要か否かは未確認

main.tf

main.tfは何もしていなくて、各リソース作成を呼び出したり引数を渡したりしているだけ。

# ##################################
# 基本設定
terraform {
  required_providers {
    aws = {
      source  = "hashicorp/aws"
      # バージョンが古いと Knowledge Bases に対応していないので注意
      version = "5.76.0"
    }
  }
}
provider "aws" {
  region = "ap-northeast-1"
}
# ##################################


# ##################################
# 各種リソース構築
# ##################################
# 作成済みのAurora内のDB構築
module "aurora_initialize_db" {
  source                      = "./modules/aurora_initialize_db"
}

# Bedrock から ベクトルDB(PostgreSQL)に繋ぐためのシークレット
module "secret" {
  source               = "./modules/secret"
}

# Bedrockで必要となるIAMロール
module "bedrock_iam" {
  source         = "./modules/bedrock_iam"
}

# Bedrock周り
module "bedrock" {
  source                  = "./modules/bedrock"
  secret_arn              = module.secret.db_conn_secret_arn
  bedrock_iam_arn         = module.bedrock_iam.bedrock_iam_role_arn
  bedrock_attach_policies = module.bedrock_iam.bedrock_attach_policies

  depends_on = [module.bedrock_iam]
}

aurora_initialize_db.tf

作成済みのAurora Serverless V2 PostgreSQL のDB, テーブル等の作成。
./modules/aurora_initialize_db.aurora_initialize_db.tf
※接続情報などは適宜外出しで変数管理すべし

ポイント

  • pgvector(PostgreSQLのベクトルストア機能)を有効化

  • Knowledge Base から利用されるベクトルストア用のテーブルを事前作成

# PostgreSQLに接続してDB作成
resource "null_resource" "create_db_sql" {
  provisioner "local-exec" {
    command = <<EOT
      psql -h <ホスト名> \
           -U <ユーザー名> \
           -d "postgres" \
           -c "
              DO \$\$
              BEGIN
                IF NOT EXISTS (SELECT 1 FROM pg_database WHERE datname = '<DB名>') THEN
                    PERFORM dblink_exec('dbname=postgres', 'CREATE DATABASE <DB名>');
                END IF;
              END
              \$\$;
           "
    EOT
    environment = {
      PGPASSWORD = <パスワード>
    }
  }
}


# 普通はアプリ専用のユーザーを作成して権限も付与すべきだけど省略


# PostgreSQLに接続してpgvector有効化
resource "null_resource" "create_extention_sql" {
  provisioner "local-exec" {
    command = <<EOT
      psql -h <ホスト名> \
           -U <ユーザー名> \
           -d "<DB名>" \
           -c "create extension if not exists vector;"
    EOT
    environment = {
      PGPASSWORD = <パスワード>
    }
  }
}

# PostgreSQLに接続して、
# BedrockからKnowledge Basesのベクトルストアとして利用するためのテーブルを作成
# テーブル名や列名は Knowledge Base 構築時にも用いる(同じ文字列を使う)必要あり
resource "null_resource" "create_kb_table" {
  provisioner "local-exec" {
    command = <<EOT
      psql -h <ホスト名> \
           -U <ユーザー名> \
           -d "<DB名>" \
           -c "
                CREATE TABLE IF NOT EXISTS bedrock_kb (
                  id         uuid PRIMARY KEY,
                  embedding  vector(1024),
                  chunks     text,
                  metadata   json
                );
           "
    EOT
    environment = {
      PGPASSWORD = <パスワード>
    }
  }
}

secret.tf

Knowledge Base から PostgreSQL に接続するための認証情報をシークレットマネージャに登録。
./modules/secret/secret.tf

# 本来はベタ書きしない(別途変数として管理すべし
locals {
  pg_conn_user_auth = {
    username   = <ユーザー名>
    password   = <パスワード>
  }
}

resource "aws_secretsmanager_secret" "db_conn_secret" {
  name = "db-conn-secret"
}

resource "aws_secretsmanager_secret_version" "latest" {
  secret_id = aws_secretsmanager_secret.db_conn_secret.id
  secret_string = jsonencode(local.pg_conn_user_auth)
}

# 後続の処理で利用する情報を output
output "db_conn_secret_arn" {
  value = aws_secretsmanager_secret.db_conn_secret.arn
}

bedrock_iam.tf

Knowledge Base で処理を実行するための専用IAMロールを作成。
※前述の構成図ではロール周りの記載を割愛
./modules/bedrock_iam/bedrock_iam.tf

# ロール作成
resource "aws_iam_role" "bedrock_exec_role_for_knowledgeBase" {
  name = "bedrock_exec_role_for_knowledgeBase"

  assume_role_policy = jsonencode({
    Version = "2012-10-17",
    Statement = [
      {
        Sid       = "AmazonBedrockKnowledgeBaseTrustPolicy",
        Effect    = "Allow",
        Principal = {
          Service = "bedrock.amazonaws.com"
        },
        Action = "sts:AssumeRole",
        Condition = {
          StringEquals = {
            # AWSアカウントID。通常は変数化して外出しすべし
            "aws:SourceAccount" = "XXXXXXXXXXXX"
          },
          ArnLike = {
            "aws:SourceArn" = "arn:aws:bedrock:ap-northeast-1:XXXXXXXXXXXX:knowledge-base/*"
          }
        }
      }
    ]
  })
}

# ポリシー作成(以下4つ分)
resource "aws_iam_policy" "bedrock_invoke_policy" {
  name   = "bedrock_invoke_policy"
  policy = jsonencode({
    Version = "2012-10-17",
    Statement = [
      {
        Sid    = "BedrockInvokeModelStatement",
        Effect = "Allow",
        Action = [
          "bedrock:InvokeModel"
        ],
        Resource = [
          "arn:aws:bedrock:ap-northeast-1::foundation-model/cohere.embed-multilingual-v3"
        ]
      }
    ]
  })
}

resource "aws_iam_policy" "bedrock_rds_policy" {
  name   = "bedrock_rds_policy"
  policy = jsonencode({
    Version = "2012-10-17",
    Statement = [
      {
        Sid    = "RdsDescribeStatementID",
        Effect = "Allow",
        Action = [
          "rds:DescribeDBClusters"
        ],
        Resource = [
          "<AuroraクラスターのARN>"
        ]
      },
      {
        Sid    = "DataAPIStatementID",
        Effect = "Allow",
        Action = [
          "rds-data:BatchExecuteStatement",
          "rds-data:ExecuteStatement"
        ],
        Resource = [
          "<AuroraクラスターのARN>"
        ]
      }
    ]
  })
}

resource "aws_iam_policy" "bedrock_s3_policy" {
  name   = "bedrock_s3_policy"
  policy = jsonencode({
    Version = "2012-10-17",
    Statement = [
      {
        Sid    = "S3ListBucketStatement",
        Effect = "Allow",
        Action = [
          "s3:ListBucket"
        ],
        Resource = [
          "<S3バケットARN>"
        ],
        Condition = {
          StringEquals = {
            # AWSアカウントID
            "aws:ResourceAccount" = ["XXXXXXXXXXXX"]
          }
        }
      },
      {
        Sid    = "S3GetObjectStatement",
        Effect = "Allow",
        Action = [
          "s3:GetObject"
        ],
        Resource = [
          "<S3バケットARN>/*"
        ],
        Condition = {
          StringEquals = {
            # AWSアカウントID
            "aws:ResourceAccount" = ["XXXXXXXXXXXX"]
          }
        }
      }
    ]
  })
}

resource "aws_iam_policy" "bedrock_secrets_manager_policy" {
  name   = "bedrock_secrets_manager_policy"
  policy = jsonencode({
    Version = "2012-10-17",
    Statement = [
      {
        Sid    = "SecretsManagerGetStatement",
        Effect = "Allow",
        Action = [
          "secretsmanager:GetSecretValue"
        ],
        Resource = [
          "<シークレットのARN>"
        ]
      }
    ]
  })
}

# ロールとポリシーを紐づけ
resource "aws_iam_role_policy_attachment" "bedrock_attach_policies" {
  for_each = {
    "bedrock_invoke_policy"          = aws_iam_policy.bedrock_invoke_policy.arn
    "bedrock_rds_policy"             = aws_iam_policy.bedrock_rds_policy.arn
    "bedrock_s3_policy"              = aws_iam_policy.bedrock_s3_policy.arn
    "bedrock_secrets_manager_policy" = aws_iam_policy.bedrock_secrets_manager_policy.arn
  }

  role       = aws_iam_role.bedrock_exec_role_for_knowledgeBase.name
  policy_arn = each.value
}

# 後続の処理で利用する情報を output
output "bedrock_iam_role_arn" {
  value = aws_iam_role.bedrock_exec_role_for_knowledgeBase.arn
}
output "bedrock_attach_policies" {
  value = aws_iam_role_policy_attachment.bedrock_attach_policies
}

bedrock.tf

いよいよKnowledge Base の作成。
./modules/bedrock/bedrock.tf

# ############################
# 前段処理で生成したリソースの値を利用
# ############################
variable "secret_arn" {}
variable "bedrock_iam_arn" {}
variable "bedrock_attach_policies" {}

########################################################
# Knowledge Base を Aurora(PostgreSQL)ベース、cohere の組み込みモデルベースで作成
########################################################
resource "aws_bedrockagent_knowledge_base" "this" {
  name     = "bedrock-kb"
  role_arn = var.bedrock_iam_arn

  knowledge_base_configuration {
    type = "VECTOR"
    vector_knowledge_base_configuration {
      # cohere が提供するEmbedding Model を使用
      embedding_model_arn = "arn:aws:bedrock:ap-northeast-1::foundation-model/cohere.embed-multilingual-v3"
    }
  }

  storage_configuration {
    type = "RDS"
    rds_configuration {
      credentials_secret_arn = var.secret_arn
      resource_arn           = "<AuroraクラスターのARN>"
      database_name          = "DB名"              # 事前作成済みのDB
      table_name             = "bedrock_kb"             # 事前作成済みのテーブル
      field_mapping {
        primary_key_field = "id"                        # テーブルの各カラム
        vector_field      = "embedding"
        text_field        = "chunks"
        metadata_field    = "metadata"
      }
    }
  }

  depends_on = [ var.bedrock_attach_policies ]
}

########################################################
# ナレッジベースが参照するデータソースのS3との紐づけ
########################################################
resource "awscc_bedrock_data_source" "this" {
  name              = "bedrock-kb-ds"
  knowledge_base_id = aws_bedrockagent_knowledge_base.this.id
  data_source_configuration = {
    type = "S3"
    s3_configuration = {
      bucket_arn = "<S3バケットのARN>"
    }
  }
}


確認

これで、前述の構成図のような構成で環境が構築される。
実際にKnowledge Base を見てみる。

AWSコンソール

ナレッジベースが作られている

IAMロールやDBやシークレット、S3、Embeddingモデルなどの情報が紐づけられている。


補足

同期処理

AWSコンソール上に「同期」ボタンがあり、このボタンを押せばS3上のファイルがナレッジベースに同期されるが、スケジュール機能やリアルタイム同期機能はない。
→追加実装が必要

解説👇️

Bedrock が使えない

AWSコンソール上で実際に動作確認しようとすると、エラーが出た。

Your request rate is too high. Reduce the frequency of requests.

各モデルを有効化していても出る。

原因はQuota。デフォルトではQuotaが0に設定されていて、有効化しても使えない。

ここで各Quotaを変更すればいいのだが、AWS側のトラブルで24年11月現在、変更ができないケースがある。
その場合は、問い合わせが必要とのこと。

チャンキング戦略

今回の構築ではChunkの設定は行っていないが、カスタマイズ可能

いいなと思ったら応援しよう!