Coverage for haystack/components/retrievers/multi_retriever.py: 99%
115 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import asyncio
6from concurrent.futures import ThreadPoolExecutor, as_completed
7from math import inf
8from typing import Any, Literal
10from haystack import component, default_from_dict, default_to_dict
11from haystack.components.retrievers.types.protocol import TextRetriever
12from haystack.core.serialization import component_from_dict, component_to_dict, import_class_by_name
13from haystack.dataclasses import Document
14from haystack.utils.async_utils import _execute_component_async, _gather_tasks_with_cancel
15from haystack.utils.experimental import _experimental
16from haystack.utils.misc import _deduplicate_documents, _reciprocal_rank_fusion
19@_experimental
20@component
21class MultiRetriever:
22 """
23 A component that accepts text retrievers and runs them in parallel, combining their results.
25 > **Note:** This component is experimental and may change or be removed in future releases without prior
26 deprecation notice.
28 All retrievers must implement the `TextRetriever` protocol. Use `TextEmbeddingRetriever` to wrap an
29 embedding-based retriever before passing it to this component.
31 Each retriever is queried concurrently using a thread pool.
32 The results are deduplicated and returned as a single list of documents.
34 ### Usage example
36 ```python
37 from haystack import Document
38 from haystack.document_stores.in_memory import InMemoryDocumentStore
39 from haystack.document_stores.types import DuplicatePolicy
40 from haystack.components.retrievers import InMemoryBM25Retriever, InMemoryEmbeddingRetriever
41 from haystack.components.retrievers import TextEmbeddingRetriever, MultiRetriever
42 from haystack.components.embedders import OpenAITextEmbedder, OpenAIDocumentEmbedder
43 from haystack.components.writers import DocumentWriter
45 documents = [
46 Document(content="Renewable energy is energy that is collected from renewable resources."),
47 Document(content="Solar energy is a type of green energy that is harnessed from the sun."),
48 Document(content="Wind energy is another type of green energy that is generated by wind turbines."),
49 ]
51 # Populate the document store
52 doc_store = InMemoryDocumentStore()
53 doc_embedder = OpenAIDocumentEmbedder()
54 doc_writer = DocumentWriter(document_store=doc_store, policy=DuplicatePolicy.SKIP)
55 doc_writer.run(documents=doc_embedder.run(documents)["documents"])
57 # Run the multi-retriever with all retrievers
58 retriever = MultiRetriever(
59 retrievers={
60 "bm25": InMemoryBM25Retriever(document_store=doc_store),
61 "embedding": TextEmbeddingRetriever(
62 retriever=InMemoryEmbeddingRetriever(document_store=doc_store),
63 text_embedder=OpenAITextEmbedder(),
64 ),
65 },
66 top_k=3,
67 )
69 # Run all retrievers
70 result = retriever.run(query="green energy sources")
72 # Run only the BM25 retriever
73 result = retriever.run(query="green energy sources", active_retrievers=["bm25"])
75 for doc in result["documents"]:
76 print(doc.content)
77 ```
78 """
80 def __init__(
81 self,
82 *,
83 retrievers: dict[str, TextRetriever],
84 filters: dict[str, Any] | None = None,
85 top_k_per_retriever: int | None = None,
86 top_k: int | None = None,
87 max_workers: int = 4,
88 join_mode: Literal["concatenate", "reciprocal_rank_fusion"] = "reciprocal_rank_fusion",
89 ) -> None:
90 """
91 Create the MultiRetriever component.
93 :param retrievers:
94 A dictionary mapping names to text retrievers (implementing the `TextRetriever` protocol) to run in
95 parallel.
96 :param filters:
97 A dictionary of filters to apply when retrieving documents.
98 :param top_k_per_retriever:
99 The maximum number of documents to return per retriever. If set, this will override the `top_k`
100 parameter for each retriever. If None, the `top_k` parameter of retrievers will be used.
101 :param top_k:
102 The maximum number of documents to return overall, extracted from the combined results of all
103 retrievers. When set, the results are always merged using reciprocal rank fusion (regardless of
104 `join_mode`) so that the combined list has a consistent global ranking before it is truncated to
105 `top_k`. If None, all results are returned.
106 :param max_workers:
107 The maximum number of threads to use for parallel retrieval.
108 :param join_mode:
109 How to merge results from multiple retrievers. Available modes:
110 - `concatenate`: Combines all results into a single list and deduplicates.
111 - `reciprocal_rank_fusion`: Deduplicates and assigns scores based on reciprocal rank fusion.
112 """
113 self.retrievers = retrievers
114 self.filters = filters
115 self.top_k_per_retriever = top_k_per_retriever
116 self.top_k = top_k
117 self.max_workers = max_workers
118 self.join_mode = join_mode
120 def _merge_results(self, document_lists: list[list[Document]], top_k: int | None = None) -> list[Document]:
121 """
122 Merge per-retriever result lists according to `join_mode`.
124 In `concatenate` mode, all lists are flattened and deduplicated. In `reciprocal_rank_fusion` mode, results
125 are deduplicated and re-scored using RRF, then returned in descending score order. When `top_k` is set, RRF
126 is always used so the combined results have a consistent global ranking, and only the top `top_k` documents
127 are returned.
128 """
129 # When top_k is set we always use reciprocal rank fusion to merge the results, regardless of join_mode,
130 # so that truncation is applied to a consistently ranked list.
131 if top_k is not None or self.join_mode == "reciprocal_rank_fusion":
132 documents = _reciprocal_rank_fusion(document_lists)
133 merged = sorted(documents, key=lambda d: d.score if d.score is not None else -inf, reverse=True)
134 return merged[:top_k] if top_k is not None else merged
135 return _deduplicate_documents([doc for docs in document_lists for doc in docs])
137 def _resolve_retrievers(self, active_retrievers: list[str] | None) -> dict[str, TextRetriever]:
138 """
139 Returns the subset of retrievers to run based on the active_retrievers list.
141 :param active_retrievers:
142 A list of retriever names to run. If None, all retrievers are returned.
144 :returns:
145 A dictionary of retriever names to retriever instances.
147 :raises ValueError:
148 If any name in `active_retrievers` does not match a retriever name.
149 """
150 if active_retrievers is None:
151 return self.retrievers
152 unknown = set(active_retrievers) - self.retrievers.keys()
153 if unknown:
154 raise ValueError(
155 f"Unknown retriever name(s): {sorted(unknown)}. Available retrievers: {sorted(self.retrievers.keys())}"
156 )
157 return {name: self.retrievers[name] for name in active_retrievers}
159 def warm_up(self) -> None:
160 """
161 Warm up the retrievers.
162 """
163 for retriever in self.retrievers.values():
164 if hasattr(retriever, "warm_up"):
165 retriever.warm_up()
167 async def warm_up_async(self) -> None:
168 """
169 Warm up the retrievers on the serving event loop.
170 """
171 for retriever in self.retrievers.values():
172 if hasattr(retriever, "warm_up_async"):
173 await retriever.warm_up_async()
174 elif hasattr(retriever, "warm_up"):
175 retriever.warm_up()
177 def close(self) -> None:
178 """
179 Release the retrievers' resources.
180 """
181 for retriever in self.retrievers.values():
182 if hasattr(retriever, "close"):
183 retriever.close()
185 async def close_async(self) -> None:
186 """
187 Release the retrievers' async resources.
188 """
189 for retriever in self.retrievers.values():
190 if hasattr(retriever, "close_async"):
191 await retriever.close_async()
192 elif hasattr(retriever, "close"):
193 retriever.close()
195 @component.output_types(documents=list[Document])
196 def run(
197 self,
198 query: str,
199 filters: dict[str, Any] | None = None,
200 top_k_per_retriever: int | None = None,
201 top_k: int | None = None,
202 *,
203 active_retrievers: list[str] | None = None,
204 ) -> dict[str, list[Document]]:
205 """
206 Runs retrievers in parallel on the given query and returns deduplicated results.
208 :param query:
209 The query to run the retrievers on.
210 :param filters:
211 Filters to apply. Defaults to the value set at initialization.
212 :param top_k_per_retriever:
213 The maximum number of documents to return per retriever. When set, this will override the `top_k`
214 parameter for each retriever. If None, the `top_k` parameter set for retrievers will be used.
215 Defaults to the value set at initialization.
216 :param top_k:
217 The maximum number of documents to return overall, extracted from the combined results of all
218 retrievers. When set, the results are always merged using reciprocal rank fusion (regardless of
219 `join_mode`) so that the combined list has a consistent global ranking before it is truncated to
220 `top_k`. If None, all results are returned. Defaults to the value set at initialization.
221 :param active_retrievers:
222 Names of retrievers to run. Defaults to all. Must match keys in the `retrievers` dictionary.
224 :returns:
225 A dictionary with the keys:
226 - "documents": A deduplicated list of retrieved documents.
228 :raises ValueError:
229 If any name in `active_retrievers` does not match a retriever name.
230 """
231 self.warm_up()
233 resolved_top_k_per_retriever = (
234 top_k_per_retriever if top_k_per_retriever is not None else self.top_k_per_retriever
235 )
236 resolved_top_k = top_k if top_k is not None else self.top_k
237 resolved_filters = filters if filters is not None else self.filters
239 retrievers_to_run = self._resolve_retrievers(active_retrievers)
241 results_by_name: dict[str, list[Document]] = {}
242 with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
243 future_to_name = {}
244 for name, retriever in retrievers_to_run.items():
245 run_kwargs: dict[str, Any] = {"query": query}
246 if resolved_top_k_per_retriever is not None:
247 run_kwargs["top_k"] = resolved_top_k_per_retriever
248 if resolved_filters is not None:
249 run_kwargs["filters"] = resolved_filters
250 future_to_name[executor.submit(retriever.run, **run_kwargs)] = name
252 for future in as_completed(future_to_name):
253 name = future_to_name[future]
254 try:
255 results_by_name[name] = future.result().get("documents", [])
256 except Exception as e:
257 raise RuntimeError(f"Retriever '{name}' failed: {e}") from e
259 document_lists = [results_by_name[name] for name in retrievers_to_run]
260 return {"documents": self._merge_results(document_lists, top_k=resolved_top_k)}
262 @component.output_types(documents=list[Document])
263 async def run_async(
264 self,
265 query: str,
266 filters: dict[str, Any] | None = None,
267 top_k_per_retriever: int | None = None,
268 top_k: int | None = None,
269 *,
270 active_retrievers: list[str] | None = None,
271 ) -> dict[str, list[Document]]:
272 """
273 Runs retrievers concurrently on the given query and returns deduplicated results.
275 Uses each retriever's `run_async` method if available, otherwise runs `run` in a thread executor.
277 :param query:
278 The query to run the retrievers on.
279 :param filters:
280 Filters to apply. Defaults to the value set at initialization.
281 :param top_k_per_retriever:
282 The maximum number of documents to return per retriever. When set, this will override the `top_k`
283 parameter for each retriever. If None, the `top_k` parameter set for retrievers will be used.
284 Defaults to the value set at initialization.
285 :param top_k:
286 The maximum number of documents to return overall, extracted from the combined results of all
287 retrievers. When set, the results are always merged using reciprocal rank fusion (regardless of
288 `join_mode`) so that the combined list has a consistent global ranking before it is truncated to
289 `top_k`. If None, all results are returned. Defaults to the value set at initialization.
290 :param active_retrievers:
291 Names of retrievers to run. Defaults to all. Must match keys in the `retrievers` dictionary.
293 :returns:
294 A dictionary with the keys:
295 - "documents": A deduplicated list of retrieved documents.
297 :raises ValueError:
298 If any name in `active_retrievers` does not match a retriever name.
299 """
300 await self.warm_up_async()
302 resolved_top_k_per_retriever = (
303 top_k_per_retriever if top_k_per_retriever is not None else self.top_k_per_retriever
304 )
305 resolved_top_k = top_k if top_k is not None else self.top_k
306 resolved_filters = filters if filters is not None else self.filters
308 retrievers_to_run = self._resolve_retrievers(active_retrievers)
310 run_kwargs: dict[str, Any] = {"query": query}
311 if resolved_top_k_per_retriever is not None:
312 run_kwargs["top_k"] = resolved_top_k_per_retriever
313 if resolved_filters is not None:
314 run_kwargs["filters"] = resolved_filters
316 async def _run_one(name: str, retriever: TextRetriever) -> list[Document]:
317 try:
318 result = await _execute_component_async(retriever, **run_kwargs)
319 return result.get("documents", [])
320 except Exception as e:
321 raise RuntimeError(f"Retriever '{name}' failed: {e}") from e
323 tasks = [asyncio.create_task(_run_one(name, retriever)) for name, retriever in retrievers_to_run.items()]
324 document_lists = await _gather_tasks_with_cancel(tasks)
325 return {"documents": self._merge_results(document_lists, top_k=resolved_top_k)}
327 def to_dict(self) -> dict[str, Any]:
328 """
329 Serializes the component to a dictionary.
331 :returns:
332 Dictionary with serialized data.
333 """
334 return default_to_dict(
335 self,
336 retrievers={name: component_to_dict(obj=r, name=name) for name, r in self.retrievers.items()},
337 filters=self.filters,
338 top_k_per_retriever=self.top_k_per_retriever,
339 top_k=self.top_k,
340 max_workers=self.max_workers,
341 join_mode=self.join_mode,
342 )
344 @classmethod
345 def from_dict(cls, data: dict[str, Any]) -> "MultiRetriever":
346 """
347 Creates an instance of the component from a dictionary.
349 :param data:
350 Dictionary with the data to create the component.
351 """
352 retrievers_data = data.get("init_parameters", {}).get("retrievers", {})
353 if retrievers_data:
354 retrievers = {}
355 for name, retriever_data in retrievers_data.items():
356 try:
357 imported_class = import_class_by_name(retriever_data["type"])
358 except ImportError as e:
359 raise ImportError(
360 f"Could not import class {retriever_data['type']} for retriever '{name}'. Error: {str(e)}"
361 ) from e
362 retrievers[name] = component_from_dict(cls=imported_class, data=retriever_data, name=name)
363 data["init_parameters"]["retrievers"] = retrievers
364 return default_from_dict(cls, data)