diff --git a/haystack/components/retrievers/multi_retriever.py b/haystack/components/retrievers/multi_retriever.py index 911fa47cb07..5c6c6b9a0fd 100644 --- a/haystack/components/retrievers/multi_retriever.py +++ b/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,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: @@ -104,21 +111,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 +172,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 +184,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. @@ -185,6 +206,9 @@ 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 @@ -192,10 +216,15 @@ def run( 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 +233,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 +254,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. @@ -239,27 +276,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]: """ @@ -272,6 +316,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, diff --git a/releasenotes/notes/fix-multiretriever-params-c893b8d7e41afde4.yaml b/releasenotes/notes/fix-multiretriever-params-c893b8d7e41afde4.yaml new file mode 100644 index 00000000000..f8aa56decaf --- /dev/null +++ b/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 diff --git a/test/components/retrievers/test_multi_retriever.py b/test/components/retrievers/test_multi_retriever.py index 6dcd2518c4c..914f148d011 100644 --- a/test/components/retrievers/test_multi_retriever.py +++ b/test/components/retrievers/test_multi_retriever.py @@ -101,7 +101,8 @@ def test_init_default_parameters(self): retriever = MultiRetriever(retrievers=retrievers) assert retriever.retrievers == retrievers assert retriever.filters is None - assert retriever.top_k == 10 + assert retriever.top_k_per_retriever is None + assert retriever.top_k is None assert retriever.max_workers == 4 assert retriever.join_mode == "reciprocal_rank_fusion" @@ -161,7 +162,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 +174,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]])} @@ -209,6 +259,7 @@ def test_to_dict(self): retriever = MultiRetriever( retrievers={"bm25": InMemoryBM25Retriever(document_store=InMemoryDocumentStore())}, filters=None, + top_k_per_retriever=3, top_k=5, max_workers=2, ) @@ -239,6 +290,7 @@ def test_to_dict(self): } }, "filters": None, + "top_k_per_retriever": 3, "top_k": 5, "max_workers": 2, "join_mode": "reciprocal_rank_fusion", @@ -272,6 +324,7 @@ def test_from_dict(self): } }, "filters": None, + "top_k_per_retriever": 3, "top_k": 5, "max_workers": 2, "join_mode": "concatenate", @@ -282,6 +335,7 @@ def test_from_dict(self): assert len(result.retrievers) == 1 assert "bm25" in result.retrievers assert isinstance(result.retrievers["bm25"], InMemoryBM25Retriever) + assert result.top_k_per_retriever == 3 assert result.top_k == 5 assert result.max_workers == 2 assert result.join_mode == "concatenate" @@ -379,7 +433,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 +445,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(