AWS Logo
Menu

How to Use Cohere Rerank on Amazon Bedrock using Langchain

This article is to show how to use cohere rerank on Amazon Bedrock using langchain framework

Agung Sidharta
Amazon Employee
Published Jan 14, 2025

Background

Prior to re-invent 2024, Amazon Web Services (AWS) announced availability of Cohere’s advanced reranking model Rerank 3.5 in Amazon Bedrock. It can be used directly through through our new Rerank API and Amazon Bedrock Knowledge Base. Rerank provides a powerful semantic boost to the search quality of any keyword or vector search system by reranking the search result. In case of generative AI (GenAI) solution with Large Language Model (LLM), it helps to provide more relevant context which can improve accuracy of Retrieval Augmented Generation (RAG) solution.

Problem Statement

There are community members (also AWS customers) building GenAI solutions using langchain and would like to continue doing so but there's lack of how-to and documentation on how to use langchain to invoke Cohere Rerank in Amazon Bedrock. Additionally, when I personally tested it, it resulted in error (more detail below).

Issue and Solution

This issue was observed in following package version (all packages are latest version during time of writing this post on 14 Jan 2025):
PackageVersion
cohere5.13.6
langchain0.3.14
langchain-cohere0.3.4
langchain-community0.3.14
langchain-core0.3.29
When I tested using Amazon Bedrock's Cohere Rerank from langchain, I used the tutorial from langchain docs but I slightly modified it in order to call Bedrock's Cohere Rerank. Below is my modified code that resulted in error.
The error is the following: TypeError: V2Client.rerank() got an unexpected keyword argument 'rank_fields'.
Upon checking and debugging, I found that cohere.BedrockClientV2 rerank method doesn't support rank_fields as argument, so I need to overwrite it with what the cohere.BedrockClientV2 rerank supports.
Workaround Solution:
  • Create a new class extending CohereRerank from langchain_cohere.
  • Use the extended class instead of CohereRerank in the implementation.
The error should disappear and the reranked results will appear.

Additional Information

As mentioned that this is a workaround, I also have noted few github issues raised in each github repo that asked to solve this permanently:
  1. Issue raised in langchain-ai/langchain-cohere : How to use aws bedrock Cohere Rerank.
  2. Issue raised in langchain-ai/langchain-aws: Support bedrock rerank API.
This post will be updated once there's already permanent solution.

Conclusion

This post showed a workaround of how to use Bedrock's Cohere Rerank from langchain. It hopes to unblock obstacle faced by community and AWS customers. Feel free to add another workaround you have and share it with the community in the comment.
 

Any opinions in this post are those of the individual author and may not reflect the opinions of AWS.

Comments