QueryExpander, merge_retrieval_results

On this page

core.query_expansion.QueryExpander, core.query_expansion.merge_retrieval_results

Overview

QueryExpander generates semantically-equivalent rephrasings of a query so that RAGSystem.retrieve() can search the vector DB with several phrasings and merge the results, improving recall for queries whose wording doesn't closely match the indexed text. It uses a two-tier strategy: an LLM (primary, higher quality) with automatic fallback to a curated synonym table (synonm_table._SYNONYM_TABLE, ~1,100+ Arabic/English keyword→synonym mappings) when the LLM is unavailable or fails.

merge_retrieval_results is a standalone helper that deduplicates and merges result lists coming from multiple query variants.

QueryExpander

Exact signature

Python · 8 lines
class QueryExpander:
    def __init__(
        self,
        llm: Optional[Any] = None,
        max_variants: int = 3,
        use_llm: bool = True,
        fallback_synonyms: bool = True,
    ) -> None: ...

Parameters

NameTypeRequiredDefaultNotes
llmOptional[Any]NoNoneAny object with .generate(prompt, max_tokens, temperature) -> str. Pass None to use the synonym table only.
max_variantsintNo3Maximum alternative queries generated (excluding the original).
use_llmboolNoTrueEffectively becomes use_llm and (llm is not None) — automatically False if no llm was supplied, regardless of the argument value.
fallback_synonymsboolNoTrueWhether to fall back to (or supplement with) the synonym table.

expand

Python
def expand(self, query: str) -> List[str]

Parameters

NameTypeRequiredNotes
querystrYesEmpty/whitespace query returns [query] (the input unchanged, wrapped in a list).

Returns

List[str] — [original_query, variant1, variant2, ...]. Always returns at least the original query as the first element.

Algorithm

  1. Empty query → return [query].
  2. If self.use_llm (LLM configured and enabled): call _llm_expand(query) inside a try block.
    • On any exception: log a warning and fall through to step 3 (do not re-raise from expand()).
  3. If self.fallback_synonyms: call the module-level _build_synonym_variants(query) (an lru_cache-wrapped function, max 1024 entries) and return up to max_variants + 1 entries (original + variants).
  4. If neither path applies: return [query].

Raises

expand() itself never raises — internal LLM failures are caught and logged, falling back to synonyms (or the bare query if fallback_synonyms=False).

Side effects

  • May call self.llm.generate() (network I/O, depending on the LLM implementation).
  • Logs at DEBUG/INFO/WARNING levels.
  • Populates the module-level _build_synonym_variants LRU cache (shared across all QueryExpander instances, since it's a plain function, not a method).

Complexity

  • LLM path: one LLM call plus O(k) parsing where k = number of output lines.
  • Synonym path: O(1) amortized after the first call for a given query (cached), else O(|_SYNONYM_TABLE| · pattern-match cost) for the first lookup of a new query.

Examples

Basic — LLM-backed expansion:

Python · 5 lines
from zerda.rag.core import QueryExpander

expander = QueryExpander(llm=my_llm)
variants = expander.expand("رقم التليفون")
# → ["رقم التليفون", "رقم الهاتف", "رقم الجوال", ...]

Synonym-table only (no LLM):

Python · 2 lines
expander = QueryExpander(llm=None)
variants = expander.expand("phone number")

Limiting variant count:

Python · 3 lines
expander = QueryExpander(llm=my_llm, max_variants=1)
variants = expander.expand("What is the refund policy?")
# at most 2 entries: [original, one_variant]

Disabling all expansion (pass-through):

Python · 2 lines
expander = QueryExpander(llm=None, fallback_synonyms=False)
expander.expand("anything")   # → ["anything"]

Error handling — LLM raises:

Python · 6 lines
class FlakyLLM:
    def generate(self, prompt, **kw): raise TimeoutError("slow")

expander = QueryExpander(llm=FlakyLLM())
variants = expander.expand("query text")
# Falls back to synonym table automatically; no exception surfaces.

Production-ready usage:

Python · 4 lines
expander = QueryExpander(llm=production_llm, max_variants=3,
                          use_llm=True, fallback_synonyms=True)
for v in expander.expand(user_query):
    results = vector_db.search(v, top_k=5)

_llm_expand (internal)

Python
def _llm_expand(self, query: str) -> List[str]
  1. Detects language via the module-level _detect_language(text) helper (>30% Arabic Unicode chars → "ar", else "en" — note this is a different threshold than PromptRouter.detect_language's 25% threshold).
  2. Formats the appropriate prompt template (_EXPANSION_PROMPT_AR or _EXPANSION_PROMPT_EN), asking for 3–4 alternative phrasings, one per line, no numbering.
  3. Calls self.llm.generate(prompt, max_tokens=300, temperature=0.3). On exception, raises QueryExpansionError(str(exc)) (chained via from exc).
  4. Parses the output via _parse_llm_variants(raw, query, self.max_variants).
  5. If fallback_synonyms is also enabled, enriches the LLM-derived variant list with any synonym-table variants not already present (up to max_variants total).
  6. Logs an INFO line with the variant count.

Raises

QueryExpansionError — wraps any exception from the LLM call itself (not parsing).

_parse_llm_variants (module-level helper)

Python
def _parse_llm_variants(raw: str, original: str, max_variants: int) -> List[str]
  1. Splits raw into non-blank lines.
  2. Strips leading bullets/numbers via regex ^[\-\d\.\)\•]+\s*.
  3. Deduplicates case-insensitively against the original and previously seen variants.
  4. Returns [original, ...cleaned_variants], capped at max_variants + 1 total entries.

_build_synonym_variants (module-level, @lru_cache(maxsize=1024))

Python
def _build_synonym_variants(query: str) -> List[str]

For each (keyword, synonyms) pair in _SYNONYM_TABLE, does a case-insensitive whole-word regex search in the query. On the first keyword match found (iteration order of the dict), it substitutes that keyword with each of its synonyms (one substitution per synonym, using count=1 so only the first occurrence is replaced) to produce new candidate queries, then stops — it does not continue trying other keywords, to avoid a combinatorial explosion of variants.

[!NOTE] Only the first matching keyword is expanded If a query contains multiple recognized keywords (e.g. both "phone" and "email"), only the first one encountered in _SYNONYM_TABLE's iteration order is used to generate variants.

merge_retrieval_results

Python · 4 lines
def merge_retrieval_results(
    results_per_query: List[List[Tuple[Any, float]]],
    top_k: int,
) -> List[Tuple[Any, float]]

Overview

Merges result lists from multiple query variants into a single deduplicated, score-sorted list — used by RAGSystem.retrieve()/aretrieve() whenever more than one query variant produced results.

Parameters

NameTypeRequiredNotes
results_per_queryList[List[Tuple[Any, float]]]YesOne (chunk, score) list per query variant.
top_kintYesMaximum number of results to return after merging.

Returns

List[Tuple[Any, float]] — deduplicated, sorted descending by score, truncated to top_k.

Algorithm

  1. For every (chunk, score) across all result lists, compute a chunk identity key: chunk.chunk_id if present, else chunk.id if present, else str(chunk) as a last resort.
  2. Keep the highest-scoring occurrence of each identity key in a dict.
  3. Sort the kept values by score descending.
  4. Truncate to top_k; log a DEBUG line if truncation occurred.

Raises

None — pure function over its inputs; malformed chunk objects fall back to str(chunk) identity rather than raising.

Complexity

Time: O(N log N) where N = total number of (chunk, score) pairs across all input lists (dominated by the final sort). Space: O(U) where U = number of unique chunk identities.

Examples

Basic:

Python · 8 lines
from zerda.rag.core import merge_retrieval_results

per_variant = [
    [(chunk_a, 0.7), (chunk_b, 0.5)],
    [(chunk_a, 0.9), (chunk_c, 0.6)],
]
merged = merge_retrieval_results(per_variant, top_k=2)
# [(chunk_a, 0.9), (chunk_c, 0.6)]  — chunk_a keeps its higher score

No duplicate chunk-id attributes (falls back to str()):

Python · 3 lines
merge_retrieval_results([[("raw string chunk", 0.8)]], top_k=5)
# identity key = str("raw string chunk") — works, but any two equal strings
# with different scores would collapse to the higher-scoring one.

Thread safety

QueryExpander instances hold no mutable state beyond constructor arguments (aside from the shared module-level LRU caches, which are safe for concurrent read/write per Python's functools.lru_cache guarantees). merge_retrieval_results is a pure function.

Best practices

  • Use merge_retrieval_results directly if you build your own multi-query retrieval loop outside of RAGSystem.
  • Be aware that _build_synonym_variants only expands on the first matched keyword — if you need multi-keyword expansion, you'll need to call it iteratively or write custom logic.

Common mistakes

  • Expecting QueryExpander.expand() to raise when the LLM fails — it degrades to synonyms/pass-through instead.
  • Assuming merge_retrieval_results is called for single-variant retrieval — RAGSystem.retrieve() skips it entirely when there's only one query variant (see RAGSystem.retrieve).

Related pages