omp
621679cee5
- C-1: corrige índice de páginas OCR'd (retorna page_number em vez de position) - C-2: adiciona tratamento para falhas de OCR no cache (não distingue erro de não processado) - Integra paralelização do OCR via subprocessos (parallel.py) - Indexador usa extração paralela quando disponível - Melhora relatório de progresso com contagem de páginas processadas
337 lines
11 KiB
Python
337 lines
11 KiB
Python
"""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)
|
|
pages_processed: int = 0
|
|
pages_total: int = 0
|
|
current_file: str | None = None
|
|
|
|
|
|
@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._indexing = False
|
|
self._pages_processed = 0
|
|
self._pages_total = 0
|
|
self._current_file: str | None = None
|
|
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 mark_indexing(self) -> None:
|
|
"""Sinaliza que uma indexação está em andamento."""
|
|
self._indexing = True
|
|
self._pages_processed = 0
|
|
self._pages_total = 0
|
|
self._current_file = None
|
|
|
|
def report_progress(
|
|
self, pdf_name: str, pages_done: int, pages_total: int
|
|
) -> None:
|
|
"""Registra o progresso da extração de um arquivo."""
|
|
self._current_file = pdf_name
|
|
self._pages_processed = pages_done
|
|
self._pages_total = pages_total
|
|
|
|
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 = []
|
|
self.mark_indexing()
|
|
pdfs = self.discover_pdfs()
|
|
|
|
try:
|
|
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()
|
|
finally:
|
|
self._indexing = False
|
|
|
|
return self.status()
|
|
|
|
def _extract_pages(self, pdf: Path, file_hash: str) -> list:
|
|
"""Extrai as páginas do PDF, usando OCR paralelo quando possível."""
|
|
extract_parallel = getattr(self.extractor, "extract_pages_parallel", None)
|
|
if extract_parallel is None:
|
|
return list(self.extractor.extract_document(pdf, file_hash=file_hash))
|
|
|
|
page_count = _pdf_page_count(pdf)
|
|
pages: list = []
|
|
batch_size = max(self.settings.ocr_workers, 1) * 4
|
|
|
|
for batch in self._page_batches(pdf, batch_size):
|
|
raw_pages = [
|
|
(page_number, text)
|
|
for page_number, text in batch
|
|
]
|
|
extracted = extract_parallel(
|
|
raw_pages,
|
|
file_hash=file_hash,
|
|
path=pdf,
|
|
workers=self.settings.ocr_workers,
|
|
)
|
|
pages.extend(extracted)
|
|
self.report_progress(pdf.name, len(pages), page_count)
|
|
|
|
return sorted(pages, key=lambda page: page.page)
|
|
|
|
def _page_batches(self, pdf: Path, batch_size: int):
|
|
"""Gera lotes de (número_da_página, texto_da_camada)."""
|
|
import pymupdf
|
|
|
|
with pymupdf.open(pdf) as document:
|
|
batch: list[tuple[int, str]] = []
|
|
for index in range(len(document)):
|
|
page = document[index]
|
|
try:
|
|
text = page.get_text("text") or ""
|
|
except Exception: # pragma: no cover - PDF malformado
|
|
text = ""
|
|
batch.append((index + 1, text))
|
|
if len(batch) >= batch_size:
|
|
yield batch
|
|
batch = []
|
|
if batch:
|
|
yield batch
|
|
|
|
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 = self._extract_pages(pdf, 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)
|
|
)
|
|
if self._indexing:
|
|
state = "indexing"
|
|
elif self._errors and indexed == 0:
|
|
state = "error"
|
|
else:
|
|
state = "ready"
|
|
return IndexStatus(
|
|
state=state,
|
|
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),
|
|
pages_processed=self._pages_processed,
|
|
pages_total=self._pages_total,
|
|
current_file=self._current_file,
|
|
)
|
|
|
|
|
|
def _pdf_page_count(pdf: Path) -> int:
|
|
"""Número de páginas do PDF, ou 0 se não for possível determinar."""
|
|
try:
|
|
import pymupdf
|
|
|
|
with pymupdf.open(pdf) as document:
|
|
return len(document)
|
|
except Exception: # pragma: no cover - PDF malformado
|
|
return 0
|
|
|
|
|
|
def _now() -> str:
|
|
return datetime.now(timezone.utc).isoformat()
|