Coverage for haystack/components/retrievers/multi_retriever.py: 99%
131 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-09 14:34 +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 in `run` and of concurrent retriever calls in `run_async`.
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.
113 :raises ValueError:
114 If `top_k` or `top_k_per_retriever` is set and is not greater than 0.
115 """
116 if top_k is not None and top_k <= 0:
117 raise ValueError(f"top_k must be greater than 0, but got {top_k}")
118 if top_k_per_retriever is not None and top_k_per_retriever <= 0:
119 raise ValueError(f"top_k_per_retriever must be greater than 0, but got {top_k_per_retriever}")
120 self.retrievers = retrievers
121 self.filters = filters
122 self.top_k_per_retriever = top_k_per_retriever
123 self.top_k = top_k
124 self.max_workers = max_workers
125 self.join_mode = join_mode
127 def _resolve_top_k(self, top_k_per_retriever: int | None, top_k: int | None) -> tuple[int | None, int | None]:
128 """
129 Resolve runtime values against the init defaults and reject negatives.
131 A resolved value of 0 is valid and means no documents are returned.
132 """
133 resolved_top_k_per_retriever = (
134 top_k_per_retriever if top_k_per_retriever is not None else self.top_k_per_retriever
135 )
136 resolved_top_k = top_k if top_k is not None else self.top_k
137 if resolved_top_k is not None and resolved_top_k < 0:
138 raise ValueError(f"top_k must be greater than or equal to 0, but got {resolved_top_k}")
139 if resolved_top_k_per_retriever is not None and resolved_top_k_per_retriever < 0:
140 raise ValueError(
141 f"top_k_per_retriever must be greater than or equal to 0, but got {resolved_top_k_per_retriever}"
142 )
143 return resolved_top_k_per_retriever, resolved_top_k
145 def _merge_results(self, document_lists: list[list[Document]], top_k: int | None = None) -> list[Document]:
146 """
147 Merge per-retriever result lists according to `join_mode`.
149 In `concatenate` mode, all lists are flattened and deduplicated. In `reciprocal_rank_fusion` mode, results
150 are deduplicated and re-scored using RRF, then returned in descending score order. When `top_k` is set, RRF
151 is always used so the combined results have a consistent global ranking, and only the top `top_k` documents
152 are returned.
153 """
154 # When top_k is set we always use reciprocal rank fusion to merge the results, regardless of join_mode,
155 # so that truncation is applied to a consistently ranked list.
156 if top_k is not None or self.join_mode == "reciprocal_rank_fusion":
157 documents = _reciprocal_rank_fusion(document_lists)
158 merged = sorted(documents, key=lambda d: d.score if d.score is not None else -inf, reverse=True)
159 return merged[:top_k] if top_k is not None else merged
160 return _deduplicate_documents([doc for docs in document_lists for doc in docs])
162 def _resolve_retrievers(self, active_retrievers: list[str] | None) -> dict[str, TextRetriever]:
163 """
164 Returns the subset of retrievers to run based on the active_retrievers list.
166 :param active_retrievers:
167 A list of retriever names to run. If None, all retrievers are returned.
169 :returns:
170 A dictionary of retriever names to retriever instances.
172 :raises ValueError:
173 If any name in `active_retrievers` does not match a retriever name.
174 """
175 if active_retrievers is None:
176 return self.retrievers
177 unknown = set(active_retrievers) - self.retrievers.keys()
178 if unknown:
179 raise ValueError(
180 f"Unknown retriever name(s): {sorted(unknown)}. Available retrievers: {sorted(self.retrievers.keys())}"
181 )
182 return {name: self.retrievers[name] for name in active_retrievers}
184 def warm_up(self) -> None:
185 """
186 Warm up the retrievers.
187 """
188 for retriever in self.retrievers.values():
189 if hasattr(retriever, "warm_up"):
190 retriever.warm_up()
192 async def warm_up_async(self) -> None:
193 """
194 Warm up the retrievers on the serving event loop.
195 """
196 for retriever in self.retrievers.values():
197 if hasattr(retriever, "warm_up_async"):
198 await retriever.warm_up_async()
199 elif hasattr(retriever, "warm_up"):
200 retriever.warm_up()
202 def close(self) -> None:
203 """
204 Release the retrievers' resources.
205 """
206 for retriever in self.retrievers.values():
207 if hasattr(retriever, "close"):
208 retriever.close()
210 async def close_async(self) -> None:
211 """
212 Release the retrievers' async resources.
213 """
214 for retriever in self.retrievers.values():
215 if hasattr(retriever, "close_async"):
216 await retriever.close_async()
217 elif hasattr(retriever, "close"):
218 retriever.close()
220 @component.output_types(documents=list[Document])
221 def run(
222 self,
223 query: str,
224 filters: dict[str, Any] | None = None,
225 top_k_per_retriever: int | None = None,
226 top_k: int | None = None,
227 *,
228 active_retrievers: list[str] | None = None,
229 ) -> dict[str, list[Document]]:
230 """
231 Runs retrievers in parallel on the given query and returns deduplicated results.
233 :param query:
234 The query to run the retrievers on.
235 :param filters:
236 Filters to apply. Defaults to the value set at initialization.
237 :param top_k_per_retriever:
238 The maximum number of documents to return per retriever. When set, this will override the `top_k`
239 parameter for each retriever. If None, the `top_k` parameter set for retrievers will be used.
240 If 0, no documents are returned. Defaults to the value set at initialization.
241 :param top_k:
242 The maximum number of documents to return overall, extracted from the combined results of all
243 retrievers. When set, the results are always merged using reciprocal rank fusion (regardless of
244 `join_mode`) so that the combined list has a consistent global ranking before it is truncated to
245 `top_k`. If None, all results are returned. If 0, no documents are returned.
246 Defaults to the value set at initialization.
247 :param active_retrievers:
248 Names of retrievers to run. Defaults to all. Must match keys in the `retrievers` dictionary.
250 :returns:
251 A dictionary with the keys:
252 - "documents": A deduplicated list of retrieved documents.
254 :raises ValueError:
255 If any name in `active_retrievers` does not match a retriever name,
256 or if the resolved `top_k` or `top_k_per_retriever` is negative.
257 """
258 self.warm_up()
260 resolved_top_k_per_retriever, resolved_top_k = self._resolve_top_k(top_k_per_retriever, top_k)
261 resolved_filters = filters if filters is not None else self.filters
263 retrievers_to_run = self._resolve_retrievers(active_retrievers)
264 if resolved_top_k == 0 or resolved_top_k_per_retriever == 0:
265 return {"documents": []}
267 results_by_name: dict[str, list[Document]] = {}
268 with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
269 future_to_name = {}
270 for name, retriever in retrievers_to_run.items():
271 run_kwargs: dict[str, Any] = {"query": query}
272 if resolved_top_k_per_retriever is not None:
273 run_kwargs["top_k"] = resolved_top_k_per_retriever
274 if resolved_filters is not None:
275 run_kwargs["filters"] = resolved_filters
276 future_to_name[executor.submit(retriever.run, **run_kwargs)] = name
278 for future in as_completed(future_to_name):
279 name = future_to_name[future]
280 try:
281 results_by_name[name] = future.result().get("documents", [])
282 except Exception as e:
283 raise RuntimeError(f"Retriever '{name}' failed: {e}") from e
285 document_lists = [results_by_name[name] for name in retrievers_to_run]
286 return {"documents": self._merge_results(document_lists, top_k=resolved_top_k)}
288 @component.output_types(documents=list[Document])
289 async def run_async(
290 self,
291 query: str,
292 filters: dict[str, Any] | None = None,
293 top_k_per_retriever: int | None = None,
294 top_k: int | None = None,
295 *,
296 active_retrievers: list[str] | None = None,
297 ) -> dict[str, list[Document]]:
298 """
299 Runs retrievers concurrently on the given query and returns deduplicated results.
301 Uses each retriever's `run_async` method if available, otherwise runs `run` in a thread executor.
303 :param query:
304 The query to run the retrievers on.
305 :param filters:
306 Filters to apply. Defaults to the value set at initialization.
307 :param top_k_per_retriever:
308 The maximum number of documents to return per retriever. When set, this will override the `top_k`
309 parameter for each retriever. If None, the `top_k` parameter set for retrievers will be used.
310 If 0, no documents are returned. Defaults to the value set at initialization.
311 :param top_k:
312 The maximum number of documents to return overall, extracted from the combined results of all
313 retrievers. When set, the results are always merged using reciprocal rank fusion (regardless of
314 `join_mode`) so that the combined list has a consistent global ranking before it is truncated to
315 `top_k`. If None, all results are returned. If 0, no documents are returned.
316 Defaults to the value set at initialization.
317 :param active_retrievers:
318 Names of retrievers to run. Defaults to all. Must match keys in the `retrievers` dictionary.
320 :returns:
321 A dictionary with the keys:
322 - "documents": A deduplicated list of retrieved documents.
324 :raises ValueError:
325 If any name in `active_retrievers` does not match a retriever name,
326 or if the resolved `top_k` or `top_k_per_retriever` is negative.
327 """
328 await self.warm_up_async()
330 resolved_top_k_per_retriever, resolved_top_k = self._resolve_top_k(top_k_per_retriever, top_k)
331 resolved_filters = filters if filters is not None else self.filters
333 retrievers_to_run = self._resolve_retrievers(active_retrievers)
334 if resolved_top_k == 0 or resolved_top_k_per_retriever == 0:
335 return {"documents": []}
337 run_kwargs: dict[str, Any] = {"query": query}
338 if resolved_top_k_per_retriever is not None:
339 run_kwargs["top_k"] = resolved_top_k_per_retriever
340 if resolved_filters is not None:
341 run_kwargs["filters"] = resolved_filters
343 # Bound concurrency to max_workers, mirroring the ThreadPoolExecutor in the sync `run`.
344 semaphore = asyncio.Semaphore(max(1, self.max_workers))
346 async def _run_one(name: str, retriever: TextRetriever) -> list[Document]:
347 async with semaphore:
348 try:
349 result = await _execute_component_async(retriever, **run_kwargs)
350 return result.get("documents", [])
351 except Exception as e:
352 raise RuntimeError(f"Retriever '{name}' failed: {e}") from e
354 tasks = [asyncio.create_task(_run_one(name, retriever)) for name, retriever in retrievers_to_run.items()]
355 document_lists = await _gather_tasks_with_cancel(tasks)
356 return {"documents": self._merge_results(document_lists, top_k=resolved_top_k)}
358 def to_dict(self) -> dict[str, Any]:
359 """
360 Serializes the component to a dictionary.
362 :returns:
363 Dictionary with serialized data.
364 """
365 return default_to_dict(
366 self,
367 retrievers={name: component_to_dict(obj=r, name=name) for name, r in self.retrievers.items()},
368 filters=self.filters,
369 top_k_per_retriever=self.top_k_per_retriever,
370 top_k=self.top_k,
371 max_workers=self.max_workers,
372 join_mode=self.join_mode,
373 )
375 @classmethod
376 def from_dict(cls, data: dict[str, Any]) -> "MultiRetriever":
377 """
378 Creates an instance of the component from a dictionary.
380 :param data:
381 Dictionary with the data to create the component.
382 """
383 retrievers_data = data.get("init_parameters", {}).get("retrievers", {})
384 if retrievers_data:
385 retrievers = {}
386 for name, retriever_data in retrievers_data.items():
387 try:
388 imported_class = import_class_by_name(retriever_data["type"])
389 except ImportError as e:
390 raise ImportError(
391 f"Could not import class {retriever_data['type']} for retriever '{name}'. Error: {str(e)}"
392 ) from e
393 retrievers[name] = component_from_dict(cls=imported_class, data=retriever_data, name=name)
394 data["init_parameters"]["retrievers"] = retrievers
395 return default_from_dict(cls, data)