rag.search.query_expansion

Query expansion for improved retrieval recall via multi-query retrieval.

class QueryExpander(llm_generate_fn=None, n_variations=3, include_original=True)[source]

Bases: object

Expand queries using LLM for multi-query retrieval.

Generates multiple reformulations of a search query to improve recall by catching different phrasings and aspects of the topic.

Parameters:

llm_generate_fn (Callable[[str], Any] | None) – Function that takes a prompt string and returns LLM response.
n_variations (int) – Number of query variations to generate (default: 3).
include_original (bool) – Whether to include original query in results (default: True).

Example

>>> def mock_llm(prompt): return "1. What is auth?\n2. How to authenticate?"
>>> expander = QueryExpander(llm_generate_fn=mock_llm, n_variations=2)
>>> queries = expander.expand("authentication methods")
>>> # Returns: ["authentication methods", "What is auth?", "How to authenticate?"]

expand(query)[source]

Generate query variations for improved retrieval.

Parameters:: query (str) – Original search query.
Returns:: List of query variations including original (if include_original=True).
Return type:: list[str]

set_llm_function(fn)[source]

Set or update the LLM generate function.

Parameters:: fn (Callable[[str], Any]) – Function that takes a prompt and returns LLM response.
Return type:: None

fuse_search_results(result_lists, n_results=5, dedupe_by='id')[source]

Fuse results from multiple searches with deduplication.

Combines results from multiple query variations, removing duplicates while preserving relevance ordering.

Parameters:

result_lists (list[list[dict[str, Any]]]) – List of search result lists from different queries.
n_results (int) – Maximum number of final results to return.
dedupe_by (str) – Field to use for deduplication (default: “id”).

Returns:

Fused and deduplicated list of results.

Return type:

list[dict[str, Any]]

Classes

class QueryExpander(llm_generate_fn=None, n_variations=3, include_original=True)[source]

Bases: object

Expand queries using LLM for multi-query retrieval.

Generates multiple reformulations of a search query to improve recall by catching different phrasings and aspects of the topic.

Parameters:

llm_generate_fn (Callable[[str], Any] | None) – Function that takes a prompt string and returns LLM response.
n_variations (int) – Number of query variations to generate (default: 3).
include_original (bool) – Whether to include original query in results (default: True).

Example

>>> def mock_llm(prompt): return "1. What is auth?\n2. How to authenticate?"
>>> expander = QueryExpander(llm_generate_fn=mock_llm, n_variations=2)
>>> queries = expander.expand("authentication methods")
>>> # Returns: ["authentication methods", "What is auth?", "How to authenticate?"]

expand(query)[source]

Generate query variations for improved retrieval.

Parameters:: query (str) – Original search query.
Returns:: List of query variations including original (if include_original=True).
Return type:: list[str]

set_llm_function(fn)[source]

Set or update the LLM generate function.

Parameters:: fn (Callable[[str], Any]) – Function that takes a prompt and returns LLM response.
Return type:: None

Functions

fuse_search_results(result_lists, n_results=5, dedupe_by='id')[source]

Fuse results from multiple searches with deduplication.

Combines results from multiple query variations, removing duplicates while preserving relevance ordering.

Parameters:

result_lists (list[list[dict[str, Any]]]) – List of search result lists from different queries.
n_results (int) – Maximum number of final results to return.
dedupe_by (str) – Field to use for deduplication (default: “id”).

Returns:

Fused and deduplicated list of results.

Return type:

list[dict[str, Any]]