diff --git a/codex/sources/arxiv.py b/codex/sources/arxiv.py new file mode 100644 index 0000000..1de3e38 --- /dev/null +++ b/codex/sources/arxiv.py @@ -0,0 +1,100 @@ +"""arXiv API client. + +Provides: +- fetch_source: download the .tar.gz source of a paper and extract the primary .tex file. +- fetch_pdf_url: return the canonical PDF URL for a given arXiv ID. +""" + +from __future__ import annotations + +import io +import logging +import tarfile + +import httpx + +logger = logging.getLogger(__name__) + +_BASE = "https://arxiv.org" + + +def fetch_source(arxiv_id: str) -> str | None: + """Download and extract the primary LaTeX source for an arXiv paper. + + Downloads the .tar.gz source bundle from ``https://arxiv.org/src/{arxiv_id}``, + locates the primary .tex file (preferring any file containing ``\\documentclass``, + falling back to the largest .tex by size), and returns its contents as a UTF-8 + string. + + Parameters + ---------- + arxiv_id: + The arXiv identifier (e.g. ``"2301.07041"``). + + Returns + ------- + str | None + Raw LaTeX source string, or None if the paper is not found or no .tex + file is present (signals Nougat fallback). + """ + url = f"{_BASE}/src/{arxiv_id}" + try: + response = httpx.get(url, timeout=60, follow_redirects=True) + except httpx.RequestError: + raise + if response.status_code == 404: + logger.debug("arXiv 404 for source id=%s", arxiv_id) + return None + if response.status_code != 200: + response.raise_for_status() + + raw = response.content + try: + with tarfile.open(fileobj=io.BytesIO(raw), mode="r:gz") as tf: + tex_members = [m for m in tf.getmembers() if m.name.endswith(".tex")] + if not tex_members: + logger.debug("No .tex files found in arXiv source for %s", arxiv_id) + return None + + # Prefer the file containing \documentclass (primary document) + primary: tarfile.TarInfo | None = None + for member in tex_members: + f = tf.extractfile(member) + if f is None: + continue + content_bytes = f.read() + if b"\\documentclass" in content_bytes: + primary = member + # Decode and return immediately — first match wins + return content_bytes.decode("utf-8", errors="replace") + + if primary is None: + # Fallback: largest .tex by size + largest = max(tex_members, key=lambda m: m.size) + f = tf.extractfile(largest) + if f is None: + return None + return f.read().decode("utf-8", errors="replace") + except tarfile.TarError as exc: + logger.warning("Failed to open tar archive for %s: %s", arxiv_id, exc) + return None + + return None # unreachable but satisfies type checker + + +def fetch_pdf_url(arxiv_id: str) -> str: + """Return the canonical PDF URL for an arXiv paper. + + This is a pure computation — no HTTP request is made. + + Parameters + ---------- + arxiv_id: + The arXiv identifier (e.g. ``"2301.07041"``). + + Returns + ------- + str + The full URL of the PDF file. + """ + return f"{_BASE}/pdf/{arxiv_id}.pdf" diff --git a/codex/sources/openalex.py b/codex/sources/openalex.py new file mode 100644 index 0000000..bd1aa2d --- /dev/null +++ b/codex/sources/openalex.py @@ -0,0 +1,136 @@ +"""OpenAlex API client. + +Provides: +- fetch_paper: resolve an arXiv ID or DOI to a Paper dataclass. +- fetch_citations: retrieve reference list as Citation dataclasses. + +All requests use the OpenAlex Polite Pool (mailto query parameter) and +are retried on 429/5xx with exponential back-off via tenacity. +""" + +from __future__ import annotations + +import logging +from typing import Any + +import httpx +from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential + +from codex.config import get_settings +from codex.models import Citation, Paper + +logger = logging.getLogger(__name__) + +_BASE = "https://api.openalex.org" + + +def _is_retryable(exc: BaseException) -> bool: + """Return True for HTTP 429 / 5xx responses wrapped in httpx.HTTPStatusError.""" + if isinstance(exc, httpx.HTTPStatusError): + return exc.response.status_code == 429 or exc.response.status_code >= 500 + return False + + +def _polite_params() -> dict[str, str]: + mailto = get_settings().openalex_mailto + if mailto: + return {"mailto": mailto} + return {} + + +@retry( + retry=retry_if_exception(_is_retryable), + stop=stop_after_attempt(5), + wait=wait_exponential(min=1, max=30), + before_sleep=lambda rs: logger.warning( + "OpenAlex retry %d after %s", + rs.attempt_number, + rs.outcome.exception(), # type: ignore[union-attr] + ), +) +def _get(url: str, params: dict[str, str] | None = None) -> httpx.Response: + merged: dict[str, str] = {**_polite_params(), **(params or {})} + response = httpx.get(url, params=merged, timeout=30) + response.raise_for_status() + return response + + +def _reconstruct_abstract(inverted_index: dict[str, list[int]] | None) -> str | None: + """Convert OpenAlex abstract_inverted_index back to plain text.""" + if not inverted_index: + return None + word_positions: list[tuple[int, str]] = [] + for word, positions in inverted_index.items(): + for pos in positions: + word_positions.append((pos, word)) + word_positions.sort(key=lambda t: t[0]) + return " ".join(w for _, w in word_positions) + + +def _map_paper(data: dict[str, Any]) -> Paper: + authors: list[str] = [ + a.get("author", {}).get("display_name") or "" for a in data.get("authorships", []) + ] + abstract = _reconstruct_abstract(data.get("abstract_inverted_index")) + return Paper( + id=data.get("doi") or data.get("id") or "", + title=data.get("title") or "", + openalex_id=data.get("id"), + authors=authors, + year=data.get("publication_year"), + abstract=abstract, + ) + + +def fetch_paper(id: str) -> Paper | None: + """Fetch a single paper by arXiv ID or DOI. + + Parameters + ---------- + id: + An arXiv ID (``"2301.07041"``) or a DOI (``"10.1145/…"``). + + Returns + ------- + Paper | None + A populated Paper instance, or None on 404. + """ + url = f"{_BASE}/works/{id}" + try: + response = _get(url) + except httpx.HTTPStatusError as exc: + if exc.response.status_code == 404: + logger.debug("OpenAlex 404 for id=%s", id) + return None + raise + return _map_paper(response.json()) + + +def fetch_citations(openalex_id: str) -> list[Citation]: + """Fetch the reference list for a paper as Citation objects. + + Parameters + ---------- + openalex_id: + The OpenAlex work ID (``"https://openalex.org/W…"`` or short ``"W…"``). + + Returns + ------- + list[Citation] + One Citation per referenced work. + """ + url = f"{_BASE}/works/{openalex_id}/references" + try: + response = _get(url) + except httpx.HTTPStatusError as exc: + if exc.response.status_code == 404: + return [] + raise + + results: list[dict[str, Any]] = response.json().get("results", []) + citations: list[Citation] = [] + for ref in results: + cited_id: str = ref.get("doi") or ref.get("id") or "" + if cited_id: + citations.append(Citation(citing_id=openalex_id, cited_id=cited_id)) + return citations diff --git a/codex/sources/semanticscholar.py b/codex/sources/semanticscholar.py new file mode 100644 index 0000000..8eb45d8 --- /dev/null +++ b/codex/sources/semanticscholar.py @@ -0,0 +1,106 @@ +"""Semantic Scholar API client. + +Provides: +- fetch_references: retrieve references for a paper as Citation dataclasses. +- fetch_recommendations: retrieve recommended paper IDs. + +Requests are rate-limited to 1 per second via tenacity wait_fixed. +""" + +from __future__ import annotations + +import logging +from typing import Any + +import httpx +from tenacity import retry, stop_after_attempt, wait_fixed + +from codex.models import Citation + +logger = logging.getLogger(__name__) + +_BASE_GRAPH = "https://api.semanticscholar.org/graph/v1" +_BASE_RECS = "https://api.semanticscholar.org/recommendations/v1" + + +@retry( + stop=stop_after_attempt(5), + wait=wait_fixed(1), + before_sleep=lambda rs: logger.warning( + "SemanticScholar retry %d after %s", + rs.attempt_number, + rs.outcome.exception(), # type: ignore[union-attr] + ), +) +def _get(url: str, params: dict[str, Any] | None = None) -> httpx.Response: + response = httpx.get(url, params=params, timeout=30) + response.raise_for_status() + return response + + +def fetch_references(paper_id: str) -> list[Citation]: + """Fetch references for a paper from Semantic Scholar. + + Parameters + ---------- + paper_id: + Semantic Scholar paper ID (or ``arXiv:…`` / ``DOI:…`` prefixed ID). + + Returns + ------- + list[Citation] + One Citation per reference, with optional context snippet. + """ + url = f"{_BASE_GRAPH}/paper/{paper_id}/references" + params: dict[str, Any] = {"fields": "externalIds,contexts"} + try: + response = _get(url, params=params) + except httpx.HTTPStatusError as exc: + if exc.response.status_code == 404: + return [] + raise + + data = response.json() + raw_refs: list[dict[str, Any]] = data.get("data", []) + citations: list[Citation] = [] + for entry in raw_refs: + cited_paper: dict[str, Any] = entry.get("citedPaper", {}) + external_ids: dict[str, str] = cited_paper.get("externalIds") or {} + contexts: list[str] = entry.get("contexts", []) + context: str | None = contexts[0] if contexts else None + + cited_id: str = ( + external_ids.get("DOI") or external_ids.get("ArXiv") or cited_paper.get("paperId") or "" + ) + if cited_id: + citations.append(Citation(citing_id=paper_id, cited_id=cited_id, context=context)) + return citations + + +def fetch_recommendations(paper_id: str, limit: int = 20) -> list[str]: + """Fetch recommended paper IDs from Semantic Scholar. + + Parameters + ---------- + paper_id: + Semantic Scholar paper ID. + limit: + Maximum number of recommendations to return. + + Returns + ------- + list[str] + List of recommended paper IDs. + """ + url = f"{_BASE_RECS}/papers/forpaper/{paper_id}" + params: dict[str, Any] = {"limit": limit} + try: + response = _get(url, params=params) + except httpx.HTTPStatusError as exc: + if exc.response.status_code == 404: + return [] + raise + + data = response.json() + recommended: list[dict[str, Any]] = data.get("recommendedPapers", []) + return [p["paperId"] for p in recommended if p.get("paperId")] diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/sources/__init__.py b/tests/sources/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/sources/test_arxiv.py b/tests/sources/test_arxiv.py new file mode 100644 index 0000000..1c08500 --- /dev/null +++ b/tests/sources/test_arxiv.py @@ -0,0 +1,140 @@ +"""Tests for codex.sources.arxiv.""" + +from __future__ import annotations + +import io +import tarfile + +import httpx +import pytest + +from codex.sources import arxiv + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_tar_gz(files: dict[str, bytes]) -> bytes: + """Build an in-memory .tar.gz archive from a dict of {name: content}.""" + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tf: + for name, content in files.items(): + info = tarfile.TarInfo(name=name) + info.size = len(content) + tf.addfile(info, io.BytesIO(content)) + return buf.getvalue() + + +# --------------------------------------------------------------------------- +# fetch_source — success: in-memory tar.gz with one .tex +# --------------------------------------------------------------------------- + + +def test_fetch_source_returns_latex(monkeypatch: pytest.MonkeyPatch) -> None: + """fetch_source extracts and returns the LaTeX content from a .tar.gz.""" + latex_content = b"\\documentclass{article}\n\\begin{document}\nHello world\n\\end{document}" + tar_bytes = _make_tar_gz({"main.tex": latex_content}) + + def mock_get( + url: str, + *, + timeout: int = 60, + follow_redirects: bool = True, + ) -> httpx.Response: + return httpx.Response(200, content=tar_bytes) + + monkeypatch.setattr(httpx, "get", mock_get) + + result = arxiv.fetch_source("2301.07041") + + assert result is not None + assert "\\documentclass" in result + assert "Hello world" in result + + +# --------------------------------------------------------------------------- +# fetch_source — 404 → None +# --------------------------------------------------------------------------- + + +def test_fetch_source_404_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: + """A 404 response should return None.""" + + def mock_get( + url: str, + *, + timeout: int = 60, + follow_redirects: bool = True, + ) -> httpx.Response: + return httpx.Response(404) + + monkeypatch.setattr(httpx, "get", mock_get) + + result = arxiv.fetch_source("nonexistent-id") + assert result is None + + +# --------------------------------------------------------------------------- +# fetch_source — no .tex in archive → None +# --------------------------------------------------------------------------- + + +def test_fetch_source_no_tex_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: + """If the archive contains no .tex files, return None.""" + tar_bytes = _make_tar_gz({"README.md": b"# Paper\nNo LaTeX here."}) + + def mock_get( + url: str, + *, + timeout: int = 60, + follow_redirects: bool = True, + ) -> httpx.Response: + return httpx.Response(200, content=tar_bytes) + + monkeypatch.setattr(httpx, "get", mock_get) + + result = arxiv.fetch_source("2301.07041") + assert result is None + + +# --------------------------------------------------------------------------- +# fetch_source — fallback to largest .tex when no \documentclass +# --------------------------------------------------------------------------- + + +def test_fetch_source_fallback_to_largest_tex(monkeypatch: pytest.MonkeyPatch) -> None: + """When no file has \\documentclass, the largest .tex is returned.""" + small_tex = b"% small file\n\\section{Intro}" + large_tex = b"% large file\n" + b"x" * 500 + + tar_bytes = _make_tar_gz({"small.tex": small_tex, "large.tex": large_tex}) + + def mock_get( + url: str, + *, + timeout: int = 60, + follow_redirects: bool = True, + ) -> httpx.Response: + return httpx.Response(200, content=tar_bytes) + + monkeypatch.setattr(httpx, "get", mock_get) + + result = arxiv.fetch_source("2301.07041") + + assert result is not None + assert "large file" in result + + +# --------------------------------------------------------------------------- +# fetch_pdf_url — pure computation, no mock needed +# --------------------------------------------------------------------------- + + +def test_fetch_pdf_url_pure_computation() -> None: + """fetch_pdf_url returns the correct URL without making any HTTP request.""" + url = arxiv.fetch_pdf_url("2301.07041") + assert url == "https://arxiv.org/pdf/2301.07041.pdf" + + url2 = arxiv.fetch_pdf_url("1234.56789") + assert url2 == "https://arxiv.org/pdf/1234.56789.pdf" diff --git a/tests/sources/test_openalex.py b/tests/sources/test_openalex.py new file mode 100644 index 0000000..54b9b8b --- /dev/null +++ b/tests/sources/test_openalex.py @@ -0,0 +1,142 @@ +"""Tests for codex.sources.openalex.""" + +from __future__ import annotations + +import httpx +import pytest + +from codex.models import Citation, Paper +from codex.sources import openalex + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +_SAMPLE_WORK = { + "id": "https://openalex.org/W2741809807", + "doi": "10.1145/3592430", + "title": "Conformal Prediction: A Review", + "publication_year": 2023, + "authorships": [ + {"author": {"display_name": "Alice Smith"}}, + {"author": {"display_name": "Bob Jones"}}, + ], + "abstract_inverted_index": { + "Conformal": [0], + "prediction": [1], + "is": [2], + "great": [3], + }, +} + +_SAMPLE_REFS = { + "results": [ + {"doi": "10.1000/ref1", "id": "https://openalex.org/W1"}, + {"doi": "10.1000/ref2", "id": "https://openalex.org/W2"}, + ] +} + + +# --------------------------------------------------------------------------- +# fetch_paper — success +# --------------------------------------------------------------------------- + + +def test_fetch_paper_success(monkeypatch: pytest.MonkeyPatch) -> None: + """Successful fetch maps all fields to Paper correctly.""" + + def mock_get(url: str, params: dict[str, str] | None = None) -> httpx.Response: + return httpx.Response(200, json=_SAMPLE_WORK) + + monkeypatch.setattr(openalex, "_get", mock_get) + + paper = openalex.fetch_paper("10.1145/3592430") + + assert paper is not None + assert isinstance(paper, Paper) + assert paper.title == "Conformal Prediction: A Review" + assert paper.year == 2023 + assert paper.authors == ["Alice Smith", "Bob Jones"] + assert paper.abstract == "Conformal prediction is great" + assert paper.openalex_id == "https://openalex.org/W2741809807" + + +# --------------------------------------------------------------------------- +# fetch_paper — 404 → None +# --------------------------------------------------------------------------- + + +def test_fetch_paper_404_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: + """A 404 response should return None without raising.""" + + def mock_get(url: str, params: dict[str, str] | None = None) -> httpx.Response: + raise httpx.HTTPStatusError( + "Not Found", + request=httpx.Request("GET", url), + response=httpx.Response(404, request=httpx.Request("GET", url)), + ) + + monkeypatch.setattr(openalex, "_get", mock_get) + + result = openalex.fetch_paper("nonexistent-id") + assert result is None + + +# --------------------------------------------------------------------------- +# fetch_paper — 429 retry then success +# --------------------------------------------------------------------------- + + +def test_fetch_paper_retry_on_429(monkeypatch: pytest.MonkeyPatch) -> None: + """On 429 x3 then 200, should return 1 Paper after 3 failed attempts.""" + attempt = [0] + + def controlled_get(url: str, params: dict[str, str] | None = None) -> httpx.Response: + attempt[0] += 1 + if attempt[0] <= 3: + resp = httpx.Response(429, request=httpx.Request("GET", url)) + raise httpx.HTTPStatusError( + "Too Many Requests", + request=httpx.Request("GET", url), + response=resp, + ) + return httpx.Response(200, json=_SAMPLE_WORK) + + monkeypatch.setattr(openalex, "_get", controlled_get) + + # Call fetch_paper in a loop to simulate retry behaviour + result: Paper | None = None + errors = 0 + for _ in range(5): + try: + result = openalex.fetch_paper("10.1145/3592430") + break + except httpx.HTTPStatusError: + errors += 1 + + assert result is not None + assert isinstance(result, Paper) + assert errors == 3 # 3 failures before success + + +# --------------------------------------------------------------------------- +# fetch_citations +# --------------------------------------------------------------------------- + + +def test_fetch_citations(monkeypatch: pytest.MonkeyPatch) -> None: + """fetch_citations returns one Citation per reference entry.""" + + def mock_get(url: str, params: dict[str, str] | None = None) -> httpx.Response: + return httpx.Response(200, json=_SAMPLE_REFS) + + monkeypatch.setattr(openalex, "_get", mock_get) + + oa_id = "https://openalex.org/W2741809807" + citations = openalex.fetch_citations(oa_id) + + assert len(citations) == 2 + assert all(isinstance(c, Citation) for c in citations) + assert citations[0].citing_id == oa_id + assert citations[0].cited_id == "10.1000/ref1" + assert citations[1].cited_id == "10.1000/ref2" diff --git a/tests/sources/test_semanticscholar.py b/tests/sources/test_semanticscholar.py new file mode 100644 index 0000000..42066d6 --- /dev/null +++ b/tests/sources/test_semanticscholar.py @@ -0,0 +1,126 @@ +"""Tests for codex.sources.semanticscholar.""" + +from __future__ import annotations + +import httpx +import pytest + +from codex.models import Citation +from codex.sources import semanticscholar + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +_SAMPLE_REFS_RESPONSE = { + "data": [ + { + "citedPaper": { + "paperId": "abc123", + "externalIds": {"DOI": "10.1000/xyz1"}, + }, + "contexts": ["This approach was first described in [1]."], + }, + { + "citedPaper": { + "paperId": "def456", + "externalIds": {"ArXiv": "2301.07041"}, + }, + "contexts": [], + }, + ] +} + +_SAMPLE_RECS_RESPONSE = { + "recommendedPapers": [ + {"paperId": "rec_id_1"}, + {"paperId": "rec_id_2"}, + {"paperId": "rec_id_3"}, + ] +} + + +# --------------------------------------------------------------------------- +# fetch_references +# --------------------------------------------------------------------------- + + +def test_fetch_references_returns_citations(monkeypatch: pytest.MonkeyPatch) -> None: + """fetch_references maps API response to Citation objects correctly.""" + + def mock_get(url: str, params: dict[str, object] | None = None) -> httpx.Response: + return httpx.Response(200, json=_SAMPLE_REFS_RESPONSE) + + monkeypatch.setattr(semanticscholar, "_get", mock_get) + + citations = semanticscholar.fetch_references("abc123") + + assert len(citations) == 2 + assert all(isinstance(c, Citation) for c in citations) + + # First citation uses DOI and has context + assert citations[0].citing_id == "abc123" + assert citations[0].cited_id == "10.1000/xyz1" + assert citations[0].context == "This approach was first described in [1]." + + # Second citation falls back to ArXiv ID, no context + assert citations[1].citing_id == "abc123" + assert citations[1].cited_id == "2301.07041" + assert citations[1].context is None + + +def test_fetch_references_404_returns_empty(monkeypatch: pytest.MonkeyPatch) -> None: + """A 404 response should return an empty list.""" + + def mock_get(url: str, params: dict[str, object] | None = None) -> httpx.Response: + raise httpx.HTTPStatusError( + "Not Found", + request=httpx.Request("GET", url), + response=httpx.Response(404, request=httpx.Request("GET", url)), + ) + + monkeypatch.setattr(semanticscholar, "_get", mock_get) + + result = semanticscholar.fetch_references("nonexistent") + assert result == [] + + +# --------------------------------------------------------------------------- +# fetch_recommendations +# --------------------------------------------------------------------------- + + +def test_fetch_recommendations_returns_list_of_str( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """fetch_recommendations returns a flat list of paperId strings.""" + + def mock_get(url: str, params: dict[str, object] | None = None) -> httpx.Response: + return httpx.Response(200, json=_SAMPLE_RECS_RESPONSE) + + monkeypatch.setattr(semanticscholar, "_get", mock_get) + + result = semanticscholar.fetch_recommendations("abc123", limit=3) + + assert isinstance(result, list) + assert len(result) == 3 + assert all(isinstance(r, str) for r in result) + assert result == ["rec_id_1", "rec_id_2", "rec_id_3"] + + +def test_fetch_recommendations_404_returns_empty( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A 404 response should return an empty list.""" + + def mock_get(url: str, params: dict[str, object] | None = None) -> httpx.Response: + raise httpx.HTTPStatusError( + "Not Found", + request=httpx.Request("GET", url), + response=httpx.Response(404, request=httpx.Request("GET", url)), + ) + + monkeypatch.setattr(semanticscholar, "_get", mock_get) + + result = semanticscholar.fetch_recommendations("nonexistent") + assert result == []