deepset-ai--haystack/haystack/components/retrievers/multi_retriever.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import asyncio
from concurrent.futures import ThreadPoolExecutor, as_completed
from math import inf
from typing import Any, Literal

from haystack import component, default_from_dict, default_to_dict
from haystack.components.retrievers.types.protocol import TextRetriever
from haystack.core.serialization import component_from_dict, component_to_dict, import_class_by_name
from haystack.dataclasses import Document
from haystack.utils.async_utils import _execute_component_async
from haystack.utils.experimental import _experimental
from haystack.utils.misc import _deduplicate_documents, _reciprocal_rank_fusion


@_experimental
@component
class MultiRetriever:
    """
    A component that accepts text retrievers and runs them in parallel, combining their results.

    > **Note:** This component is experimental and may change or be removed in future releases without prior
    deprecation notice.

    All retrievers must implement the `TextRetriever` protocol. Use `TextEmbeddingRetriever` to wrap an
    embedding-based retriever before passing it to this component.

    Each retriever is queried concurrently using a thread pool.
    The results are deduplicated and returned as a single list of documents.

    ### Usage example

    ```python
    from haystack import Document
    from haystack.document_stores.in_memory import InMemoryDocumentStore
    from haystack.document_stores.types import DuplicatePolicy
    from haystack.components.retrievers import InMemoryBM25Retriever, InMemoryEmbeddingRetriever
    from haystack.components.retrievers import TextEmbeddingRetriever, MultiRetriever
    from haystack.components.embedders import OpenAITextEmbedder, OpenAIDocumentEmbedder
    from haystack.components.writers import DocumentWriter

    documents = [
        Document(content="Renewable energy is energy that is collected from renewable resources."),
        Document(content="Solar energy is a type of green energy that is harnessed from the sun."),
        Document(content="Wind energy is another type of green energy that is generated by wind turbines."),
    ]

    # Populate the document store
    doc_store = InMemoryDocumentStore()
    doc_embedder = OpenAIDocumentEmbedder()
    doc_writer = DocumentWriter(document_store=doc_store, policy=DuplicatePolicy.SKIP)
    doc_writer.run(documents=doc_embedder.run(documents)["documents"])

    # Run the multi-retriever with all retrievers
    retriever = MultiRetriever(
        retrievers={
            "bm25": InMemoryBM25Retriever(document_store=doc_store),
            "embedding": TextEmbeddingRetriever(
                retriever=InMemoryEmbeddingRetriever(document_store=doc_store),
                text_embedder=OpenAITextEmbedder(),
            ),
        },
        top_k=3,
    )

    # Run all retrievers
    result = retriever.run(query="green energy sources")

    # Run only the BM25 retriever
    result = retriever.run(query="green energy sources", active_retrievers=["bm25"])

    for doc in result["documents"]:
        print(doc.content)
    ```
    """

    def __init__(
        self,
        *,
        retrievers: dict[str, TextRetriever],
        filters: dict[str, Any] | None = None,
        top_k_per_retriever: int | None = None,
        top_k: int | None = None,
        max_workers: int = 4,
        join_mode: Literal["concatenate", "reciprocal_rank_fusion"] = "reciprocal_rank_fusion",
    ) -> None:
        """
        Create the MultiRetriever component.

        :param retrievers:
            A dictionary mapping names to text retrievers (implementing the `TextRetriever` protocol) to run in
            parallel.
        :param filters:
            A dictionary of filters to apply when retrieving documents.
        :param top_k_per_retriever:
            The maximum number of documents to return per retriever. If set, this will override the `top_k`
            parameter for each retriever. If None, the `top_k` parameter of retrievers will be used.
        :param top_k:
            The maximum number of documents to return overall, extracted from the combined results of all
            retrievers. When set, the results are always merged using reciprocal rank fusion (regardless of
            `join_mode`) so that the combined list has a consistent global ranking before it is truncated to
            `top_k`. If None, all results are returned.
        :param max_workers:
            The maximum number of threads to use for parallel retrieval.
        :param join_mode:
            How to merge results from multiple retrievers. Available modes:
            - `concatenate`: Combines all results into a single list and deduplicates.
            - `reciprocal_rank_fusion`: Deduplicates and assigns scores based on reciprocal rank fusion.
        """
        self.retrievers = retrievers
        self.filters = filters
        self.top_k_per_retriever = top_k_per_retriever
        self.top_k = top_k
        self.max_workers = max_workers
        self.join_mode = join_mode

    def _merge_results(self, document_lists: list[list[Document]], top_k: int | None = None) -> list[Document]:
        """
        Merge per-retriever result lists according to `join_mode`.

        In `concatenate` mode, all lists are flattened and deduplicated. In `reciprocal_rank_fusion` mode, results
        are deduplicated and re-scored using RRF, then returned in descending score order. When `top_k` is set, RRF
        is always used so the combined results have a consistent global ranking, and only the top `top_k` documents
        are returned.
        """
        # When top_k is set we always use reciprocal rank fusion to merge the results, regardless of join_mode,
        # so that truncation is applied to a consistently ranked list.
        if top_k is not None or self.join_mode == "reciprocal_rank_fusion":
            documents = _reciprocal_rank_fusion(document_lists)
            merged = sorted(documents, key=lambda d: d.score if d.score is not None else -inf, reverse=True)
            return merged[:top_k] if top_k is not None else merged
        return _deduplicate_documents([doc for docs in document_lists for doc in docs])

    def _resolve_retrievers(self, active_retrievers: list[str] | None) -> dict[str, TextRetriever]:
        """
        Returns the subset of retrievers to run based on the active_retrievers list.

        :param active_retrievers:
            A list of retriever names to run. If None, all retrievers are returned.

        :returns:
            A dictionary of retriever names to retriever instances.

        :raises ValueError:
            If any name in `active_retrievers` does not match a retriever name.
        """
        if active_retrievers is None:
            return self.retrievers
        unknown = set(active_retrievers) - self.retrievers.keys()
        if unknown:
            raise ValueError(
                f"Unknown retriever name(s): {sorted(unknown)}. Available retrievers: {sorted(self.retrievers.keys())}"
            )
        return {name: self.retrievers[name] for name in active_retrievers}

    def warm_up(self) -> None:
        """
        Warm up the retrievers.
        """
        for retriever in self.retrievers.values():
            if hasattr(retriever, "warm_up"):
                retriever.warm_up()

    async def warm_up_async(self) -> None:
        """
        Warm up the retrievers on the serving event loop.
        """
        for retriever in self.retrievers.values():
            if hasattr(retriever, "warm_up_async"):
                await retriever.warm_up_async()
            elif hasattr(retriever, "warm_up"):
                retriever.warm_up()

    def close(self) -> None:
        """
        Release the retrievers' resources.
        """
        for retriever in self.retrievers.values():
            if hasattr(retriever, "close"):
                retriever.close()

    async def close_async(self) -> None:
        """
        Release the retrievers' async resources.
        """
        for retriever in self.retrievers.values():
            if hasattr(retriever, "close_async"):
                await retriever.close_async()
            elif hasattr(retriever, "close"):
                retriever.close()

    @component.output_types(documents=list[Document])
    def run(
        self,
        query: str,
        filters: dict[str, Any] | None = None,
        top_k_per_retriever: int | None = None,
        top_k: int | None = None,
        *,
        active_retrievers: list[str] | None = None,
    ) -> dict[str, list[Document]]:
        """
        Runs retrievers in parallel on the given query and returns deduplicated results.

        :param query:
            The query to run the retrievers on.
        :param filters:
            Filters to apply. Defaults to the value set at initialization.
        :param top_k_per_retriever:
            The maximum number of documents to return per retriever. When set, this will override the `top_k`
            parameter for each retriever. If None, the `top_k` parameter set for retrievers will be used.
            Defaults to the value set at initialization.
        :param top_k:
            The maximum number of documents to return overall, extracted from the combined results of all
            retrievers. When set, the results are always merged using reciprocal rank fusion (regardless of
            `join_mode`) so that the combined list has a consistent global ranking before it is truncated to
            `top_k`. If None, all results are returned. Defaults to the value set at initialization.
        :param active_retrievers:
            Names of retrievers to run. Defaults to all. Must match keys in the `retrievers` dictionary.

        :returns:
            A dictionary with the keys:
                - "documents": A deduplicated list of retrieved documents.

        :raises ValueError:
            If any name in `active_retrievers` does not match a retriever name.
        """
        self.warm_up()

        resolved_top_k_per_retriever = (
            top_k_per_retriever if top_k_per_retriever is not None else self.top_k_per_retriever
        )
        resolved_top_k = top_k if top_k is not None else self.top_k
        resolved_filters = filters if filters is not None else self.filters

        retrievers_to_run = self._resolve_retrievers(active_retrievers)

        results_by_name: dict[str, list[Document]] = {}
        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            future_to_name = {}
            for name, retriever in retrievers_to_run.items():
                run_kwargs: dict[str, Any] = {"query": query}
                if resolved_top_k_per_retriever is not None:
                    run_kwargs["top_k"] = resolved_top_k_per_retriever
                if resolved_filters is not None:
                    run_kwargs["filters"] = resolved_filters
                future_to_name[executor.submit(retriever.run, **run_kwargs)] = name

            for future in as_completed(future_to_name):
                name = future_to_name[future]
                try:
                    results_by_name[name] = future.result().get("documents", [])
                except Exception as e:
                    raise RuntimeError(f"Retriever '{name}' failed: {e}") from e

        document_lists = [results_by_name[name] for name in retrievers_to_run]
        return {"documents": self._merge_results(document_lists, top_k=resolved_top_k)}

    @component.output_types(documents=list[Document])
    async def run_async(
        self,
        query: str,
        filters: dict[str, Any] | None = None,
        top_k_per_retriever: int | None = None,
        top_k: int | None = None,
        *,
        active_retrievers: list[str] | None = None,
    ) -> dict[str, list[Document]]:
        """
        Runs retrievers concurrently on the given query and returns deduplicated results.

        Uses each retriever's `run_async` method if available, otherwise runs `run` in a thread executor.

        :param query:
            The query to run the retrievers on.
        :param filters:
            Filters to apply. Defaults to the value set at initialization.
        :param top_k_per_retriever:
            The maximum number of documents to return per retriever. When set, this will override the `top_k`
            parameter for each retriever. If None, the `top_k` parameter set for retrievers will be used.
            Defaults to the value set at initialization.
        :param top_k:
            The maximum number of documents to return overall, extracted from the combined results of all
            retrievers. When set, the results are always merged using reciprocal rank fusion (regardless of
            `join_mode`) so that the combined list has a consistent global ranking before it is truncated to
            `top_k`. If None, all results are returned. Defaults to the value set at initialization.
        :param active_retrievers:
            Names of retrievers to run. Defaults to all. Must match keys in the `retrievers` dictionary.

        :returns:
            A dictionary with the keys:
                - "documents": A deduplicated list of retrieved documents.

        :raises ValueError:
            If any name in `active_retrievers` does not match a retriever name.
        """
        await self.warm_up_async()

        resolved_top_k_per_retriever = (
            top_k_per_retriever if top_k_per_retriever is not None else self.top_k_per_retriever
        )
        resolved_top_k = top_k if top_k is not None else self.top_k
        resolved_filters = filters if filters is not None else self.filters

        retrievers_to_run = self._resolve_retrievers(active_retrievers)

        run_kwargs: dict[str, Any] = {"query": query}
        if resolved_top_k_per_retriever is not None:
            run_kwargs["top_k"] = resolved_top_k_per_retriever
        if resolved_filters is not None:
            run_kwargs["filters"] = resolved_filters

        async def _run_one(name: str, retriever: TextRetriever) -> list[Document]:
            try:
                result = await _execute_component_async(retriever, **run_kwargs)
                return result.get("documents", [])
            except Exception as e:
                raise RuntimeError(f"Retriever '{name}' failed: {e}") from e

        document_lists = list(await asyncio.gather(*[_run_one(name, r) for name, r in retrievers_to_run.items()]))
        return {"documents": self._merge_results(document_lists, top_k=resolved_top_k)}

    def to_dict(self) -> dict[str, Any]:
        """
        Serializes the component to a dictionary.

        :returns:
            Dictionary with serialized data.
        """
        return default_to_dict(
            self,
            retrievers={name: component_to_dict(obj=r, name=name) for name, r in self.retrievers.items()},
            filters=self.filters,
            top_k_per_retriever=self.top_k_per_retriever,
            top_k=self.top_k,
            max_workers=self.max_workers,
            join_mode=self.join_mode,
        )

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "MultiRetriever":
        """
        Creates an instance of the component from a dictionary.

        :param data:
            Dictionary with the data to create the component.
        """
        retrievers_data = data.get("init_parameters", {}).get("retrievers", {})
        if retrievers_data:
            retrievers = {}
            for name, retriever_data in retrievers_data.items():
                try:
                    imported_class = import_class_by_name(retriever_data["type"])
                except ImportError as e:
                    raise ImportError(
                        f"Could not import class {retriever_data['type']} for retriever '{name}'. Error: {str(e)}"
                    ) from e
                retrievers[name] = component_from_dict(cls=imported_class, data=retriever_data, name=name)
            data["init_parameters"]["retrievers"] = retrievers
        return default_from_dict(cls, data)