diff --git a/codex/ingest.py b/codex/ingest.py new file mode 100644 index 0000000..3c243b2 --- /dev/null +++ b/codex/ingest.py @@ -0,0 +1,202 @@ +"""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, + ) diff --git a/tests/ingest/__init__.py b/tests/ingest/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/ingest/test_ingest.py b/tests/ingest/test_ingest.py new file mode 100644 index 0000000..1f54db2 --- /dev/null +++ b/tests/ingest/test_ingest.py @@ -0,0 +1,290 @@ +"""Tests for codex.ingest — end-to-end idempotent ingest pipeline. + +All external calls (DB, OpenAlex, S2, embed, parsing) are mocked so +the suite runs offline in milliseconds. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import Any +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from codex.ingest import IngestResult, ingest_paper +from codex.models import Citation, Paper + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_conn_cm(mock_conn: MagicMock) -> Any: + """Return a context-manager factory that yields *mock_conn*.""" + + @contextmanager # type: ignore[arg-type] + def _cm() -> Any: + yield mock_conn + + return _cm + + +def _make_paper( + paper_id: str = "2301.07041", + openalex_id: str | None = "W123", + abstract: str | None = "Test abstract.", +) -> Paper: + return Paper( + id=paper_id, + title="A Test Paper", + openalex_id=openalex_id, + authors=["Alice", "Bob"], + year=2023, + abstract=abstract, + ) + + +def _fake_embedder(dim: int = 1024) -> MagicMock: + """Return a mock Embedder whose encode_dense returns deterministic output.""" + embedder = MagicMock() + embedder.dim = dim + + def encode_dense(texts: list[str]) -> np.ndarray: + return np.ones((len(texts), dim), dtype=np.float32) + + embedder.encode_dense.side_effect = encode_dense + return embedder + + +# --------------------------------------------------------------------------- +# 1. test_ingest_paper_basic +# --------------------------------------------------------------------------- + + +def test_ingest_paper_basic() -> None: + """fetch_paper returns a Paper; DB upsert is called; IngestResult has correct counts.""" + paper = _make_paper() + mock_conn = MagicMock() + mock_conn.execute = MagicMock() + mock_conn.executemany = MagicMock() + mock_conn.commit = MagicMock() + + api_citations = [ + Citation(citing_id=paper.openalex_id or "", cited_id="W999"), + ] + + with ( + patch("codex.ingest.openalex.fetch_paper", return_value=paper) as mock_fetch, + patch("codex.ingest.openalex.fetch_citations", return_value=api_citations), + patch("codex.ingest.get_embedder", return_value=_fake_embedder()), + patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)), + ): + result = ingest_paper(paper.id) + + mock_fetch.assert_called_once_with(paper.id) + + # Paper upsert must have been issued + assert mock_conn.execute.call_count >= 1 + first_sql: str = mock_conn.execute.call_args_list[0][0][0] + assert "INSERT INTO papers" in first_sql + assert "ON CONFLICT" in first_sql + + assert isinstance(result, IngestResult) + assert result.paper_id == paper.id + assert result.chunks_upserted == 0 + assert result.citations_upserted == len(api_citations) + + +# --------------------------------------------------------------------------- +# 2. test_ingest_paper_source_tex +# --------------------------------------------------------------------------- + + +def test_ingest_paper_source_tex(tmp_path: Any) -> None: + """``.tex`` source → latex_to_text, chunk_text, embed, chunks inserted in DB.""" + tex_file = tmp_path / "paper.tex" + tex_file.write_text(r"\section{Introduction} Hello world. This is a test paper.") + + paper = _make_paper() + mock_conn = MagicMock() + mock_conn.execute = MagicMock() + mock_conn.executemany = MagicMock() + mock_conn.commit = MagicMock() + + fake_chunks = ["chunk one text here.", "chunk two text here."] + + with ( + patch("codex.ingest.openalex.fetch_paper", return_value=paper), + patch("codex.ingest.openalex.fetch_citations", return_value=[]), + patch("codex.ingest.get_embedder", return_value=_fake_embedder()), + patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)), + patch("codex.parsing.tex.latex_to_text", return_value="Hello world. Intro text."), + patch("codex.parsing.tex.chunk_text", return_value=fake_chunks), + ): + result = ingest_paper(paper.id, source_path=str(tex_file)) + + # DELETE + INSERT for chunks + sql_calls = [str(c[0][0]) for c in mock_conn.execute.call_args_list] + assert any("DELETE FROM chunks" in sql for sql in sql_calls) + + # executemany is called on the cursor, not on the connection directly + mock_cursor = mock_conn.cursor.return_value.__enter__.return_value + assert mock_cursor.executemany.call_count >= 1 + chunk_insert_call = mock_cursor.executemany.call_args_list[0] + chunk_sql: str = chunk_insert_call[0][0] + assert "INSERT INTO chunks" in chunk_sql + + assert result.chunks_upserted == len(fake_chunks) + + +# --------------------------------------------------------------------------- +# 3. test_ingest_paper_source_pdf +# --------------------------------------------------------------------------- + + +def test_ingest_paper_source_pdf(tmp_path: Any) -> None: + """``.pdf`` source → pdf_to_markdown, chunk_text, embed, grobid refs, chunks + citations.""" + pdf_file = tmp_path / "paper.pdf" + pdf_file.write_bytes(b"%PDF-1.4 fake content") + + paper = _make_paper() + mock_conn = MagicMock() + mock_conn.execute = MagicMock() + mock_conn.executemany = MagicMock() + mock_conn.commit = MagicMock() + + fake_chunks = ["pdf chunk one.", "pdf chunk two."] + grobid_refs = [ + {"title": "Ref A", "authors": "X", "year": "2020", "doi": "10.1/ref_a", "arxiv_id": ""}, + {"title": "Ref B", "authors": "Y", "year": "2021", "doi": "", "arxiv_id": "2101.00001"}, + ] + + with ( + patch("codex.ingest.openalex.fetch_paper", return_value=paper), + patch("codex.ingest.openalex.fetch_citations", return_value=[]), + patch("codex.ingest.get_embedder", return_value=_fake_embedder()), + patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)), + patch("codex.parsing.nougat.pdf_to_markdown", return_value="pdf markdown text."), + patch("codex.parsing.tex.chunk_text", return_value=fake_chunks), + patch("codex.parsing.grobid.extract_references", return_value=grobid_refs), + ): + result = ingest_paper(paper.id, source_path=str(pdf_file)) + + # Chunks were inserted + assert result.chunks_upserted == len(fake_chunks) + + # GROBID refs were merged into citations + # Two grobid refs → both have either doi or arxiv_id → 2 citations + assert result.citations_upserted == 2 + + # executemany is called on the cursor for chunks + citations (at least 2 calls total) + mock_cursor = mock_conn.cursor.return_value.__enter__.return_value + assert mock_cursor.executemany.call_count >= 2 + + +# --------------------------------------------------------------------------- +# 4. test_ingest_paper_not_found +# --------------------------------------------------------------------------- + + +def test_ingest_paper_not_found() -> None: + """OpenAlex returns None for unknown ID → ValueError raised.""" + with ( + patch("codex.ingest.openalex.fetch_paper", return_value=None), + patch("codex.ingest.get_embedder", return_value=_fake_embedder()), + pytest.raises(ValueError, match="Paper not found"), + ): + ingest_paper("10.9999/unknown-doi") + + +# --------------------------------------------------------------------------- +# 5. test_ingest_paper_idempotent +# --------------------------------------------------------------------------- + + +def test_ingest_paper_idempotent() -> None: + """Second call with the same paper_id overwrites without error (no unique violations).""" + paper = _make_paper() + mock_conn = MagicMock() + mock_conn.execute = MagicMock() + mock_conn.executemany = MagicMock() + mock_conn.commit = MagicMock() + + with ( + patch("codex.ingest.openalex.fetch_paper", return_value=paper), + patch("codex.ingest.openalex.fetch_citations", return_value=[]), + patch("codex.ingest.get_embedder", return_value=_fake_embedder()), + patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)), + ): + result1 = ingest_paper(paper.id) + result2 = ingest_paper(paper.id) + + assert result1.paper_id == result2.paper_id == paper.id + + # ON CONFLICT … DO UPDATE means no error on second call + all_execute_calls = mock_conn.execute.call_args_list + paper_upserts = [c for c in all_execute_calls if "INSERT INTO papers" in str(c[0][0])] + # Two calls → two paper upserts + assert len(paper_upserts) == 2 + + +# --------------------------------------------------------------------------- +# 6. test_ingest_paper_arxiv_s2_fallback +# --------------------------------------------------------------------------- + + +def test_ingest_paper_arxiv_s2_fallback() -> None: + """If OpenAlex returns None for an arXiv ID, S2 fallback produces a stub Paper.""" + paper_id = "2301.07041" + mock_conn = MagicMock() + mock_conn.execute = MagicMock() + mock_conn.executemany = MagicMock() + mock_conn.commit = MagicMock() + + with ( + patch("codex.ingest.openalex.fetch_paper", return_value=None), + patch("codex.ingest.semanticscholar.fetch_references", return_value=[]) as mock_s2, + patch("codex.ingest.get_embedder", return_value=_fake_embedder()), + patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)), + ): + result = ingest_paper(paper_id) + + # S2 was consulted as fallback + mock_s2.assert_called() + assert result.paper_id == paper_id + + +# --------------------------------------------------------------------------- +# 7. test_ingest_paper_no_abstract_uses_zero_vector +# --------------------------------------------------------------------------- + + +def test_ingest_paper_no_abstract_uses_zero_vector() -> None: + """A paper without an abstract gets a zero embedding vector (no model call).""" + paper = _make_paper(abstract=None) + mock_conn = MagicMock() + mock_conn.execute = MagicMock() + mock_conn.executemany = MagicMock() + mock_conn.commit = MagicMock() + + fake_emb = _fake_embedder(dim=1024) + + with ( + patch("codex.ingest.openalex.fetch_paper", return_value=paper), + patch("codex.ingest.openalex.fetch_citations", return_value=[]), + patch("codex.ingest.get_embedder", return_value=fake_emb), + patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)), + ): + ingest_paper(paper.id) + + # encode_dense must NOT have been called (no abstract → zero vector shortcut) + fake_emb.encode_dense.assert_not_called() + + # Verify zero vector was passed in the paper upsert + upsert_call = mock_conn.execute.call_args_list[0] + params: dict[str, Any] = upsert_call[0][1] + emb: list[float] = params["abstract_emb"] + assert len(emb) == 1024 + assert all(v == 0.0 for v in emb)