feat(backend): add incremental indexer with manifest and partial states
This commit is contained in:
omp committed
1 parent
abf7d802ab
commit
1748b561bd
2 files changed
+453
No files matched your search
@@ -0,0 +1,244 @@
|
||||
"""Indexador: orquestra extração, parsing, embeddings e persistência vetorial.
|
||||
|
||||
O manifesto (`manifest.json`) registra, por arquivo, o SHA-256 do conteúdo, o
|
||||
estado (`complete`, `partial`, `error`) e as páginas com falha. Arquivos cujo
|
||||
hash não mudou não são reprocessados.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
||||
from .bible_parser import BibleParser, parse_book_index
|
||||
from .config import Settings
|
||||
from .embedding import Embedder
|
||||
from .extraction import PageText, file_sha256
|
||||
from .vector_store import VectorStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
STATE_COMPLETE = "complete"
|
||||
STATE_PARTIAL = "partial"
|
||||
STATE_ERROR = "error"
|
||||
|
||||
|
||||
@dataclass
|
||||
class IndexStatus:
|
||||
"""Estado público da indexação."""
|
||||
|
||||
state: str = "ready"
|
||||
pdfs_total: int = 0
|
||||
pdfs_indexed: int = 0
|
||||
chunks: int = 0
|
||||
model: str = ""
|
||||
last_indexed_at: str | None = None
|
||||
errors: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _ManifestEntry:
|
||||
sha256: str
|
||||
state: str
|
||||
chunks: int = 0
|
||||
indexed_at: str = ""
|
||||
failed_pages: list[int] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"sha256": self.sha256,
|
||||
"state": self.state,
|
||||
"chunks": self.chunks,
|
||||
"indexed_at": self.indexed_at,
|
||||
"failed_pages": self.failed_pages,
|
||||
}
|
||||
|
||||
|
||||
class Indexer:
|
||||
"""Coordena a indexação dos PDFs presentes em `settings.pdf_dir`."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
settings: Settings,
|
||||
extractor,
|
||||
parser: BibleParser,
|
||||
embedder: Embedder,
|
||||
store: VectorStore,
|
||||
):
|
||||
self.settings = settings
|
||||
self.extractor = extractor
|
||||
self.parser = parser
|
||||
self.embedder = embedder
|
||||
self.store = store
|
||||
self._manifest: dict[str, _ManifestEntry] = {}
|
||||
self._lock = threading.Lock()
|
||||
self._last_indexed_at: str | None = None
|
||||
self._errors: list[str] = []
|
||||
self._load_manifest()
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Manifesto
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def _load_manifest(self) -> None:
|
||||
path = self.settings.manifest_path
|
||||
if not path.exists():
|
||||
return
|
||||
try:
|
||||
raw = json.loads(path.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError) as exc:
|
||||
logger.warning("Manifesto ilegível em %s: %s", path, exc)
|
||||
return
|
||||
for name, entry in raw.items():
|
||||
self._manifest[name] = _ManifestEntry(
|
||||
sha256=entry.get("sha256", ""),
|
||||
state=entry.get("state", STATE_ERROR),
|
||||
chunks=int(entry.get("chunks", 0)),
|
||||
indexed_at=entry.get("indexed_at", ""),
|
||||
failed_pages=list(entry.get("failed_pages", [])),
|
||||
)
|
||||
|
||||
def _save_manifest(self) -> None:
|
||||
path = self.settings.manifest_path
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
payload = {name: entry.to_dict() for name, entry in self._manifest.items()}
|
||||
temporary = path.with_suffix(".json.tmp")
|
||||
temporary.write_text(
|
||||
json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
|
||||
)
|
||||
temporary.replace(path)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Índice de livros
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def _book_starts(self, pages: list[PageText]) -> dict[str, int]:
|
||||
"""Índice `{livro: página}` do PDF, memoizado em disco.
|
||||
|
||||
As primeiras páginas do documento contêm o índice de livros. Se ele não
|
||||
puder ser lido, devolve um dicionário vazio: o parser cairá para o
|
||||
cabeçalho de cada página.
|
||||
"""
|
||||
cached_path = self.settings.book_index_path
|
||||
if cached_path.exists():
|
||||
try:
|
||||
cached = json.loads(cached_path.read_text(encoding="utf-8"))
|
||||
if cached:
|
||||
return {str(key): int(value) for key, value in cached.items()}
|
||||
except (OSError, json.JSONDecodeError, ValueError) as exc:
|
||||
logger.warning("Índice de livros ilegível em %s: %s", cached_path, exc)
|
||||
|
||||
starts: dict[str, int] = {}
|
||||
for page in pages[:12]:
|
||||
if not page.text.strip():
|
||||
continue
|
||||
starts.update(parse_book_index(page.text))
|
||||
if len(starts) >= 20:
|
||||
break
|
||||
|
||||
if starts:
|
||||
try:
|
||||
cached_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
cached_path.write_text(
|
||||
json.dumps(starts, ensure_ascii=False, indent=2), encoding="utf-8"
|
||||
)
|
||||
except OSError as exc: # pragma: no cover
|
||||
logger.warning("Não foi possível gravar o índice de livros: %s", exc)
|
||||
return starts
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Indexação
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def discover_pdfs(self) -> list[Path]:
|
||||
"""Lista os PDFs do diretório configurado, ordenados por nome."""
|
||||
directory = self.settings.pdf_dir
|
||||
if not directory.exists():
|
||||
return []
|
||||
return sorted(path for path in directory.glob("*.pdf") if path.is_file())
|
||||
|
||||
def reindex(self, force: bool = False) -> IndexStatus:
|
||||
"""Indexa os PDFs novos ou alterados."""
|
||||
with self._lock:
|
||||
self._errors = []
|
||||
pdfs = self.discover_pdfs()
|
||||
|
||||
for pdf in pdfs:
|
||||
try:
|
||||
self._index_file(pdf, force=force)
|
||||
except Exception as exc: # pragma: no cover - falha inesperada
|
||||
logger.exception("Falha ao indexar %s", pdf.name)
|
||||
self._manifest[pdf.name] = _ManifestEntry(
|
||||
sha256="", state=STATE_ERROR, indexed_at=_now()
|
||||
)
|
||||
self._errors.append(f"{pdf.name}: {exc}")
|
||||
|
||||
self._save_manifest()
|
||||
self._last_indexed_at = _now()
|
||||
return self.status()
|
||||
|
||||
def _index_file(self, pdf: Path, force: bool) -> None:
|
||||
file_hash = file_sha256(pdf)
|
||||
entry = self._manifest.get(pdf.name)
|
||||
|
||||
if entry is not None and entry.sha256 == file_hash and not force:
|
||||
logger.info("%s já está indexado; pulando", pdf.name)
|
||||
return
|
||||
|
||||
logger.info("Indexando %s", pdf.name)
|
||||
pages = list(self.extractor.extract_document(pdf, file_hash=file_hash))
|
||||
failed_pages = [page.page for page in pages if page.failed]
|
||||
|
||||
book_starts = self._book_starts(pages)
|
||||
chunks = self.parser.parse(pages, pdf.name, book_starts)
|
||||
|
||||
if not chunks:
|
||||
state = STATE_ERROR
|
||||
self._errors.append(f"{pdf.name}: nenhum versículo pôde ser extraído")
|
||||
elif failed_pages:
|
||||
state = STATE_PARTIAL
|
||||
self._errors.append(
|
||||
f"{pdf.name}: {len(failed_pages)} página(s) falharam na extração"
|
||||
)
|
||||
else:
|
||||
state = STATE_COMPLETE
|
||||
|
||||
# Remove os chunks anteriores deste arquivo antes de gravar os novos.
|
||||
self.store.delete_file(pdf.name)
|
||||
|
||||
if chunks:
|
||||
vectors = self.embedder.encode([chunk.text for chunk in chunks])
|
||||
self.store.upsert(chunks, vectors)
|
||||
|
||||
self._manifest[pdf.name] = _ManifestEntry(
|
||||
sha256=file_hash,
|
||||
state=state,
|
||||
chunks=len(chunks),
|
||||
indexed_at=_now(),
|
||||
failed_pages=failed_pages,
|
||||
)
|
||||
|
||||
def status(self) -> IndexStatus:
|
||||
"""Estado atual da indexação."""
|
||||
indexed = sum(
|
||||
1
|
||||
for entry in self._manifest.values()
|
||||
if entry.state in (STATE_COMPLETE, STATE_PARTIAL)
|
||||
)
|
||||
return IndexStatus(
|
||||
state="error" if self._errors and indexed == 0 else "ready",
|
||||
pdfs_total=len(self.discover_pdfs()),
|
||||
pdfs_indexed=indexed,
|
||||
chunks=self.store.count(),
|
||||
model=self.settings.embedding_model,
|
||||
last_indexed_at=self._last_indexed_at,
|
||||
errors=list(self._errors),
|
||||
)
|
||||
|
||||
|
||||
def _now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
Reference in new issue
Block a user