Kendra retrieveAPIで返戻のページ数のデフォルト値を変える方法

結論

デフォルトでは10件が返戻されますが、以下のように、PageSizeを指定することで最大値100個までパッセージを取得することが出来ます。

以下の例ではPageSize＝30とすることで、最大30件分のパッセージをresponseに含めるように設定しています。

import boto3
(省略)

kendra = boto3.client('kendra', region_name="ap-northeast-1")

(省略)
    response = kendra.retrieve(
        QueryText=query_text,
        IndexId=index_id,
        AttributeFilter={
            "EqualsTo": {
                "Key": "_language_code",
                "Value": {"StringValue": "ja"},
            },
        },
        PageSize=30,
    )
    # Kendraの応答から上位30件の結果を抽出 
    results = response['ResultItems'][:30] if response['ResultItems'] else []
    
(省略)

背景

KendraとBedrockを利用したRAG実装を試みた際のこと。
RAGの実装に際して、AWSのKendraに対してLambdaから検索をかけるRetrieve部分の処理を実装する際に、検索結果のパッセージを何件取得するかプログラム側で制御するのですが、当初返戻が10件以上得られず上限値を変える方法を少し調べることとなりました。
ドキュメントを読み込むことも重要かとは思いますが、情報を見つけにくい部分もあったため、誰かの助けとなればと思い本記事を執筆しました。

Query APIとRetrieve API について

Kendraでは、検索結果を取得する方法としてQueryAPIとRetrieve APIの二通りがあります。
これらの二つのAPIでは、出来ることや検索結果の取得方法が異なります。詳細についてはこちらの記事が参考になりました。
https://dev.classmethod.jp/articles/rag-knowledge-on-real-projects/

RAGの実装では、基本的に Retrieve APIの利用が推奨されています。Retrieve APIについては返戻件数のデフォルトの値は10件となっており、上記の通りPageSizeを指定することで10件以上の返戻を得ることができます。

※Query API のドキュメントは以下。
https://docs.aws.amazon.com/kendra/latest/APIReference/API_Query.html

※Retrieve APIのドキュメントは以下。
https://docs.aws.amazon.com/kendra/latest/APIReference/API_Retrieve.html

まとめ

KendraではRetrieveAPIを利用することでデフォルトでは10件が返戻されますが、リクエストのPageSizeパラメータを指定することで最大値100個までパッセージを取得することが出来ました。
実装する上で公式ドキュメントのRequest Syntax の各項目は把握しておくと良いですね。

以上、最後までお読みいただきありがとうございました。

その他参考記事