scripts/ingest.py

"""
ByteAstra — Textbook Ingestion Script.

Reads source documents (plain text or markdown files), chunks them, and
indexes the chunks into ChromaDB under the specified domain's collection.

Usage:
  python scripts/ingest.py --domain ayurveda --source-dir ./data/ayurveda

Arguments:
  --domain       Domain key (must match a YAML in app/domains/)
  --source-dir   Directory containing .txt or .md source files
  --chunk-size   Token-approximate chunk size in characters (default: 800)
  --overlap      Character overlap between consecutive chunks (default: 100)
  --batch-size   Upsert batch size for ChromaDB (default: 64)
  --dry-run      Parse and chunk without writing to ChromaDB
"""
import argparse
import hashlib
import logging
import sys
from pathlib import Path

# Allow running from the backend/ directory
sys.path.insert(0, str(Path(__file__).parent.parent))

from app.config import get_settings
from app.services.rag import index_chunks, get_or_create_collection

logging.basicConfig(level=logging.INFO, format="%(levelname)s | %(message)s")
logger = logging.getLogger(__name__)


# ── Chunking ───────────────────────────────────────────────────────────────────

def chunk_text(
    text: str,
    chunk_size: int = 800,
    overlap: int = 100,
) -> list[str]:
    """
    Simple character-window chunker.
    Splits on paragraph boundaries where possible to preserve semantic units.
    """
    paragraphs = text.split("\n\n")
    chunks: list[str] = []
    current = ""

    for para in paragraphs:
        para = para.strip()
        if not para:
            continue
        if len(current) + len(para) + 2 <= chunk_size:
            current += ("\n\n" if current else "") + para
        else:
            if current:
                chunks.append(current)
            # If paragraph itself exceeds chunk_size, hard-split it
            if len(para) > chunk_size:
                for i in range(0, len(para), chunk_size - overlap):
                    chunks.append(para[i : i + chunk_size])
            else:
                current = para

    if current:
        chunks.append(current)

    return chunks


def parse_frontmatter(text: str) -> tuple[dict, str]:
    """
    Parse optional YAML-style frontmatter from a markdown file.
    Returns (metadata_dict, body_text).

    Frontmatter format (optional):
    ---
    source: Charaka Samhita
    chapter: Sutrasthana
    section: 1.1
    ---
    """
    meta: dict = {}
    if text.startswith("---"):
        lines = text.splitlines()
        end_idx = None
        for i, line in enumerate(lines[1:], 1):
            if line.strip() == "---":
                end_idx = i
                break
        if end_idx:
            import yaml
            try:
                meta = yaml.safe_load("\n".join(lines[1:end_idx])) or {}
            except Exception:
                pass
            text = "\n".join(lines[end_idx + 1:]).strip()
    return meta, text


# ── Main ───────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Ingest textbooks into ByteAstra ChromaDB")
    parser.add_argument("--domain", required=True, help="Domain key, e.g. 'ayurveda'")
    parser.add_argument("--source-dir", required=True, type=Path, help="Directory of .txt/.md files")
    parser.add_argument("--chunk-size", type=int, default=800)
    parser.add_argument("--overlap", type=int, default=100)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--dry-run", action="store_true", help="Parse without writing to DB")
    args = parser.parse_args()

    source_dir: Path = args.source_dir
    if not source_dir.is_dir():
        logger.error("Source directory not found: %s", source_dir)
        sys.exit(1)

    files = list(source_dir.glob("**/*.txt")) + list(source_dir.glob("**/*.md"))
    if not files:
        logger.warning("No .txt or .md files found in %s", source_dir)
        sys.exit(0)

    collection_name = f"domain_{args.domain}"
    logger.info("Ingesting into collection: %s", collection_name)
    logger.info("Found %d source files", len(files))

    all_chunks: list[dict] = []

    for file_path in sorted(files):
        logger.info("Processing: %s", file_path.name)
        raw = file_path.read_text(encoding="utf-8", errors="ignore")
        meta, body = parse_frontmatter(raw)

        source = meta.get("source", file_path.stem)
        chapter = meta.get("chapter", "")
        section = meta.get("section", "")

        text_chunks = chunk_text(body, chunk_size=args.chunk_size, overlap=args.overlap)
        logger.info("  → %d chunks", len(text_chunks))

        for i, chunk_text_str in enumerate(text_chunks):
            # Deterministic ID based on content hash — safe to re-run (upsert)
            chunk_hash = hashlib.md5(chunk_text_str.encode()).hexdigest()[:12]
            chunk_id = f"{args.domain}_{file_path.stem}_{i:04d}_{chunk_hash}"

            all_chunks.append({
                "id": chunk_id,
                "text": chunk_text_str,
                "source": source,
                "chapter": str(chapter),
                "section": str(section),
            })

    logger.info("Total chunks prepared: %d", len(all_chunks))

    if args.dry_run:
        logger.info("DRY RUN — no data written to ChromaDB.")
        for c in all_chunks[:3]:
            print(f"\n[{c['id']}]\n{c['text'][:200]}...")
        return

    # Batch upsert
    for i in range(0, len(all_chunks), args.batch_size):
        batch = all_chunks[i : i + args.batch_size]
        index_chunks(collection_name, batch)
        logger.info("Indexed batch %d–%d", i + 1, i + len(batch))

    logger.info("✓ Ingestion complete. %d chunks in '%s'", len(all_chunks), collection_name)


if __name__ == "__main__":
    main()