エムオーテックス株式会社が運営するテックブログです。

CloudFormation で Amazon Bedrock Knowledge Bases を構築してみた - Auth0 Changelog 分析システムを事例に

CloudFormation で Amazon Bedrock Knowledge Bases を構築してみた - Auth0 Changelog 分析システムを事例に

はじめに

こんにちは、アプリケーションチームの辻です。

LANSCOPE エンドポイントマネージャー クラウド版(以下、クラウド版)では認証基盤として Auth0 を利用しています。
Auth0 の機能変更や仕様変更がクラウド版に影響を与える可能性があるため、定期的に Changelog を確認する必要があります。
しかし、Auth0 の Changelog を手動で調査するのは大変なので自動化したいと考えていました。

いつもの調査では、Changelog の日本語翻訳・要約・クラウド版への影響を確認しています。
これらを自動化するために Amazon Bedrock Knowledge Bases(以下、Knowledge Base)を活用して「Auth0 Changelog 分析システム」を構築しました。

本記事では、このシステムの核となる Knowledge Base を CloudFormation で構築した話を中心に紹介します。

本題に入る前に「なぜ Knowledge Base を選択したのか」と「Auth0 Changelog 分析システムの全体アーキテクチャ」を話しておきます。

なぜ Knowledge Base を選択したのか

日本語翻訳・要約するだけなら Bedrock だけで十分です。
しかし、クラウド版への影響まで調査するとなると、Auth0 の機能ごとの使用有無や設定情報と照らし合わせる必要があります。

kiro cli のような対話型ツールも検討しました。
しかし、Changelog が更新されたときに自動で調査してほしかったため、手動での質問が必要なツールでは要件を満たせませんでした。

RAG システムを構築する方法として、OpenSearch + Bedrock や RDS + Bedrock での組み合わせも検討しました。
これらの場合はベクトル検索と LLM との統合部分を自分で実装する必要があり、開発・運用コストが高くなります。
Knowledge Base なら RAG 全体がマネージドサービスとして提供されるため、開発・運用コストを大幅に削減できることから採用しました。

また、ベクトルストアには S3 Vectors を選択しました。
Auth0 のドキュメントデータは各機能の使用有無を記録したシンプルな CSV 形式です。
OpenSearch Serverless と比較して、S3 Vectors はコスト効率が良く、データの規模に対しても検索性能で十分であるため、最適だと判断し採用しました。

Auth0 Changelog 分析システムの全体アーキテクチャ

構築した Knowledge Base は、Auth0 Changelog 分析システム全体の中でどのような位置づけになるのかを紹介します。

実際の分析システムでは、Auth0 Changelog の RSS 監視、API Gateway、SQS、Slack 通知など多くのコンポーネントで構成されています。
その中で、今回 CloudFormation で実装した Knowledge Base が中核的な役割を担っています。

Auth0 Changelog 分析システムの全体アーキテクチャ図

次章からは、この Knowledge Base 部分の CloudFormation 実装について詳しく解説していきます。

CloudFormation での Knowledge Base 実装

ここからが本記事のメイン部分です。
CloudFormation を使って Knowledge Base を構築する際のポイントを詳しく解説します。

S3 Vectors の設定

Auth0 の機能ごとの使用有無や設定情報を CSV 形式で作成します。
今回は「カテゴリ、サブカテゴリ、項目名、使用有無、使用箇所」など、Auth0 の各機能がクラウド版で実際に利用されているかを記録した CSV データを用意しました。
これらをドキュメント格納用の S3 バケットに配置し、Knowledge Base のデータソースとして活用します。

ここでは、そのドキュメント格納用 S3 バケットとベクトルデータを格納するための S3 Vectors を設定します。

# Knowledge Baseが読み込むソースドキュメントを配置するためのS3バケット
DocumentsBucket:
  Type: AWS::S3::Bucket
  Properties:
    BucketName: 'my-documents-bucket'
    PublicAccessBlockConfiguration:
      BlockPublicAcls: true
      BlockPublicPolicy: true
      IgnorePublicAcls: true
      RestrictPublicBuckets: true

# Knowledge Baseが生成したベクトル(埋め込み)を格納するためのS3ベクトルバケット
VectorBucket:
  Type: AWS::S3Vectors::VectorBucket
  Properties:
    VectorBucketName: 'my-vectors-bucket'
    EncryptionConfiguration:
      # AES256: AWS管理キーによる暗号化を採用
      # 作成後の暗号化方式変更は不可のため慎重に選択
      SseType: AES256

# ベクトル検索用のインデックス
VectorIndex:
  Type: AWS::S3Vectors::Index
  Properties:
    IndexName: 'my-vector-index'
    VectorBucketArn: !GetAtt VectorBucket.VectorBucketArn
    # 精度とストレージサイズのバランスを考慮
    DataType: float32
    # Titan Embed Text v2の出力次元数に合わせる必須設定
    Dimension: 1024
    # cosine: テキスト検索に最適な類似度計算方式を採用
    DistanceMetric: 'cosine'
    EncryptionConfiguration:
      SseType: AES256
    MetadataConfiguration:
      # Bedrockが自動付与するメタデータキーをフィルタリング対象外に設定
      # 2KBのメタデータ制限を回避するため、ほぼ必須の設定
      NonFilterableMetadataKeys:
        - AMAZON_BEDROCK_TEXT      # テキストチャンク本体
        - AMAZON_BEDROCK_METADATA  # システムメタデータ

IAM ロールの設計

Knowledge Base が必要なリソースにアクセスできるよう、適切な権限を持つ IAM ロールを作成します。

# Knowledge BaseがS3 Vectorsにアクセスするための権限を定義
BedrockS3VectorsPolicyForKB:
  Type: AWS::IAM::ManagedPolicy
  Properties:
    ManagedPolicyName: 'BedrockS3VectorsPolicy'
    PolicyDocument:
      Version: '2012-10-17'
      Statement:
        # S3 Vectorsへのアクセス権限
        - Effect: Allow
          Action:
            - 's3vectors:GetIndex'
            - 's3vectors:QueryVectors'
            - 's3vectors:PutVectors' 
            - 's3vectors:GetVectors'
            - 's3vectors:DeleteVectors'
          Resource: !Ref VectorIndex

# Knowledge Baseが各種リソースにアクセスするための権限を定義
KnowledgeBaseRole:
  Type: AWS::IAM::Role
  Properties:
    RoleName: 'KnowledgeBaseRole'
    AssumeRolePolicyDocument:
      Version: '2012-10-17'
      Statement:
        - Effect: Allow
          Principal:
            Service: bedrock.amazonaws.com
          Action: sts:AssumeRole
    Policies:
      # S3データソースへのアクセス権限
      - PolicyName: S3DataSourcePolicy
        PolicyDocument:
          Version: '2012-10-17'
          Statement:
            - Effect: Allow
              Action:
                - s3:GetObject
                - s3:ListBucket
                - s3:GetBucketLocation
              Resource:
                - !GetAtt DocumentsBucket.Arn
                - !Sub '${DocumentsBucket.Arn}/*'
      # 埋め込みモデルへのアクセス権限
      - PolicyName: BedrockModelPolicy
        PolicyDocument:
          Version: '2012-10-17'
          Statement:
            - Effect: Allow
              Action:
                - bedrock:InvokeModel
              Resource: !Sub 'arn:aws:bedrock:${AWS::Region}::foundation-model/amazon.titan-embed-text-v2:0'
    ManagedPolicyArns:
      # S3 Vectorsへのアクセス権限(上記で定義したポリシーを参照)
      - !Ref BedrockS3VectorsPolicyForKB

Knowledge Base 本体の定義

作成した IAM ロールと S3 Vectors を使用して、Knowledge Base 本体を定義します。

# Knowledge Base本体
KnowledgeBase:
  Type: AWS::Bedrock::KnowledgeBase
  # IAMポリシーの作成完了を明示的に待つ
  DependsOn: BedrockS3VectorsPolicyForKB
  Properties:
    Name: 'my-knowledge-base'
    Description: 'Knowledge base for document analysis and impact assessment'
    # 上記で作成したIAMロールを指定
    RoleArn: !GetAtt KnowledgeBaseRole.Arn
    KnowledgeBaseConfiguration:
      # ベクトル検索ベースのKnowledge Baseを採用
      Type: VECTOR
      VectorKnowledgeBaseConfiguration:
        # 埋め込みモデルを明示的に指定(Titan Embed Text v2)
        EmbeddingModelArn: !Sub 'arn:aws:bedrock:${AWS::Region}::foundation-model/amazon.titan-embed-text-v2:0'
    StorageConfiguration:
      # S3 Vectorsをベクトルストアとして使用
      Type: S3_VECTORS
      S3VectorsConfiguration:
        # 上記で作成したベクトルインデックスとバケットを指定
        IndexArn: !Ref VectorIndex
        VectorBucketArn: !GetAtt VectorBucket.VectorBucketArn

Data Source の設定

最後に、S3 バケットからドキュメントを読み込むための Data Source を設定します。

# Data Source: Knowledge BaseがS3バケットからドキュメントを読み込むための設定
DataSource:
  Type: AWS::Bedrock::DataSource
  Properties:
    # 上記で作成したKnowledge BaseのIDを指定
    KnowledgeBaseId: !Ref KnowledgeBase
    Name: 'my-knowledge-base-data-source'
    Description: 'S3 data source for system documents'
    DataSourceConfiguration:
      # S3: S3バケットをデータソースとして使用
      Type: S3
      S3Configuration:
        # ドキュメント格納用バケットを指定
        BucketArn: !GetAtt DocumentsBucket.Arn
        InclusionPrefixes:
          # 特定のプレフィックス配下のファイルのみを対象
          - 'documents/'
    VectorIngestionConfiguration:
      ChunkingConfiguration:
        # 一旦固定サイズでのチャンキング戦略を採用
        ChunkingStrategy: FIXED_SIZE
        FixedSizeChunkingConfiguration:
          # 検索精度と処理効率のバランスを考慮
          MaxTokens: 800
          # チャンク間の重複により文脈の連続性を保持
          OverlapPercentage: 15

Lambda から Knowledge Base への質問

CloudFormation で Knowledge Base を構築した後、Lambda から実際に質問する方法を紹介します。

通常の Bedrock モデルを使用する場合はinvoke_modelconverse API を使用しますが、Knowledge Base の場合はretrieve_and_generate API を使用します。
この API で RAG(検索 + 生成)を一度に実行でき、Knowledge Base に格納された Auth0 設定情報を参照しながら分析が可能になります。

import boto3

client = boto3.client('bedrock-agent-runtime', region_name='ap-northeast-1')

def query_knowledge_base(query_text: str, knowledge_base_id: str) -> str:
    response = client.retrieve_and_generate(
        input={'text': query_text},
        retrieveAndGenerateConfiguration={
            'type': 'KNOWLEDGE_BASE',
            'knowledgeBaseConfiguration': {
                'knowledgeBaseId': knowledge_base_id,
                'modelArn': 'arn:aws:bedrock:ap-northeast-1::foundation-model/anthropic.claude-3-5-sonnet-20240620-v1:0'
            }
        }
    )
    return response['output']['text']

ちなみに: Knowledge Base の自動同期も実装

Knowledge Base の構築だけでなく、ドキュメントの自動同期も実装しました。

S3 のドキュメントが更新された後、Knowledge Base のベクトルデータを S3 Vectors に同期する必要があります。
手動で Ingestion Job を実行するのは面倒なので、buildspec.yaml 内に書いて自動化しました。

pre_build:
  commands:
    # S3にドキュメントをアップロード
    - aws s3 cp ./documents s3://my-documents-bucket/documents --recursive
    
    # CloudFormation ExportからIDを取得
    - |
      KNOWLEDGE_BASE_ID=$(aws cloudformation list-exports \
        --query "Exports[?Name=='my-knowledge-base-id'].Value" \
        --output text)
      DATA_SOURCE_ID=$(aws cloudformation list-exports \
        --query "Exports[?Name=='my-data-source-id'].Value" \
        --output text)
      
      # Ingestion Jobを実行してKnowledge Baseを更新
      aws bedrock-agent start-ingestion-job \
        --knowledge-base-id $KNOWLEDGE_BASE_ID \
        --data-source-id $DATA_SOURCE_ID

これにより、デプロイ時に自動的にドキュメントが更新され、Knowledge Base に反映されます。

動かしてみた

Knowledge Base を活用することで、以下のようにクラウド版の影響有無の分析もできるようになりました。
(※セキュリティ上の理由により、分析結果の詳細部分にはモザイク処理を施しています)

Before(単純な Bedrock のみ)

Auth0 の機能ごとの使用有無や設定情報が分からないため、日本語翻訳・要約が限界です。

単純な Bedrock 使用時のSlackキャプチャ

After(Knowledge Base 使用)

Auth0 の機能ごとの使用有無や設定情報と照らし合わせて、影響有無の分析まで実施できるようになりました。

Knowledge Base 使用時のSlackキャプチャ

おわりに

Amazon Bedrock Knowledge Base を活用することで、単純な翻訳・要約を超えた、文脈を考慮した影響分析の自動化を実現できました。
また、CloudFormation で Infrastructure as Code として管理することで、環境間での設定統一やバージョン管理も実現できています。

今後の改善予定として、プロンプトの改善やチャンキング戦略の見直しなど、さらなるチューニングを行って精度向上を目指したいと考えています。

この仕組みは他の SaaS サービスにも応用可能です。
AWS Service Health Dashboard などの AWS サービスやその他の外部サービスの変更監視にも同様のアプローチが活用できそうなので、今後の改善ネタとして取り組んでみたいと思います。

最後までお読みいただき、ありがとうございました。
本記事がお役に立てれば幸いです。