Add to langchain_chroma.Chroma the method max_marginal_relevance_search_with_relevance_scores

[MaximeSongIdris]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

We already have the following method maximal_marginal_relevance, but it is missing the relevance score like in similarity_search_with_relevance_scores.

Here is a version of the method written as a function :

from typing import List, Tuple

import numpy as np
from langchain_chroma import Chroma
from langchain_chroma.vectorstores import _results_to_docs_and_scores, maximal_marginal_relevance
from langchain_core.documents import Document

def max_marginal_relevance_search_with_relevance_scores(
    vdb: Chroma,
    query: str,
    k: int,
    fetch_k: int
) -> List[Tuple[Document, float]]:
    """
    Perform Maximal Marginal Relevance (MMR) retrieval with relevance scores.

    This function retrieves a set of candidate documents from a Chroma vectorstore,
    ranks them using Maximal Marginal Relevance (MMR) to maximize both relevance and diversity,
    and returns the top-k documents along with their relevance scores.

    References: https://github.com/langchain-ai/langchain/blob/master/libs/partners/chroma/langchain_chroma/vectorstores.py#L144

    Parameters:
        vdb (Chroma): The Chroma vectorstore object.
        query (str): The natural language query string.
        k (int): The number of final documents to return.
        fetch_k (int): The number of initial similar documents to retrieve before applying MMR.
                       Must be greater than k.

    Returns:
        List[Tuple[Document, float]]: A list of (Document, relevance_score) tuples.
                                      The relevance_score is a similarity score in [0, 1].
    """
    assert fetch_k > k, 'fetch_k should be greater than k to allow MMR to select a diverse subset.'

    # Embed the query using the vectorstore's embedding function
    query_embed = vdb._embedding_function.embed_query(query)

    # Retrieve top `fetch_k` most similar documents using raw Chroma query
    results = vdb._Chroma__query_collection(
        query_embeddings=query_embed,
        n_results=fetch_k,
        where=None,
        where_document=None,
        include=["metadatas", "documents", "distances", "embeddings"],
    )

    # Apply MMR to select a diverse subset of `k` documents
    mmr_selected_doc_index = maximal_marginal_relevance(
        np.array(query_embed, dtype=np.float32),
        results["embeddings"][0],
        k=k,
        lambda_mult=0.5,
    )
    selected_results = [r for i, r in enumerate(_results_to_docs_and_scores(results)) if i in mmr_selected_doc_index]

    # Convert distance (lower is better) into similarity (higher is better)
    relevance_score_fn = vdb._select_relevance_score_fn()
    selected_results = [(doc, relevance_score_fn(score)) for doc, score in selected_results]

    return selected_results

Motivation

Useful after MMR retrieval to print relevance score and compute metrics based on that.

Proposal (If applicable)

No response