Knowledge base for Amazon Bedrock と LangChainでRAG実装

Knowledge base for Amazon BedrockがGAしたので、KendraではなくKnowledge base for Amazon Bedrockを使用したRAGをLangChainを使って実装してみます。
インデックスしてしまえばLangChainからはKendraとほぼ同じ使用感で使用でき、チャンクサイズの自由度が高いのと、Agents for Amazon Bedrockから「も」呼べるのが良いところかなと思います。あと料金ですかね

必要ライブラリのアップデート

boto3とLangChainを最新にします。
LangChainを動作確認したバージョンは0.0.345です。

pip install -U boto3
pip install -U langchain==0.0.345

Knowledge base の作成

マネコンのAmazon BedrockのOrchestrationのKnowledge baseのCreate knowledge baseから作成します。
名前を変えるかどうか程度でわりとデフォルトでいけるのですが、Set up data sourceでドキュメントが格納されたS3を指定します。それとAddional settings - optionalから、ベクトルDBのチャンクサイズとチャンク間でオーバーラップさせるパーセンテージを指定できます。チャンクサイズがどうあるべきかはユースケースとストアするドキュメント次第かなと思いますが、Claudeの入力可能トークンの大きさを頼ってとりあえず大きめにします（しました）。
おおまかに1000トークン弱でPDF1ページぐらいかな？という感じなので、この設定であれば3-4ページ分の大きさでチャンクを作ってくれるはずです。後はデフォルトで作成します（OpenSearchにベクトルDBが作成されます。不要になればマネコンのOpenSearchから削除してください）。

作成が終わったらSyncを実行する事でS3がクロールされてベクトル化されます。
私はテスト用に以下のPDFをS3に格納、Syncしました。

作成したKnowledge baseのKnowledge base overviewの右上から、Knowledge base IDをコピーしておきます。

RAGアプリの構築

まず最小のRAGを作ってみます。
knowledge_base_idには各自のKnowledge base IDを設定します。

実際には上述のチャンクサイズやベクトルDBからの取得件数はチューニング必要かなと思います。今の設定ではベクトルの類似度が高いチャンク（PDFの3-4ページごと）を10件、つまり、PDFから類似度が高いページを30-40ページ程度引っ張ってきて、参考ドキュメントとしてプロンプトに渡している動きになります。金額が嵩みそうなのでLLMはClaude Instantを使用しています。

from langchain.llms import Bedrock
from langchain.chains import RetrievalQA
from langchain.retrievers.bedrock import AmazonKnowledgeBasesRetriever

# クエリ内容を設定
user_input="bedrockで使用可能なモデルIDは"

#各自のKnowledge base ID
knowledge_base_id = "XXXXXXXXXX"

# ベクトルDBからの取得件数
retrieval_config={"vectorSearchConfiguration": {"numberOfResults": 10}}

# LLMのモデルID
model_id="anthropic.claude-instant-v1" #Claude Instant
#model_id="anthropic.claude-v2:1" #Claude2.1

# 最大出力トークン数
model_kwargs={"max_tokens_to_sample": 1000} #Claude系の場合

# Retriever(KnowledgeBase)の定義
retriever = AmazonKnowledgeBasesRetriever(knowledge_base_id=knowledge_base_id,retrieval_config=retrieval_config)

# LLMの定義
llm = Bedrock(model_id=model_id,model_kwargs=model_kwargs)

# Chainの定義
qa = RetrievalQA.from_chain_type(retriever=retriever,llm=llm)

# chainの実行
result = qa.run(user_input)
print(result)

実行結果

Amazon Bedrock で現在使用可能な主なモデルとそのモデルIDは以下の通りです。

- Titan Text G1-Express 8K (アマゾン): titan-text-express-v1:0:8キロ
- Titan Embeddings G1-Text (アマゾン): titan-embed-text-v1:2:8キロ
- Anthropic Claude V2 18K: anthropic.claude-v2:0:18k
- Anthropic Claude V2 100K: anthropic.claude-v2:0:100k
- Anthropic Claude Instant V1 100K: anthropic.claude-instant-v1:2:100メートル
- Stable Diffusion XL 1.0: stable-diffusion-xl-v1:0
- Meta Llama 2 Chat 13B: meta.llama2-13b-chat-v1:0:4キロバイト

これらのモデルIDは、プロビジョニングされたスループットを作成する際に使用します。
また、ベースモデルやカスタムモデルを指定する際の参考情報となります。

新しいモデルがサポートされる場合は、Amazon Bedrock ドキュメントやコンソールからその情報を
確認ください。

元のドキュメントの日本語が怪しめなのでRAG結果の日本語も怪しめですが、きちんと元PDFから取得されています。

プロンプトを指定する場合は以下のように書きます。プロンプトを指定した方が生成内容が安定するような気がします。

from langchain.llms import Bedrock
from langchain.chains import RetrievalQA
from langchain.retrievers.bedrock import AmazonKnowledgeBasesRetriever
from langchain.prompts import PromptTemplate

# クエリ内容を設定
user_input="bedrockで使用可能なモデルIDは"

#各自のKnowledge base ID
knowledge_base_id = "XXXXXXXXXX"

# ベクトルDBからの取得件数
retrieval_config={"vectorSearchConfiguration": {"numberOfResults": 10}}

# LLMのモデルID
model_id="anthropic.claude-instant-v1" #Claude Instant
#model_id="anthropic.claude-v2:1" #Claude2.1

# 最大出力トークン数
model_kwargs={"max_tokens_to_sample": 1000} #Claude系の場合

# Retriever(KnowledgeBase)の定義
retriever = AmazonKnowledgeBasesRetriever(knowledge_base_id=knowledge_base_id,retrieval_config=retrieval_config)

# LLMの定義
llm = Bedrock(model_id=model_id,model_kwargs=model_kwargs)

# promptの定義
prompt_template = """
  <documents>タグには参考文書が書かれています。
  <documents>{context}</documents>
  \n\nHuman: 上記参考文書を元に、<question>に対して1000文字以内で説明してください。言語の指定が無い場合は日本語で答えてください。但し固有名詞と思われる単語は英語のままとしてください。
    もし<question>の内容が参考文書に無かった場合は「文書にありません」と答えてください。
  <question>{question}</question>
  \n\nAssistant:"""
PROMPT = PromptTemplate(template=prompt_template, input_variables=["context", "question"])

# Chainの定義
chain_type_kwargs = {"prompt": PROMPT}
qa = RetrievalQA.from_chain_type(retriever=retriever,llm=llm,chain_type_kwargs=chain_type_kwargs)

# chainの実行
result = qa.run(user_input)
print(result)

Streamlitでガワを被せるとこんな感じです。

from langchain.llms import Bedrock
from langchain.chains import RetrievalQA
from langchain.retrievers.bedrock import AmazonKnowledgeBasesRetriever
from langchain.prompts import PromptTemplate

#各自のKnowledge base ID
knowledge_base_id = "XXXXXXXXXX"

# ベクトルDBからの取得件数
retrieval_config={"vectorSearchConfiguration": {"numberOfResults": 10}}

# LLMのモデルID
model_id="anthropic.claude-instant-v1" #Claude Instant
#model_id="anthropic.claude-v2:1" #Claude2.1

# 最大出力トークン数
model_kwargs={"max_tokens_to_sample": 1000} #Claude系の場合

# Retriever(KnowledgeBase)の定義
retriever = AmazonKnowledgeBasesRetriever(knowledge_base_id=knowledge_base_id,retrieval_config=retrieval_config)

# LLMの定義
llm = Bedrock(model_id=model_id,model_kwargs=model_kwargs)

# promptの定義
prompt_template = """
  <documents>タグには参考文書が書かれています。
  <documents>{context}</documents>
  \n\nHuman: 上記参考文書を元に、<question>に対して1000文字以内で説明してください。言語の指定が無い場合は日本語で答えてください。但し固有名詞と思われる単語は英語のままとしてください。
    もし<question>の内容が参考文書に無かった場合は「文書にありません」と答えてください。
  <question>{question}</question>
  \n\nAssistant:"""
PROMPT = PromptTemplate(template=prompt_template, input_variables=["context", "question"])

# Chainの定義
chain_type_kwargs = {"prompt": PROMPT}
qa = RetrievalQA.from_chain_type(retriever=retriever,llm=llm,chain_type_kwargs=chain_type_kwargs)

# Streamlitでガワを被せる
import streamlit as st

st.title("LangChainからのKnowledge base for Amazon Bedrock呼出し")
input_text = st.text_input("このテキストをBedrockに送信します")
send_button = st.button("送信")

if send_button:

  # chainの実行
  result = qa.run(input_text)

  # 結果をStreamlitに出力
  st.write(result)