fix: update Multiretriever params (#11823)

Co-authored-by: bogdankostic <bogdankostic@web.de>

Author: Amnah199

haystack/components/retrievers/multi_retriever.py

@@ -83,7 +83,8 @@ def __init__(
         *,
         retrievers: dict[str, TextRetriever],
         filters: dict[str, Any] | None = None,
-        top_k: int = 10,
+        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:
@@ -95,8 +96,14 @@ def __init__(
             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 per retriever.
+            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:
@@ -106,21 +113,27 @@ def __init__(
         """
         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
         self._is_warmed_up = False
 
-    def _merge_results(self, document_lists: list[list[Document]]) -> list[Document]:
+    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.
+        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.
         """
-        if self.join_mode == "reciprocal_rank_fusion":
+        # 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)
-            return sorted(documents, key=lambda d: d.score if d.score is not None else -inf, reverse=True)
+            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]:
@@ -161,6 +174,7 @@ 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,
@@ -172,8 +186,15 @@ def run(
             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:
-            Maximum documents to return per retriever. Defaults to the value set at initialization.
+            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.
 
@@ -187,17 +208,25 @@ def run(
         if not self._is_warmed_up:
             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 = {
-                executor.submit(retriever.run, query=query, filters=resolved_filters, top_k=resolved_top_k): name
-                for name, retriever in retrievers_to_run.items()
-            }
+            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:
@@ -206,13 +235,14 @@ def run(
                     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)}
+        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,
@@ -226,8 +256,15 @@ async def run_async(
             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:
-            Maximum documents to return per retriever. Defaults to the value set at initialization.
+            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.
 
@@ -241,27 +278,34 @@ async def run_async(
         if not self._is_warmed_up:
             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)
 
+        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
+
         loop = asyncio.get_running_loop()
 
         async def _run_one(name: str, retriever: TextRetriever) -> list[Document]:
             try:
                 if hasattr(retriever, "run_async") and callable(retriever.run_async):
-                    result = await retriever.run_async(query=query, filters=resolved_filters, top_k=resolved_top_k)
+                    result = await retriever.run_async(**run_kwargs)
                 else:
-                    result = await loop.run_in_executor(
-                        None, lambda: retriever.run(query=query, filters=resolved_filters, top_k=resolved_top_k)
-                    )
+                    result = await loop.run_in_executor(None, lambda: retriever.run(**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)}
+        return {"documents": self._merge_results(document_lists, top_k=resolved_top_k)}
 
     def to_dict(self) -> dict[str, Any]:
         """
@@ -274,6 +318,7 @@ def to_dict(self) -> dict[str, Any]:
             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,

releasenotes/notes/fix-multiretriever-params-c893b8d7e41afde4.yaml

@@ -0,0 +1,6 @@
+---
+fixes:
+  - |
+    Update parameters in MultiRetriever to allow for more flexible retrieval of documents. Changes include:
+    - ``top_k_per_retriever``: Allows specifying the maximum number of documents to return per retriever
+    - ``top_k``: Allows specifying the maximum number of documents to be retrieved from combined results of all retrievers