LangChain AmazonKendraRetriever から必要な情報を返してもらえない人へ

## はじめに
こんにちは、KDDIアジャイル開発センターのはしもと（仮名）です。

LLM を使用した独自アプリケーション開発において、**Retrieval-augmented Generation（以下、RAG）** はもはや必須の技術となりました。

:::note info
**RAGとは？**

LLM 内部で保有していない情報（最新のウェブ検索結果や社内文書検索結果）を関連する文脈として与えることで、モデルの応答精度や関連性を向上させる技術のこと
:::

ここでは、Amazon Kendra (以下、Kendra) と Amazon Bedrock (以下、Bedrock) で簡易的な RAG アーキテクチャの実装を試した際、自分がハマった箇所を中心に備忘も兼ねて簡単に解説します。

## 使用するサンプルドキュメント
検索対象として準備したファイルの内容は以下のとおりです
- ファイル名：`PTA集会のお知らせ.pdf`
- ファイル形式：PDF
- S3 URI：`s3://<BUCKET_NAME>/PTA集会のお知らせ.pdf`

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/463957/169da201-4951-2821-0ee6-d37a6f12e2d4.png)

このドキュメントの内容について検索を実行します。

## つまづいたこと
RAG を実装する際は、検索精度を向上させることも必要ですが、提示された回答が正しいかを利用者が判断できるような仕組みがあるべきです。

後述する `RetrievalQA` において、Kendra を retriever に指定した際、関連ドキュメントのタイトル (`Title`) や抜粋 (`Excerpt`) は取得できましたが、保存場所や URL といった情報は応答に含まれていませんでした。

<details><summary>うまくいかなかったコードと実行結果</summary><div>

```python
from langchain.prompts import PromptTemplate
from langchain.llms.bedrock import Bedrock

import langchain

langchain.verbose = True

AWS_REGION = "ap-northeast-1"
MAX_TOKENS = 400
MODEL_NAME = "anthropic.claude-instant-v1"

LLM = Bedrock(
    region_name=AWS_REGION,
    model_id=MODEL_NAME,
    model_kwargs={
        "temperature": 0,
        "max_tokens_to_sample": MAX_TOKENS,
    },
    verbose=True,
)


def _create_prompt() -> PromptTemplate:
    PROMPT_TEMPLATE = """

    Human:
    ユーザーからの質問に回答してください。
    回答生成時に必要な情報が不足している場合は、以下の情報を使用してください。
    出力には、質問への回答を含めた上で、追加の情報を使用した場合は [ページタイトル](URL) の形式でこれらを含めてください。

    Given Information:
    {context}

    Question: {question}

    Assistant:
    """
    prompt = PromptTemplate(
        input_variables=["context", "question"], template=PROMPT_TEMPLATE
    )
    return prompt
```

```python
from langchain.retrievers import AmazonKendraRetriever
from langchain.retrievers.kendra import clean_excerpt


INDEX_ID = "xxxxxxx"
retriever = AmazonKendraRetriever(
    index_id=INDEX_ID,
    attribute_filter={
        "EqualsTo": {
            "Key": "_language_code",
            "Value": {"StringValue": "ja"},
        },
    },
)
```

```python
from langchain.chains import RetrievalQA

PROMPT = _create_prompt()
chain_type_kwargs = {"prompt": PROMPT}
qa = RetrievalQA.from_chain_type(
    llm=LLM,
    chain_type="stuff",
    retriever=retriever,
    chain_type_kwargs=chain_type_kwargs,
)
```

```python
query = "PTA集会に参加したい場合はどうすればいい？"
qa.run(query)

"""
// ここにはファイルの場所が `DocumentId`/`DocumentURI` として含まれる
> Entering new RetrievalQA chain...
{'Id': 'xxxxxxx', 'DocumentId': 's3://<BUCKET_NAME>/<DOC_NAME>.pdf', 'DocumentURI': 'https://<BUCKET_NAME>.s3.ap-northeast-1.amazonaws.com/<DOC_NAME>.pdf', 'DocumentAttributes': [DocumentAttribute(Key='_source_uri', Value=DocumentAttributeValue(DateValue=None, LongValue=None, StringListValue=None, StringValue='https://<BUCKET_NAME>.s3.ap-northeast-1.amazonaws.com/<DOC_NAME>.pdf')), DocumentAttribute(Key='s3_document_id', Value=DocumentAttributeValue(DateValue=None, LongValue=None, StringListValue=None, StringValue='PTA集会のお知らせ.pdf')), DocumentAttribute(Key='_excerpt_page_number', Value=DocumentAttributeValue(DateValue=None, LongValue=1, StringListValue=None, StringValue=None))], 'DocumentTitle': '<DOC_NAME>.pdf', 'Content': '架空小学校 PTA集会のお知らせ 日時:\t2023 年 11 月 5 日\t(土)\t10:00〜12:00 場所:\t架空小学校\t多目的ホール 内容: 1. 校長先生より学校の現状報告 2. PTA 活動の報告と今後の方針説明 3. 各クラス代表からの情報共有 4. 質疑応答 参加を希望される方は、2023 年 10 月 30 日までに担任の先生へ回答してください。 皆様のご参加を心よりお待ちしております。', 'ScoreAttributes': {'ScoreConfidence': 'NOT_AVAILABLE'}}

> Entering new StuffDocumentsChain chain...

// ここで失われる
> Entering new LLMChain chain...
Prompt after formatting:
    
    Human:
    ユーザーからの質問に回答してください。
    回答生成時に必要な情報が不足している場合は、以下の情報を使用してください。
    出力には、質問への回答を含めた上で、追加の情報を使用した場合は [ページタイトル](URL) の形式でこれらを含めてください。

    Given Information:
    Document Title: PTA集会のお知らせ
Document Excerpt: 
架空小学校 PTA集会のお知らせ 日時: 2023 年 11 月 5 日 (土) 10:00〜12:00 場所: 架空小学校 多目的ホール 内容: 1. 校長先生より学校の現状報告 2. PTA 活動の報告と今後の方針説明 3. 各クラス代表からの情報共有 4. 質疑応答 参加を希望される方は、2023 年 10 月 30 日までに担任の先生へ回答してください。 皆様のご参加を心よりお待ちしております。


    Question: PTA集会に参加したい場合はどうすればいい？

...

> Finished chain.
 PTA集会に参加したい場合は、文書から以下の通りです。
[PTA集会のお知らせ](https://example.com/pta-meeting)に記載されている通り、2023年10月30日までに担任の先生へ参加の意向を回答してください。参加を希望される方は、この期日までに担任の先生に参加する旨を伝えてください。

"""
```
</div></details>

LLM を通して得られた出力は、**質問に対する回答としては正しいです。**
一方、参照したドキュメントへのリンクとして出力された URL は間違っており、ステップごとに見るとこの情報は途中で破棄されていることが分かります。
これでは、**回答した内容が正しいのか間違っているのか、ユーザーはすぐに判断できません。**

`RetrievalQA` を用いる場合でも、参照したドキュメントへのリンクを出力したい。
どうすればよいでしょうか。

## Retriever の使用方法による出力結果の違い
解決策を探る前に、Retriever の挙動を確認します。
ここでは、Retriever 単体で検索を行うケースと、`RetrievalQA` などを使って処理フローの1ステップとして利用する2つのパターンを考えます

### `AmazonKendraRetriever` について
LLM を使用するアプリにおいて、Kendra を Retriever（ユーザーからの入力文やタスクの実行に必要な関連情報を大量のデータから抽出する部分）として使用するには、LangChain の `AmazonKendraRetriever` が便利です。

https://python.langchain.com/docs/integrations/retrievers/amazon_kendra_retriever

### Retriever 単独での検索 
まずはこれを使用して関連するドキュメントの検索を行ってみます。
検索実行時に設定できるパラメータもいくつかありますが、最低限のものならば、以下のような少ないコード量で簡単に実現できます。

```python
import boto3
from langchain.retrievers import AmazonKendraRetriever

retriever = AmazonKendraRetriever(index_id="xxxxxxxxx")

retriever.get_relevant_documents("<QUESTION>")

"""
{'page_content': 'Document Title: xxx\nDocument Excerpt: \nxxx\n',
 'metadata': {'result_id': 'xxx',
  'document_id': 's3://xxx/xxx.pdf',
  'source': 'https://xxx/xxx.pdf',
  'title': 'xxx',
  'excerpt': 'xxx',
  'document_attributes': {'_source_uri': 'https://xxx',
   '_excerpt_page_number': x}},
 'type': 'Document'}
"""
```

Kendra のインデックスに、S3上のドキュメントを同期した状態で上記のコマンドを実行すると、関連するドキュメントのタイトル・本文からの抜粋以外に、`metadata`としてドキュメントの S3 URI や抜粋が含まれたページ番号などの情報も含まれています。

期待した検索が行えることを知りたいだけであれば、このように retriever を直接呼び出せばよいのですが、最終的に行いたいことは、**質問文に対する適切な回答を人間に分かりやすい文章で返す**ことです。

そのため、実際のアプリ開発においては、retriever 単独で使用する機会はあまりなく、後述の **RetrievalQA** など、なにかしらの chain と組み合わせて利用することが多いイメージです。

### Chain 内での Retriever の利用
先ほど説明した一連の処理をフローとして実行するのに便利なのが `RetrievalQA` です。
名前のとおり、インデックスへの検索結果に基いた質問応答を簡単に実装できるクラスです。

https://python.langchain.com/docs/use_cases/question_answering/vector_db_qa#custom-prompts

`RetrievalQA` ではプロンプトをカスタマイズして使用することもできます。
用途の一例ですが、出力が特定のフォーマットになるように指示を与えたり、質問への回答だけでなく参照したドキュメントへの導線も出力させる、といったことが可能になります。

しかし前述のように、`RetrievalQA` を使用した場合には、`source` などのメタデータはコンテキストとして LLM に与えられていません。


## 回避策
### AmazonKendraRetriever の Docs を読んだ

`AmazonKendraRetriever` は、裏側では Kendra の API を呼び出しているだけのようなので、公式ドキュメントに記載がないかを探したところ見つかりました。

どうやら `page_content_formatter` というプロパティが応答にどのような属性を含むかを制御しているようです。
デフォルトだとタイトルと抜粋のみが使用されると書いてあるので間違いなさそうです。

> page_content_formatter – generates the Document page_content allowing access to all result item attributes. By default, it uses the item’s title and excerpt.

https://api.python.langchain.com/en/latest/retrievers/langchain.retrievers.kendra.AmazonKendraRetriever.html#:~:text=kendra/latest/APIReference-,page_content_formatter,-%E2%80%93%20generates%20the%20Document

### AmazonKendraRetriever のコードを読んだ
念のためソースコードも確認します。

`page_content_formatter` の値はデフォルト値が設定されており、通常は、
```
Document Title: <TITLE>\n"Document Excerpt: \n<EXCERPT>\n
```
のような文字列になることが分かります。

https://github.com/langchain-ai/langchain/blob/master/libs/langchain/langchain/retrievers/kendra.py#L27-L44

### `page_content_formatter` に指定する関数の作成

全体のフォーマットを変更することも可能ですが、今回は単純に、もともとの出力の後ろに、`Document URI: <DOCUMENT_URI>\n` を連結するだけとしました。

```python
def original_content_formatter(item):
    text = ""
    title = item.get_title()
    if title:
        text += f"Document Title: {title}\n"
    excerpt = clean_excerpt(item.get_excerpt())
    if excerpt:
        text += f"Document Excerpt: \n{excerpt}\n"
    // 以下2行を追加    
    if item.DocumentURI:
        text += f"Document URI: {item.DocumentURI}\n"
    return text
```


## 修正後の実行結果
先ほど作成した `original_content_formatter` を使用するようにした上で、再度 chain を実行してみます。

```python
retriever = AmazonKendraRetriever(
    index_id=INDEX_ID,
    attribute_filter={
        "EqualsTo": {
            "Key": "_language_code",
            "Value": {"StringValue": "ja"},
        },
    },
    page_content_formatter=original_content_formatter, // 追加
)

qa = RetrievalQA.from_chain_type(
    llm=LLM,
    chain_type="stuff",
    retriever=retriever,
    chain_type_kwargs=chain_type_kwargs,
)

query = "PTA集会に参加したい場合はどうすればいい？"
qa.run(query)


"""
> Entering new RetrievalQA chain...


> Entering new StuffDocumentsChain chain...


> Entering new LLMChain chain...
Prompt after formatting:


    Human:
    ユーザーからの質問に回答してください。
    回答生成時に必要な情報が不足している場合は、以下の情報を使用してください。
    出力には、質問への回答を含めた上で、追加の情報を使用した場合は [ページタイトル](URL) の形式でこれらを含めてください。

    Given Information:
    Document Title: PTA集会のお知らせ
Document Excerpt: 
架空小学校 PTA集会のお知らせ 日時: 2023 年 11 月 5 日 (土) 10:00〜12:00 場所: 架空小学校 多目的ホール 内容: 1. 校長先生より学校の現状報告 2. PTA 活動の報告と今後の方針説明 3. 各クラス代表からの情報共有 4. 質疑応答 参加を希望される方は、2023 年 10 月 30 日までに担任の先生へ回答してください。 皆様のご参加を心よりお待ちしております。
Document URI: https://s3.ap-northeast-1.amazonaws.com/<S3_BUCKET_NAME>/<DOC_NAME>.pdf


    Question: PTA集会に参加したい場合はどうすればいい？
...

> Finished chain.
PTA集会に参加したい場合は、文書から以下の通り参加を希望する方は2023年10月30日までに担任の先生へ回答してください、とあるため、担任の先生に参加希望の意思表示を行う必要があります。
[PTA集会のお知らせ](https://s3.ap-northeast-1.amazonaws.com/<S3_BUCKET_NAME>/<DOC_NAME>.pdf)

```

正しいドキュメント URL が結果に含まれるようになりました！

まだまだプロンプトに工夫できる箇所はたくさんあるので、引き続き改善していきます。

## 費用について
Amazon Kendra、個人で利用するにしてはなかなかいいお値段です。

費用の大部分を占めるであろう「インデックス」については、30日間で最大750時間分を無料で利用することができますが、
検証用途向けに設定された「Developer Edition」モードでも、一ヶ月放置するだけで **810 USD** 発生します。

実際に同様の検証を行われる際はご注意ください。

https://aws.amazon.com/jp/kendra/pricing/


## まとめ

RAG の実装には、Kendra のような全文検索エンジンを使う以外にも、OpenSearch Service や Pinecone といったベクトル検索に対応したデータベースをインデックスとして使用するなどいくつかパターンがあります。

主観ですが、後者はオープンソースで公開されているものも多く、導入にあたってのハードルは低いものの、ベクトル検索だけでは十分な検索精度が得られないケースがあります。

前者は、費用はかかる（まじでつらい）ものの、大手クラウド事業者が1サービスとして提供していることも多く、多機能・高性能（機械学習アルゴリズムの組み込み、外部データソースとのコネクタ、一般的なファイル形式のネイティブ対応）で、それぞれにメリット・デメリットはあります。

費用対効果を見極めて、適切なアーキテクチャ設計をしたいですね！ハイ！

----

<blockquote class="twitter-tweet"><p lang="ja" dir="ltr">我が家の和室には「Kendra 一日 一時間」と書かれた掛け軸があります</p>&mdash; はしもと（仮名） (@s3kzk) <a href="https://twitter.com/s3kzk/status/1714251914723090723?ref_src=twsrc%5Etfw">October 17, 2023</a></blockquote> <script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script>

「Kendra 一日 一時間」
昔はそんなこともあったねと、仲間たちと肩を組んで笑い会える日が来るといいな。

はしもと（仮名）でした。