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から削除してください）。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3573242/8c6a0f22-0e91-fbe8-0d45-f653edaf01ab.png)

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

https://docs.aws.amazon.com/ja_jp/bedrock/latest/userguide/bedrock-ug.pdf#what-is-bedrock

作成した`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を使用しています。

```Python
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から取得されています。

プロンプトを指定する場合は以下のように書きます。プロンプトを指定した方が生成内容が安定するような気がします。
```Python
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でガワを被せるとこんな感じです。
```Python
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)
```

https://qiita.com/cyberBOSE/items/fe66c2049aaa04568d61