scripts/build_db.py

"""
build_db.py — Pre-build the ChromaDB vector store from all 5 bundled company docs.

Run this ONCE locally before pushing to HuggingFace Spaces:

    cd hf_space
    python scripts/build_db.py

Then commit the resulting db/chroma/ directory and push:

    git add db/chroma/
    git commit -m "chore: add pre-built ChromaDB"
    git push origin main

HF Space will now start instantly — no re-ingestion on every restart.
"""
from __future__ import annotations

import logging
import sys
from pathlib import Path

# Ensure project root is on sys.path
_ROOT = Path(__file__).resolve().parents[1]
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)-8s  %(name)s — %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("build_db")

BUNDLED_DOCS: list[tuple[Path, str]] = [
    (_ROOT / "data" / "notion"  / "notion_support_docs.txt",  "Notion"),
    (_ROOT / "data" / "slack"   / "slack_support_docs.txt",   "Slack"),
    (_ROOT / "data" / "github"  / "github_support_docs.txt",  "GitHub"),
    (_ROOT / "data" / "zoom"    / "zoom_support_docs.txt",    "Zoom"),
    (_ROOT / "data" / "shopify" / "shopify_support_docs.txt", "Shopify"),
]


def main() -> None:
    logger.info("=" * 60)
    logger.info("ProdAssist RAG — Pre-building ChromaDB")
    logger.info("=" * 60)

    # Validate all source files exist before doing anything heavy
    missing = [(p, name) for p, name in BUNDLED_DOCS if not p.exists()]
    if missing:
        for p, name in missing:
            logger.error("Missing: %s  →  %s", name, p)
        sys.exit(1)

    logger.info("All %d source files found.", len(BUNDLED_DOCS))

    # Build the pipeline (loads embedding model — ~30 s on first run)
    logger.info("Loading pipeline (downloading embedding model if needed)…")
    from app.pipeline.rag_pipeline import build_pipeline
    pipeline = build_pipeline()

    # If the DB already has data, wipe it so we get a clean, reproducible build
    if pipeline.document_count > 0:
        logger.warning(
            "Existing DB has %d chunks — wiping and rebuilding from scratch.",
            pipeline.document_count,
        )
        pipeline.reset_database()

    # Ingest each company's docs
    total_chunks = 0
    for path, name in BUNDLED_DOCS:
        logger.info("Ingesting %s (%s)…", name, path.name)
        stored = pipeline.ingest([path])
        total_chunks += stored
        logger.info("  → %d chunks stored (running total: %d)", stored, total_chunks)

    logger.info("=" * 60)
    logger.info("Build complete!  Total chunks in DB: %d", pipeline.document_count)
    logger.info("=" * 60)

    db_path = _ROOT / "db" / "chroma"
    logger.info("ChromaDB written to: %s", db_path)
    logger.info("")
    logger.info("Next steps — commit and push the pre-built DB:")
    logger.info("  cd %s", _ROOT)
    logger.info("  git add db/chroma/")
    logger.info('  git commit -m "chore: add pre-built ChromaDB"')
    logger.info("  git push origin main")


if __name__ == "__main__":
    main()