Knowledge Base for Amazon Bedrock を🦜️🔗LangChain から呼び出す

この記事は セゾン情報システムズ Advent Calendar 2023 21 日目の記事です。

Knowledge Base for Amazon Bedrock とは

AWS re:Invent 2023 で Knowledge Base for Amazon Bedrock の一般提供開始がアナウンスされました🎉

大規模言語モデルに対し、外部検索によって得られた知識を組み合わせて回答精度の向上をおこなう RAG (Retrieval-Augmented Generation) という手法が一般化してきています。

Knowledge Base for Amazon Bedrock を使用すると、ユーザーは任意の S3 バケットに検索対象とさせたいオブジェクトを配置するだけで、Amazon Bedrcok 上でマネージドな RAG の構成を簡単に実装できます。Azure OpenAI Service で提供されている Add your data に近い機能です。

詳しい動作イメージはドキュメントの How it works 参照ください。

作成した Knowledge Base を利用する方法としては同日に一般的提供開始となった Agents for Amazon Bedrock を使用することが多いかと思いますが、本記事では LangChain から利用する方法についてまとめています。

Knowledge base の作成

最初に今回の検証で使用する Knowledge base を作成します。Bedrock コンソールの Orchestration から Knowledge base を選択し、Create knowledge base へ進みます。

Step 1 で Knowledge base の名前および詳細を入力します。ここではすべてデフォルト値のまま次へ進みます。サービスロールも新規作成を選択します。

Step 2 でデータソースの配置先となる S3 バケットを指定します。オブジェクトがカスタマー管理の KMS キーで暗号化されている場合は Add customer-managed KMS key for S3 data にチェックを入れます。

画面下部の Advanced settings では Knowledge Base for Amazon Bedrock がデータを Embedding 処理をする際に一時的にデータを暗号化するための KMS キーの指定および、ドキュメントのチャンク分割戦略 (Chunking strategy) を指定できます。

Chunking strategy は以下の 3 つから選択することができます。指定した戦略に従ってドキュメントのチャンク分割を行い、チャンクごとに Embedding が生成されます。

チャンクはベクトル検索結果の精度にも影響しますが、データソースの追加後に Chunking strategy を変更することはできないため注意する必要があります。

• Default chunking
  各チャンクに最大 200 個のトークンが含まれるように自動的にチャンクを作成します。

• Fiexed size chunking
  ユーザーが指定した Max tokenss サイズを超えない範囲でチャンクを作成します。

• No chunking
  Knowledge Base for Amazon Bedrock はチャンクを作成しません。ユーザー独自の前処理でドキュメント分割をおこなっている場合に選択します。

今回はデフォルトのまま次へ進みます。

Step 3 でベクターストアの設定を行います。Embedding 用のモデルは 2023/12/21 時点で Amazon Titan の Embeddings G1 - Text のみ利用できます。Vector database セクションで新しいベクターストア (Amazon OpenSearch Serverless) をクイック作成するか、Knowledge Base for Amazon Bedrock がサポートするユーザー管理のベクターストアを選択できます。

ここでは Quick create a new vector store - Recommended を選択して次に進みます。作成される OpenSearch Serverless のコレクションをカスタマー管理の KMS キーで暗号化したい場合は、Add customer-managed KMS key for Amazon OpenSearch Serverless vector にチェックをいれます。

ここで作成される Amazon OpenSearch Serverless のコレクションは Knowledge base の削除を行っても自動で削除されることはありません。費用を抑えたい場合はユーザー自身でコレクションの削除を行う必要があります。

Step 4 でこれまでの入力内容を確認し、画面下部の Create knowledge base をクリックします。作成の完了までにはしばらくかかります。

Knowledge base の作成が完了したら最初に S3 バケット上に配置したデータを同期する必要があるため、Sync をクリックします。

今回はあらかじめ PDF 版 のAmazon Bedrock のユーザーガイド (英語) を S3 バケットに格納しています。

新規作成されるサービスロールにはTitan Embbedings モデルを Invoke するためのポリシー、OpenSearch Serverless コレクションにアクセスするためのポリシー、データソースの S3 バケットにアクセスするためのポリシーがそれぞれ作成、設定されていました。

OpenSearch Serverless のコレクションもベクトル検索タイプのものが作成され、自動で埋め込みされたベクトルデータが保存されていることを確認できます。

また Serverless コレクションに対して、サービスロールおよびリソース作成者をプリンシパルとしてデータアクセスポリシーが設定されています。

サービスロールを新規作成せず、独自のものを指定した場合は自身でデータアクセスポリシーを編集する必要があります

コンソールの Test Knowledge base から動作確認をおこなうことも可能です。

以上で Knowledge base の作成および設定の確認が完了しました。

🦜️🔗LangChain から呼び出す

AmazonKnowledgeBasesRetriever

LangChain は LLM を活用したアプリケーションを開発支援する OSS のフレームワークです。LangChain が提供する機能の一つに Retriever と呼ばれるユーザーの入力などの非構造化クエリを指定してドキュメントを返すインターフェースがあります。

コミュニティによって様々なサービスの Retriever が開発されており、AWS のサービスではAmazon Kendra に対応した AmazonKendraRetriever と Knowledge Base for Amazon Bedrock に対応した AmazonKnowledgeBasesRetriever を利用できます。

Knowledge base を使用する際は RetrieveAndGenerate API を使用するか Knowledge base を Agents for Amazon Bedrock に追加してドキュメントの検索～回答の生成までを Bedrock で完結せることができます。

一方 AmazonKnowledgeBasesRetriever では Retrieve API を使用してドキュメントの検索のみを行います。

あえて LangChain から Knowledge base を使用することのメリットとしては以下のようなものが考えられます。

• Amazon Bedrock で提供されていない大規模言語モデルを使用して回答を生成できる
• LangChain の豊富な機能やエコシステムを利用できる

もちろん逆もしかりで、本来 Amazon Bedrock の機能のみで完結できる内容をわざわざ LangChain で実装する必要はなく、運用管理などのデメリットも発生します。

サンプルコード

以下のバージョンで動作確認しています

• boto3: v1.34.4
• LangChain: v0.0.352

以下のような少量のコードで Knowledge Base の情報をもとに回答を生成することができます。knowledge_base_id には作成した Knowledge base ID を指定します。

from langchain.chains import RetrievalQA
from langchain.chat_models import BedrockChat
from langchain.retrievers import AmazonKnowledgeBasesRetriever

def main():
    llm = BedrockChat(
        model_id="anthropic.claude-v2:1",
        model_kwargs={
            "temperature":0,
            "max_tokens_to_sample":1024
        }
    )
    
    retriever = AmazonKnowledgeBasesRetriever(
        knowledge_base_id="SNXTFPISRP",
        retrieval_config={
            "vectorSearchConfiguration": {
                "numberOfResults": 4
            }
        }
    )

    qa = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type='stuff',
        retriever=retriever,
        verbose=True
    )
    query = "Does Amazon Bedrock support PrivateLink?"
    result = qa.run(query)
    print(result)

if __name__ == "__main__":
    main()

実行結果は以下です。「Amazon Bedrock は PrivateLink をサポートしていますか？」という質問に対してそれっぽい回答を生成できています。

 Yes, Amazon Bedrock supports PrivateLink. Here are some key points about using PrivateLink with Amazon Bedrock:

- You can create an interface endpoint in your VPC to privately access Amazon Bedrock APIs. This allows you to access Bedrock without needing public IP addresses or going over the public internet.

- The supported service names for Amazon Bedrock interface endpoints are:
  - com.amazonaws.region.bedrock
  - com.amazonaws.region.bedrock-runtime 
  - com.amazonaws.region.bedrock-agent-runtime

- Interface endpoints are created in specific Availability Zones, so you need to list the supported AZs before creating the endpoint. 

- You can attach an endpoint policy to control access to Bedrock through the interface endpoint.

- Traffic to Bedrock flows through endpoint network interfaces created in your subnets, allowing private access from your VPC.

So in summary, yes PrivateLink via interface endpoints allows private connectivity from your VPC to Amazon Bedrock without traversing the internet.

Chainlit でサンプルアプリを作ってみる

せっかくなら UI も欲しいというこで、Chainlit を使ったサンプルアプリを作ってみます。Chainlit は Chat GPT ライクなアプリケーションを簡単に構築できるオープンソースの Python パッケージです。

以下がサンプルコードです。ConversationalRetrievalChain を使用しているので会話の履歴も踏まえて質問の言い換えを行い、ドキュメントを検索した結果で回答を生成するはずです。

• chainlit: v0.7.700 で動作確認しています
• chainlit run app.py の実行でアプリケーションが起動します

app.py

import os
from langchain.chat_models import BedrockChat
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ChatMessageHistory, ConversationBufferMemory
from langchain.retrievers import AmazonKnowledgeBasesRetriever
import chainlit as cl

aws_region = os.environ["AWS_REGION"]

@cl.on_chat_start
async def main():
    llm = BedrockChat(
        model_id="anthropic.claude-v2:1",
        model_kwargs={
            "temperature":0,
            "max_tokens_to_sample":1024
        }
    )

    message_history = ChatMessageHistory()
    
    retriever = AmazonKnowledgeBasesRetriever(
        knowledge_base_id="SNXTFPISRP",
        retrieval_config={
            "vectorSearchConfiguration": {
                "numberOfResults": 4
            }
        }
    )

    memory = ConversationBufferMemory(
        memory_key="chat_history",
        output_key="answer",
        chat_memory=message_history,
        return_messages=True,
    )
    
    chain = ConversationalRetrievalChain.from_llm(
        llm,
        chain_type="stuff",
        retriever=retriever,
        memory=memory,
        return_source_documents=True,
        verbose=True
    )
    # Store the chain in the user session
    cl.user_session.set("chain", chain)

@cl.on_message
async def main(message: cl.Message):
    # Retrieve the chain from the user session
    chain = cl.user_session.get("chain")

    # Call the chain asynchronously
    res = await chain.acall(
        message.content, 
        callbacks=[cl.AsyncLangchainCallbackHandler()]
    )
    answer = res["answer"]
    source_documents = res["source_documents"]  # type: List[Document]

    text_elements = []  # type: List[cl.Text]

    if source_documents:
        for source_idx, source_doc in enumerate(source_documents):
            source_name = f"source_{source_idx}"
            # Create the text element referenced in the message
            text_elements.append(
                cl.Text(content=source_doc.page_content, name=source_name)
            )
        source_names = [text_el.name for text_el in text_elements]

        if source_names:
            answer += f"\nSources: {', '.join(source_names)}"
        else:
            answer += "\nNo sources found"

    await cl.Message(content=answer, elements=text_elements).send()

以下が実行イメージです。

今度は「Titan Image Generator が、ユーザーがアップロードした画像を元にして新しい画像を生成できるか」という最新かつ、ちょっと難易度高めの質問をぶつけてみました。モデルは結局 Bedrock 上 の Claude 2.1 を使用しており、Claude に適したプロンプトのカスタマイズは行っていませんが、ドキュメントの内容をもとに的確に回答しているようにみえます。

以上です。
参考になれば幸いです。