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