Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 63 additions & 18 deletions haystack/components/retrievers/multi_retriever.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand All @@ -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.
Comment on lines +125 to +127

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's document this in public facing docstrings, private methods are not in shown in our API reference.

"""
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]:
Expand Down Expand Up @@ -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,
Expand All @@ -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.

Expand All @@ -185,17 +206,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:
Expand All @@ -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,
Expand All @@ -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.

Expand All @@ -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]:
"""
Expand All @@ -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,
Expand Down
Original file line number Diff line number Diff line change
@@ -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
Loading
Loading