Fix params

Author: Amnah199

haystack/components/retrievers/multi_retriever.py

@@ -81,7 +81,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:
@@ -93,8 +94,12 @@ 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. When set, this will extract the top_k documents
+            from the combined results of all retrievers. If None, all results are returned.
         :param max_workers:
             The maximum number of threads to use for parallel retrieval.
         :param join_mode:
@@ -104,21 +109,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]:
@@ -159,6 +170,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,
@@ -170,8 +182,14 @@ 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. When set, this will extract the top_k documents
+            from the combined results of all retrievers. 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.
 
@@ -185,17 +203,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:
@@ -204,13 +230,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,
@@ -224,8 +251,14 @@ 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. When set, this will extract the top_k documents
+            from the combined results of all retrievers. 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.
 
@@ -239,27 +272,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 r=retriever: r.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]:
         """

test/components/retrievers/test_multi_retriever.py

@@ -161,7 +161,7 @@ def test_run_deduplicates_results(self, sample_documents):
         ids = [doc.id for doc in result["documents"]]
         assert ids.count("doc1") == 1
 
-    def test_run_resolves_filters_and_top_k(self):
+    def test_run_resolves_filters_and_top_k_per_retriever(self):
         received: dict = {}
 
         @component
@@ -173,19 +173,68 @@ def run(self, query: str, filters: dict[str, Any] | None = None, top_k: int | No
                 return {"documents": []}
 
         retriever = MultiRetriever(
-            retrievers={"capturing": CapturingRetriever()}, filters={"field": "meta.category"}, top_k=5
+            retrievers={"capturing": CapturingRetriever()}, filters={"field": "meta.category"}, top_k_per_retriever=5
         )
 
-        # Should use init-time values when not overridden
+        # Should use init-time values when not overridden (top_k_per_retriever is forwarded as the retriever's top_k)
         retriever.run(query="energy")
         assert received["filters"] == {"field": "meta.category"}
         assert received["top_k"] == 5
 
         # Should prefer run-time values when provided
-        retriever.run(query="energy", filters={"field": "meta.other"}, top_k=2)
+        retriever.run(query="energy", filters={"field": "meta.other"}, top_k_per_retriever=2)
         assert received["filters"] == {"field": "meta.other"}
         assert received["top_k"] == 2
 
+    def test_run_forwards_top_k_per_retriever_not_overall_top_k(self):
+        received: dict = {}
+
+        @component
+        class CapturingRetriever:
+            @component.output_types(documents=list[Document])
+            def run(self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None):
+                received["top_k"] = top_k
+                return {"documents": []}
+
+        retriever = MultiRetriever(retrievers={"capturing": CapturingRetriever()})
+
+        # top_k_per_retriever is forwarded to each retriever as its top_k
+        retriever.run(query="energy", top_k_per_retriever=3)
+        assert received["top_k"] == 3
+
+        # the overall top_k is applied at merge-time only, not forwarded to retrievers
+        received.clear()
+        retriever.run(query="energy", top_k=5)
+        assert received.get("top_k") is None
+
+    def test_run_top_k_truncates_merged_results(self, sample_documents):
+        retriever = MultiRetriever(
+            retrievers={
+                "a": MockRetriever(documents=sample_documents[:3]),
+                "b": MockRetriever(documents=sample_documents[2:5]),
+            },
+            max_workers=2,
+        )
+        result = retriever.run(query="energy", top_k=2)
+        assert len(result["documents"]) == 2
+        scores = [doc.score for doc in result["documents"]]
+        assert all(score is not None for score in scores)
+        assert scores == sorted(scores, reverse=True)
+
+    def test_run_top_k_forces_rrf_in_concatenate_mode(self, sample_documents):
+        # In concatenate mode there is no global ranking, so setting top_k falls back to RRF to truncate consistently
+        retriever = MultiRetriever(
+            retrievers={
+                "a": MockRetriever(documents=sample_documents[:3]),
+                "b": MockRetriever(documents=sample_documents[1:4]),
+            },
+            join_mode="concatenate",
+            max_workers=2,
+        )
+        result = retriever.run(query="energy", top_k=2)
+        assert len(result["documents"]) == 2
+        assert all(doc.score is not None for doc in result["documents"])
+
     def test_run_with_active_retrievers(self, sample_documents):
         retriever = MultiRetriever(
             retrievers={"a": MockRetriever([sample_documents[0]]), "b": MockRetriever([sample_documents[1]])}
@@ -379,7 +428,7 @@ async def test_run_async_rrf_assigns_scores_and_sorts(self, sample_documents):
         assert ids.index("doc1") < ids.index("doc3")
 
     @pytest.mark.asyncio
-    async def test_run_async_resolves_filters_and_top_k(self):
+    async def test_run_async_resolves_filters_and_top_k_per_retriever(self):
         received: dict = {}
 
         @component
@@ -391,17 +440,68 @@ def run(self, query: str, filters: dict[str, Any] | None = None, top_k: int | No
                 return {"documents": []}
 
         retriever = MultiRetriever(
-            retrievers={"capturing": CapturingRetriever()}, filters={"field": "meta.category"}, top_k=5
+            retrievers={"capturing": CapturingRetriever()}, filters={"field": "meta.category"}, top_k_per_retriever=5
         )
 
+        # top_k_per_retriever is forwarded as the retriever's top_k
         await retriever.run_async(query="energy")
         assert received["filters"] == {"field": "meta.category"}
         assert received["top_k"] == 5
 
-        await retriever.run_async(query="energy", filters={"field": "meta.other"}, top_k=2)
+        await retriever.run_async(query="energy", filters={"field": "meta.other"}, top_k_per_retriever=2)
         assert received["filters"] == {"field": "meta.other"}
         assert received["top_k"] == 2
 
+    @pytest.mark.asyncio
+    async def test_run_async_forwards_top_k_per_retriever_not_overall_top_k(self):
+        received: dict = {}
+
+        @component
+        class CapturingRetriever:
+            @component.output_types(documents=list[Document])
+            def run(self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None):
+                received["top_k"] = top_k
+                return {"documents": []}
+
+        retriever = MultiRetriever(retrievers={"capturing": CapturingRetriever()})
+
+        # top_k_per_retriever is forwarded to each retriever as its top_k
+        await retriever.run_async(query="energy", top_k_per_retriever=3)
+        assert received["top_k"] == 3
+
+        # the overall top_k is applied at merge-time only, not forwarded to retrievers
+        received.clear()
+        await retriever.run_async(query="energy", top_k=5)
+        assert received.get("top_k") is None
+
+    @pytest.mark.asyncio
+    async def test_run_async_top_k_truncates_merged_results(self, sample_documents):
+        retriever = MultiRetriever(
+            retrievers={
+                "a": MockRetriever(documents=sample_documents[:3]),
+                "b": MockRetriever(documents=sample_documents[2:5]),
+            }
+        )
+        result = await retriever.run_async(query="energy", top_k=2)
+        assert len(result["documents"]) == 2
+        scores = [doc.score for doc in result["documents"]]
+        assert all(score is not None for score in scores)
+        assert scores == sorted(scores, reverse=True)
+
+    @pytest.mark.asyncio
+    async def test_run_async_top_k_forces_rrf_in_concatenate_mode(self, sample_documents):
+        # In concatenate mode there is no global ranking, so setting top_k falls back to RRF to truncate consistently
+        retriever = MultiRetriever(
+            retrievers={
+                "a": MockRetriever(documents=sample_documents[:3]),
+                "b": MockRetriever(documents=sample_documents[1:4]),
+            },
+            join_mode="concatenate",
+        )
+        result = await retriever.run_async(query="energy", top_k=2)
+        assert len(result["documents"]) == 2
+        assert all(doc.score is not None for doc in result["documents"])
+
     @pytest.mark.asyncio
     async def test_run_async_with_active_retrievers(self, sample_documents):
         retriever = MultiRetriever(