Add rerank document compressor

[jpfcabral]

Fixes #298

Added:

• BedrockRerank based on BaseDocumentCompressor so we can use ContextualCompressionRetriever
• Import from root import i.e: from langchain_aws import BedrockRerank
• Unit tests

Some snippets:

• Example 1 (from documents):

from langchain_core.documents import Document
from langchain_aws import BedrockRerank

# Initialize the class
reranker = BedrockRerank(model_arn=model_arn)

# List of documents to rerank
documents = [
    Document(page_content="LangChain is a powerful library for LLMs."),
    Document(page_content="AWS Bedrock enables access to AI models."),
    Document(page_content="Artificial intelligence is transforming the world."),
]

# Query for reranking
query = "What is AWS Bedrock?"

# Call the rerank method
results = reranker.compress_documents(documents, query)

# Display the most relevant documents
for doc in results:
    print(f"Content: {doc.page_content}")
    print(f"Score: {doc.metadata['relevance_score']}")

• Example 2 (with contextual compression retriever):

from langchain_aws import BedrockEmbeddings
from langchain.retrievers.contextual_compression import ContextualCompressionRetriever
from langchain.vectorstores import FAISS
from langchain_core.documents import Document
from langchain_aws import BedrockRerank

# Create a vector store using FAISS with Bedrock embeddings
documents = [
    Document(page_content="LangChain integrates LLM models."),
    Document(page_content="AWS Bedrock provides cloud-based AI models."),
    Document(page_content="Machine learning can be used for predictions."),
]
embeddings = BedrockEmbeddings()
vectorstore = FAISS.from_documents(documents, embeddings)

# Create the document compressor using BedrockRerank
reranker = BedrockRerank(model_arn=model_arn)

# Create the retriever with contextual compression
retriever = ContextualCompressionRetriever(
    base_compressor=reranker,
    base_retriever=vectorstore.as_retriever(),
)

# Execute a query
query = "How does AWS Bedrock work?"
retrieved_docs = retriever.invoke(query)

# Display the most relevant documents
for doc in retrieved_docs:
    print(f"Content: {doc.page_content}")
    print(f"Score: {doc.metadata.get('relevance_score', 'N/A')}")

• Example 3 (from list):

from langchain_aws import BedrockRerank

# Initialize BedrockRerank
reranker = BedrockRerank(model_arn=model_arn)

# Unstructured documents
documents = [
    "LangChain is used to integrate LLM models.",
    "AWS Bedrock provides access to cloud-based models.",
    "Machine learning is revolutionizing the world.",
]

# Query
query = "What is the role of AWS Bedrock?"

# Rerank the documents
results = reranker.rerank(query=query, documents=documents)

# Display the results
for res in results:
    print(f"Index: {res['index']}, Score: {res['relevance_score']}")
    print(f"Document: {documents[res['index']]}")

[mgvalverde]

Hi @jpfcabral, interesting contribution!

I noticed that the default region is set to aws_region="us-west-2". From what I understand, if you have a region configured using a profile, not specifying the region should use the one from the profile. However, it defaults always to "us-west-2" instead.
Would it make sense for you to replace that value with aws_region=None?

[jpfcabral]

Fair point, @mgvalverde , I just changed on bbc0243

[3coins]

@jpfcabral
It seems like the model is called directly via invoke here, but this rerank API seems to have more options for controlling the re-ranking. Is this the right API to implement?
https://docs.aws.amazon.com/bedrock/latest/userguide/rerank-use.html

[jpfcabral]

@3coins
I was using bedrock-runtime client to hit the model so we couldn't use more options for controlling the re-ranking.

Following on the documentation you sent me, I need to call it by bedrock-agent-runtime so we can configure the request with rerankingConfiguration attribute. Based on that, the user can add additional information, if needed, in additional_model_request_fields parameter.

Note: bedrock-agent-runtime rerank only support requests with model arn but does not have a method to get the arn from a model id. To surpass this, I created a funcion that instantiate a simple bedrock client to get session information + model id and return model arn.

I just added a commit 13ad88f with the changes above, let me know if need some fixes on that.

[3coins]

@jpfcabral
Thanks for submitting this PR and a quick turnaround on the updates.
Great job on adding the examples in the PR description, it might be more useful to add a notebook with those samples in the samples/document_compressors directory.

Also, to keep the module organization consistent with community, does it sound better to put this under document_compressors rather than rerank?

[jpfcabral]

Based on the recommendations:

• Samples added (5c8aca3)
• Rename module from rerank to document_compressors (66f0d64)

[3coins]

@jpfcabral
Code looks good, can you fix the lint errors. The only other thing would be that this code is missing the nextTokens implementation, which will not handle paged results from the rerank API; we can tackle that in a separate PR.

[jpfcabral]

@3coins The lint issues have been resolved. I also refactored d69a71f the initialize_client method to align with the pattern used in other langchain libraries, following this reference. I agree that the nextTokens implementation should be handled in a separate PR—let me know if I can help with that in any way. Thanks for the review!

[3coins]

@jpfcabral
Great job on getting the updates done. LGTM! 🚀

[3coins]

@jpfcabral
A small suggestion that I missed previously, can we change the extra_model_fields to additional_model_request_fields that aligns with the naming in the rerank API. I have the changes ready locally, but did not have permissions to push to your main.

[jpfcabral]

Done!