🐡

【Bedrock】ナレッジベースへの直接取り込みを試してみた(IngestKnowledgeBaseDocuments)

2025/01/17に公開

はじめに

2024年12月1日のre:Inventで、Amazon Bedrockにおける新しいAPI、IngestKnowledgeBaseDocumentsorに関する発表がありました。
このアップデートにより、Bedrock APIを使って、ナレッジの追加から同期までを一貫して行えるようになりました。
https://aws.amazon.com/jp/blogs/aws/new-apis-in-amazon-bedrock-to-enhance-rag-applications-now-available/

従来はデータソース全体をナレッジベースと同期する必要がありましたが、
このアップデートで、特定のデータのみをナレッジベースに反映することが可能になりました。

今回は、データソースタイプS3で直接取り込みを試してみました。(表で言う右上)

https://docs.aws.amazon.com/ja_jp/bedrock/latest/userguide/kb-direct-ingestion-add.html

ということで早速試してみました。

前提条件

以下の設定がすでに完了していることを前提とします。

  • ナレッジベースのデータソースとしてS3を使用している
  • Pythonで開発したLambda関数を利用

書いたコード

ドキュメントをS3からナレッジベースへ取り込む関数

boto3のingest_knowledge_base_documentsを使用します。
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent/client/ingest_knowledge_base_documents.html

ingest_knowledge_base_documents.py
import boto3

client = boto3.client('bedrock')

BEDROCK_KNOWLEDGE_BASE_ID = "your-knowledge-base-id"
GOOD_KNOWLEDGE_DATA_SOURCE_ID = "your-data-source-id"

def ingest_document_s3(document_s3_uri):
    """
    ナレッジベースにドキュメントを取り込む
    """
    response = client.ingest_knowledge_base_documents(
        knowledgeBaseId=BEDROCK_KNOWLEDGE_BASE_ID,
        dataSourceId=GOOD_KNOWLEDGE_DATA_SOURCE_ID,
        documents=[
            {
                'content': {
                    'dataSourceType': 'S3',
                    's3': {
                        's3Location': {
                            'uri': document_s3_uri
                        }
                    }
                },
            },
        ],
    )
    return response

ドキュメントのステータスを確認する関数

取り込んだドキュメントのステータスを確認するための関数です。
boto3のget_knowledge_base_documentsを使用しました。
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent/client/get_knowledge_base_documents.html

get_knowledge_base_documents.py
def check_status(document_s3_uri):
    """
    ドキュメントのステータスを確認する
    """
    response = client.get_knowledge_base_documents(
        knowledgeBaseId=BEDROCK_KNOWLEDGE_BASE_ID,
        dataSourceId=GOOD_KNOWLEDGE_DATA_SOURCE_ID,
        documentIdentifiers=[
            {
                'dataSourceType': 'S3',
                's3': {
                    'uri': document_s3_uri,
                },
            },
        ],
    )
    return response['documentDetails'][0]['status']

Lambda関数の権限設定
Bedrock APIを利用するには、適切な権限をLambda関数に付与する必要があります。
以下のコードはCDKでLambda関数に権限を追加しているコード例です。

cdk.ts
lambda.addToRolePolicy(
  new iam.PolicyStatement({
    effect: iam.Effect.ALLOW,
    actions: [
      "bedrock:IngestKnowledgeBaseDocuments",
      "bedrock:GetKnowledgeBaseDocuments",
    ],
    resources: ["arn:aws:bedrock:region:account-id:resource-id"],
  })
);

試してみた結果

データソースタイプ "S3" でドキュメントを直接取り込むことができました。
取り込み後はcheck_status関数でステータスを確認することで、正常にナレッジベースに反映されたことを確認しました。

詰まったところ

新しい機能ですので、Lambdaがデフォルトでimportしたboto3のバージョンでは、ingest_knowledge_base_documentsやget_knowledge_base_documentsを使用できませんでした。
Lambda Layerを用意してboto3のバージョンを指定することで、解決できました。

Custom Connectorの使い所について

実際に直接取り込みをやってみたところ、ナレッジベースへの反映は4,5秒程度で終わっていました。
ナレッジベース全体を同期するより速さの面でメリットがありそうです。

また、先日Bedrock勉強会に参加し、Bedrock本[1]著者のみのるんさん[2]にも聞いてみました。
IOTなどリアルタイムでデータ収集して、ナレッジベースに即反映したい場合などに、インラインで直接取り込みをするというようなユースケースを教えていただきました。
確かに、S3等に保存する必要がなければ、インラインでの直接取り込みが有用そうです。

脚注
  1. https://www.amazon.co.jp/Amazon-Bedrock-生成AIアプリ開発入門-AWS深掘りガイド-御田/dp/4815626448 ↩︎

  2. https://zenn.dev/minorun365 ↩︎

NCDCエンジニアブログ

Discussion