LlamaIndex QueryFusionRetriever — fusion errors and BM25/vector imbalance

from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.retrievers.bm25 import BM25Retriever
from llama_index.core import VectorStoreIndex

# ... build vector_index and get nodes ready ...

# --- Vector retriever ---
vector_retriever = vector_index.as_retriever(similarity_top_k=20)

# --- BM25 retriever ---
bm25_retriever = BM25Retriever.from_defaults(
    nodes=all_nodes,
    similarity_top_k=20,
)

# --- Fusion retriever ---
fusion_retriever = QueryFusionRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    similarity_top_k=10,          # final top-k after fusion
    num_queries=1,                # no query expansion (see next fix)
    mode="reciprocal_rerank",     # RRF — robust to score scale differences
    use_async=True,
    verbose=True,
)

# Use it in a query engine
from llama_index.core.query_engine import RetrieverQueryEngine
query_engine = RetrieverQueryEngine.from_args(fusion_retriever, llm=llm)

response = query_engine.query("What is the return window for damaged items?")
print(response)

# --- Inspect fusion behaviour ---
retrieved = fusion_retriever.retrieve("What is the return window?")
for i, node in enumerate(retrieved):
    print(f"{i}: score={node.score:.3f} source={node.metadata.get('source')}")

# RRF scores are always in the [0,1] range because they are computed as
#   sum over retrievers of 1 / (k + rank)
# with k typically 60. Small values but comparable across queries.

from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.llms.openai import OpenAI

llm = OpenAI(model="gpt-4o-mini")

# --- Query expansion + hybrid fusion ---
fusion_retriever = QueryFusionRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    similarity_top_k=10,
    num_queries=4,               # LLM generates 4 query variants
    mode="reciprocal_rerank",
    llm=llm,                     # for query generation
    verbose=True,                # print generated queries
    query_gen_prompt=(
        # Optional custom prompt — default is decent
        "You are helpful. Given the user's query, generate {num_queries} "
        "paraphrased versions that capture different phrasing of the same intent.\n"
        "Original query: {query}\n\n"
        "Generate {num_queries} paraphrases (one per line):"
    ),
    use_async=True,               # crucial — otherwise VERY slow
)

# Execution: 4 queries × 2 retrievers = 8 parallel retrieval calls,
# fused via RRF into top-10

response = query_engine.query("How can I return a damaged product?")
# LLM generates variants like:
#   "What is the procedure for returning a damaged item?"
#   "How do I initiate a return for a broken product?"
#   "Return process for defective goods?"

# Each variant is retrieved separately, boosting recall on paraphrased queries

from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.core.schema import NodeWithScore
from typing import List

class NormalizingRetriever:
    """Wraps a retriever and normalizes its scores to [0, 1]."""
    def __init__(self, retriever, weight: float = 1.0):
        self.retriever = retriever
        self.weight = weight

    async def aretrieve(self, query: str) -> List[NodeWithScore]:
        results = await self.retriever.aretrieve(query)
        if not results:
            return results
        max_score = max(n.score or 0 for n in results)
        min_score = min(n.score or 0 for n in results)
        rng = max_score - min_score
        if rng > 0:
            for n in results:
                # Normalize to [0,1], then apply weight
                normalized = (n.score - min_score) / rng
                n.score = normalized * self.weight
        return results

# --- Weighted fusion favouring vector over BM25 ---
weighted_vector = NormalizingRetriever(vector_retriever, weight=0.7)
weighted_bm25 = NormalizingRetriever(bm25_retriever, weight=0.3)

fusion = QueryFusionRetriever(
    retrievers=[weighted_vector, weighted_bm25],
    similarity_top_k=10,
    mode="relative_score",         # score-based since we've normalized
    use_async=True,
)

# --- Simpler alternative: post-hoc reranking with a cross-encoder ---
# Often better than any fusion tuning — rerank fused top-30 to top-5

from llama_index.postprocessor.cohere_rerank import CohereRerank

fusion_broad = QueryFusionRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    similarity_top_k=30,       # broad candidate pool
    mode="reciprocal_rerank",
    use_async=True,
)

reranker = CohereRerank(top_n=5, model="rerank-english-v3.0", api_key="...")

from llama_index.core.query_engine import RetrieverQueryEngine
query_engine = RetrieverQueryEngine.from_args(
    fusion_broad,
    node_postprocessors=[reranker],
    llm=llm,
)