【AWS/Bedrock/RAG】PythonからKnowledge BaseをAPIで利用する

やりたいこと

Cloud9からRAG構成のAmazon Bedrockを利用する。尚、KnowledgeBaseによるS3及びベクトルDBは事前に作成/配置済みとする。

環境情報

$ python3 --version
Python 3.8.16

やり方

作成したKnowledge Baseから以下を値を確認した上で

• KnowledgeBaseId
• DataSourceId="XX"

下記のコードを実行する。

import boto3

KnowledgeBaseId="XX"
DataSourceId="XX"
Text = "hogeの好きな食べ物は？"
ModelArn = "arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-v2"
my_session = boto3.Session(profile_name='default')
client = my_session.client('bedrock-agent-runtime', region_name="us-east-1")
response = client.retrieve_and_generate(
        input={
            'text': Text
        },
        retrieveAndGenerateConfiguration={
            'knowledgeBaseConfiguration': {
                'knowledgeBaseId': KnowledgeBaseId,
                'modelArn': ModelArn,
                'retrievalConfiguration': {
                    'vectorSearchConfiguration': {
                        'overrideSearchType': 'SEMANTIC'
                    }
                }
            },
            'type': 'KNOWLEDGE_BASE'
        }
    )
print(response["output"]["text"],flush=False)

すると標準出力に、S3に登録した文書を参考にして

hogeの好きな食べ物は肉です

といった文字列が出力される。

以下、補足です。

補足

S3のデータソースには以下のようなテキストファイルを配置しました。

好きな食べ物
hoge : 肉
fuga : 魚
piyo : 牛乳

上記に記載の出力部分をprint(response)のみとすると、以下のようにソースに当たる箇所も出力されます。

{'ResponseMetadata': {'RequestId': 'xx', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'xx', 'content-type': 'application/json', 'content-length': '579', 'connection': 'keep-alive', 'x-amzn-requestid': 'xx'}, 'RetryAttempts': 0}, 'citations': [{'generatedResponsePart': {'textResponsePart': {'span': {'end': 17, 'start': 0}, 'text': 'hogeの好>きな食べ物は肉です。'}}, 'retrievedReferences': [{'content': {'text': '好きな食べ物\r hoge : 肉\r fuga : 魚\r piyo : 牛乳'}, 'location': {'s3Location': {'uri': 'xx'}, 'type': 'S3'}, 'metadata': {'x-amz-bedrock-kb-source-uri': 'xx', 'x-amz-bedrock-kb-data-source-id': 'xx'}}]}], 'output': {'text': 'hogeの好きな食べ物は肉です。'}, 'sessionId': 'xx'}

以上、ご参考ください。