chore: import upstream snapshot with attribution

haystack/components/joiners/document_joiner.py

@@ -0,0 +1,271 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import itertools
+from collections import defaultdict
+from dataclasses import replace
+from enum import Enum
+from math import inf
+from typing import Any
+
+from haystack import Document, component, default_from_dict, default_to_dict, logging
+from haystack.core.component.types import Variadic
+from haystack.utils.misc import _reciprocal_rank_fusion
+
+logger = logging.getLogger(__name__)
+
+
+class JoinMode(Enum):
+    """
+    Enum for join mode.
+    """
+
+    CONCATENATE = "concatenate"
+    MERGE = "merge"
+    RECIPROCAL_RANK_FUSION = "reciprocal_rank_fusion"
+    DISTRIBUTION_BASED_RANK_FUSION = "distribution_based_rank_fusion"
+
+    def __str__(self) -> str:
+        return self.value
+
+    @staticmethod
+    def from_str(string: str) -> "JoinMode":
+        """
+        Convert a string to a JoinMode enum.
+        """
+        enum_map = {e.value: e for e in JoinMode}
+        mode = enum_map.get(string)
+        if mode is None:
+            msg = f"Unknown join mode '{string}'. Supported modes in DocumentJoiner are: {list(enum_map.keys())}"
+            raise ValueError(msg)
+        return mode
+
+
+@component
+class DocumentJoiner:
+    """
+    Joins multiple lists of documents into a single list.
+
+    It supports different join modes:
+    - concatenate: Keeps the highest-scored document in case of duplicates.
+    - merge: Calculates a weighted sum of scores for duplicates and merges them.
+    - reciprocal_rank_fusion: Merges and assigns scores based on reciprocal rank fusion.
+    - distribution_based_rank_fusion: Merges and assigns scores based on scores distribution in each Retriever.
+
+    ### Usage example:
+
+    ```python
+    from haystack import Pipeline, Document
+    from haystack.components.embedders import OpenAITextEmbedder, OpenAIDocumentEmbedder
+    from haystack.components.joiners import DocumentJoiner
+    from haystack.components.retrievers import InMemoryBM25Retriever
+    from haystack.components.retrievers import InMemoryEmbeddingRetriever
+    from haystack.document_stores.in_memory import InMemoryDocumentStore
+
+    document_store = InMemoryDocumentStore()
+    docs = [Document(content="Paris"), Document(content="Berlin"), Document(content="London")]
+    embedder = OpenAIDocumentEmbedder()
+    docs_embeddings = embedder.run(docs)
+    document_store.write_documents(docs_embeddings['documents'])
+
+    p = Pipeline()
+    p.add_component(instance=InMemoryBM25Retriever(document_store=document_store), name="bm25_retriever")
+    p.add_component(
+            instance=OpenAITextEmbedder(),
+            name="text_embedder",
+        )
+    p.add_component(instance=InMemoryEmbeddingRetriever(document_store=document_store), name="embedding_retriever")
+    p.add_component(instance=DocumentJoiner(), name="joiner")
+    p.connect("bm25_retriever", "joiner")
+    p.connect("embedding_retriever", "joiner")
+    p.connect("text_embedder", "embedding_retriever")
+    query = "What is the capital of France?"
+    p.run(data={"query": query, "text": query, "top_k": 1})
+    ```
+    """
+
+    def __init__(
+        self,
+        join_mode: str | JoinMode = JoinMode.CONCATENATE,
+        weights: list[float] | None = None,
+        top_k: int | None = None,
+        sort_by_score: bool = True,
+    ) -> None:
+        """
+        Creates a DocumentJoiner component.
+
+        :param join_mode:
+            Specifies the join mode to use. Available modes:
+            - `concatenate`: Keeps the highest-scored document in case of duplicates.
+            - `merge`: Calculates a weighted sum of scores for duplicates and merges them.
+            - `reciprocal_rank_fusion`: Merges and assigns scores based on reciprocal rank fusion.
+            - `distribution_based_rank_fusion`: Merges and assigns scores based on scores
+            distribution in each Retriever.
+        :param weights:
+            Assign importance to each list of documents to influence how they're joined.
+            This parameter is ignored for
+            `concatenate` or `distribution_based_rank_fusion` join modes.
+            Weight for each list of documents must match the number of inputs.
+        :param top_k:
+            The maximum number of documents to return.
+        :param sort_by_score:
+            If `True`, sorts the documents by score in descending order.
+            If a document has no score, it is handled as if its score is -infinity.
+        """
+        if isinstance(join_mode, str):
+            join_mode = JoinMode.from_str(join_mode)
+        join_mode_functions = {
+            JoinMode.CONCATENATE: DocumentJoiner._concatenate,
+            JoinMode.MERGE: self._merge,
+            JoinMode.RECIPROCAL_RANK_FUSION: self._rrf,
+            JoinMode.DISTRIBUTION_BASED_RANK_FUSION: DocumentJoiner._distribution_based_rank_fusion,
+        }
+        self.join_mode_function = join_mode_functions[join_mode]
+        self.join_mode = join_mode
+        if weights:
+            weight_sum = sum(weights)
+            if weight_sum == 0:
+                raise ValueError("The provided `weights` must not sum to zero.")
+            self.weights: list[float] | None = [float(i) / weight_sum for i in weights]
+        else:
+            self.weights = None
+        self.top_k = top_k
+        self.sort_by_score = sort_by_score
+
+    @component.output_types(documents=list[Document])
+    def run(self, documents: Variadic[list[Document]], top_k: int | None = None) -> dict[str, Any]:
+        """
+        Joins multiple lists of Documents into a single list depending on the `join_mode` parameter.
+
+        :param documents:
+            List of list of documents to be merged.
+        :param top_k:
+            The maximum number of documents to return. Overrides the instance's `top_k` if provided.
+
+        :returns:
+            A dictionary with the following keys:
+            - `documents`: Merged list of Documents
+        """
+        documents = list(documents)
+        output_documents = self.join_mode_function(documents)
+
+        if self.sort_by_score:
+            output_documents = sorted(
+                output_documents, key=lambda doc: doc.score if doc.score is not None else -inf, reverse=True
+            )
+            if any(doc.score is None for doc in output_documents):
+                logger.info(
+                    "Some of the Documents DocumentJoiner got have score=None. It was configured to sort Documents by "
+                    "score, so those with score=None were sorted as if they had a score of -infinity."
+                )
+
+        if top_k:
+            output_documents = output_documents[:top_k]
+        elif self.top_k:
+            output_documents = output_documents[: self.top_k]
+
+        return {"documents": output_documents}
+
+    @staticmethod
+    def _concatenate(document_lists: list[list[Document]]) -> list[Document]:
+        """
+        Concatenate multiple lists of Documents and return only the Document with the highest score for duplicates.
+        """
+        output = []
+        docs_per_id = defaultdict(list)
+        for doc in itertools.chain.from_iterable(document_lists):
+            docs_per_id[doc.id].append(doc)
+        for docs in docs_per_id.values():
+            doc_with_best_score = max(docs, key=lambda doc: doc.score if doc.score is not None else -inf)
+            output.append(doc_with_best_score)
+        return output
+
+    def _merge(self, document_lists: list[list[Document]]) -> list[Document]:
+        """
+        Merge multiple lists of Documents and calculate a weighted sum of the scores of duplicate Documents.
+        """
+        # This check prevents a division by zero when no documents are passed
+        if not document_lists:
+            return []
+
+        scores_map: dict = defaultdict(int)
+        documents_map = {}
+        weights = self.weights if self.weights else [1 / len(document_lists)] * len(document_lists)
+
+        for documents, weight in zip(document_lists, weights, strict=True):
+            for doc in documents:
+                scores_map[doc.id] += (doc.score if doc.score is not None else 0) * weight
+                documents_map[doc.id] = doc
+
+        return [replace(doc, score=scores_map[doc.id]) for doc in documents_map.values()]
+
+    def _rrf(self, document_lists: list[list[Document]]) -> list[Document]:
+        """
+        Merge multiple lists of Documents and assign scores based on reciprocal rank fusion.
+        """
+        return _reciprocal_rank_fusion(document_lists, weights=self.weights)
+
+    @staticmethod
+    def _distribution_based_rank_fusion(document_lists: list[list[Document]]) -> list[Document]:
+        """
+        Merge multiple lists of Documents and assign scores based on Distribution-Based Score Fusion.
+
+        (https://medium.com/plain-simple-software/distribution-based-score-fusion-dbsf-a-new-approach-to-vector-search-ranking-f87c37488b18)
+        If a Document is in more than one retriever, the one with the highest score is used.
+        """
+        rescaled_lists: list[list[Document]] = []
+        for documents in document_lists:
+            if len(documents) == 0:
+                rescaled_lists.append(documents)
+                continue
+
+            scores_list = [doc.score if doc.score is not None else 0 for doc in documents]
+
+            mean_score = sum(scores_list) / len(scores_list)
+            std_dev = (sum((x - mean_score) ** 2 for x in scores_list) / len(scores_list)) ** 0.5
+            min_score = mean_score - 3 * std_dev
+            max_score = mean_score + 3 * std_dev
+            delta_score = max_score - min_score
+
+            # if all docs have the same score delta_score is 0, the docs are uninformative for the query
+            rescaled_lists.append(
+                [
+                    replace(
+                        doc,
+                        score=((doc.score if doc.score is not None else 0) - min_score) / delta_score
+                        if delta_score != 0.0
+                        else 0.0,
+                    )
+                    for doc in documents
+                ]
+            )
+
+        return DocumentJoiner._concatenate(document_lists=rescaled_lists)
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Serializes the component to a dictionary.
+
+        :returns:
+            Dictionary with serialized data.
+        """
+        return default_to_dict(
+            self,
+            join_mode=str(self.join_mode),
+            weights=self.weights,
+            top_k=self.top_k,
+            sort_by_score=self.sort_by_score,
+        )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "DocumentJoiner":
+        """
+        Deserializes the component from a dictionary.
+
+        :param data:
+            The dictionary to deserialize from.
+        :returns:
+            The deserialized component.
+        """
+        return default_from_dict(cls, data)