Coverage for haystack/components/query/query_expander.py: 93%
142 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
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_responses import OpenAIResponsesChatGenerator
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 `OpenAIResponsesChatGenerator` with the gpt-4.1-mini model and `store` set to `False`
102 is used.
103 :param prompt_template: Custom [PromptBuilder](https://docs.haystack.deepset.ai/docs/promptbuilder)
104 template for query expansion. The template should instruct the LLM to return a JSON response with the
105 structure: `{"queries": ["query1", "query2", "query3"]}`. The template should include 'query' and
106 'n_expansions' variables.
107 :param n_expansions: Number of alternative queries to generate (default: 4).
108 :param include_original_query: Whether to include the original query in the output.
109 """
110 if n_expansions <= 0:
111 raise ValueError("n_expansions must be positive")
113 self.n_expansions = n_expansions
114 self.include_original_query = include_original_query
116 if chat_generator is None:
117 self.chat_generator: ChatGenerator = OpenAIResponsesChatGenerator(
118 model="gpt-4.1-mini",
119 generation_kwargs={
120 "temperature": 0.7,
121 "store": False,
122 "text": {
123 "format": {
124 "type": "json_schema",
125 "name": "query_expansion",
126 "schema": {
127 "type": "object",
128 "properties": {"queries": {"type": "array", "items": {"type": "string"}}},
129 "required": ["queries"],
130 "additionalProperties": False,
131 },
132 }
133 },
134 },
135 )
136 else:
137 self.chat_generator = chat_generator
139 self.prompt_template = prompt_template or DEFAULT_PROMPT_TEMPLATE
141 # Check if required variables are present in the template
142 if "query" not in self.prompt_template:
143 logger.warning(
144 "The prompt template does not contain the 'query' variable. This may cause issues during execution."
145 )
146 if "n_expansions" not in self.prompt_template:
147 logger.warning(
148 "The prompt template does not contain the 'n_expansions' variable. "
149 "This may cause issues during execution."
150 )
152 self._prompt_builder = PromptBuilder(
153 template=self.prompt_template, required_variables=["n_expansions", "query"]
154 )
156 def to_dict(self) -> dict[str, Any]:
157 """
158 Serializes the component to a dictionary.
160 :return: Dictionary with serialized data.
161 """
162 return default_to_dict(
163 self,
164 chat_generator=component_to_dict(self.chat_generator, name="chat_generator"),
165 prompt_template=self.prompt_template,
166 n_expansions=self.n_expansions,
167 include_original_query=self.include_original_query,
168 )
170 @classmethod
171 def from_dict(cls, data: dict[str, Any]) -> "QueryExpander":
172 """
173 Deserializes the component from a dictionary.
175 :param data: Dictionary with serialized data.
176 :return: Deserialized component.
177 """
178 init_params = data.get("init_parameters", {})
180 deserialize_chatgenerator_inplace(init_params, key="chat_generator")
182 return default_from_dict(cls, data)
184 @component.output_types(queries=list[str])
185 def run(self, query: str, n_expansions: int | None = None) -> dict[str, list[str]]:
186 """
187 Expand the input query into multiple semantically similar queries.
189 The language of the original query is preserved in the expanded queries.
191 :param query: The original query to expand.
192 :param n_expansions: Number of additional queries to generate (not including the original).
193 If None, uses the value from initialization. Must be a positive integer.
194 :return: Dictionary with "queries" key containing the list of expanded queries.
195 If include_original_query=True, the original query will be included in addition
196 to the n_expansions alternative queries.
197 :raises ValueError: If n_expansions is not positive (less than or equal to 0).
198 """
200 self.warm_up()
202 if not isinstance(query, str) or not query.strip():
203 logger.warning("Empty query provided to QueryExpander. {fallback}", fallback=self._fallback_message())
204 return {"queries": [query] if self.include_original_query and isinstance(query, str) else []}
206 response = {"queries": [query] if self.include_original_query else []}
208 expansion_count = n_expansions if n_expansions is not None else self.n_expansions
209 if expansion_count <= 0:
210 raise ValueError("n_expansions must be positive")
212 try:
213 prompt_result = self._prompt_builder.run(query=query.strip(), n_expansions=expansion_count)
214 messages = [ChatMessage.from_user(prompt_result["prompt"])]
215 with _trace_chat_generator_run(self.chat_generator, {"messages": messages}) as span:
216 generator_result = self.chat_generator.run(messages=messages)
217 span.set_content_tag("haystack.component.output", generator_result)
219 if not generator_result.get("replies") or len(generator_result["replies"]) == 0:
220 logger.warning(
221 "ChatGenerator returned no replies for query: {query}. {fallback}",
222 query=query,
223 fallback=self._fallback_message(),
224 )
225 return response
227 expanded_text = generator_result["replies"][0].text.strip()
228 expanded_queries = self._parse_expanded_queries(expanded_text)
230 # Limit the number of expanded queries to the requested amount
231 if len(expanded_queries) > expansion_count:
232 logger.warning(
233 "Generated {generated_count} queries but only {requested_count} were requested. "
234 "Truncating to the first {requested_count} queries. ",
235 generated_count=len(expanded_queries),
236 requested_count=expansion_count,
237 )
238 expanded_queries = expanded_queries[:expansion_count]
240 # Add original query if not already present
241 if self.include_original_query:
242 expanded_queries_lower = [q.lower() for q in expanded_queries]
243 if query.lower() not in expanded_queries_lower:
244 expanded_queries.append(query)
246 response["queries"] = expanded_queries
247 return response
249 except Exception as e:
250 # Fallback: return original query to maintain pipeline functionality
251 logger.exception(
252 "Failed to expand query {query}. {fallback} Error: {error}",
253 query=query,
254 fallback=self._fallback_message(),
255 error=str(e),
256 )
257 return response
259 @component.output_types(queries=list[str])
260 async def run_async(self, query: str, n_expansions: int | None = None) -> dict[str, list[str]]:
261 """
262 Asynchronously expand the input query into multiple semantically similar queries.
264 The language of the original query is preserved in the expanded queries.
266 This is the asynchronous version of the `run` method. It has the same parameters and return values
267 but can be used with `await` in an async code. If the chat generator only implements a synchronous
268 `run` method, it is executed in a thread to avoid blocking the event loop.
270 :param query: The original query to expand.
271 :param n_expansions: Number of additional queries to generate (not including the original).
272 If None, uses the value from initialization. Must be a positive integer.
273 :return: Dictionary with "queries" key containing the list of expanded queries.
274 If include_original_query=True, the original query will be included in addition
275 to the n_expansions alternative queries.
276 :raises ValueError: If n_expansions is not positive (less than or equal to 0).
277 """
279 await self.warm_up_async()
281 if not isinstance(query, str) or not query.strip():
282 logger.warning("Empty query provided to QueryExpander. {fallback}", fallback=self._fallback_message())
283 return {"queries": [query] if self.include_original_query and isinstance(query, str) else []}
285 response = {"queries": [query] if self.include_original_query else []}
287 expansion_count = n_expansions if n_expansions is not None else self.n_expansions
288 if expansion_count <= 0:
289 raise ValueError("n_expansions must be positive")
291 try:
292 prompt_result = self._prompt_builder.run(query=query.strip(), n_expansions=expansion_count)
293 messages = [ChatMessage.from_user(prompt_result["prompt"])]
294 with _trace_chat_generator_run(self.chat_generator, {"messages": messages}) as span:
295 generator_result = await _execute_component_async(self.chat_generator, messages=messages)
296 span.set_content_tag("haystack.component.output", generator_result)
298 if not generator_result.get("replies") or len(generator_result["replies"]) == 0:
299 logger.warning(
300 "ChatGenerator returned no replies for query: {query}. {fallback}",
301 query=query,
302 fallback=self._fallback_message(),
303 )
304 return response
306 expanded_text = generator_result["replies"][0].text.strip()
307 expanded_queries = self._parse_expanded_queries(expanded_text)
309 # Limit the number of expanded queries to the requested amount
310 if len(expanded_queries) > expansion_count:
311 logger.warning(
312 "Generated {generated_count} queries but only {requested_count} were requested. "
313 "Truncating to the first {requested_count} queries. ",
314 generated_count=len(expanded_queries),
315 requested_count=expansion_count,
316 )
317 expanded_queries = expanded_queries[:expansion_count]
319 # Add original query if not already present
320 if self.include_original_query:
321 expanded_queries_lower = [q.lower() for q in expanded_queries]
322 if query.lower() not in expanded_queries_lower:
323 expanded_queries.append(query)
325 response["queries"] = expanded_queries
326 return response
328 except Exception as e:
329 # Fallback: return original query to maintain pipeline functionality
330 logger.exception(
331 "Failed to expand query {query}. {fallback} Error: {error}",
332 query=query,
333 fallback=self._fallback_message(),
334 error=str(e),
335 )
336 return response
338 def warm_up(self) -> None:
339 """
340 Warm up the underlying chat generator.
341 """
342 if hasattr(self.chat_generator, "warm_up"):
343 self.chat_generator.warm_up()
345 async def warm_up_async(self) -> None:
346 """
347 Warm up the underlying chat generator on the serving event loop.
348 """
349 if hasattr(self.chat_generator, "warm_up_async"):
350 await self.chat_generator.warm_up_async()
351 elif hasattr(self.chat_generator, "warm_up"):
352 self.chat_generator.warm_up()
354 def close(self) -> None:
355 """
356 Release the underlying chat generator's resources.
357 """
358 if hasattr(self.chat_generator, "close"):
359 self.chat_generator.close()
361 async def close_async(self) -> None:
362 """
363 Release the underlying chat generator's async resources.
364 """
365 if hasattr(self.chat_generator, "close_async"):
366 await self.chat_generator.close_async()
367 elif hasattr(self.chat_generator, "close"):
368 self.chat_generator.close()
370 def _fallback_message(self) -> str:
371 return "Returning only the original query." if self.include_original_query else "Returning no queries."
373 def _parse_expanded_queries(self, generator_response: str) -> list[str]:
374 """
375 Parse the generator response to extract individual expanded queries.
377 :param generator_response: The raw text response from the generator.
378 :return: List of parsed expanded queries, deduplicated in first-seen order.
379 """
380 try:
381 parsed = _parse_dict_from_json(generator_response, expected_keys=["queries"], raise_on_failure=True)
382 except ValueError as e:
383 logger.warning(
384 "QueryExpander failed to parse the ChatGenerator response: {response}. {fallback} Error: {error}",
385 response=generator_response,
386 fallback=self._fallback_message(),
387 error=e,
388 )
389 return []
391 if not isinstance(parsed["queries"], list):
392 logger.warning(
393 "Expected 'queries' to be a list but got {type}. {fallback}",
394 type=type(parsed["queries"]).__name__,
395 fallback=self._fallback_message(),
396 )
397 return []
399 queries = []
400 seen: set[str] = set()
401 for item in parsed["queries"]:
402 if isinstance(item, str) and item.strip():
403 stripped = item.strip()
404 if stripped not in seen:
405 seen.add(stripped)
406 queries.append(stripped)
407 else:
408 logger.warning("Skipping non-string or empty query in response: {item}", item=item)
410 if not queries:
411 logger.warning(
412 "QueryExpander found no valid queries in the ChatGenerator response. {fallback}",
413 fallback=self._fallback_message(),
414 )
416 return queries