feat: merge F-05 ingest pipeline (idempotent upsert, sources+parsing+embed)
This commit is contained in:
Tarik Moussa
202
codex/ingest.py
Normal file
202
codex/ingest.py
Normal file
@@ -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,
|
||||
)
|
||||
0
tests/ingest/__init__.py
Normal file
0
tests/ingest/__init__.py
Normal file
290
tests/ingest/test_ingest.py
Normal file
290
tests/ingest/test_ingest.py
Normal file
@@ -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)
|
||||
Reference in New Issue
Block a user