【Bedrock】ナレッジベースへの直接取り込みを試してみた(IngestKnowledgeBaseDocuments)
はじめに
2024年12月1日のre:Inventで、Amazon Bedrockにおける新しいAPI、IngestKnowledgeBaseDocumentsorに関する発表がありました。
このアップデートにより、Bedrock APIを使って、ナレッジの追加から同期までを一貫して行えるようになりました。
従来はデータソース全体をナレッジベースと同期する必要がありましたが、
このアップデートで、特定のデータのみをナレッジベースに反映することが可能になりました。
今回は、データソースタイプS3で直接取り込みを試してみました。(表で言う右上)
ということで早速試してみました。
前提条件
以下の設定がすでに完了していることを前提とします。
- ナレッジベースのデータソースとしてS3を使用している
- Pythonで開発したLambda関数を利用
書いたコード
ドキュメントをS3からナレッジベースへ取り込む関数
boto3のingest_knowledge_base_documentsを使用します。
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を使用しました。
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関数に権限を追加しているコード例です。
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等に保存する必要がなければ、インラインでの直接取り込みが有用そうです。

NCDC株式会社( ncdc.co.jp/ )のエンジニアチームです。 募集中のエンジニアのポジションや、採用している技術スタックの紹介などはこちら( github.com/ncdcdev/recruitment )をご覧ください! ※エンジニア以外も記事を投稿することがあります
Discussion