Coverage for haystack/components/query/query_expander.py: 93%
137 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
5from typing import Any
7from haystack import default_from_dict, default_to_dict, logging
8from haystack.components.builders.prompt_builder import PromptBuilder
9from haystack.components.generators.chat.openai import OpenAIChatGenerator
10from haystack.components.generators.chat.types import ChatGenerator
11from haystack.components.generators.utils import _trace_chat_generator_run
12from haystack.core.component import component
13from haystack.core.serialization import component_to_dict
14from haystack.dataclasses.chat_message import ChatMessage
15from haystack.utils import deserialize_chatgenerator_inplace
16from haystack.utils.async_utils import _execute_component_async
17from haystack.utils.misc import _parse_dict_from_json
19logger = logging.getLogger(__name__)
22DEFAULT_PROMPT_TEMPLATE = """
23You are part of an information system that processes user queries for retrieval.
24You have to expand a given query into {{ n_expansions }} queries that are
25semantically similar to improve retrieval recall.
27Structure:
28Follow the structure shown below in examples to generate expanded queries.
30Examples:
311. Query: "climate change effects"
32 {"queries": ["impact of climate change", "consequences of global warming", "effects of environmental changes"]}
342. Query: "machine learning algorithms"
35 {"queries": ["neural networks", "clustering techniques", "supervised learning methods", "deep learning models"]}
373. Query: "open source NLP frameworks"
38 {"queries": ["natural language processing tools", "free nlp libraries", "open-source NLP platforms"]}
40Guidelines:
41- Generate queries that use different words and phrasings
42- Include synonyms and related terms
43- Maintain the same core meaning and intent
44- Make queries that are likely to retrieve relevant information the original might miss
45- Focus on variations that would work well with keyword-based search
46- Respond in the same language as the input query
48Your Task:
49Query: "{{ query }}"
51You *must* respond with a JSON object containing a "queries" array with the expanded queries.
52Example: {"queries": ["query1", "query2", "query3"]}"""
55@component
56class QueryExpander:
57 """
58 A component that returns a list of semantically similar queries to improve retrieval recall in RAG systems.
60 The component uses a chat generator to expand queries. The chat generator is expected to return a JSON response
61 with the following structure:
62 ```json
63 {"queries": ["expanded query 1", "expanded query 2", "expanded query 3"]}
64 ```
66 ### Usage example
68 ```python
69 from haystack.components.generators.chat.openai import OpenAIChatGenerator
70 from haystack.components.query import QueryExpander
72 expander = QueryExpander(
73 chat_generator=OpenAIChatGenerator(model="gpt-4.1-mini"),
74 n_expansions=3
75 )
77 result = expander.run(query="green energy sources")
78 print(result["queries"])
79 # Output: ['alternative query 1', 'alternative query 2', 'alternative query 3', 'green energy sources']
80 # Note: Up to 3 additional queries + 1 original query (if include_original_query=True)
82 # To control total number of queries:
83 expander = QueryExpander(n_expansions=2, include_original_query=True) # Up to 3 total
84 # or
85 expander = QueryExpander(n_expansions=3, include_original_query=False) # Exactly 3 total
86 ```
87 """
89 def __init__(
90 self,
91 *,
92 chat_generator: ChatGenerator | None = None,
93 prompt_template: str | None = None,
94 n_expansions: int = 4,
95 include_original_query: bool = True,
96 ) -> None:
97 """
98 Initialize the QueryExpander component.
100 :param chat_generator: The chat generator component to use for query expansion.
101 If None, a default OpenAIChatGenerator with gpt-4.1-mini model is used.
102 :param prompt_template: Custom [PromptBuilder](https://docs.haystack.deepset.ai/docs/promptbuilder)
103 template for query expansion. The template should instruct the LLM to return a JSON response with the
104 structure: `{"queries": ["query1", "query2", "query3"]}`. The template should include 'query' and
105 'n_expansions' variables.
106 :param n_expansions: Number of alternative queries to generate (default: 4).
107 :param include_original_query: Whether to include the original query in the output.
108 """
109 if n_expansions <= 0:
110 raise ValueError("n_expansions must be positive")
112 self.n_expansions = n_expansions
113 self.include_original_query = include_original_query
115 if chat_generator is None:
116 self.chat_generator: ChatGenerator = OpenAIChatGenerator(
117 model="gpt-4.1-mini",
118 generation_kwargs={
119 "temperature": 0.7,
120 "response_format": {
121 "type": "json_schema",
122 "json_schema": {
123 "name": "query_expansion",
124 "schema": {
125 "type": "object",
126 "properties": {"queries": {"type": "array", "items": {"type": "string"}}},
127 "required": ["queries"],
128 "additionalProperties": False,
129 },
130 },
131 },
132 "seed": 42,
133 },
134 )
135 else:
136 self.chat_generator = chat_generator
138 self.prompt_template = prompt_template or DEFAULT_PROMPT_TEMPLATE
140 # Check if required variables are present in the template
141 if "query" not in self.prompt_template:
142 logger.warning(
143 "The prompt template does not contain the 'query' variable. This may cause issues during execution."
144 )
145 if "n_expansions" not in self.prompt_template:
146 logger.warning(
147 "The prompt template does not contain the 'n_expansions' variable. "
148 "This may cause issues during execution."
149 )
151 self._prompt_builder = PromptBuilder(
152 template=self.prompt_template, required_variables=["n_expansions", "query"]
153 )
155 def to_dict(self) -> dict[str, Any]:
156 """
157 Serializes the component to a dictionary.
159 :return: Dictionary with serialized data.
160 """
161 return default_to_dict(
162 self,
163 chat_generator=component_to_dict(self.chat_generator, name="chat_generator"),
164 prompt_template=self.prompt_template,
165 n_expansions=self.n_expansions,
166 include_original_query=self.include_original_query,
167 )
169 @classmethod
170 def from_dict(cls, data: dict[str, Any]) -> "QueryExpander":
171 """
172 Deserializes the component from a dictionary.
174 :param data: Dictionary with serialized data.
175 :return: Deserialized component.
176 """
177 init_params = data.get("init_parameters", {})
179 deserialize_chatgenerator_inplace(init_params, key="chat_generator")
181 return default_from_dict(cls, data)
183 @component.output_types(queries=list[str])
184 def run(self, query: str, n_expansions: int | None = None) -> dict[str, list[str]]:
185 """
186 Expand the input query into multiple semantically similar queries.
188 The language of the original query is preserved in the expanded queries.
190 :param query: The original query to expand.
191 :param n_expansions: Number of additional queries to generate (not including the original).
192 If None, uses the value from initialization. Must be a positive integer.
193 :return: Dictionary with "queries" key containing the list of expanded queries.
194 If include_original_query=True, the original query will be included in addition
195 to the n_expansions alternative queries.
196 :raises ValueError: If n_expansions is not positive (less than or equal to 0).
197 """
199 self.warm_up()
201 response = {"queries": [query] if self.include_original_query else []}
203 if not query.strip():
204 logger.warning("Empty query provided to QueryExpander")
205 return response
207 expansion_count = n_expansions if n_expansions is not None else self.n_expansions
208 if expansion_count <= 0:
209 raise ValueError("n_expansions must be positive")
211 try:
212 prompt_result = self._prompt_builder.run(query=query.strip(), n_expansions=expansion_count)
213 messages = [ChatMessage.from_user(prompt_result["prompt"])]
214 with _trace_chat_generator_run(self.chat_generator, {"messages": messages}) as span:
215 generator_result = self.chat_generator.run(messages=messages)
216 span.set_content_tag("haystack.component.output", generator_result)
218 if not generator_result.get("replies") or len(generator_result["replies"]) == 0:
219 logger.warning("ChatGenerator returned no replies for query: {query}", query=query)
220 return response
222 expanded_text = generator_result["replies"][0].text.strip()
223 expanded_queries = self._parse_expanded_queries(expanded_text)
225 # Limit the number of expanded queries to the requested amount
226 if len(expanded_queries) > expansion_count:
227 logger.warning(
228 "Generated {generated_count} queries but only {requested_count} were requested. "
229 "Truncating to the first {requested_count} queries. ",
230 generated_count=len(expanded_queries),
231 requested_count=expansion_count,
232 )
233 expanded_queries = expanded_queries[:expansion_count]
235 # Add original query if not already present
236 if self.include_original_query:
237 expanded_queries_lower = [q.lower() for q in expanded_queries]
238 if query.lower() not in expanded_queries_lower:
239 expanded_queries.append(query)
241 response["queries"] = expanded_queries
242 return response
244 except Exception as e:
245 # Fallback: return original query to maintain pipeline functionality
246 logger.exception("Failed to expand query {query}: {error}", query=query, error=str(e))
247 return response
249 @component.output_types(queries=list[str])
250 async def run_async(self, query: str, n_expansions: int | None = None) -> dict[str, list[str]]:
251 """
252 Asynchronously expand the input query into multiple semantically similar queries.
254 The language of the original query is preserved in the expanded queries.
256 This is the asynchronous version of the `run` method. It has the same parameters and return values
257 but can be used with `await` in an async code. If the chat generator only implements a synchronous
258 `run` method, it is executed in a thread to avoid blocking the event loop.
260 :param query: The original query to expand.
261 :param n_expansions: Number of additional queries to generate (not including the original).
262 If None, uses the value from initialization. Must be a positive integer.
263 :return: Dictionary with "queries" key containing the list of expanded queries.
264 If include_original_query=True, the original query will be included in addition
265 to the n_expansions alternative queries.
266 :raises ValueError: If n_expansions is not positive (less than or equal to 0).
267 """
269 await self.warm_up_async()
271 response = {"queries": [query] if self.include_original_query else []}
273 if not query.strip():
274 logger.warning("Empty query provided to QueryExpander")
275 return response
277 expansion_count = n_expansions if n_expansions is not None else self.n_expansions
278 if expansion_count <= 0:
279 raise ValueError("n_expansions must be positive")
281 try:
282 prompt_result = self._prompt_builder.run(query=query.strip(), n_expansions=expansion_count)
283 messages = [ChatMessage.from_user(prompt_result["prompt"])]
284 with _trace_chat_generator_run(self.chat_generator, {"messages": messages}) as span:
285 generator_result = await _execute_component_async(self.chat_generator, messages=messages)
286 span.set_content_tag("haystack.component.output", generator_result)
288 if not generator_result.get("replies") or len(generator_result["replies"]) == 0:
289 logger.warning("ChatGenerator returned no replies for query: {query}", query=query)
290 return response
292 expanded_text = generator_result["replies"][0].text.strip()
293 expanded_queries = self._parse_expanded_queries(expanded_text)
295 # Limit the number of expanded queries to the requested amount
296 if len(expanded_queries) > expansion_count:
297 logger.warning(
298 "Generated {generated_count} queries but only {requested_count} were requested. "
299 "Truncating to the first {requested_count} queries. ",
300 generated_count=len(expanded_queries),
301 requested_count=expansion_count,
302 )
303 expanded_queries = expanded_queries[:expansion_count]
305 # Add original query if not already present
306 if self.include_original_query:
307 expanded_queries_lower = [q.lower() for q in expanded_queries]
308 if query.lower() not in expanded_queries_lower:
309 expanded_queries.append(query)
311 response["queries"] = expanded_queries
312 return response
314 except Exception as e:
315 # Fallback: return original query to maintain pipeline functionality
316 logger.exception("Failed to expand query {query}: {error}", query=query, error=str(e))
317 return response
319 def warm_up(self) -> None:
320 """
321 Warm up the underlying chat generator.
322 """
323 if hasattr(self.chat_generator, "warm_up"):
324 self.chat_generator.warm_up()
326 async def warm_up_async(self) -> None:
327 """
328 Warm up the underlying chat generator on the serving event loop.
329 """
330 if hasattr(self.chat_generator, "warm_up_async"):
331 await self.chat_generator.warm_up_async()
332 elif hasattr(self.chat_generator, "warm_up"):
333 self.chat_generator.warm_up()
335 def close(self) -> None:
336 """
337 Release the underlying chat generator's resources.
338 """
339 if hasattr(self.chat_generator, "close"):
340 self.chat_generator.close()
342 async def close_async(self) -> None:
343 """
344 Release the underlying chat generator's async resources.
345 """
346 if hasattr(self.chat_generator, "close_async"):
347 await self.chat_generator.close_async()
348 elif hasattr(self.chat_generator, "close"):
349 self.chat_generator.close()
351 @staticmethod
352 def _parse_expanded_queries(generator_response: str) -> list[str]:
353 """
354 Parse the generator response to extract individual expanded queries.
356 :param generator_response: The raw text response from the generator.
357 :return: List of parsed expanded queries, deduplicated in first-seen order.
358 """
359 parsed = _parse_dict_from_json(generator_response, expected_keys=["queries"], raise_on_failure=False)
361 if parsed is None:
362 return []
364 if not isinstance(parsed["queries"], list):
365 logger.warning(
366 "Expected 'queries' to be a list but got {type}. Returning no expanded queries.",
367 type=type(parsed["queries"]).__name__,
368 )
369 return []
371 queries = []
372 seen: set[str] = set()
373 for item in parsed["queries"]:
374 if isinstance(item, str) and item.strip():
375 stripped = item.strip()
376 if stripped not in seen:
377 seen.add(stripped)
378 queries.append(stripped)
379 else:
380 logger.warning("Skipping non-string or empty query in response: {item}", item=item)
382 return queries