"""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