Coverage for haystack/components/retrievers/multi_query_text_retriever.py: 96%
72 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
7from typing import Any
9from haystack import Document, component, default_from_dict, default_to_dict
10from haystack.components.retrievers.types import TextRetriever
11from haystack.core.serialization import component_to_dict
12from haystack.utils.async_utils import _execute_component_async, _gather_tasks_with_cancel
13from haystack.utils.misc import _deduplicate_documents
16@component
17class MultiQueryTextRetriever:
18 """
19 A component that retrieves documents using multiple queries in parallel with a text-based retriever.
21 This component takes a list of text queries and uses a text-based retriever to find relevant documents for each
22 query in parallel, using a thread pool to manage concurrent execution. The results are combined and sorted by
23 relevance score.
25 You can use this component in combination with QueryExpander component to enhance the retrieval process.
27 ### Usage example
28 ```python
29 from haystack import Document
30 from haystack.components.writers import DocumentWriter
31 from haystack.document_stores.in_memory import InMemoryDocumentStore
32 from haystack.document_stores.types import DuplicatePolicy
33 from haystack.components.retrievers import InMemoryBM25Retriever
34 from haystack.components.query import QueryExpander
35 from haystack.components.retrievers.multi_query_text_retriever import MultiQueryTextRetriever
37 documents = [
38 Document(content="Renewable energy is energy that is collected from renewable resources."),
39 Document(content="Solar energy is a type of green energy that is harnessed from the sun."),
40 Document(content="Wind energy is another type of green energy that is generated by wind turbines."),
41 Document(content="Hydropower is a form of renewable energy using the flow of water to generate electricity."),
42 Document(content="Geothermal energy is heat that comes from the sub-surface of the earth.")
43 ]
45 document_store = InMemoryDocumentStore()
46 doc_writer = DocumentWriter(document_store=document_store, policy=DuplicatePolicy.SKIP)
47 doc_writer.run(documents=documents)
49 in_memory_retriever = InMemoryBM25Retriever(document_store=document_store, top_k=1)
50 multiquery_retriever = MultiQueryTextRetriever(retriever=in_memory_retriever)
51 results = multiquery_retriever.run(queries=["renewable energy?", "Geothermal", "Hydropower"])
52 for doc in results["documents"]:
53 print(f"Content: {doc.content}, Score: {doc.score}")
54 # >>
55 # >> Content: Geothermal energy is heat that comes from the sub-surface of the earth., Score: 1.6474448833731097
56 # >> Content: Hydropower is a form of renewable energy using the flow of water to generate electricity., Score: 1.615
57 # >> Content: Renewable energy is energy that is collected from renewable resources., Score: 1.5255309812344944
58 ```
59 """ # noqa E501
61 def __init__(self, *, retriever: TextRetriever, max_workers: int = 3) -> None:
62 """
63 Initialize MultiQueryTextRetriever.
65 :param retriever: The text-based retriever to use for document retrieval.
66 :param max_workers: Maximum number of worker threads in `run` and of concurrent retriever calls in
67 `run_async`. Default is 3.
68 """
69 self.retriever = retriever
70 self.max_workers = max_workers
72 def warm_up(self) -> None:
73 """
74 Warm up the retriever.
75 """
76 if hasattr(self.retriever, "warm_up"):
77 self.retriever.warm_up()
79 async def warm_up_async(self) -> None:
80 """
81 Warm up the retriever on the serving event loop.
82 """
83 if hasattr(self.retriever, "warm_up_async"):
84 await self.retriever.warm_up_async()
85 elif hasattr(self.retriever, "warm_up"):
86 self.retriever.warm_up()
88 def close(self) -> None:
89 """
90 Release the retriever's resources.
91 """
92 if hasattr(self.retriever, "close"):
93 self.retriever.close()
95 async def close_async(self) -> None:
96 """
97 Release the retriever's async resources.
98 """
99 if hasattr(self.retriever, "close_async"):
100 await self.retriever.close_async()
101 elif hasattr(self.retriever, "close"):
102 self.retriever.close()
104 @component.output_types(documents=list[Document])
105 def run(self, queries: list[str], retriever_kwargs: dict[str, Any] | None = None) -> dict[str, list[Document]]:
106 """
107 Retrieve documents using multiple queries in parallel.
109 :param queries: List of text queries to process.
110 :param retriever_kwargs: Optional dictionary of arguments to pass to the retriever's run method.
111 :returns:
112 A dictionary containing:
113 `documents`: List of retrieved documents sorted by relevance score.
114 """
115 docs: list[Document] = []
116 retriever_kwargs = retriever_kwargs or {}
118 self.warm_up()
120 with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
121 queries_results = executor.map(lambda query: self._run_on_thread(query, retriever_kwargs), queries)
122 for result in queries_results:
123 if not result:
124 continue
125 docs.extend(result)
127 # de-duplicate and sort
128 docs = _deduplicate_documents(docs)
129 docs.sort(key=lambda x: x.score or 0.0, reverse=True)
130 return {"documents": docs}
132 @component.output_types(documents=list[Document])
133 async def run_async(
134 self, queries: list[str], retriever_kwargs: dict[str, Any] | None = None
135 ) -> dict[str, list[Document]]:
136 """
137 Retrieve documents using multiple queries concurrently.
139 Uses the retriever's `run_async` method if available, otherwise falls back to running `run`
140 in a thread executor. Queries are processed concurrently using asyncio.gather.
142 :param queries: List of text queries to process.
143 :param retriever_kwargs: Optional dictionary of arguments to pass to the retriever's run method.
144 :returns:
145 A dictionary containing:
146 `documents`: List of retrieved documents sorted by relevance score.
147 """
148 retriever_kwargs = retriever_kwargs or {}
150 await self.warm_up_async()
152 # Bound concurrency to max_workers, mirroring the ThreadPoolExecutor in the sync `run`.
153 semaphore = asyncio.Semaphore(max(1, self.max_workers))
155 async def _bounded_run_one(query: str) -> list[Document] | None:
156 async with semaphore:
157 return await self._run_one_async(query, retriever_kwargs)
159 tasks = [asyncio.create_task(_bounded_run_one(query)) for query in queries]
160 results = await _gather_tasks_with_cancel(tasks)
161 docs: list[Document] = [doc for result in results if result for doc in result]
162 docs = _deduplicate_documents(docs)
163 docs.sort(key=lambda x: x.score or 0.0, reverse=True)
164 return {"documents": docs}
166 def _run_on_thread(self, query: str, retriever_kwargs: dict[str, Any] | None = None) -> list[Document] | None:
167 """
168 Process a single query on a separate thread.
170 :param query: The text query to process.
171 :param retriever_kwargs: Optional dictionary of arguments to pass to the retriever's run method.
172 :returns:
173 List of retrieved documents or None if no results.
174 """
175 result = self.retriever.run(query=query, **(retriever_kwargs or {}))
176 if result and "documents" in result:
177 return result["documents"]
178 return None
180 async def _run_one_async(self, query: str, retriever_kwargs: dict[str, Any]) -> list[Document] | None:
181 """
182 Process a single query asynchronously.
184 :param query: The text query to process.
185 :param retriever_kwargs: Arguments to pass to the retriever's run method.
186 :returns:
187 List of retrieved documents or None if no results.
188 """
189 result = await _execute_component_async(self.retriever, query=query, **retriever_kwargs)
191 if result and "documents" in result:
192 return result["documents"]
193 return None
195 def to_dict(self) -> dict[str, Any]:
196 """
197 Serializes the component to a dictionary.
199 :returns:
200 The serialized component as a dictionary.
201 """
202 return default_to_dict(
203 self, retriever=component_to_dict(obj=self.retriever, name="retriever"), max_workers=self.max_workers
204 )
206 @classmethod
207 def from_dict(cls, data: dict[str, Any]) -> "MultiQueryTextRetriever":
208 """
209 Deserializes the component from a dictionary.
211 :param data: The dictionary to deserialize from.
212 :returns:
213 The deserialized component.
214 """
215 return default_from_dict(cls, data)