Add rerank document compressor

[jpfcabral]

This PR resolves #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?

[3coins]

Rather than relying on the API to fetch the model arn, I would suggest, we keep this consistent with the rerank API and take model_arn as input rather than the model_id.

Suggested change

	def _get_model_arn(self) -> str:
	"""Fetch the ARN of the reranker model using the model ID."""
	session = self._get_session()
	client = session.client("bedrock", self.aws_region)
	response = client.get_foundation_model(modelIdentifier=self.model_id)
	return response["modelDetails"]["modelArn"]

[jpfcabral]

Solved on f45340b

[3coins]

Update this to model_arn.

Suggested change

	model_id: Optional[str] = "amazon.rerank-v1:0"
	model_arn: str

[jpfcabral]

Solved on f45340b

[3coins]

Though I like the short name here, we have been using region_name in other places, and would prefer we keep it consistent here as well.

Suggested change

	aws_region: str = Field(
	default_factory=from_env("AWS_DEFAULT_REGION", default=None)
	)
	region_name: str = Field(
	default_factory=from_env("AWS_DEFAULT_REGION", default=None)
	)

[jpfcabral]

Solved on 0489125

[3coins]

Similar to the region, would like this to be consistent with other implementations in langchain_aws.

Suggested change

	aws_profile: Optional[str] = Field(
	default_factory=from_env("AWS_PROFILE", default=None)
	)
	credentials_profile_name: Optional[str] = Field(
	default_factory=from_env("AWS_PROFILE", default=None)
	)

[jpfcabral]

Solved on 0489125

[jpfcabral]

Based on the recommendations:

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