Coverage for haystack/components/query/query_expander.py: 93%
142 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +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.
29These examples show the output when 3 queries are requested.
31Examples:
321. Query: "climate change effects"
33 {"queries": ["impact of climate change", "consequences of global warming", "effects of environmental changes"]}
352. Query: "machine learning algorithms"
36 {"queries": ["neural networks", "clustering techniques", "supervised learning methods"]}
383. Query: "open source NLP frameworks"
39 {"queries": ["natural language processing tools", "free nlp libraries", "open-source NLP platforms"]}
41Guidelines:
42- Generate queries that use different words and phrasings
43- Include synonyms and related terms
44- Maintain the same core meaning and intent
45- Make queries that are likely to retrieve relevant information the original might miss
46- Focus on variations that would work well with keyword-based search
47- Respond in the same language as the input query
49Your Task:
50Query: "{{ query }}"
52You *must* respond with a JSON object containing a "queries" array with exactly {{ n_expansions }} expanded queries."""
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 import OpenAIResponsesChatGenerator
70 from haystack.components.query import QueryExpander
72 chat_generator = OpenAIResponsesChatGenerator(
73 model="gpt-5.6-luna",
74 generation_kwargs={
75 "text": {
76 "format": {
77 "type": "json_schema",
78 "name": "query_expansion",
79 "strict": True,
80 "schema": {
81 "type": "object",
82 "properties": {"queries": {"type": "array", "items": {"type": "string"}}},
83 "required": ["queries"],
84 "additionalProperties": False,
85 },
86 }
87 }
88 },
89 )
91 expander = QueryExpander(chat_generator=chat_generator, n_expansions=3)
93 result = expander.run(query="green energy sources")
94 print(result["queries"])
95 # Output: ['alternative query 1', 'alternative query 2', 'alternative query 3', 'green energy sources']
96 # Note: Up to 3 additional queries + 1 original query (if include_original_query=True)
98 # To control total number of queries:
99 expander = QueryExpander(n_expansions=2, include_original_query=True) # Up to 3 total
100 # or
101 expander = QueryExpander(n_expansions=3, include_original_query=False) # Exactly 3 total
102 ```
103 """
105 def __init__(
106 self,
107 *,
108 chat_generator: ChatGenerator | None = None,
109 prompt_template: str | None = None,
110 n_expansions: int = 4,
111 include_original_query: bool = True,
112 ) -> None:
113 """
114 Initialize the QueryExpander component.
116 :param chat_generator: The chat generator component to use for query expansion.
117 If None, a default `OpenAIResponsesChatGenerator` with the gpt-4.1-mini model and `store` set to `False`
118 is used.
119 :param prompt_template: Custom [PromptBuilder](https://docs.haystack.deepset.ai/docs/promptbuilder)
120 template for query expansion. The template should instruct the LLM to return a JSON response with the
121 structure: `{"queries": ["query1", "query2", "query3"]}`. The template should include 'query' and
122 'n_expansions' variables.
123 :param n_expansions: Number of alternative queries to generate (default: 4).
124 :param include_original_query: Whether to include the original query in the output.
125 """
126 if n_expansions <= 0:
127 raise ValueError("n_expansions must be positive")
129 self.n_expansions = n_expansions
130 self.include_original_query = include_original_query
132 if chat_generator is None:
133 self.chat_generator: ChatGenerator = OpenAIResponsesChatGenerator(
134 model="gpt-4.1-mini",
135 generation_kwargs={
136 "temperature": 0.7,
137 "store": False,
138 "text": {
139 "format": {
140 "type": "json_schema",
141 "name": "query_expansion",
142 "schema": {
143 "type": "object",
144 "properties": {"queries": {"type": "array", "items": {"type": "string"}}},
145 "required": ["queries"],
146 "additionalProperties": False,
147 },
148 }
149 },
150 },
151 )
152 else:
153 self.chat_generator = chat_generator
155 self.prompt_template = prompt_template or DEFAULT_PROMPT_TEMPLATE
157 # Check if required variables are present in the template
158 if "query" not in self.prompt_template:
159 logger.warning(
160 "The prompt template does not contain the 'query' variable. This may cause issues during execution."
161 )
162 if "n_expansions" not in self.prompt_template:
163 logger.warning(
164 "The prompt template does not contain the 'n_expansions' variable. "
165 "This may cause issues during execution."
166 )
168 self._prompt_builder = PromptBuilder(
169 template=self.prompt_template, required_variables=["n_expansions", "query"]
170 )
172 def to_dict(self) -> dict[str, Any]:
173 """
174 Serializes the component to a dictionary.
176 :return: Dictionary with serialized data.
177 """
178 return default_to_dict(
179 self,
180 chat_generator=component_to_dict(self.chat_generator, name="chat_generator"),
181 prompt_template=self.prompt_template,
182 n_expansions=self.n_expansions,
183 include_original_query=self.include_original_query,
184 )
186 @classmethod
187 def from_dict(cls, data: dict[str, Any]) -> "QueryExpander":
188 """
189 Deserializes the component from a dictionary.
191 :param data: Dictionary with serialized data.
192 :return: Deserialized component.
193 """
194 init_params = data.get("init_parameters", {})
196 deserialize_chatgenerator_inplace(init_params, key="chat_generator")
198 return default_from_dict(cls, data)
200 @component.output_types(queries=list[str])
201 def run(self, query: str, n_expansions: int | None = None) -> dict[str, list[str]]:
202 """
203 Expand the input query into multiple semantically similar queries.
205 The language of the original query is preserved in the expanded queries.
207 :param query: The original query to expand.
208 :param n_expansions: Number of additional queries to generate (not including the original).
209 If None, uses the value from initialization. Must be a positive integer.
210 :return: Dictionary with "queries" key containing the list of expanded queries.
211 If include_original_query=True, the original query will be included in addition
212 to the n_expansions alternative queries.
213 :raises ValueError: If n_expansions is not positive (less than or equal to 0).
214 """
216 self.warm_up()
218 if not isinstance(query, str) or not query.strip():
219 logger.warning("Empty query provided to QueryExpander. {fallback}", fallback=self._fallback_message())
220 return {"queries": [query] if self.include_original_query and isinstance(query, str) else []}
222 response = {"queries": [query] if self.include_original_query else []}
224 expansion_count = n_expansions if n_expansions is not None else self.n_expansions
225 if expansion_count <= 0:
226 raise ValueError("n_expansions must be positive")
228 try:
229 prompt_result = self._prompt_builder.run(query=query.strip(), n_expansions=expansion_count)
230 messages = [ChatMessage.from_user(prompt_result["prompt"])]
231 with _trace_chat_generator_run(self.chat_generator, {"messages": messages}) as span:
232 generator_result = self.chat_generator.run(messages=messages)
233 span.set_content_tag("haystack.component.output", generator_result)
235 if not generator_result.get("replies") or len(generator_result["replies"]) == 0:
236 logger.warning(
237 "ChatGenerator returned no replies for query: {query}. {fallback}",
238 query=query,
239 fallback=self._fallback_message(),
240 )
241 return response
243 expanded_text = generator_result["replies"][0].text.strip()
244 expanded_queries = self._parse_expanded_queries(expanded_text)
246 # Limit the number of expanded queries to the requested amount
247 if len(expanded_queries) > expansion_count:
248 logger.warning(
249 "Generated {generated_count} queries but only {requested_count} were requested. "
250 "Truncating to the first {requested_count} queries. ",
251 generated_count=len(expanded_queries),
252 requested_count=expansion_count,
253 )
254 expanded_queries = expanded_queries[:expansion_count]
256 # Add original query if not already present
257 if self.include_original_query:
258 expanded_queries_lower = [q.lower() for q in expanded_queries]
259 if query.lower() not in expanded_queries_lower:
260 expanded_queries.append(query)
262 response["queries"] = expanded_queries
263 return response
265 except Exception as e:
266 # Fallback: return original query to maintain pipeline functionality
267 logger.exception(
268 "Failed to expand query {query}. {fallback} Error: {error}",
269 query=query,
270 fallback=self._fallback_message(),
271 error=str(e),
272 )
273 return response
275 @component.output_types(queries=list[str])
276 async def run_async(self, query: str, n_expansions: int | None = None) -> dict[str, list[str]]:
277 """
278 Asynchronously expand the input query into multiple semantically similar queries.
280 The language of the original query is preserved in the expanded queries.
282 This is the asynchronous version of the `run` method. It has the same parameters and return values
283 but can be used with `await` in an async code. If the chat generator only implements a synchronous
284 `run` method, it is executed in a thread to avoid blocking the event loop.
286 :param query: The original query to expand.
287 :param n_expansions: Number of additional queries to generate (not including the original).
288 If None, uses the value from initialization. Must be a positive integer.
289 :return: Dictionary with "queries" key containing the list of expanded queries.
290 If include_original_query=True, the original query will be included in addition
291 to the n_expansions alternative queries.
292 :raises ValueError: If n_expansions is not positive (less than or equal to 0).
293 """
295 await self.warm_up_async()
297 if not isinstance(query, str) or not query.strip():
298 logger.warning("Empty query provided to QueryExpander. {fallback}", fallback=self._fallback_message())
299 return {"queries": [query] if self.include_original_query and isinstance(query, str) else []}
301 response = {"queries": [query] if self.include_original_query else []}
303 expansion_count = n_expansions if n_expansions is not None else self.n_expansions
304 if expansion_count <= 0:
305 raise ValueError("n_expansions must be positive")
307 try:
308 prompt_result = self._prompt_builder.run(query=query.strip(), n_expansions=expansion_count)
309 messages = [ChatMessage.from_user(prompt_result["prompt"])]
310 with _trace_chat_generator_run(self.chat_generator, {"messages": messages}) as span:
311 generator_result = await _execute_component_async(self.chat_generator, messages=messages)
312 span.set_content_tag("haystack.component.output", generator_result)
314 if not generator_result.get("replies") or len(generator_result["replies"]) == 0:
315 logger.warning(
316 "ChatGenerator returned no replies for query: {query}. {fallback}",
317 query=query,
318 fallback=self._fallback_message(),
319 )
320 return response
322 expanded_text = generator_result["replies"][0].text.strip()
323 expanded_queries = self._parse_expanded_queries(expanded_text)
325 # Limit the number of expanded queries to the requested amount
326 if len(expanded_queries) > expansion_count:
327 logger.warning(
328 "Generated {generated_count} queries but only {requested_count} were requested. "
329 "Truncating to the first {requested_count} queries. ",
330 generated_count=len(expanded_queries),
331 requested_count=expansion_count,
332 )
333 expanded_queries = expanded_queries[:expansion_count]
335 # Add original query if not already present
336 if self.include_original_query:
337 expanded_queries_lower = [q.lower() for q in expanded_queries]
338 if query.lower() not in expanded_queries_lower:
339 expanded_queries.append(query)
341 response["queries"] = expanded_queries
342 return response
344 except Exception as e:
345 # Fallback: return original query to maintain pipeline functionality
346 logger.exception(
347 "Failed to expand query {query}. {fallback} Error: {error}",
348 query=query,
349 fallback=self._fallback_message(),
350 error=str(e),
351 )
352 return response
354 def warm_up(self) -> None:
355 """
356 Warm up the underlying chat generator.
357 """
358 if hasattr(self.chat_generator, "warm_up"):
359 self.chat_generator.warm_up()
361 async def warm_up_async(self) -> None:
362 """
363 Warm up the underlying chat generator on the serving event loop.
364 """
365 if hasattr(self.chat_generator, "warm_up_async"):
366 await self.chat_generator.warm_up_async()
367 elif hasattr(self.chat_generator, "warm_up"):
368 self.chat_generator.warm_up()
370 def close(self) -> None:
371 """
372 Release the underlying chat generator's resources.
373 """
374 if hasattr(self.chat_generator, "close"):
375 self.chat_generator.close()
377 async def close_async(self) -> None:
378 """
379 Release the underlying chat generator's async resources.
380 """
381 if hasattr(self.chat_generator, "close_async"):
382 await self.chat_generator.close_async()
383 elif hasattr(self.chat_generator, "close"):
384 self.chat_generator.close()
386 def _fallback_message(self) -> str:
387 return "Returning only the original query." if self.include_original_query else "Returning no queries."
389 def _parse_expanded_queries(self, generator_response: str) -> list[str]:
390 """
391 Parse the generator response to extract individual expanded queries.
393 :param generator_response: The raw text response from the generator.
394 :return: List of parsed expanded queries, deduplicated in first-seen order.
395 """
396 try:
397 parsed = _parse_dict_from_json(generator_response, expected_keys=["queries"], raise_on_failure=True)
398 except ValueError as e:
399 logger.warning(
400 "QueryExpander failed to parse the ChatGenerator response: {response}. {fallback} Error: {error}",
401 response=generator_response,
402 fallback=self._fallback_message(),
403 error=e,
404 )
405 return []
407 if not isinstance(parsed["queries"], list):
408 logger.warning(
409 "Expected 'queries' to be a list but got {type}. {fallback}",
410 type=type(parsed["queries"]).__name__,
411 fallback=self._fallback_message(),
412 )
413 return []
415 queries = []
416 seen: set[str] = set()
417 for item in parsed["queries"]:
418 if isinstance(item, str) and item.strip():
419 stripped = item.strip()
420 if stripped not in seen:
421 seen.add(stripped)
422 queries.append(stripped)
423 else:
424 logger.warning("Skipping non-string or empty query in response: {item}", item=item)
426 if not queries:
427 logger.warning(
428 "QueryExpander found no valid queries in the ChatGenerator response. {fallback}",
429 fallback=self._fallback_message(),
430 )
432 return queries