diff --git a/codex/discover.py b/codex/discover.py new file mode 100644 index 0000000..8f1e0f2 --- /dev/null +++ b/codex/discover.py @@ -0,0 +1,64 @@ +"""Discovery queries over the citation graph stored in PostgreSQL.""" + +from __future__ import annotations + +from codex.db import get_conn + + +def discovery_leads(limit: int = 20) -> list[dict[str, int | str]]: + """Return papers referenced by already-ingested papers but not yet collected. + + Returns list of dicts with keys: + cited_id (str) — arXiv ID, DOI, or OpenAlex W-ID of the target + pull (int) — how many already-ingested papers cite this target + + Ordered by pull DESC, limited to `limit` rows. + """ + sql = """ + SELECT cited_id, count(*) AS pull + FROM citations + WHERE cited_id NOT IN (SELECT id FROM papers) + GROUP BY cited_id + ORDER BY pull DESC + LIMIT %(limit)s + """ + with get_conn() as conn: + rows = conn.execute(sql, {"limit": limit}).fetchall() + return [{"cited_id": row["cited_id"], "pull": row["pull"]} for row in rows] + + +def citing_papers(paper_id: str) -> list[str]: + """Return IDs of all papers that cite `paper_id` (in-graph reverse citations).""" + sql = "SELECT citing_id FROM citations WHERE cited_id = %(paper_id)s" + with get_conn() as conn: + rows = conn.execute(sql, {"paper_id": paper_id}).fetchall() + return [row["citing_id"] for row in rows] + + +def cited_by(paper_id: str) -> list[str]: + """Return IDs of all papers that `paper_id` cites (forward citations).""" + sql = "SELECT cited_id FROM citations WHERE citing_id = %(paper_id)s" + with get_conn() as conn: + rows = conn.execute(sql, {"paper_id": paper_id}).fetchall() + return [row["cited_id"] for row in rows] + + +def cocited_papers(paper_id: str, limit: int = 10) -> list[dict[str, int | str]]: + """Return papers frequently co-cited with `paper_id`. + + A paper X is co-cited with `paper_id` if some paper cites both. + Returns list of dicts: {paper_id: str, co_citations: int}, ordered DESC. + """ + sql = """ + SELECT c2.cited_id AS paper_id, count(*) AS co_citations + FROM citations c1 + JOIN citations c2 ON c1.citing_id = c2.citing_id + WHERE c1.cited_id = %(paper_id)s + AND c2.cited_id != %(paper_id)s + GROUP BY c2.cited_id + ORDER BY co_citations DESC + LIMIT %(limit)s + """ + with get_conn() as conn: + rows = conn.execute(sql, {"paper_id": paper_id, "limit": limit}).fetchall() + return [{"paper_id": row["paper_id"], "co_citations": row["co_citations"]} for row in rows] diff --git a/codex/provenance.py b/codex/provenance.py new file mode 100644 index 0000000..2a56c5a --- /dev/null +++ b/codex/provenance.py @@ -0,0 +1,184 @@ +"""Provenance tracking: @cite-scan, code_links CRUD, references.bib export.""" + +from __future__ import annotations + +import re +from pathlib import Path + +from codex.db import get_conn +from codex.models import CodeLink + +# --------------------------------------------------------------------------- +# @cite scanning +# --------------------------------------------------------------------------- + +_CITE_RE = re.compile(r"@cite\s+(\S+)", re.IGNORECASE) +_FUNC_RE = re.compile(r"(\w+)\s*\(") +_CLASS_RE = re.compile(r"class\s+(\w+)") + +_CPP_SUFFIXES = {".cpp", ".h", ".cxx", ".hpp"} + + +def scan_cite_tags(root_dir: str) -> list[dict[str, str]]: + """Recursively scan C++ source files for @cite tags (Doxygen-style). + + Returns one dict per hit: + file (str) — relative path to the source file + line (int as str) — 1-based line number + symbol (str) — function or class name (extracted from context, best-effort) + bibkey (str) — the @cite argument + + Scans files matching: *.cpp, *.h, *.cxx, *.hpp + Pattern matched: ``@cite `` (case-insensitive, optional whitespace) + """ + root = Path(root_dir) + results: list[dict[str, str]] = [] + + for path in sorted(root.rglob("*")): + if path.suffix.lower() not in _CPP_SUFFIXES: + continue + try: + lines = path.read_text(encoding="utf-8", errors="replace").splitlines() + except OSError: + continue + + rel_path = str(path.relative_to(root)) + + for lineno, line in enumerate(lines, start=1): + for match in _CITE_RE.finditer(line): + bibkey = match.group(1) + # Best-effort symbol extraction: scan backwards for func or class + symbol = "" + context_lines = lines[: lineno - 1] + for prev_line in reversed(context_lines): + func_match = _FUNC_RE.search(prev_line) + class_match = _CLASS_RE.search(prev_line) + if func_match: + symbol = func_match.group(1) + break + if class_match: + symbol = class_match.group(1) + break + + results.append( + { + "file": rel_path, + "line": str(lineno), + "symbol": symbol, + "bibkey": bibkey, + } + ) + + return results + + +# --------------------------------------------------------------------------- +# code_links CRUD +# --------------------------------------------------------------------------- + + +def add_code_link( + symbol: str, + paper_id: str | None = None, + role: str | None = None, + note: str | None = None, +) -> CodeLink: + """Insert a new code_link row and return the populated CodeLink (with db-assigned id).""" + sql = """ + INSERT INTO code_links (symbol, paper_id, role, note) + VALUES (%(symbol)s, %(paper_id)s, %(role)s, %(note)s) + RETURNING id, added_at + """ + with get_conn() as conn: + row = conn.execute( + sql, + {"symbol": symbol, "paper_id": paper_id, "role": role, "note": note}, + ).fetchone() + conn.commit() + assert row is not None + return CodeLink( + symbol=symbol, + paper_id=paper_id, + role=role, + note=note, + id=row["id"], + added_at=row["added_at"], + ) + + +def list_code_links(paper_id: str | None = None) -> list[CodeLink]: + """Return all code_links, optionally filtered by paper_id.""" + if paper_id is not None: + sql = ( + "SELECT id, symbol, paper_id, role, note, added_at" + " FROM code_links WHERE paper_id = %(paper_id)s" + ) + params: dict[str, str] = {"paper_id": paper_id} + else: + sql = "SELECT id, symbol, paper_id, role, note, added_at FROM code_links" + params = {} + + with get_conn() as conn: + rows = conn.execute(sql, params).fetchall() + + return [ + CodeLink( + symbol=row["symbol"], + paper_id=row["paper_id"], + role=row["role"], + note=row["note"], + id=row["id"], + added_at=row["added_at"], + ) + for row in rows + ] + + +# --------------------------------------------------------------------------- +# references.bib export +# --------------------------------------------------------------------------- + + +def export_bib(paper_ids: list[str] | None = None) -> str: + """Export papers as a BibTeX string. + + If paper_ids is None, exports ALL papers with a non-NULL bibkey. + Only papers with a non-NULL bibkey are included (bibkey is the BibTeX key). + + Entry format (article): + @article{, + title = {}, + author = {<authors joined with " and ">}, + year = {<year>}, + note = {<id>}, + } + """ + if paper_ids is not None: + sql = ( + "SELECT id, bibkey, title, authors, year" + " FROM papers" + " WHERE bibkey IS NOT NULL AND id = ANY(%(paper_ids)s)" + ) + params: dict[str, list[str]] = {"paper_ids": paper_ids} + else: + sql = "SELECT id, bibkey, title, authors, year FROM papers WHERE bibkey IS NOT NULL" + params = {} + + with get_conn() as conn: + rows = conn.execute(sql, params).fetchall() + + entries: list[str] = [] + for row in rows: + authors: list[str] = row["authors"] or [] + author_str = " and ".join(authors) if authors else "" + entry = ( + f"@article{{{row['bibkey']},\n" + f" title = {{{row['title']}}},\n" + f" author = {{{author_str}}},\n" + f" year = {{{row['year']}}},\n" + f" note = {{{row['id']}}},\n" + f"}}" + ) + entries.append(entry) + + return "\n\n".join(entries) diff --git a/tests/discover/__init__.py b/tests/discover/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/discover/test_discover.py b/tests/discover/test_discover.py new file mode 100644 index 0000000..2284942 --- /dev/null +++ b/tests/discover/test_discover.py @@ -0,0 +1,145 @@ +"""Tests for codex.discover — graph-based discovery queries. + +All DB access is mocked; no real PostgreSQL server required. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import Any +from unittest.mock import MagicMock, patch + +# --------------------------------------------------------------------------- +# 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 _dict_row(**kwargs: Any) -> dict[str, Any]: + return dict(kwargs) + + +# --------------------------------------------------------------------------- +# 1. test_discovery_leads_returns_correct_format +# --------------------------------------------------------------------------- + + +def test_discovery_leads_returns_correct_format() -> None: + """Mock conn returns rows; verify output has correct shape and ordering.""" + from codex.discover import discovery_leads + + mock_conn = MagicMock() + mock_conn.execute.return_value.fetchall.return_value = [ + _dict_row(cited_id="W111", pull=5), + _dict_row(cited_id="W222", pull=3), + _dict_row(cited_id="W333", pull=1), + ] + + with patch("codex.discover.get_conn", side_effect=_make_conn_cm(mock_conn)): + result = discovery_leads(limit=10) + + assert len(result) == 3 + assert result[0] == {"cited_id": "W111", "pull": 5} + assert result[1] == {"cited_id": "W222", "pull": 3} + assert result[2] == {"cited_id": "W333", "pull": 1} + + # Verify the SQL and limit param were passed + call_args = mock_conn.execute.call_args + sql: str = call_args[0][0] + params: dict[str, Any] = call_args[0][1] + assert "cited_id NOT IN" in sql + assert "pull DESC" in sql + assert params["limit"] == 10 + + +# --------------------------------------------------------------------------- +# 2. test_citing_papers +# --------------------------------------------------------------------------- + + +def test_citing_papers() -> None: + """Mock returns citing_id list; verify flat list of IDs is returned.""" + from codex.discover import citing_papers + + mock_conn = MagicMock() + mock_conn.execute.return_value.fetchall.return_value = [ + _dict_row(citing_id="2301.00001"), + _dict_row(citing_id="2302.00002"), + ] + + with patch("codex.discover.get_conn", side_effect=_make_conn_cm(mock_conn)): + result = citing_papers("W999") + + assert result == ["2301.00001", "2302.00002"] + + call_args = mock_conn.execute.call_args + sql: str = call_args[0][0] + params: dict[str, Any] = call_args[0][1] + assert "cited_id" in sql + assert params["paper_id"] == "W999" + + +# --------------------------------------------------------------------------- +# 3. test_cited_by +# --------------------------------------------------------------------------- + + +def test_cited_by() -> None: + """Mock returns cited_id list; verify flat list of IDs is returned.""" + from codex.discover import cited_by + + mock_conn = MagicMock() + mock_conn.execute.return_value.fetchall.return_value = [ + _dict_row(cited_id="10.1000/ref_a"), + _dict_row(cited_id="10.1000/ref_b"), + ] + + with patch("codex.discover.get_conn", side_effect=_make_conn_cm(mock_conn)): + result = cited_by("2301.00001") + + assert result == ["10.1000/ref_a", "10.1000/ref_b"] + + call_args = mock_conn.execute.call_args + sql: str = call_args[0][0] + params: dict[str, Any] = call_args[0][1] + assert "citing_id" in sql + assert params["paper_id"] == "2301.00001" + + +# --------------------------------------------------------------------------- +# 4. test_cocited_papers +# --------------------------------------------------------------------------- + + +def test_cocited_papers() -> None: + """Mock returns co-citation list; verify shape and ordering.""" + from codex.discover import cocited_papers + + mock_conn = MagicMock() + mock_conn.execute.return_value.fetchall.return_value = [ + _dict_row(paper_id="W444", co_citations=7), + _dict_row(paper_id="W555", co_citations=2), + ] + + with patch("codex.discover.get_conn", side_effect=_make_conn_cm(mock_conn)): + result = cocited_papers("W111", limit=5) + + assert len(result) == 2 + assert result[0] == {"paper_id": "W444", "co_citations": 7} + assert result[1] == {"paper_id": "W555", "co_citations": 2} + + call_args = mock_conn.execute.call_args + sql: str = call_args[0][0] + params: dict[str, Any] = call_args[0][1] + assert "co_citations DESC" in sql + assert params["paper_id"] == "W111" + assert params["limit"] == 5 diff --git a/tests/provenance/__init__.py b/tests/provenance/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/provenance/test_provenance.py b/tests/provenance/test_provenance.py new file mode 100644 index 0000000..228552a --- /dev/null +++ b/tests/provenance/test_provenance.py @@ -0,0 +1,209 @@ +"""Tests for codex.provenance — @cite scan, code_links CRUD, bib export. + +All DB access is mocked; no real PostgreSQL server required. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from datetime import UTC, datetime +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock, patch + +from codex.models import CodeLink + +# --------------------------------------------------------------------------- +# 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 _dict_row(**kwargs: Any) -> dict[str, Any]: + return dict(kwargs) + + +# --------------------------------------------------------------------------- +# 1. test_scan_cite_tags +# --------------------------------------------------------------------------- + + +def test_scan_cite_tags(tmp_path: Path) -> None: + """A .cpp file with @cite produces correct dict entries.""" + from codex.provenance import scan_cite_tags + + src_file = tmp_path / "algo.cpp" + src_file.write_text( + "// Helper function\n" + "void compute_hodge() {\n" + " // @cite somepaper2023\n" + " do_stuff();\n" + "}\n", + encoding="utf-8", + ) + + results = scan_cite_tags(str(tmp_path)) + + assert len(results) == 1 + hit = results[0] + assert hit["bibkey"] == "somepaper2023" + assert hit["line"] == "3" + assert hit["file"] == "algo.cpp" + # symbol extracted from "compute_hodge()" on the line above + assert hit["symbol"] == "compute_hodge" + + +# --------------------------------------------------------------------------- +# 2. test_scan_cite_tags_no_files +# --------------------------------------------------------------------------- + + +def test_scan_cite_tags_no_files(tmp_path: Path) -> None: + """Empty directory returns an empty list.""" + from codex.provenance import scan_cite_tags + + result = scan_cite_tags(str(tmp_path)) + assert result == [] + + +# --------------------------------------------------------------------------- +# 3. test_add_code_link +# --------------------------------------------------------------------------- + + +def test_add_code_link() -> None: + """Mock conn, verify INSERT is issued and returned CodeLink has id from RETURNING.""" + from codex.provenance import add_code_link + + now = datetime(2026, 6, 5, tzinfo=UTC) + mock_conn = MagicMock() + mock_conn.execute.return_value.fetchone.return_value = _dict_row(id=42, added_at=now) + + with patch("codex.provenance.get_conn", side_effect=_make_conn_cm(mock_conn)): + link = add_code_link( + symbol="dec::hodge_star", + paper_id="2301.07041", + role="implements Thm 3.2", + note="core algorithm", + ) + + assert isinstance(link, CodeLink) + assert link.id == 42 + assert link.added_at == now + assert link.symbol == "dec::hodge_star" + assert link.paper_id == "2301.07041" + assert link.role == "implements Thm 3.2" + assert link.note == "core algorithm" + + call_args = mock_conn.execute.call_args + sql: str = call_args[0][0] + assert "INSERT INTO code_links" in sql + assert "RETURNING" in sql + + mock_conn.commit.assert_called_once() + + +# --------------------------------------------------------------------------- +# 4. test_list_code_links_all +# --------------------------------------------------------------------------- + + +def test_list_code_links_all() -> None: + """No filter → all rows returned as CodeLink list.""" + from codex.provenance import list_code_links + + now = datetime(2026, 6, 5, tzinfo=UTC) + mock_conn = MagicMock() + mock_conn.execute.return_value.fetchall.return_value = [ + _dict_row(id=1, symbol="foo::bar", paper_id="W111", role="uses", note=None, added_at=now), + _dict_row(id=2, symbol="baz::qux", paper_id=None, role=None, note="ref", added_at=now), + ] + + with patch("codex.provenance.get_conn", side_effect=_make_conn_cm(mock_conn)): + links = list_code_links() + + assert len(links) == 2 + assert all(isinstance(lnk, CodeLink) for lnk in links) + assert links[0].symbol == "foo::bar" + assert links[0].id == 1 + assert links[1].symbol == "baz::qux" + assert links[1].paper_id is None + + # No WHERE clause in the SQL + call_args = mock_conn.execute.call_args + sql: str = call_args[0][0] + assert "WHERE" not in sql + + +# --------------------------------------------------------------------------- +# 5. test_list_code_links_filtered +# --------------------------------------------------------------------------- + + +def test_list_code_links_filtered() -> None: + """Filter by paper_id → WHERE clause applied and only matching rows returned.""" + from codex.provenance import list_code_links + + now = datetime(2026, 6, 5, tzinfo=UTC) + mock_conn = MagicMock() + mock_conn.execute.return_value.fetchall.return_value = [ + _dict_row( + id=3, + symbol="some::func", + paper_id="2301.07041", + role=None, + note=None, + added_at=now, + ), + ] + + with patch("codex.provenance.get_conn", side_effect=_make_conn_cm(mock_conn)): + links = list_code_links(paper_id="2301.07041") + + assert len(links) == 1 + assert links[0].paper_id == "2301.07041" + + call_args = mock_conn.execute.call_args + sql: str = call_args[0][0] + params: dict[str, Any] = call_args[0][1] + assert "WHERE paper_id" in sql + assert params["paper_id"] == "2301.07041" + + +# --------------------------------------------------------------------------- +# 6. test_export_bib_format +# --------------------------------------------------------------------------- + + +def test_export_bib_format() -> None: + """Mock returns one paper row; verify valid BibTeX string output.""" + from codex.provenance import export_bib + + mock_conn = MagicMock() + mock_conn.execute.return_value.fetchall.return_value = [ + _dict_row( + id="2301.07041", + bibkey="smith2023conformal", + title="Conformal Approaches to Laplacian Eigenmaps", + authors=["Alice Smith", "Bob Jones"], + year=2023, + ), + ] + + with patch("codex.provenance.get_conn", side_effect=_make_conn_cm(mock_conn)): + bib = export_bib() + + assert "@article{smith2023conformal," in bib + assert "title = {Conformal Approaches to Laplacian Eigenmaps}" in bib + assert "author = {Alice Smith and Bob Jones}" in bib + assert "year = {2023}" in bib + assert "note = {2301.07041}" in bib