src/retrieval/fusion.py

"""Fusion algorithms for combining multiple retrieval results."""

import logging
from collections import defaultdict
from typing import Any

from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever

logger = logging.getLogger(__name__)


def reciprocal_rank_fusion(
    results_list: list[list[Document]],
    k: int = 60,
    weights: list[float] | None = None,
) -> list[Document]:
    """Combine multiple ranked lists using Reciprocal Rank Fusion (RRF).

    RRF is a simple but effective method for combining rankings from
    multiple retrieval systems. It's based on the formula:
        RRF(d) = sum(1 / (k + rank(d)))

    where k is a constant (typically 60) and rank(d) is the position
    of document d in each ranking (1-indexed).

    Args:
        results_list: List of ranked document lists from different retrievers
        k: RRF constant (higher = less aggressive re-ranking, default: 60)
        weights: Optional weights for each result list (must sum to 1.0)

    Returns:
        Combined and re-ranked list of documents

    Reference:
        Cormack, G. V., Clarke, C. L., & Buettcher, S. (2009).
        Reciprocal rank fusion outperforms condorcet and individual
        rank learning methods. SIGIR.
    """
    if not results_list:
        return []

    # Normalize weights if provided
    if weights is None:
        weights = [1.0] * len(results_list)
    else:
        if len(weights) != len(results_list):
            raise ValueError("Number of weights must match number of result lists")
        # Normalize weights to sum to 1
        total_weight = sum(weights)
        if total_weight > 0:
            weights = [w / total_weight for w in weights]

    # Calculate RRF scores
    # Use page_content as document identifier (could also use metadata)
    doc_scores: dict[str, float] = defaultdict(float)
    doc_map: dict[str, Document] = {}

    for result_idx, results in enumerate(results_list):
        weight = weights[result_idx]
        for rank, doc in enumerate(results, start=1):
            # Create a unique key for the document
            doc_key = doc.page_content

            # RRF formula with weight
            score = weight * (1.0 / (k + rank))
            doc_scores[doc_key] += score

            # Keep the document object (prefer first occurrence)
            if doc_key not in doc_map:
                doc_map[doc_key] = doc

    # Sort by RRF score (descending)
    sorted_keys = sorted(doc_scores.keys(), key=lambda x: doc_scores[x], reverse=True)

    # Return sorted documents
    return [doc_map[key] for key in sorted_keys]


def weighted_fusion(
    results_list: list[list[tuple[Document, float]]],
    weights: list[float] | None = None,
) -> list[Document]:
    """Combine multiple scored result lists using weighted scoring.

    This method is useful when retrievers provide relevance scores.
    Documents are combined by weighted sum of their scores.

    Args:
        results_list: List of (document, score) tuples from different retrievers
        weights: Weights for each result list (normalized internally)

    Returns:
        Combined and re-ranked list of documents
    """
    if not results_list:
        return []

    # Normalize weights
    if weights is None:
        weights = [1.0] * len(results_list)
    else:
        total_weight = sum(weights)
        if total_weight > 0:
            weights = [w / total_weight for w in weights]

    # Combine scores
    doc_scores: dict[str, float] = defaultdict(float)
    doc_map: dict[str, Document] = {}

    for result_idx, results in enumerate(results_list):
        weight = weights[result_idx]
        for doc, score in results:
            doc_key = doc.page_content
            doc_scores[doc_key] += weight * score
            if doc_key not in doc_map:
                doc_map[doc_key] = doc

    # Sort by combined score
    sorted_keys = sorted(doc_scores.keys(), key=lambda x: doc_scores[x], reverse=True)

    return [doc_map[key] for key in sorted_keys]


class FusionRetriever(BaseRetriever):
    """A retriever that combines results from multiple retrievers using fusion.

    This retriever implements the Ensemble pattern, allowing multiple
    retrieval strategies to be combined for better recall and precision.

    Attributes:
        retrievers: List of retrievers to combine
        weights: Weights for each retriever (optional)
        fusion_algorithm: Algorithm to use ('rrf' or 'weighted')
        rrf_k: Constant for RRF algorithm
        final_k: Number of documents to return after fusion
    """

    retrievers: list[BaseRetriever]
    weights: list[float] | None = None
    fusion_algorithm: str = "rrf"
    rrf_k: int = 60
    final_k: int = 4

    model_config = {"arbitrary_types_allowed": True}

    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun | None = None,  # noqa: ARG002
    ) -> list[Document]:
        """Get documents from all retrievers and fuse results.

        Args:
            query: Query string to search for
            run_manager: Callback manager

        Returns:
            Fused and re-ranked list of documents
        """
        # Collect results from all retrievers
        all_results: list[list[Document]] = []

        for retriever in self.retrievers:
            try:
                results = retriever.invoke(query)
                all_results.append(results)
                logger.debug(f"Retriever returned {len(results)} documents")
            except Exception as e:
                logger.warning(f"Retriever failed: {e}")
                all_results.append([])

        # Apply fusion algorithm
        if self.fusion_algorithm == "rrf":
            fused = reciprocal_rank_fusion(
                all_results,
                k=self.rrf_k,
                weights=self.weights,
            )
        else:
            # For non-RRF, we don't have scores, so use RRF as fallback
            logger.warning(
                f"Fusion algorithm '{self.fusion_algorithm}' not supported "
                "for rank-only results, falling back to RRF"
            )
            fused = reciprocal_rank_fusion(all_results, k=self.rrf_k, weights=self.weights)

        # Return top k results
        result = fused[: self.final_k]
        logger.debug(f"Fusion returned {len(result)} documents from {len(fused)} total")

        return result

    def get_retriever_info(self) -> dict[str, Any]:
        """Get information about the fusion configuration.

        Returns:
            Dictionary with fusion configuration details
        """
        return {
            "num_retrievers": len(self.retrievers),
            "weights": self.weights,
            "fusion_algorithm": self.fusion_algorithm,
            "rrf_k": self.rrf_k,
            "final_k": self.final_k,
        }