chore: import upstream snapshot with attribution

haystack/components/retrievers/multi_retriever.py

@@ -0,0 +1,363 @@
+# 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)