omp
957c73b09e
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.
297 lines
10 KiB
Python
297 lines
10 KiB
Python
"""Banco vetorial persistente (ChromaDB) e implementação em memória para testes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import math
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Iterable, Sequence
|
|
|
|
from .corpus import VerseChunk
|
|
from .embedding import Embedder, HashEmbedder
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SearchResult:
|
|
"""Trecho recuperado da busca.
|
|
|
|
`text` é a janela de contexto que a busca encontrou (e que vai para o LLM);
|
|
`verse_text` é só o versículo citado, que é o que o usuário vê.
|
|
"""
|
|
|
|
book: str
|
|
chapter: int
|
|
verse: int
|
|
file: str
|
|
text: str
|
|
verse_text: str = ""
|
|
distance: float = 0.0
|
|
|
|
@property
|
|
def reference(self) -> str:
|
|
return f"{self.book} {self.chapter}:{self.verse}"
|
|
|
|
@property
|
|
def cited_text(self) -> str:
|
|
"""Texto a exibir como citação: o versículo, ou a janela se não houver."""
|
|
return self.verse_text or self.text
|
|
|
|
@property
|
|
def score(self) -> float:
|
|
"""Similaridade em escala 0..1 (1 = idêntico)."""
|
|
return max(0.0, min(1.0, 1.0 - self.distance))
|
|
|
|
|
|
def cosine_similarity(left: Sequence[float], right: Sequence[float]) -> float:
|
|
"""Similaridade de cosseno entre dois vetores."""
|
|
if not left or not right:
|
|
return 0.0
|
|
size = min(len(left), len(right))
|
|
dot = sum(left[i] * right[i] for i in range(size))
|
|
left_norm = math.sqrt(sum(value * value for value in left[:size]))
|
|
right_norm = math.sqrt(sum(value * value for value in right[:size]))
|
|
if left_norm == 0 or right_norm == 0:
|
|
return 0.0
|
|
return dot / (left_norm * right_norm)
|
|
|
|
|
|
class VectorStore:
|
|
"""Fachada sobre o ChromaDB persistente."""
|
|
|
|
def __init__(
|
|
self,
|
|
persist_dir: Path | str,
|
|
collection_name: str = "biblia",
|
|
embedder: Embedder | None = None,
|
|
):
|
|
import chromadb
|
|
from chromadb.config import Settings as ChromaSettings
|
|
|
|
self.persist_dir = Path(persist_dir)
|
|
self.persist_dir.mkdir(parents=True, exist_ok=True)
|
|
self.collection_name = collection_name
|
|
self.embedder = embedder or HashEmbedder()
|
|
self._client = chromadb.PersistentClient(
|
|
path=str(self.persist_dir),
|
|
settings=ChromaSettings(anonymized_telemetry=False, allow_reset=True),
|
|
)
|
|
self._collection = self._client.get_or_create_collection(
|
|
name=collection_name,
|
|
metadata={"hnsw:space": "cosine"},
|
|
)
|
|
|
|
def upsert(self, chunks: Iterable[VerseChunk], vectors: Sequence[Sequence[float]]) -> None:
|
|
"""Grava chunks e vetores, substituindo os de mesmo id."""
|
|
chunk_list = list(chunks)
|
|
if not chunk_list:
|
|
return
|
|
if len(vectors) != len(chunk_list):
|
|
raise ValueError("vectors e chunks devem ter o mesmo tamanho")
|
|
|
|
self._collection.upsert(
|
|
ids=[chunk.chunk_id for chunk in chunk_list],
|
|
embeddings=[list(vector) for vector in vectors],
|
|
# O documento gravado é a janela de contexto: é ela que foi embeddada
|
|
# e é ela que serve de evidência ao LLM. O versículo citado vai no
|
|
# metadado, para a UI mostrar a referência exata.
|
|
documents=[chunk.embedding_text for chunk in chunk_list],
|
|
metadatas=[
|
|
{
|
|
"book": chunk.book,
|
|
"chapter": chunk.chapter,
|
|
"verse": chunk.verse,
|
|
"file": chunk.file,
|
|
"verse_text": chunk.text,
|
|
"reference": f"{chunk.book} {chunk.chapter}:{chunk.verse}",
|
|
}
|
|
for chunk in chunk_list
|
|
],
|
|
)
|
|
|
|
def delete_file(self, file_name: str) -> None:
|
|
"""Remove todos os chunks de um arquivo."""
|
|
existing = self._collection.get(where={"file": file_name})
|
|
ids = existing.get("ids") or []
|
|
if ids:
|
|
self._collection.delete(ids=ids)
|
|
|
|
def query(
|
|
self, vector: Sequence[float], n_results: int = 6, min_similarity: float = 0.35
|
|
) -> list[SearchResult]:
|
|
"""Busca os chunks mais próximos, descartando os abaixo do limiar."""
|
|
total = self._collection.count()
|
|
if total == 0:
|
|
return []
|
|
|
|
raw = self._collection.query(
|
|
query_embeddings=[list(vector)],
|
|
n_results=min(n_results, total),
|
|
include=["documents", "metadatas", "distances"],
|
|
)
|
|
|
|
documents = (raw.get("documents") or [[]])[0]
|
|
metadatas = (raw.get("metadatas") or [[]])[0]
|
|
distances = (raw.get("distances") or [[]])[0]
|
|
|
|
results: list[SearchResult] = []
|
|
for document, metadata, distance in zip(documents, metadatas, distances):
|
|
similarity = 1.0 - float(distance)
|
|
if similarity < min_similarity:
|
|
continue
|
|
results.append(
|
|
SearchResult(
|
|
book=str(metadata.get("book", "")),
|
|
chapter=int(metadata.get("chapter", 0)),
|
|
verse=int(metadata.get("verse", 0)),
|
|
file=str(metadata.get("file", "")),
|
|
text=document or "",
|
|
verse_text=str(metadata.get("verse_text", "") or ""),
|
|
distance=float(distance),
|
|
)
|
|
)
|
|
return results
|
|
|
|
def get_by_reference(
|
|
self, book: str, chapter: int, verse: int
|
|
) -> list[SearchResult]:
|
|
"""Busca determinística por livro, capítulo e versículo."""
|
|
raw = self._collection.get(
|
|
where={
|
|
"$and": [
|
|
{"book": book},
|
|
{"chapter": int(chapter)},
|
|
{"verse": int(verse)},
|
|
]
|
|
},
|
|
include=["documents", "metadatas"],
|
|
)
|
|
documents = raw.get("documents") or []
|
|
metadatas = raw.get("metadatas") or []
|
|
return [
|
|
SearchResult(
|
|
book=str(metadata.get("book", book)),
|
|
chapter=int(metadata.get("chapter", chapter)),
|
|
verse=int(metadata.get("verse", verse)),
|
|
file=str(metadata.get("file", "")),
|
|
text=document or "",
|
|
verse_text=str(metadata.get("verse_text", "") or ""),
|
|
distance=0.0,
|
|
)
|
|
for document, metadata in zip(documents, metadatas)
|
|
]
|
|
|
|
def snapshot(self) -> list[SearchResult]:
|
|
"""Inventário de todos os trechos indexados, sem pontuação.
|
|
|
|
Serve para (re)construir o índice de palavras-chave depois de uma
|
|
reinicialização, quando o Chroma persistiu mas o BM25 em memória não.
|
|
`distance=1.0` deixa `score` em 0: snapshot não é ranking e não pode
|
|
fingir similaridade.
|
|
"""
|
|
raw = self._collection.get(include=["documents", "metadatas"])
|
|
documents = raw.get("documents") or []
|
|
metadatas = raw.get("metadatas") or []
|
|
return [
|
|
SearchResult(
|
|
book=str(metadata.get("book", "")),
|
|
chapter=int(metadata.get("chapter", 0)),
|
|
verse=int(metadata.get("verse", 0)),
|
|
file=str(metadata.get("file", "")),
|
|
text=document or "",
|
|
verse_text=str(metadata.get("verse_text", "") or ""),
|
|
distance=1.0,
|
|
)
|
|
for document, metadata in zip(documents, metadatas)
|
|
]
|
|
|
|
def count(self) -> int:
|
|
"""Total de chunks indexados."""
|
|
return int(self._collection.count())
|
|
|
|
|
|
class InMemoryVectorStore:
|
|
"""Banco vetorial em memória, usado nos testes.
|
|
|
|
Implementa a mesma interface de `VectorStore`, com busca por similaridade
|
|
de cosseno e filtros por metadados.
|
|
"""
|
|
|
|
def __init__(self, embedder: Embedder | None = None):
|
|
self.embedder = embedder or HashEmbedder()
|
|
self._vectors: dict[str, list[float]] = {}
|
|
self._chunks: dict[str, VerseChunk] = {}
|
|
|
|
def upsert(self, chunks: Iterable[VerseChunk], vectors: Sequence[Sequence[float]]) -> None:
|
|
chunk_list = list(chunks)
|
|
if len(vectors) != len(chunk_list):
|
|
raise ValueError("vectors e chunks devem ter o mesmo tamanho")
|
|
for chunk, vector in zip(chunk_list, vectors):
|
|
self._vectors[chunk.chunk_id] = list(vector)
|
|
self._chunks[chunk.chunk_id] = chunk
|
|
|
|
def delete_file(self, file_name: str) -> None:
|
|
for chunk_id in [cid for cid, chunk in self._chunks.items() if chunk.file == file_name]:
|
|
self._vectors.pop(chunk_id, None)
|
|
self._chunks.pop(chunk_id, None)
|
|
|
|
def query(
|
|
self, vector: Sequence[float], n_results: int = 6, min_similarity: float = 0.35
|
|
) -> list[SearchResult]:
|
|
scored: list[SearchResult] = []
|
|
for chunk_id, stored in self._vectors.items():
|
|
similarity = cosine_similarity(vector, stored)
|
|
if similarity < min_similarity:
|
|
continue
|
|
chunk = self._chunks[chunk_id]
|
|
scored.append(
|
|
SearchResult(
|
|
book=chunk.book,
|
|
chapter=chunk.chapter,
|
|
verse=chunk.verse,
|
|
file=chunk.file,
|
|
text=chunk.embedding_text,
|
|
verse_text=chunk.text,
|
|
distance=1.0 - similarity,
|
|
)
|
|
)
|
|
scored.sort(key=lambda result: result.distance)
|
|
return scored[:n_results]
|
|
|
|
def get_by_reference(
|
|
self, book: str, chapter: int, verse: int
|
|
) -> list[SearchResult]:
|
|
return [
|
|
SearchResult(
|
|
book=chunk.book,
|
|
chapter=chunk.chapter,
|
|
verse=chunk.verse,
|
|
file=chunk.file,
|
|
text=chunk.embedding_text,
|
|
verse_text=chunk.text,
|
|
distance=0.0,
|
|
)
|
|
for chunk in self._chunks.values()
|
|
if chunk.book == book and chunk.chapter == chapter and chunk.verse == verse
|
|
]
|
|
|
|
def snapshot(self) -> list[SearchResult]:
|
|
"""Inventário de todos os trechos indexados, sem pontuação."""
|
|
return [
|
|
SearchResult(
|
|
book=chunk.book,
|
|
chapter=chunk.chapter,
|
|
verse=chunk.verse,
|
|
file=chunk.file,
|
|
text=chunk.embedding_text,
|
|
verse_text=chunk.text,
|
|
distance=1.0,
|
|
)
|
|
for chunk in self._chunks.values()
|
|
]
|
|
|
|
def count(self) -> int:
|
|
return len(self._chunks)
|