backend/scripts/ingest_data.py

"""
Data ingestion script — loads jamescalam/ai-arxiv2-chunks into PostgreSQL.

Usage:
    python -m scripts.ingest_data [--limit N] [--batch-size 256] [--role researcher]

The script:
  1. Streams the HuggingFace dataset (no full download required for large corpora).
  2. Upserts documents (deduplication by arxiv_id).
  3. Embeds chunk content locally with sentence-transformers.
  4. Bulk-inserts chunks with their 384-dim vectors via COPY (fast).
  5. Grants access to the specified role for all ingested documents.

NOTE: 241k chunks × 384 floats ≈ 370 MB of vectors.  Embedding on CPU takes
      ~2-4 hours for the full dataset; use --limit for quick smoke tests.
"""
import argparse
import logging
import time
import uuid
from typing import Iterator

import numpy as np
import psycopg2
import psycopg2.extras
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from tqdm import tqdm

from app.config import get_settings

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

DATASET_NAME = "jamescalam/ai-arxiv2-chunks"
EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
EMBED_BATCH = 64        # chunks per embedding batch
DB_BATCH = 256          # rows per DB insert


# ── Dataset helpers ───────────────────────────────────────────────────────────

def _iter_dataset(limit: int | None) -> Iterator[dict]:
    """Stream dataset rows; respect optional limit."""
    ds = load_dataset(DATASET_NAME, split="train", streaming=True)
    for i, row in enumerate(ds):
        if limit and i >= limit:
            break
        yield row


def _arxiv_url(arxiv_id: str) -> str:
    return f"https://arxiv.org/abs/{arxiv_id}"


def _full_citation(row: dict) -> str:
    title = row.get("title", "")
    arxiv_id = row.get("doi", row.get("arxiv_id", ""))
    return f"{title}. arXiv:{arxiv_id}" if arxiv_id else title


# ── Database helpers ──────────────────────────────────────────────────────────

def _get_or_create_document(cur, row: dict) -> str:
    """Upsert a document row and return its UUID."""
    arxiv_id = row.get("doi", row.get("arxiv_id", str(uuid.uuid4())))
    cur.execute(
        """
        INSERT INTO documents (arxiv_id, title, source_url, full_citation, metadata)
        VALUES (%s, %s, %s, %s, %s)
        ON CONFLICT (arxiv_id) DO UPDATE SET title = EXCLUDED.title
        RETURNING id
        """,
        (
            arxiv_id,
            row.get("title"),
            _arxiv_url(arxiv_id),
            _full_citation(row),
            psycopg2.extras.Json({"source": row.get("source", "")}),
        ),
    )
    return cur.fetchone()[0]


def _get_role_id(cur, role_name: str) -> str:
    cur.execute("SELECT id FROM roles WHERE name = %s", (role_name,))
    result = cur.fetchone()
    if result is None:
        raise ValueError(f"Role '{role_name}' not found. Run init_db.py first.")
    return result[0]


def _grant_access(cur, role_id: str, doc_ids: list[str]) -> None:
    psycopg2.extras.execute_values(
        cur,
        "INSERT INTO role_document_access (role_id, document_id) VALUES %s ON CONFLICT DO NOTHING",
        [(role_id, doc_id) for doc_id in doc_ids],
    )


def _get_existing_chunk_indices(cur, doc_id: str) -> set:
    """Return the set of chunk_index values already stored for this document."""
    cur.execute("SELECT chunk_index FROM chunks WHERE document_id = %s", (doc_id,))
    return {row[0] for row in cur.fetchall()}


def _bulk_insert_chunks(cur, rows: list[tuple]) -> None:
    """rows: (doc_id, content, embedding_str, chunk_index, metadata_json)"""
    psycopg2.extras.execute_values(
        cur,
        """
        INSERT INTO chunks (document_id, content, embedding, chunk_index, metadata)
        VALUES %s
        ON CONFLICT (document_id, chunk_index) DO NOTHING
        """,
        rows,
        template="(%s, %s, %s::vector, %s, %s)",
    )


# ── Main ingestion loop ────────────────────────────────────────────────────────

def ingest(limit: int | None, batch_size: int, role_name: str) -> None:
    settings = get_settings()
    embedder = SentenceTransformer(EMBED_MODEL)

    conn = psycopg2.connect(settings.database_sync_url)
    conn.autocommit = False

    with conn.cursor() as cur:
        role_id = _get_role_id(cur, role_name)

    logger.info(
        "Starting ingestion — dataset=%s  limit=%s  role=%s",
        DATASET_NAME,
        limit or "ALL",
        role_name,
    )

    pending_chunks: list[tuple] = []
    doc_ids_seen: set[str] = set()
    doc_chunk_cache: dict[str, set] = {}  # doc_id -> set of already-stored chunk_index
    total_chunks = 0
    skipped_chunks = 0
    failed_chunks = 0
    texts_buf: list[str] = []
    meta_buf: list[dict] = []

    def _flush(force: bool = False):
        nonlocal total_chunks, failed_chunks
        if not texts_buf:
            return
        if not force and len(texts_buf) < EMBED_BATCH:
            return

        embeddings = embedder.encode(texts_buf, batch_size=EMBED_BATCH, show_progress_bar=False)
        for emb, meta in zip(embeddings, meta_buf):
            pending_chunks.append((
                meta["doc_id"],
                meta["content"],
                str(emb.tolist()),
                meta["chunk_index"],
                psycopg2.extras.Json(meta.get("metadata", {})),
            ))

        texts_buf.clear()
        meta_buf.clear()

        if force or len(pending_chunks) >= batch_size:
            try:
                with conn.cursor() as cur:
                    _bulk_insert_chunks(cur, pending_chunks)
                    if doc_ids_seen:
                        _grant_access(cur, role_id, list(doc_ids_seen))
                conn.commit()
                total_chunks += len(pending_chunks)
                logger.info("Inserted %d chunks (total: %d)", len(pending_chunks), total_chunks)
            except Exception as batch_exc:
                conn.rollback()
                logger.warning("Batch insert failed (%s) — retrying row-by-row", batch_exc)
                inserted = 0
                for chunk_row in pending_chunks:
                    try:
                        with conn.cursor() as cur:
                            _bulk_insert_chunks(cur, [chunk_row])
                        conn.commit()
                        inserted += 1
                    except Exception as row_exc:
                        conn.rollback()
                        logger.error(
                            "Skipping chunk (doc=%s idx=%s): %s",
                            chunk_row[0], chunk_row[3], row_exc,
                        )
                        failed_chunks += 1
                total_chunks += inserted
                logger.info("Row-by-row retry: %d inserted, %d failed", inserted, failed_chunks)
            finally:
                pending_chunks.clear()
                doc_ids_seen.clear()

    try:
        for row in tqdm(_iter_dataset(limit), desc="Ingesting", unit="chunk"):
            with conn.cursor() as cur:
                doc_id = _get_or_create_document(cur, row)
                if doc_id not in doc_chunk_cache:
                    doc_chunk_cache[doc_id] = _get_existing_chunk_indices(cur, doc_id)
            conn.commit()
            doc_ids_seen.add(doc_id)

            chunk_index = row.get("chunk-id", row.get("chunk_index", 0))
            content = row.get("chunk", row.get("text", "")).replace("\x00", "")
            if not content.strip():
                continue

            if chunk_index in doc_chunk_cache[doc_id]:
                skipped_chunks += 1
                continue

            texts_buf.append(content)
            meta_buf.append({
                "doc_id": doc_id,
                "content": content,
                "chunk_index": chunk_index,
                "metadata": {"source": row.get("source", "")},
            })
            _flush()

        _flush(force=True)  # flush remainder

    finally:
        conn.close()

    logger.info(
        "Ingestion complete. Inserted: %d  Skipped (already exist): %d  Failed: %d",
        total_chunks,
        skipped_chunks,
        failed_chunks,
    )


# ── CLI ───────────────────────────────────────────────────────────────────────

def _parse_args():
    parser = argparse.ArgumentParser(description="Ingest ArXiv chunks into pgvector")
    parser.add_argument("--limit", type=int, default=None, help="Max chunks to ingest")
    parser.add_argument("--batch-size", type=int, default=DB_BATCH, help="DB insert batch size")
    parser.add_argument("--role", default="researcher", help="Role to grant document access")
    return parser.parse_args()


if __name__ == "__main__":
    args = _parse_args()
    ingest(limit=args.limit, batch_size=args.batch_size, role_name=args.role)