🍣

RAG実装①：Amazon Bedrock,Kendra(LangChain不使用)

2023/11/26に公開

AWS

はじめに

BedrockとKendraをつかってRAG(Retrieval-augmented Generation)を作ってみる。
今回は、LangChainを使っていません。

構成図

このような構成にします。

S3バケットには格納したサンプルデータはgithubにあります。

ファイル構成

以下のように作ります

iam0001:~/environment/rag $ tree
.
├── Dockerfile_RAG
├── dockerbuild-run.sh
├── dockerstop-imgdell.sh
└── work
    └── rag-NotUseLangchein.py

1 directory, 4 files

スクリプト

スクリプトは以下です。
Kendraの実行結果でヒットしたドキュメントの中の0番目のドキュメントの内容をBedrockで要約します。

rag-NotUseLangchein.py

# Bedrock + Kendra を組み合わせる

import boto3
import json
import datetime

query = 'おばあさんの杖'  # クエリ
index_id = 'YOUR_INDEX_ID'  # KendraのインデックスID
modelId = "anthropic.claude-v2"  # BedrockのModel
region_name = 'us-east-1'  # リージョンの指定

# Kendraクライアントを作成
kendra = boto3.client('kendra', region_name=region_name)

# Boto3 セッションの作成
session = boto3.Session()

# Bedrock クライアントの作成
bedrock_client = session.client(service_name='bedrock')
bedrock_runtime_client = session.client(service_name='bedrock-runtime')

def search_kendra(indexId, query):
    kendra = boto3.client('kendra')

    response = kendra.retrieve(
        QueryText=query,
        IndexId=indexId,
        AttributeFilter={
            "EqualsTo": {
                "Key": "_language_code",
                "Value": {"StringValue": "ja"},
            },
        },
    )
    
    return response

def generate_response(modelId, documents):
    # モデルの同期的呼び出し（Anthropic Claude-v2 の場合）
    prompt = f"Human:『${documents}』の『から』までの内容を100文字に要約して\nAssistant:"

    body = json.dumps({
        "prompt": prompt,
        "max_tokens_to_sample": 300,
        "temperature": 0.1,
        "top_k": 1,
        "stop_sequences": ["\n\nHuman"]
    })
    
    accept = "application/json"
    content_type = "application/json"
    
    response = bedrock_runtime_client.invoke_model(
        body=body,
        modelId=modelId,
        accept=accept,
        contentType=content_type
    )
    
    response_body = json.loads(response.get('body').read())

    return response_body


def log(message):
    # 現在の日時を取得
    now = datetime.datetime.now()
    # メッセージと日時を表示
    print(f'{now}: {message}')


def rag(indexId, modelId, query):
    log(f"query = {query}")
    log(f"--- Kendra Start---")
    documents = search_kendra(indexId, query)  # Kendraの処理
    log(f"--- Kendra End---")
    log(f"Kendra Result : {documents}\n")
    targetString = documents['ResultItems'][0]['Content']  # Kendraからの応答結果から一部抜き取る
    log(f"--- Bedrock Start---")
    response = generate_response(modelId, targetString)  # Kendra応答結果(抜粋版)をBedrockで処理する
    log(f"--- Bedrock End---")
    log(f"Bedrock Result : {response}\n")

    return response

response = rag(indexId, modelId, query)
print(response)

実行結果

クエリ『おばあさんの杖』で実行した結果です。それっぽい結果が出力されました。

root@d531ad4fa825:/work# python rag-NotUseLangchein.py 
2023-11-26 04:16:43.120501: query = おばあさんの杖
2023-11-26 04:16:43.120559: --- Kendra Start---
2023-11-26 04:16:43.488952: --- Kendra End---
2023-11-26 04:16:43.490458: Kendra Result : {'QueryId': '11111111-2222-3333-4444-555555555555', 'ResultItems': [{'Id': '11111111-2222-3333-4444-555555555555-66666666-7777-8888-9999-000000000000', 'DocumentId': 's3://hogefugapiyo-us-east-1/momo-taro/桃太郎物語_アイテム_おばあさんの杖.txt', 'DocumentTitle': '桃太郎物語_アイテム_おばあさんの杖.txt', 'Content': 'アイテム名：おばあさんの杖 クラス: 癒し手専用 レベル: 4 攻撃力: 1/10 魔法　　※省略※
2023-11-26 04:16:43.490957: --- Bedrock Start---
2023-11-26 04:16:44.655657: --- Bedrock End---
2023-11-26 04:16:44.655713: Bedrock Result : {'completion': ' おばあさんの杖は、物理攻撃力はほとんどないが、魔法力が高く、特に癒しの魔法を効果的に使えるレベル4の癒し手専用アイテム。おばあさんの優しさと癒しの力を象徴している。', 'stop_reason': 'stop_sequence', 'stop': '\n\nHuman'}

{'completion': ' おばあさんの杖は、物理攻撃力はほとんどないが、魔法力が高く、特に癒しの魔法を効果的に使えるレベル4の癒し手専用アイテム。おばあさんの優しさと癒しの力を象徴している。', 'stop_reason': 'stop_sequence', 'stop': '\n\nHuman'}
root@d531ad4fa825:/work#

はじめに

構成図

ファイル構成

スクリプト

実行結果

Discussion