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

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
        )


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,
    ):
        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._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}
        self.keyword_index.build([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.top_k]

        fused = reciprocal_rank_fusion(
            [[_key(hit) for hit in vector_hits], [_key(hit) for hit in keyword_hits]],
            limit=self.top_k,
            k=self.rrf_k,
        )

        # O lado vetorial tem precedência: 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)

        return [best[key] for key, _ in fused if key in best]

    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