codex-py/codex/ingest.py

"""End-to-end idempotent ingest pipeline for a single paper."""

from __future__ import annotations

import contextlib
import logging
from dataclasses import dataclass
from pathlib import Path

import numpy as np

from codex.db import get_conn
from codex.embed import get_embedder
from codex.models import Citation, Paper
from codex.sources import openalex, semanticscholar

logger = logging.getLogger(__name__)


@dataclass
class IngestResult:
    paper_id: str
    chunks_upserted: int
    citations_upserted: int


def ingest_paper(
    paper_id: str,
    source_path: str | None = None,
) -> IngestResult:
    """Idempotent ingest of one paper.

    Parameters
    ----------
    paper_id:
        An arXiv ID (``"2301.07041"``), DOI (``"10.1145/…"``), or
        OpenAlex W-ID (``"W2741809807"``).
    source_path:
        Optional path to a local ``.tex`` or ``.pdf`` file.
        If given, the file is parsed into text chunks and stored.
        Supports:
        - ``.tex``  → :func:`codex.parsing.tex.latex_to_text` + chunk
        - ``.pdf``  → :func:`codex.parsing.nougat.pdf_to_markdown` + chunk
                      + :func:`codex.parsing.grobid.extract_references` for refs

    Returns
    -------
    IngestResult
        Counts of upserted chunks and citations.
    """
    # ---------------------------------------------------------------
    # 1. Fetch metadata (OpenAlex primary, SemanticScholar fallback)
    # ---------------------------------------------------------------
    paper: Paper | None = openalex.fetch_paper(paper_id)

    if paper is None:
        # Detect arXiv IDs: numeric pattern like "2301.07041" or "arxiv:..." prefix
        pid_lower = paper_id.lower()
        looks_like_arxiv = pid_lower.startswith("arxiv:") or (
            len(paper_id) > 4 and paper_id[4:5] == "." and paper_id[:4].isdigit()
        )
        if looks_like_arxiv:
            arxiv_key = paper_id if pid_lower.startswith("arxiv:") else f"arXiv:{paper_id}"
            with contextlib.suppress(Exception):
                semanticscholar.fetch_references(arxiv_key)
            paper = Paper(id=paper_id, title="")
        else:
            raise ValueError(f"Paper not found: {paper_id}")

    if paper is None:
        raise ValueError(f"Paper not found: {paper_id}")

    # ---------------------------------------------------------------
    # 2. Embed abstract (dense only — schema has one vector column)
    # ---------------------------------------------------------------
    embedder = get_embedder()
    abstract_text = paper.abstract or ""
    if abstract_text:
        abstract_emb_arr: np.ndarray = embedder.encode_dense([abstract_text])
        abstract_emb: list[float] = abstract_emb_arr[0].tolist()
    else:
        dim = embedder.dim
        abstract_emb = [0.0] * dim

    # ---------------------------------------------------------------
    # 3. Upsert paper  (INSERT … ON CONFLICT (id) DO UPDATE SET …)
    # ---------------------------------------------------------------
    with get_conn() as conn:
        conn.execute(
            """
            INSERT INTO papers (id, openalex_id, bibkey, title, authors, year, abstract,
                                source_path, abstract_emb)
            VALUES (%(id)s, %(openalex_id)s, %(bibkey)s, %(title)s, %(authors)s, %(year)s,
                    %(abstract)s, %(source_path)s, %(abstract_emb)s)
            ON CONFLICT (id) DO UPDATE SET
                openalex_id  = EXCLUDED.openalex_id,
                title        = EXCLUDED.title,
                authors      = EXCLUDED.authors,
                year         = EXCLUDED.year,
                abstract     = EXCLUDED.abstract,
                source_path  = COALESCE(EXCLUDED.source_path, papers.source_path),
                abstract_emb = EXCLUDED.abstract_emb
            """,
            {
                "id": paper.id,
                "openalex_id": paper.openalex_id,
                "bibkey": paper.bibkey,
                "title": paper.title,
                "authors": paper.authors,
                "year": paper.year,
                "abstract": paper.abstract,
                "source_path": source_path,
                "abstract_emb": abstract_emb,
            },
        )

        # ---------------------------------------------------------------
        # 4. Parse + chunk + embed source file (if provided)
        #    DELETE existing chunks for this paper first, then bulk INSERT
        # ---------------------------------------------------------------
        chunks_upserted = 0
        pdf_citations: list[Citation] = []

        if source_path is not None:
            from codex.parsing.grobid import extract_references
            from codex.parsing.nougat import pdf_to_markdown
            from codex.parsing.tex import chunk_text, latex_to_text

            suffix = Path(source_path).suffix.lower()
            text = ""

            if suffix == ".tex":
                text = latex_to_text(Path(source_path).read_text())
            elif suffix == ".pdf":
                text = pdf_to_markdown(source_path)
                # Also extract GROBID refs
                grobid_refs = extract_references(source_path)
                for ref in grobid_refs:
                    cited_id = ref.get("doi") or ref.get("arxiv_id")
                    if cited_id:
                        pdf_citations.append(Citation(citing_id=paper.id, cited_id=cited_id))
            else:
                logger.warning("Unbekannter Dateityp: %s — kein Text-Parsing", source_path)

            chunks_text = chunk_text(text) if text else []

            # Delete existing chunks for this paper
            conn.execute("DELETE FROM chunks WHERE paper_id = %s", (paper.id,))

            if chunks_text:
                # Embed all chunks in one batch
                chunk_embeddings = embedder.encode_dense(chunks_text)
                chunk_rows = [
                    (paper.id, ord_idx, content, chunk_embeddings[ord_idx].tolist())
                    for ord_idx, content in enumerate(chunks_text)
                ]
                with conn.cursor() as cur:
                    cur.executemany(
                        "INSERT INTO chunks (paper_id, ord, content, embedding)"
                        " VALUES (%s, %s, %s, %s)",
                        chunk_rows,
                    )
                chunks_upserted = len(chunk_rows)

        # ---------------------------------------------------------------
        # 5. Fetch citations  (OpenAlex if openalex_id, else S2)
        #    Insert with ON CONFLICT (citing_id, cited_id) DO NOTHING
        # ---------------------------------------------------------------
        api_citations: list[Citation]
        if paper.openalex_id:
            api_citations = openalex.fetch_citations(paper.openalex_id)
        else:
            api_citations = semanticscholar.fetch_references(paper_id)

        # Merge API citations and GROBID PDF citations; dedup via set
        all_citations_set: set[tuple[str, str]] = set()
        merged_citations: list[Citation] = []
        for cit in api_citations + pdf_citations:
            key = (cit.citing_id, cit.cited_id)
            if key not in all_citations_set:
                all_citations_set.add(key)
                merged_citations.append(cit)

        if merged_citations:
            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 merged_citations],
                )

        citations_upserted = len(merged_citations)
        conn.commit()

    return IngestResult(
        paper_id=paper.id,
        chunks_upserted=chunks_upserted,
        citations_upserted=citations_upserted,
    )