BedrockでRerankモデルが使えるようになったのでLangChainに組み込んでみた！

def rerank(inp: dict[str, Any], top_n: int = 3) -> list[Document]:
    """
    この関数はmazon.rerank-v1:0モデルを用いてリランキングします。

    Args:
        inp (dict[str, Any]): リランクするための入力データを含む辞書。
            - "question" (str): ドキュメントをリランクする基準となる質問。
            - "documents" (list[Document]): リランク対象となるドキュメントのリスト。
        top_n (int): リランク後に上位から返されるドキュメントの数。デフォルトは3。

    Returns:
        list[Document]: 与えられた質問に基づいてリランクされたドキュメントのリスト。
    """

    question = inp["question"]
    documents = inp["documents"]

    # AWSクライアントの設定
    client = boto3.client('bedrock-runtime', region_name='ap-northeast-1')

    # documentsからpage_contentを抽出
    document_contents = [doc.page_content for doc in documents]
    
    # リランク
    response = client.invoke_model(
        modelId="amazon.rerank-v1:0",
        body=json.dumps({
            "query": question,
            "documents": document_contents,
            "top_n": top_n
        })
    )

    # レスポンス処理
    response_body = json.loads(response['body'].read())
    ranked_results = response_body.get('results', [])

    # 結果をDocumentのリストに変換
    ranked_documents = [
        Document(
            page_content=document_contents[result['index']],
            metadata={"relevance_score": result['relevance_score']}
        )
        for result in ranked_results
    ]

    return ranked_documents

rerank_rag_chain = (
    {
        "question": RunnablePassthrough(),
        "documents": retriever,
    }
    | RunnablePassthrough.assign(context=rerank)
    | prompt | model | StrOutputParser()
)
output = rerank_rag_chain.invoke("LangChainの概要")
print(output)