Hybrid RAG with RRF

# Retrieval settings
k_bm25_candidates = 15      # BM25 candidate pool size
k_semantic_candidates = 15  # Semantic candidate pool size
k_rrf_pool = 10            # Documents after RRF fusion
k_final = 5                # Final documents after MMR
rrf_k = 60                 # RRF constant (typically 60)
mmr_lambda = 0.7           # MMR diversity factor (0.7 = 70% relevance, 30% diversity)

def reciprocal_rank_fusion(
    rankings: List[List[Document]], 
    k_constant: int = 60, 
    top_k: int = 5
) -> List[Document]:
    """Fuses multiple ranked lists using RRF and returns top-k documents."""
    scores: Dict[str, float] = {}
    documents_by_id: Dict[str, Document] = {}
    
    for ranked_docs in rankings:
        for rank, doc in enumerate(ranked_docs, start=1):
            doc_id = _document_unique_id(doc)
            # RRF formula: 1 / (k + rank)
            scores[doc_id] = scores.get(doc_id, 0.0) + (1.0 / (k_constant + rank))
            if doc_id not in documents_by_id:
                documents_by_id[doc_id] = doc
    
    # Sort by RRF score descending
    sorted_doc_ids = sorted(scores.keys(), key=lambda doc_id: scores[doc_id], reverse=True)
    return [documents_by_id[doc_id] for doc_id in sorted_doc_ids[:top_k]]

def mmr_select(
    query: str, 
    candidate_docs: List[Document], 
    top_k: int, 
    lambda_mult: float = 0.7
) -> List[Document]:
    """
    Selects top-k documents using MMR to balance relevance and diversity.
    
    MMR objective: λ * relevance(d, q) - (1-λ) * max_similarity(d, selected)
    """
    query_embedding = embeddings.embed_query(query)
    doc_embeddings = embeddings.embed_documents([doc.page_content for doc in candidate_docs])
    
    relevance_scores = [_cosine_similarity(query_embedding, doc_vec) 
                       for doc_vec in doc_embeddings]
    
    selected_indices: List[int] = []
    remaining_indices = list(range(len(candidate_docs)))
    
    # First pick: most relevant to query
    first_idx = max(remaining_indices, key=lambda idx: relevance_scores[idx])
    selected_indices.append(first_idx)
    remaining_indices.remove(first_idx)
    
    # Next picks: maximize MMR objective
    while remaining_indices and len(selected_indices) < top_k:
        best_idx = None
        best_score = float("-inf")
        
        for idx in remaining_indices:
            # Find max similarity to already selected docs
            max_similarity_to_selected = max(
                _cosine_similarity(doc_embeddings[idx], doc_embeddings[selected_idx])
                for selected_idx in selected_indices
            )
            # MMR score: balance relevance and diversity
            mmr_score = (lambda_mult * relevance_scores[idx]) - \
                       ((1.0 - lambda_mult) * max_similarity_to_selected)
            if mmr_score > best_score:
                best_score = mmr_score
                best_idx = idx
        
        if best_idx is None:
            break
        
        selected_indices.append(best_idx)
        remaining_indices.remove(best_idx)
    
    return [candidate_docs[idx] for idx in selected_indices]

def retrieve_hybrid_rrf(query: str) -> List[Document]:
    """Retrieves candidates, fuses with RRF, then applies MMR diversification."""
    # 1. Retrieve candidates from both retrievers
    bm25_docs = bm25_retriever.invoke(query)           # 15 candidates
    semantic_docs = semantic_retriever.invoke(query)   # 15 candidates
    
    # 2. Fuse using RRF
    rrf_ranked_docs = reciprocal_rank_fusion(
        [bm25_docs, semantic_docs],
        k_constant=rrf_k,      # 60
        top_k=k_rrf_pool,      # 10 documents
    )
    
    # 3. Diversify using MMR
    return mmr_select(
        query=query, 
        candidate_docs=rrf_ranked_docs, 
        top_k=k_final,         # 5 final documents
        lambda_mult=mmr_lambda # 0.7 (70% relevance, 30% diversity)
    )

from src.rag.hybrid_rrf import query_for_evaluation

# Basic usage with default model (gpt-4o)
result = query_for_evaluation(
    question="¿Cuáles son los síntomas del parto prematuro?"
)

# With custom model
result = query_for_evaluation(
    question="¿Qué es la diabetes gestacional?",
    llm_model="gpt-4o-mini"
)

# With custom LLM instance
from langchain_openai import ChatOpenAI
custom_llm = ChatOpenAI(model_name="gpt-4o", temperature=0)
result = query_for_evaluation(
    question="¿Qué cuidados necesito en el embarazo?",
    custom_llm=custom_llm
)

{
    "question": str,
    "answer": str,
    "contexts": List[str],
    "source_documents": List,
    "metadata": {
        "num_contexts": 5,
        "retrieval_method": "hybrid_bm25_semantic_rrf_mmr",
        "rrf_k": 60,
        "k_bm25_candidates": 15,
        "k_semantic_candidates": 15,
        "k_rrf_pool": 10,
        "k_final": 5,
        "mmr_lambda": 0.7,
        "llm_model": "gpt-4o",
        "execution_time": 3.21,
        "total_cost": 0.002678
    }
}

# Conservative (faster, lower recall)
k_bm25_candidates = 10
k_semantic_candidates = 10
k_rrf_pool = 7
k_final = 5

# Default (balanced)
k_bm25_candidates = 15
k_semantic_candidates = 15
k_rrf_pool = 10
k_final = 5

# Aggressive (slower, maximum recall)
k_bm25_candidates = 20
k_semantic_candidates = 20
k_rrf_pool = 15
k_final = 5