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()
|
||||
@@ -0,0 +1,209 @@
|
||||
"""Testes do indexador: manifesto incremental, reprocessamento e falhas parciais."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
|
||||
from app.bible_parser import BibleParser # noqa: E402
|
||||
from app.embedding import HashEmbedder # noqa: E402
|
||||
from app.extraction import PageText # noqa: E402
|
||||
from app.indexer import Indexer, IndexStatus # noqa: E402
|
||||
from app.vector_store import InMemoryVectorStore # noqa: E402
|
||||
|
||||
|
||||
class FakeExtractor:
|
||||
"""Extrator controlado: devolve as páginas pré-definidas por arquivo."""
|
||||
|
||||
def __init__(self, pages_by_file: dict[str, list[PageText]]):
|
||||
self.pages_by_file = pages_by_file
|
||||
self.calls: list[str] = []
|
||||
|
||||
def extract_document(self, path: Path, file_hash: str | None = None):
|
||||
self.calls.append(Path(path).name)
|
||||
yield from self.pages_by_file.get(Path(path).name, [])
|
||||
|
||||
|
||||
GENESIS_PAGE = (
|
||||
"GÊNESIS\n"
|
||||
"[1]\n"
|
||||
"1 No princípio criou Deus os céus e a terra.\n"
|
||||
"2 A terra era sem forma e vazia.\n"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def settings(tmp_path):
|
||||
from app.config import Settings
|
||||
|
||||
return Settings(
|
||||
pdf_dir=tmp_path / "pdfs",
|
||||
vector_db_dir=tmp_path / "vector_db",
|
||||
embedding_model="hash-test",
|
||||
)
|
||||
|
||||
|
||||
def build_indexer(settings, extractor, embedder=None, store=None):
|
||||
embedder = embedder or HashEmbedder(dimensions=32)
|
||||
store = store or InMemoryVectorStore(embedder=embedder)
|
||||
indexer = Indexer(
|
||||
settings=settings,
|
||||
extractor=extractor,
|
||||
parser=BibleParser(),
|
||||
embedder=embedder,
|
||||
store=store,
|
||||
)
|
||||
return indexer, store
|
||||
|
||||
|
||||
def write_pdf(directory: Path, name: str, content: bytes = b"%PDF-1.3") -> Path:
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
path = directory / name
|
||||
path.write_bytes(content)
|
||||
return path
|
||||
|
||||
|
||||
def test_reindex_indexes_pdfs_and_reports_counts(settings):
|
||||
write_pdf(settings.pdf_dir, "biblia.pdf")
|
||||
extractor = FakeExtractor({"biblia.pdf": [PageText(page=6, text=GENESIS_PAGE)]})
|
||||
indexer, store = build_indexer(settings, extractor)
|
||||
|
||||
status = indexer.reindex()
|
||||
|
||||
assert status.state == "ready"
|
||||
assert status.pdfs_total == 1
|
||||
assert status.pdfs_indexed == 1
|
||||
assert status.chunks == 2
|
||||
assert store.count() == 2
|
||||
|
||||
|
||||
def test_reindex_is_incremental_when_hash_is_unchanged(settings):
|
||||
write_pdf(settings.pdf_dir, "biblia.pdf")
|
||||
extractor = FakeExtractor({"biblia.pdf": [PageText(page=6, text=GENESIS_PAGE)]})
|
||||
indexer, _ = build_indexer(settings, extractor)
|
||||
|
||||
indexer.reindex()
|
||||
indexer.reindex()
|
||||
|
||||
assert extractor.calls == ["biblia.pdf"], "o PDF não deveria ser reprocessado"
|
||||
|
||||
|
||||
def test_reindex_reprocesses_file_when_content_changes(settings):
|
||||
pdf_path = write_pdf(settings.pdf_dir, "biblia.pdf", b"%PDF-1.3 original")
|
||||
extractor = FakeExtractor({"biblia.pdf": [PageText(page=6, text=GENESIS_PAGE)]})
|
||||
indexer, _ = build_indexer(settings, extractor)
|
||||
|
||||
indexer.reindex()
|
||||
pdf_path.write_bytes(b"%PDF-1.3 modified content")
|
||||
indexer.reindex()
|
||||
|
||||
assert extractor.calls == ["biblia.pdf", "biblia.pdf"]
|
||||
|
||||
|
||||
def test_reindex_with_force_reprocesses_everything(settings):
|
||||
write_pdf(settings.pdf_dir, "biblia.pdf")
|
||||
extractor = FakeExtractor({"biblia.pdf": [PageText(page=6, text=GENESIS_PAGE)]})
|
||||
indexer, _ = build_indexer(settings, extractor)
|
||||
|
||||
indexer.reindex()
|
||||
indexer.reindex(force=True)
|
||||
|
||||
assert extractor.calls == ["biblia.pdf", "biblia.pdf"]
|
||||
|
||||
|
||||
def test_manifest_is_persisted(settings):
|
||||
write_pdf(settings.pdf_dir, "biblia.pdf")
|
||||
extractor = FakeExtractor({"biblia.pdf": [PageText(page=6, text=GENESIS_PAGE)]})
|
||||
indexer, _ = build_indexer(settings, extractor)
|
||||
|
||||
indexer.reindex()
|
||||
|
||||
manifest = json.loads(settings.manifest_path.read_text(encoding="utf-8"))
|
||||
assert "biblia.pdf" in manifest
|
||||
assert manifest["biblia.pdf"]["state"] == "complete"
|
||||
assert manifest["biblia.pdf"]["chunks"] == 2
|
||||
assert manifest["biblia.pdf"]["sha256"]
|
||||
|
||||
|
||||
def test_partial_file_is_recorded_as_partial(settings):
|
||||
write_pdf(settings.pdf_dir, "biblia.pdf")
|
||||
extractor = FakeExtractor(
|
||||
{
|
||||
"biblia.pdf": [
|
||||
PageText(page=6, text=GENESIS_PAGE),
|
||||
PageText(page=7, text="", failed=True),
|
||||
]
|
||||
}
|
||||
)
|
||||
indexer, store = build_indexer(settings, extractor)
|
||||
|
||||
status = indexer.reindex()
|
||||
|
||||
manifest = json.loads(settings.manifest_path.read_text(encoding="utf-8"))
|
||||
assert manifest["biblia.pdf"]["state"] == "partial"
|
||||
assert manifest["biblia.pdf"]["failed_pages"] == [7]
|
||||
assert status.errors
|
||||
|
||||
|
||||
def test_failed_file_is_recorded_as_error(settings):
|
||||
write_pdf(settings.pdf_dir, "biblia.pdf")
|
||||
extractor = FakeExtractor({"biblia.pdf": [PageText(page=6, text="", failed=True)]})
|
||||
indexer, store = build_indexer(settings, extractor)
|
||||
|
||||
indexer.reindex()
|
||||
|
||||
manifest = json.loads(settings.manifest_path.read_text(encoding="utf-8"))
|
||||
assert manifest["biblia.pdf"]["state"] == "error"
|
||||
assert store.count() == 0
|
||||
|
||||
|
||||
def test_reindex_replaces_old_chunks_of_the_same_file(settings):
|
||||
write_pdf(settings.pdf_dir, "biblia.pdf")
|
||||
extractor = FakeExtractor({"biblia.pdf": [PageText(page=6, text=GENESIS_PAGE)]})
|
||||
indexer, store = build_indexer(settings, extractor)
|
||||
|
||||
indexer.reindex()
|
||||
indexer.reindex(force=True)
|
||||
|
||||
assert store.count() == 2, "não deve haver duplicação após reindexação"
|
||||
|
||||
|
||||
def test_status_reports_chunks_from_the_store(settings):
|
||||
write_pdf(settings.pdf_dir, "biblia.pdf")
|
||||
extractor = FakeExtractor({"biblia.pdf": [PageText(page=6, text=GENESIS_PAGE)]})
|
||||
indexer, store = build_indexer(settings, extractor)
|
||||
|
||||
indexer.reindex()
|
||||
status = indexer.status()
|
||||
|
||||
assert isinstance(status, IndexStatus)
|
||||
assert status.chunks == store.count()
|
||||
|
||||
|
||||
def test_empty_pdf_directory_is_ready_with_zero_chunks(settings):
|
||||
settings.pdf_dir.mkdir(parents=True, exist_ok=True)
|
||||
extractor = FakeExtractor({})
|
||||
indexer, _ = build_indexer(settings, extractor)
|
||||
|
||||
status = indexer.reindex()
|
||||
|
||||
assert status.state == "ready"
|
||||
assert status.pdfs_total == 0
|
||||
assert status.chunks == 0
|
||||
|
||||
|
||||
def test_non_pdf_files_are_ignored(settings):
|
||||
settings.pdf_dir.mkdir(parents=True, exist_ok=True)
|
||||
(settings.pdf_dir / "notas.txt").write_text("ignorado", encoding="utf-8")
|
||||
extractor = FakeExtractor({})
|
||||
indexer, _ = build_indexer(settings, extractor)
|
||||
|
||||
status = indexer.reindex()
|
||||
|
||||
assert status.pdfs_total == 0
|
||||
assert extractor.calls == []
|
||||
Reference in new issue
Block a user