Coverage for haystack/components/retrievers/multi_query_embedding_retriever.py: 98%
82 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.embedders.types.protocol import TextEmbedder
11from haystack.components.retrievers.types import EmbeddingRetriever
12from haystack.core.serialization import component_to_dict
13from haystack.utils.async_utils import _execute_component_async, _gather_tasks_with_cancel
14from haystack.utils.misc import _deduplicate_documents
17@component
18class MultiQueryEmbeddingRetriever:
19 """
20 A component that retrieves documents using multiple queries in parallel with an embedding-based retriever.
22 This component takes a list of text queries, converts them to embeddings using a query embedder,
23 and then uses an embedding-based retriever to find relevant documents for each query in parallel.
24 The results are combined and sorted by relevance score.
26 ### Usage example
28 ```python
29 from haystack import Document
30 from haystack.document_stores.in_memory import InMemoryDocumentStore
31 from haystack.document_stores.types import DuplicatePolicy
32 from haystack.components.embedders import OpenAITextEmbedder
33 from haystack.components.embedders import OpenAIDocumentEmbedder
34 from haystack.components.retrievers import InMemoryEmbeddingRetriever
35 from haystack.components.writers import DocumentWriter
36 from haystack.components.retrievers import MultiQueryEmbeddingRetriever
38 documents = [
39 Document(content="Renewable energy is energy that is collected from renewable resources."),
40 Document(content="Solar energy is a type of green energy that is harnessed from the sun."),
41 Document(content="Wind energy is another type of green energy that is generated by wind turbines."),
42 Document(content="Geothermal energy is heat that comes from the sub-surface of the earth."),
43 Document(content="Biomass energy is produced from organic materials, such as plant and animal waste."),
44 Document(content="Fossil fuels, such as coal, oil, and natural gas, are non-renewable energy sources."),
45 ]
47 # Populate the document store
48 doc_store = InMemoryDocumentStore()
49 doc_embedder = OpenAIDocumentEmbedder()
50 doc_writer = DocumentWriter(document_store=doc_store, policy=DuplicatePolicy.SKIP)
51 documents = doc_embedder.run(documents)["documents"]
52 doc_writer.run(documents=documents)
54 # Run the multi-query retriever
55 in_memory_retriever = InMemoryEmbeddingRetriever(document_store=doc_store, top_k=1)
56 query_embedder = OpenAITextEmbedder()
58 multi_query_retriever = MultiQueryEmbeddingRetriever(
59 retriever=in_memory_retriever,
60 query_embedder=query_embedder,
61 max_workers=3
62 )
64 queries = ["Geothermal energy", "natural gas", "turbines"]
65 result = multi_query_retriever.run(queries=queries)
66 for doc in result["documents"]:
67 print(f"Content: {doc.content}, Score: {doc.score}")
68 # >> Content: Geothermal energy is heat that comes from the sub-surface of the earth., Score: 0.8509603046266574
69 # >> Content: Renewable energy is energy that is collected from renewable resources., Score: 0.42763211298893034
70 # >> Content: Solar energy is a type of green energy that is harnessed from the sun., Score: 0.40077417016494354
71 # >> Content: Fossil fuels, such as coal, oil, and natural gas, are non-renewable energy sources., Score: 0.3774863680
72 # >> Content: Wind energy is another type of green energy that is generated by wind turbines., Score: 0.30914239725622
73 # >> Content: Biomass energy is produced from organic materials, such as plant and animal waste., Score: 0.25173074243
74 ```
75 """ # noqa E501
77 def __init__(self, *, retriever: EmbeddingRetriever, query_embedder: TextEmbedder, max_workers: int = 3) -> None:
78 """
79 Initialize MultiQueryEmbeddingRetriever.
81 :param retriever: The embedding-based retriever to use for document retrieval.
82 :param query_embedder: The query embedder to convert text queries to embeddings.
83 :param max_workers: Maximum number of worker threads in `run` and of concurrent retriever calls in
84 `run_async`.
85 """
86 self.retriever = retriever
87 self.query_embedder = query_embedder
88 self.max_workers = max_workers
90 def warm_up(self) -> None:
91 """
92 Warm up the query embedder and the retriever.
93 """
94 for inner in (self.query_embedder, self.retriever):
95 if hasattr(inner, "warm_up"):
96 inner.warm_up()
98 async def warm_up_async(self) -> None:
99 """
100 Warm up the query embedder and the retriever on the serving event loop.
101 """
102 for inner in (self.query_embedder, self.retriever):
103 if hasattr(inner, "warm_up_async"):
104 await inner.warm_up_async()
105 elif hasattr(inner, "warm_up"):
106 inner.warm_up()
108 def close(self) -> None:
109 """
110 Release the query embedder's and the retriever's resources.
111 """
112 for inner in (self.query_embedder, self.retriever):
113 if hasattr(inner, "close"):
114 inner.close()
116 async def close_async(self) -> None:
117 """
118 Release the query embedder's and the retriever's async resources.
119 """
120 for inner in (self.query_embedder, self.retriever):
121 if hasattr(inner, "close_async"):
122 await inner.close_async()
123 elif hasattr(inner, "close"):
124 inner.close()
126 @component.output_types(documents=list[Document])
127 def run(self, queries: list[str], retriever_kwargs: dict[str, Any] | None = None) -> dict[str, list[Document]]:
128 """
129 Retrieve documents using multiple queries in parallel.
131 :param queries: List of text queries to process.
132 :param retriever_kwargs: Optional dictionary of arguments to pass to the retriever's run method.
133 :returns:
134 A dictionary containing:
135 - `documents`: List of retrieved documents sorted by relevance score.
136 """
137 docs: list[Document] = []
138 retriever_kwargs = retriever_kwargs or {}
140 self.warm_up()
142 with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
143 queries_results = executor.map(lambda query: self._run_on_thread(query, retriever_kwargs), queries)
144 for result in queries_results:
145 if not result:
146 continue
147 docs.extend(result)
149 # de-duplicate and sort
150 docs = _deduplicate_documents(docs)
151 docs.sort(key=lambda x: x.score or 0.0, reverse=True)
152 return {"documents": docs}
154 @component.output_types(documents=list[Document])
155 async def run_async(
156 self, queries: list[str], retriever_kwargs: dict[str, Any] | None = None
157 ) -> dict[str, list[Document]]:
158 """
159 Retrieve documents using multiple queries concurrently.
161 Uses each component's `run_async` method if available, otherwise falls back to running `run`
162 in a thread executor. Queries are processed concurrently using asyncio.gather.
164 :param queries: List of text queries to process.
165 :param retriever_kwargs: Optional dictionary of arguments to pass to the retriever's run method.
166 :returns:
167 A dictionary containing:
168 - `documents`: List of retrieved documents sorted by relevance score.
169 """
170 retriever_kwargs = retriever_kwargs or {}
172 await self.warm_up_async()
174 # Bound concurrency to max_workers, mirroring the ThreadPoolExecutor in the sync `run`.
175 semaphore = asyncio.Semaphore(max(1, self.max_workers))
177 async def _bounded_run_one(query: str) -> list[Document] | None:
178 async with semaphore:
179 return await self._run_one_async(query, retriever_kwargs)
181 tasks = [asyncio.create_task(_bounded_run_one(query)) for query in queries]
182 results = await _gather_tasks_with_cancel(tasks)
183 docs: list[Document] = [doc for result in results if result for doc in result]
184 docs = _deduplicate_documents(docs)
185 docs.sort(key=lambda x: x.score or 0.0, reverse=True)
186 return {"documents": docs}
188 def _run_on_thread(self, query: str, retriever_kwargs: dict[str, Any] | None = None) -> list[Document] | None:
189 """
190 Process a single query on a separate thread.
192 :param query: The text query to process.
193 :param retriever_kwargs: Arguments to pass to the retriever's run method.
194 :returns:
195 List of retrieved documents or None if no results.
196 """
197 embedding_result = self.query_embedder.run(text=query)
198 query_embedding = embedding_result["embedding"]
199 result = self.retriever.run(query_embedding=query_embedding, **(retriever_kwargs or {}))
200 if result and "documents" in result:
201 return result["documents"]
202 return None
204 async def _run_one_async(self, query: str, retriever_kwargs: dict[str, Any]) -> list[Document] | None:
205 """
206 Process a single query asynchronously.
208 :param query: The text query to process.
209 :param retriever_kwargs: Arguments to pass to the retriever's run method.
210 :returns:
211 List of retrieved documents or None if no results.
212 """
213 embedding_result = await _execute_component_async(self.query_embedder, text=query)
215 query_embedding = embedding_result["embedding"]
217 result = await _execute_component_async(self.retriever, query_embedding=query_embedding, **retriever_kwargs)
219 if result and "documents" in result:
220 return result["documents"]
221 return None
223 def to_dict(self) -> dict[str, Any]:
224 """
225 Serializes the component to a dictionary.
227 :returns:
228 A dictionary representing the serialized component.
229 """
230 return default_to_dict(
231 self,
232 retriever=component_to_dict(obj=self.retriever, name="retriever"),
233 query_embedder=component_to_dict(obj=self.query_embedder, name="query_embedder"),
234 max_workers=self.max_workers,
235 )
237 @classmethod
238 def from_dict(cls, data: dict[str, Any]) -> "MultiQueryEmbeddingRetriever":
239 """
240 Deserializes the component from a dictionary.
242 :param data: The dictionary to deserialize from.
243 :returns:
244 The deserialized component.
245 """
246 return default_from_dict(cls, data)