RAG Fusion - Multi-Query Generation with Reciprocal Rank Fusion

from typing import List, Dict

def generate_queries(original_query: str, llm, n: int = 4) -> List[str]:
    """Use an LLM to generate N reworded/expanded variants of the original query."""
    prompt = f"""You are a query expansion assistant. Given a user question, generate {n - 1}
different ways to ask the same question. Vary the vocabulary, perspective, and framing
to maximize document recall. Return one query per line, no numbering.

Original question: {original_query}

Alternative queries:"""

    response = llm.invoke(prompt)
    variants = [q.strip() for q in response.strip().split('\n') if q.strip()]
    # Always include the original first
    return [original_query] + variants[: n - 1]


def retrieve_for_query(query: str, retriever, top_k: int = 10) -> List[Dict]:
    """Retrieve documents for a single query and tag each result with its rank."""
    results = retriever.search(query, top_k=top_k)
    return [{'doc': doc, 'rank': i + 1, 'query': query} for i, doc in enumerate(results)]


def rrf_merge(ranked_lists: List[List[Dict]], k: int = 60, top_n: int = 10) -> List[Dict]:
    """
    Reciprocal Rank Fusion across multiple ranked lists.

    score(d) = sum of 1 / (k + rank_i(d))  for every list i where d appears.

    k=60 is the standard default from the original RRF paper (Cormack 2009).
    Documents appearing in more lists, and ranked higher within each list, score highest.
    """
    scores: Dict[str, float] = {}
    doc_map: Dict[str, Dict] = {}

    for ranked_list in ranked_lists:
        for item in ranked_list:
            doc_id = item['doc']['id']
            rank = item['rank']
            rrf_score = 1.0 / (k + rank)

            if doc_id not in scores:
                scores[doc_id] = 0.0
                doc_map[doc_id] = item['doc']
            scores[doc_id] += rrf_score

    sorted_ids = sorted(scores, key=lambda x: scores[x], reverse=True)
    return [
        {'doc': doc_map[doc_id], 'rrf_score': scores[doc_id]}
        for doc_id in sorted_ids[:top_n]
    ]


def rag_fusion(query: str, retriever, llm, n_queries: int = 4, top_k: int = 10) -> Dict:
    """Full RAG Fusion pipeline: generate → retrieve → fuse → generate."""
    # Step 1: Generate query variants
    queries = generate_queries(query, llm, n=n_queries)

    # Step 2: Retrieve for each query independently
    all_ranked_lists = []
    for q in queries:
        ranked = retrieve_for_query(q, retriever, top_k=top_k)
        all_ranked_lists.append(ranked)

    # Step 3: Fuse ranked lists with RRF
    fused_docs = rrf_merge(all_ranked_lists, k=60, top_n=top_k)

    # Step 4: Generate answer from fused context
    context = "\n\n".join([item['doc']['content'] for item in fused_docs])
    answer = llm.invoke(f"Context:\n{context}\n\nQuestion: {query}\n\nAnswer:")

    return {
        'answer': answer,
        'queries_used': queries,
        'fused_docs': fused_docs,
        'total_candidates': len(queries) * top_k,
    }