ingest.py

#!/usr/bin/env python3
"""
Pipeline de ingestão: lê PDFs de docs/ e indexa no ChromaDB.

Uso:
    python ingest.py           # indexa (pula se já indexado)
    python ingest.py --rebuild # apaga e re-indexa tudo
"""
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

from src.labdaps.config import DOCS_DIR
from src.labdaps.ingestion.pdf_extractor import extract_all_pdfs
from src.labdaps.ingestion.chunker import chunk_pages
from src.labdaps.ingestion.embedder import Embedder
from src.labdaps.retrieval.vector_store import build_index, collection_count


def main():
    rebuild = "--rebuild" in sys.argv

    print("=" * 60)
    print("  LABDAPS RAG -- Ingestão do Livro Sinal de Alerta")
    print("=" * 60)

    if not rebuild:
        count = collection_count()
        if count > 0:
            print(f"[INFO] Coleção já contém {count} chunks. Pulando ingestão.")
            print("[INFO] Use --rebuild para forçar re-indexação.")
            return

    if not DOCS_DIR.exists():
        print(f"[ERROR] Pasta de documentos não encontrada: {DOCS_DIR}")
        sys.exit(1)

    pdf_files = list(DOCS_DIR.glob("*.pdf"))
    print(f"[INFO] Encontrados {len(pdf_files)} PDF(s) em {DOCS_DIR}")

    print("\n[1/3] Extraindo texto dos PDFs...")
    t0 = time.time()
    pages = extract_all_pdfs(DOCS_DIR)
    print(f"[INFO] {len(pages)} páginas extraídas em {time.time()-t0:.1f}s")

    print("\n[2/3] Dividindo em chunks...")
    t0 = time.time()
    chunks = chunk_pages(pages)
    print(f"[INFO] {len(chunks)} chunks gerados em {time.time()-t0:.1f}s")

    print("\n[3/3] Gerando embeddings e indexando no ChromaDB...")
    embedder = Embedder()
    t0 = time.time()
    build_index(chunks, embedder, rebuild=rebuild)
    print(f"[INFO] Indexação concluída em {time.time()-t0:.1f}s")

    print(f"\n[OK] Ingestão concluída! Total de chunks: {collection_count()}")


if __name__ == "__main__":
    main()