pastor-online/backend/app/retrieval.py

"""Estratégias de recuperação: vetorial e híbrida (vetorial + BM25).

A recuperação híbrida existe porque as duas buscas erram de modos diferentes. A
vetorial acha "não temas" para "estou com medo" só se o embedding captar a
paráfrase — medido: não capta. A lexical acha Mateus 5:44 para "amar os
inimigos" pela palavra, mas não entende sentimento. Fundir os dois rankings
cobre os dois casos.

A fusão é por **posição** (RRF), não por pontuação: cosseno vive em 0..1 e BM25
não tem teto, então normalizar um contra o outro seria arbitrário.
"""

from __future__ import annotations

import logging
from typing import Protocol, Sequence

from .embedding import Embedder
from .keyword import BM25Index, reciprocal_rank_fusion
from .vector_store import SearchResult

logger = logging.getLogger(__name__)

# Identidade de um trecho para fusão. Não pode ser o `SearchResult` inteiro: o
# mesmo versículo chega das duas buscas com `distance` diferente, e como o
# dataclass compara todos os campos as duas ocorrências pareceriam itens
# distintos — o RRF nunca premiaria o que ambas confirmam.
CandidateKey = tuple[str, str, int, int]


def _key(result: SearchResult) -> CandidateKey:
    return (result.file, result.book, result.chapter, result.verse)


class Retriever(Protocol):
    """Interface de recuperação de trechos."""

    def search(self, message: str) -> list[SearchResult]:
        """Trechos que sustentam a resposta, em ordem de relevância."""
        ...


class VectorRetriever:
    """Recuperação só por similaridade vetorial."""

    def __init__(
        self,
        store,
        embedder: Embedder,
        top_k: int = 6,
        min_similarity: float = 0.35,
    ):
        self.store = store
        self.embedder = embedder
        self.top_k = top_k
        self.min_similarity = min_similarity

    def search(self, message: str) -> list[SearchResult]:
        if not message or not message.strip():
            return []
        vector = self.embedder.encode([message])[0]
        return self.store.query(
            vector, n_results=self.top_k, min_similarity=self.min_similarity
        )


def fuse_rankings(
    vector_hits: Sequence[SearchResult],
    keyword_hits: Sequence[SearchResult],
    limit: int,
    rrf_k: int = 60,
    max_keyword_only: int = 2,
) -> list[SearchResult]:
    """Funde os dois rankings por RRF, com cota para achados só-lexicais.

    A cota existe porque a **nota** BM25 não separa acerto de ruído — medido no
    corpus real: para "como posso perdoar quem me feriu?", Deuteronômio 28:27
    ("o Senhor te ferirá") pontua na razão 1,000 do topo enquanto o correto
    Lucas 7:47 fica em 0,716. Um piso de nota cortaria o certo e manteria o
    errado. Limitar quantas posições o lado lexical pode ocupar sozinho preserva
    os achados fortes (Jó 3:13 para "não consigo dormir") sem deixar o ruído
    dominar a resposta.

    O que os dois lados confirmam nunca conta contra a cota: é a evidência mais
    forte disponível.
    """
    vector_keys = [_key(hit) for hit in vector_hits]
    keyword_keys = [_key(hit) for hit in keyword_hits]

    # O lado vetorial tem precedência no dicionário: ele traz a similaridade
    # real. Um trecho achado só por palavra-chave fica com score 0, que é
    # honesto — nenhuma similaridade vetorial foi medida para ele.
    best: dict[CandidateKey, SearchResult] = {}
    for hit in vector_hits:
        best.setdefault(_key(hit), hit)
    for hit in keyword_hits:
        best.setdefault(_key(hit), hit)

    confirmed = set(vector_keys)
    fused = reciprocal_rank_fusion(
        [vector_keys, keyword_keys],
        limit=len(best),
        k=rrf_k,
    )

    results: list[SearchResult] = []
    keyword_only_used = 0
    for key, _score in fused:
        if len(results) >= limit:
            break
        hit = best.get(key)
        if hit is None:
            continue
        if key not in confirmed:
            if keyword_only_used >= max_keyword_only:
                continue
            keyword_only_used += 1
        results.append(hit)
    return results


class HybridRetriever(VectorRetriever):
    """Vetorial + BM25 fundidos por RRF."""

    def __init__(
        self,
        store,
        embedder: Embedder,
        top_k: int = 6,
        min_similarity: float = 0.35,
        keyword_index: BM25Index | None = None,
        candidate_multiplier: int = 4,
        rrf_k: int = 60,
        max_keyword_only: int = 2,
    ):
        super().__init__(
            store=store,
            embedder=embedder,
            top_k=top_k,
            min_similarity=min_similarity,
        )
        self.keyword_index = keyword_index or BM25Index()
        self.candidate_multiplier = max(1, candidate_multiplier)
        self.rrf_k = rrf_k
        self.max_keyword_only = max(0, max_keyword_only)
        self._entries: dict[CandidateKey, SearchResult] = {}
        self._order: list[CandidateKey] = []

    def rebuild(self) -> None:
        """Reconstrói o índice lexical a partir do que está no banco vetorial.

        Chamado após cada indexação e no arranque: o Chroma persiste em disco,
        mas o BM25 vive em memória e nasce vazio a cada processo novo.
        """
        entries = self.store.snapshot()
        self._order = [_key(entry) for entry in entries]
        self._entries = {_key(entry): entry for entry in entries}
        # BM25 sobre o versículo, não sobre a janela de contexto. Os dois lados
        # querem granularidades opostas: a janela dá tema ao embedding, mas
        # multiplica os termos incidentais que o BM25 pontua. Medido no corpus
        # real, para "o que Jesus disse sobre amar os inimigos?" Mateus 5:44 fica
        # em 12º indexando a janela e em 1º indexando o versículo.
        self.keyword_index.build(
            [entry.verse_text or entry.text for entry in entries]
        )
        logger.info("Índice de palavras-chave construído com %s trechos", len(entries))

    def search(self, message: str) -> list[SearchResult]:
        if not message or not message.strip():
            return []

        candidates = self.top_k * self.candidate_multiplier
        vector_hits = self.store.query(
            self.embedder.encode([message])[0],
            n_results=candidates,
            min_similarity=self.min_similarity,
        )
        keyword_hits = self._keyword_search(message, candidates)

        if not keyword_hits:
            return vector_hits[: self.top_k]
        if not vector_hits:
            return keyword_hits[: self.max_keyword_only or self.top_k]

        return fuse_rankings(
            vector_hits,
            keyword_hits,
            limit=self.top_k,
            rrf_k=self.rrf_k,
            max_keyword_only=self.max_keyword_only,
        )

    def _keyword_search(self, message: str, limit: int) -> list[SearchResult]:
        if not len(self.keyword_index):
            return []
        hits: list[SearchResult] = []
        for doc_index, _score in self.keyword_index.search(message, k=limit):
            if 0 <= doc_index < len(self._order):
                entry = self._entries.get(self._order[doc_index])
                if entry is not None:
                    hits.append(entry)
        return hits