feat: busca híbrida (vetorial + BM25 com fusão RRF) e janela de contexto
Recuperação: - keyword.py: tokenizador PT (acentos, stopwords, dígitos), BM25 Okapi com índice invertido, e reciprocal_rank_fusion com desempate determinístico - retrieval.py: VectorRetriever e HybridRetriever; fusão por posição, não por nota (cosseno 0..1 e BM25 sem teto são escalas incomparáveis) - identidade de fusão é a coordenada (arquivo,livro,capítulo,versículo): o mesmo versículo chega das duas buscas com distance diferente e o dataclass nunca os igualaria, então o RRF não premiaria o que ambas confirmam - ChatService recebe retriever injetado; referência explícita mantém precedência Contexto: - VerseChunk ganha context (janela de CONTEXT_WINDOW vizinhos por lado, dentro do livro, atravessando capítulos) e embedding_text - store grava a janela como documento e o versículo em verse_text; a UI mostra o versículo citado e o modelo recebe a janela - snapshot() nos dois stores: o BM25 é em memória e o Chroma persiste, então o arranque precisa reconstruir o índice lexical ou a busca degrada em silêncio Motivação medida em 5703 versículos reais: a busca vetorial não liga 'estou com medo do futuro' a 'não temas' (0,19, nível de par sem relação), e e5-small tem faixa de similaridade de só 0,02 — inutilizaria o piso de evidência.
This commit is contained in:
omp committed
1 parent
2043d5ddcf
commit
957c73b09e
16 files changed
+1085
-45
No files matched your search
@@ -0,0 +1,147 @@
|
||||
"""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
|
||||
Reference in new issue
Block a user