Coverage for haystack/components/retrievers/multi_query_text_retriever.py: 96%
68 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +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 for parallel processing. Default is 3.
67 """
68 self.retriever = retriever
69 self.max_workers = max_workers
71 def warm_up(self) -> None:
72 """
73 Warm up the retriever.
74 """
75 if hasattr(self.retriever, "warm_up"):
76 self.retriever.warm_up()
78 async def warm_up_async(self) -> None:
79 """
80 Warm up the retriever on the serving event loop.
81 """
82 if hasattr(self.retriever, "warm_up_async"):
83 await self.retriever.warm_up_async()
84 elif hasattr(self.retriever, "warm_up"):
85 self.retriever.warm_up()
87 def close(self) -> None:
88 """
89 Release the retriever's resources.
90 """
91 if hasattr(self.retriever, "close"):
92 self.retriever.close()
94 async def close_async(self) -> None:
95 """
96 Release the retriever's async resources.
97 """
98 if hasattr(self.retriever, "close_async"):
99 await self.retriever.close_async()
100 elif hasattr(self.retriever, "close"):
101 self.retriever.close()
103 @component.output_types(documents=list[Document])
104 def run(self, queries: list[str], retriever_kwargs: dict[str, Any] | None = None) -> dict[str, list[Document]]:
105 """
106 Retrieve documents using multiple queries in parallel.
108 :param queries: List of text queries to process.
109 :param retriever_kwargs: Optional dictionary of arguments to pass to the retriever's run method.
110 :returns:
111 A dictionary containing:
112 `documents`: List of retrieved documents sorted by relevance score.
113 """
114 docs: list[Document] = []
115 retriever_kwargs = retriever_kwargs or {}
117 self.warm_up()
119 with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
120 queries_results = executor.map(lambda query: self._run_on_thread(query, retriever_kwargs), queries)
121 for result in queries_results:
122 if not result:
123 continue
124 docs.extend(result)
126 # de-duplicate and sort
127 docs = _deduplicate_documents(docs)
128 docs.sort(key=lambda x: x.score or 0.0, reverse=True)
129 return {"documents": docs}
131 @component.output_types(documents=list[Document])
132 async def run_async(
133 self, queries: list[str], retriever_kwargs: dict[str, Any] | None = None
134 ) -> dict[str, list[Document]]:
135 """
136 Retrieve documents using multiple queries concurrently.
138 Uses the retriever's `run_async` method if available, otherwise falls back to running `run`
139 in a thread executor. Queries are processed concurrently using asyncio.gather.
141 :param queries: List of text queries to process.
142 :param retriever_kwargs: Optional dictionary of arguments to pass to the retriever's run method.
143 :returns:
144 A dictionary containing:
145 `documents`: List of retrieved documents sorted by relevance score.
146 """
147 retriever_kwargs = retriever_kwargs or {}
149 await self.warm_up_async()
151 tasks = [asyncio.create_task(self._run_one_async(query, retriever_kwargs)) for query in queries]
152 results = await _gather_tasks_with_cancel(tasks)
153 docs: list[Document] = [doc for result in results if result for doc in result]
154 docs = _deduplicate_documents(docs)
155 docs.sort(key=lambda x: x.score or 0.0, reverse=True)
156 return {"documents": docs}
158 def _run_on_thread(self, query: str, retriever_kwargs: dict[str, Any] | None = None) -> list[Document] | None:
159 """
160 Process a single query on a separate thread.
162 :param query: The text query to process.
163 :param retriever_kwargs: Optional dictionary of arguments to pass to the retriever's run method.
164 :returns:
165 List of retrieved documents or None if no results.
166 """
167 result = self.retriever.run(query=query, **(retriever_kwargs or {}))
168 if result and "documents" in result:
169 return result["documents"]
170 return None
172 async def _run_one_async(self, query: str, retriever_kwargs: dict[str, Any]) -> list[Document] | None:
173 """
174 Process a single query asynchronously.
176 :param query: The text query to process.
177 :param retriever_kwargs: Arguments to pass to the retriever's run method.
178 :returns:
179 List of retrieved documents or None if no results.
180 """
181 result = await _execute_component_async(self.retriever, query=query, **retriever_kwargs)
183 if result and "documents" in result:
184 return result["documents"]
185 return None
187 def to_dict(self) -> dict[str, Any]:
188 """
189 Serializes the component to a dictionary.
191 :returns:
192 The serialized component as a dictionary.
193 """
194 return default_to_dict(
195 self, retriever=component_to_dict(obj=self.retriever, name="retriever"), max_workers=self.max_workers
196 )
198 @classmethod
199 def from_dict(cls, data: dict[str, Any]) -> "MultiQueryTextRetriever":
200 """
201 Deserializes the component from a dictionary.
203 :param data: The dictionary to deserialize from.
204 :returns:
205 The deserialized component.
206 """
207 return default_from_dict(cls, data)