# SPDX-FileCopyrightText: 2022-present deepset GmbH # # 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)