codex-py/scripts/ra_grobid_backfill.py

"""R-A — GROBID reference backfill (references-only, idempotent).

Adds GROBID-parsed citations to the live corpus **without** re-chunking the body,
so the DQ-3-verified-clean ``.txt`` chunks are preserved. Complements the DQ-1
Semantic-Scholar supplement with references the APIs lack (notably the two
TU-Berlin ``depositonce`` theses, which neither OpenAlex nor S2 indexes).

Why references-only instead of ``ingest_paper(source_path=pdf)``
----------------------------------------------------------------
The PDF branch of :func:`codex.ingest.ingest_paper` runs Nougat OCR and then
``DELETE FROM chunks`` + re-INSERT — it would replace the clean ``.txt`` chunks
(DQ-3) with unmeasured PDF-OCR text and additionally requires the Nougat server.
This backfill touches only the ``citations`` table (``ON CONFLICT DO NOTHING``),
so it is safe, idempotent, and needs only a reachable GROBID server.

Prerequisite
------------
A reachable GROBID server. The corpus' ``GROBID_URL`` points at the Jetson
(``http://jetson.local:8070``). GROBID now runs natively on that aarch64 host
via the multi-arch CRF-only image (``grobid/grobid:0.9.0-crf``) — see
``docs/infra/grobid-jetson.md`` for the deploy + docker-permission fix and the
``GET /api/isalive`` check. (The old ``grobid/grobid:0.8.2`` pin was amd64-only,
which is why GROBID never came up; pass ``--grobid-url`` to override the default.)
Use the SSH DB tunnel for ``DATABASE_URL`` as documented in
``docs/audit/DATA-QUALITY-2026-06-15.md``.

Usage
-----
    # dry run (no writes) over the two zero-edge theses
    PYTHONPATH=. python scripts/ra_grobid_backfill.py --zero-edge-only

    # write every paper's GROBID-parsed refs to the live DB
    PYTHONPATH=. python scripts/ra_grobid_backfill.py --write

Default is **dry-run**; pass ``--write`` to mutate the live DB.
"""

from __future__ import annotations

import argparse
import logging
from pathlib import Path

import psycopg
import psycopg.rows

from codex.config import get_settings
from codex.db import get_conn
from codex.models import Citation
from codex.parsing.grobid import extract_references

logger = logging.getLogger(__name__)

DEFAULT_PDF_DIR = "/Users/tarikmoussa/Desktop/ConformalLabpp/papers"


def _normalize_cited_id(doi: str = "", arxiv_id: str = "") -> str | None:
    """Normalize a GROBID-extracted reference id to the corpus' canonical form.

    DOIs → bare, lower-cased (strip ``https://doi.org/`` / ``http://`` / ``doi:``),
    matching ``papers.id`` and ``codex.sources.openalex._normalize_doi`` (DQ-5).
    arXiv ids → strip a leading ``arXiv:`` prefix, leaving the bare id
    (``2305.10988``, ``math/0603097``). DOI takes precedence when both exist.
    Kept self-contained rather than reusing ``ingest._norm_cited_id`` because it
    additionally strips the ``arXiv:`` prefix (which GROBID can emit but that
    helper leaves untouched), and to avoid importing a private cross-module name.
    """
    doi = (doi or "").strip()
    if doi:
        s = doi.lower()
        for prefix in ("https://doi.org/", "http://doi.org/", "doi:"):
            if s.startswith(prefix):
                return s[len(prefix) :]
        return s
    arx = (arxiv_id or "").strip()
    if arx:
        return arx[len("arxiv:") :].strip() if arx.lower().startswith("arxiv:") else arx
    return None


def build_grobid_citations(
    paper_id: str,
    pdf_path: str | Path,
    grobid_url: str | None = None,
    timeout: float = 60.0,
) -> list[Citation]:
    """Extract a paper's references via GROBID and return normalized Citations.

    ``citing_id`` is the canonical ``paper_id``; ``cited_id`` is normalized.
    Self-citations and within-paper duplicates are dropped. Refs with neither a
    DOI nor an arXiv id are skipped (they cannot join the graph).
    """
    refs = extract_references(str(pdf_path), grobid_url=grobid_url, timeout=timeout)
    seen: set[str] = set()
    citations: list[Citation] = []
    for ref in refs:
        cited = _normalize_cited_id(ref.get("doi", ""), ref.get("arxiv_id", ""))
        if not cited or cited == paper_id or cited in seen:
            continue
        seen.add(cited)
        citations.append(Citation(citing_id=paper_id, cited_id=cited))
    return citations


def _coverage(conn: psycopg.Connection[psycopg.rows.DictRow]) -> tuple[int, int, int]:
    """Return (n_papers, n_papers_with_out_edges, n_citations) from the live DB."""
    row = conn.execute(
        """
        SELECT (SELECT count(*) FROM papers)                    AS papers,
               (SELECT count(DISTINCT citing_id) FROM citations) AS citing,
               (SELECT count(*) FROM citations)                  AS edges
        """
    ).fetchone()
    assert row is not None
    return int(row["papers"]), int(row["citing"]), int(row["edges"])


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--pdf-dir", default=DEFAULT_PDF_DIR, help="directory of source PDFs")
    parser.add_argument("--grobid-url", default=None, help="override Settings().grobid_url")
    parser.add_argument("--only", default=None, help="restrict to a single paper id")
    parser.add_argument(
        "--zero-edge-only",
        action="store_true",
        help="only papers with 0 out-edges (the R-A primary targets)",
    )
    parser.add_argument("--limit", type=int, default=None, help="cap number of papers")
    parser.add_argument(
        "--timeout",
        type=float,
        default=300.0,
        help="GROBID HTTP timeout (s); theses can need >60s, esp. emulated",
    )
    parser.add_argument(
        "--write", action="store_true", help="WRITE to the live DB (default: dry-run)"
    )
    args = parser.parse_args(argv)

    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
    grobid_url = args.grobid_url or get_settings().grobid_url
    pdf_dir = Path(args.pdf_dir)
    pdfs = {p.stem: p for p in pdf_dir.glob("*.pdf")}

    # Read the target list in a short-lived connection, then release it: the GROBID
    # extraction below is a slow per-paper network loop and must NOT hold an
    # idle-in-transaction connection open across it (a mid-loop SSH-tunnel drop would
    # otherwise abort the whole batch — see docs/infra/jetson-ssh-stability.md).
    with get_conn() as conn:
        rows = conn.execute(
            """
            SELECT p.id, p.source_path,
                   (SELECT count(*) FROM citations ci WHERE ci.citing_id = p.id) AS out_edges
            FROM papers p
            ORDER BY out_edges, p.id
            """
        ).fetchall()

    targets: list[tuple[str, Path]] = []
    for r in rows:
        if args.only and r["id"] != args.only:
            continue
        if args.zero_edge_only and r["out_edges"] != 0:
            continue
        stem = Path(r["source_path"]).stem if r["source_path"] else None
        pdf = pdfs.get(stem) if stem else None
        if pdf is None:
            logger.warning("no PDF for %s (stem=%s) — skipping", r["id"], stem)
            continue
        targets.append((r["id"], pdf))
    if args.limit is not None:
        targets = targets[: args.limit]

    logger.info(
        "GROBID=%s | papers to process=%d | mode=%s",
        grobid_url,
        len(targets),
        "WRITE" if args.write else "DRY-RUN",
    )

    # GROBID extraction — no DB connection held during this slow network loop.
    all_citations: list[Citation] = []
    for paper_id, pdf in targets:
        try:
            cits = build_grobid_citations(
                paper_id, pdf, grobid_url=grobid_url, timeout=args.timeout
            )
        except Exception as exc:
            logger.warning("GROBID extraction failed for %s (%s): %s", paper_id, pdf.name, exc)
            continue
        logger.info("%-38s %3d refs  (%s)", paper_id, len(cits), pdf.name)
        all_citations.extend(cits)

    logger.info("total candidate citations: %d", len(all_citations))

    if not args.write:
        logger.info("DRY-RUN — no rows written. Re-run with --write to apply.")
        return 0

    # Fresh connection for the write — the read connection was already released.
    with get_conn() as conn:
        before = _coverage(conn)
        with conn.cursor() as cur:
            cur.executemany(
                """
                INSERT INTO citations (citing_id, cited_id, context)
                VALUES (%s, %s, %s)
                ON CONFLICT (citing_id, cited_id) DO NOTHING
                """,
                [(c.citing_id, c.cited_id, c.context) for c in all_citations],
            )
        conn.commit()
        after = _coverage(conn)
    logger.info(
        "coverage papers=%d  citing %d->%d  edges %d->%d (+%d)",
        after[0],
        before[1],
        after[1],
        before[2],
        after[2],
        after[2] - before[2],
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())