feat: merge F-06 discover + provenance (graph queries, @cite scan, bib export)

This commit is contained in:
Tarik Moussa
2026-06-05 07:15:58 +02:00
6 changed files with 602 additions and 0 deletions

64
codex/discover.py Normal file
View File

@@ -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]

184
codex/provenance.py Normal file
View File

@@ -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 <bibkey>`` (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{<bibkey>,
title = {<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)

View File

View File

@@ -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

View File

View File

@@ -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