fusion.py

"""
Retrieval fusion, reranking, and source-focused filtering.

Phase 4 upgrade: merge_retrieved_docs now uses Reciprocal Rank Fusion (RRF)
instead of simple concatenation so vector and keyword ranks are balanced.
"""
import logging
from langchain_core.documents import Document

from retrieval.query import (
    normalize_for_match,
    infer_source_hints,
    source_matches_hints,
)

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Document identity
# ---------------------------------------------------------------------------

def doc_identity_key(doc: Document) -> tuple:
    return (
        doc.metadata.get("source", "unknown"),
        doc.metadata.get("page"),
        doc.page_content[:200],
    )


# ---------------------------------------------------------------------------
# Reciprocal Rank Fusion
# ---------------------------------------------------------------------------

def reciprocal_rank_fusion(
    ranked_lists: list[list[Document]],
    k: int = 60,
    limit: int | None = None,
) -> list[Document]:
    """
    Merge multiple ranked lists using Reciprocal Rank Fusion.

    Each document's RRF score is the sum of 1/(k+rank) across all lists
    it appears in.  k=60 is the standard default.
    """
    scores: dict[tuple, float] = {}
    docs_map: dict[tuple, Document] = {}

    for ranked_list in ranked_lists:
        for rank, doc in enumerate(ranked_list, 1):
            key = doc_identity_key(doc)
            scores[key] = scores.get(key, 0.0) + 1.0 / (k + rank)
            if key not in docs_map:
                docs_map[key] = doc

    sorted_keys = sorted(scores, key=lambda x: scores[x], reverse=True)
    if limit:
        sorted_keys = sorted_keys[:limit]
    return [docs_map[k] for k in sorted_keys]


# ---------------------------------------------------------------------------
# Legacy simple merge (kept for fallback / compatibility)
# ---------------------------------------------------------------------------

def merge_retrieved_docs(
    vector_docs: list[Document],
    keyword_docs: list[Document],
    limit: int,
) -> list[Document]:
    """
    Merge using RRF; falls back to dedup-by-order if lists are empty.
    Vector docs are list 1, keyword docs are list 2.
    """
    if vector_docs and keyword_docs:
        return reciprocal_rank_fusion([vector_docs, keyword_docs], limit=limit)

    # Dedup-by-insertion-order when only one source has results
    merged: list[Document] = []
    seen: set[tuple] = set()
    for doc in (vector_docs or []) + (keyword_docs or []):
        key = doc_identity_key(doc)
        if key in seen:
            continue
        seen.add(key)
        merged.append(doc)
        if len(merged) >= limit:
            break
    return merged


# ---------------------------------------------------------------------------
# Source and type filtering
# ---------------------------------------------------------------------------

def prioritize_non_image_docs(docs: list[Document], limit: int) -> list[Document]:
    """For non-image queries, push text/PDF/Excel results ahead of image docs."""
    if not docs:
        return []
    non_image = [d for d in docs if d.metadata.get("source_type") != "image"]
    image = [d for d in docs if d.metadata.get("source_type") == "image"]
    return (non_image + image)[:limit]


def search_non_image_context(vectorstore, question: str, limit: int) -> list[Document]:
    """Retrieve text-only context without making multiple vector queries."""
    collected: list[Document] = []
    seen: set[tuple] = set()
    try:
        docs = vectorstore.similarity_search(question, k=max(limit * 6, 18))
    except Exception:
        docs = []

    for doc in docs:
        if doc.metadata.get("source_type") == "image":
            continue
        key = doc_identity_key(doc)
        if key in seen:
            continue
        seen.add(key)
        collected.append(doc)
        if len(collected) >= limit:
            return collected
    return collected


def rank_text_docs(question: str, docs: list[Document]) -> list[Document]:
    """Re-rank text docs by token overlap with query (source name weighted 4×)."""
    q_tokens = set(normalize_for_match(question).split())
    if not q_tokens:
        return docs

    def score(doc: Document) -> int:
        source = normalize_for_match(doc.metadata.get("source", ""))
        content = normalize_for_match(doc.page_content[:800])
        source_overlap = len(q_tokens.intersection(set(source.split())))
        content_overlap = len(q_tokens.intersection(set(content.split())))
        return source_overlap * 4 + content_overlap

    return sorted(docs, key=score, reverse=True)


def apply_source_focus(
    vectorstore,
    question: str,
    docs: list[Document],
    limit: int,
) -> list[Document]:
    """
    When a product term is detected, restrict results to documents from the
    matching source before falling back to the full retrieved set.
    """
    hints = infer_source_hints(question)
    if not hints:
        return docs

    current_hits = [
        d for d in (docs or [])
        if d.metadata.get("source_type") != "image"
        and source_matches_hints(d.metadata.get("source", ""), hints)
    ]
    if current_hits:
        return rank_text_docs(question, current_hits)[:limit]

    try:
        expanded = vectorstore.similarity_search(question, k=max(24, limit * 6))
    except Exception:
        expanded = []

    focused: list[Document] = []
    seen: set[tuple] = set()
    for doc in expanded:
        source = doc.metadata.get("source", "")
        if doc.metadata.get("source_type") == "image":
            continue
        if not source_matches_hints(source, hints):
            continue
        key = doc_identity_key(doc)
        if key in seen:
            continue
        seen.add(key)
        focused.append(doc)
        if len(focused) >= limit:
            break

    if focused:
        return rank_text_docs(question, focused)[:limit]
    return docs