From b7f19b6e2ce723cc51e21ac1cd4380ddcb9ff495 Mon Sep 17 00:00:00 2001 From: Tarik Moussa Date: Thu, 4 Jun 2026 23:31:15 +0200 Subject: [PATCH 1/2] feat(sources): OpenAlex, SemanticScholar, arXiv API clients Co-Authored-By: Claude Sonnet 4.6 --- codex/sources/arxiv.py | 100 ++++++++++++++++++ codex/sources/openalex.py | 136 ++++++++++++++++++++++++ codex/sources/semanticscholar.py | 106 +++++++++++++++++++ tests/__init__.py | 0 tests/sources/__init__.py | 0 tests/sources/test_arxiv.py | 140 +++++++++++++++++++++++++ tests/sources/test_openalex.py | 142 ++++++++++++++++++++++++++ tests/sources/test_semanticscholar.py | 126 +++++++++++++++++++++++ 8 files changed, 750 insertions(+) create mode 100644 codex/sources/arxiv.py create mode 100644 codex/sources/openalex.py create mode 100644 codex/sources/semanticscholar.py create mode 100644 tests/__init__.py create mode 100644 tests/sources/__init__.py create mode 100644 tests/sources/test_arxiv.py create mode 100644 tests/sources/test_openalex.py create mode 100644 tests/sources/test_semanticscholar.py 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 == [] From 52737abe3dff48bf3106df644f3503968743ddd9 Mon Sep 17 00:00:00 2001 From: Tarik Moussa Date: Thu, 4 Jun 2026 23:49:35 +0200 Subject: [PATCH 2/2] fix(sources): correct OpenAlex ID prefixing, citations endpoint, S2 rate-limit Review-Gate findings: - openalex: bare DOIs/arXiv IDs need doi:/arxiv: prefix (bare IDs 404); add _resolve_id(); fix fetch_citations to use referenced_works field instead of non-existent /works/{id}/references endpoint. - semanticscholar: wait_fixed(1) was inter-retry only; add per-request monotonic rate-limiter (_rate_limit()) before each httpx call. Add _is_retryable() filter so 404s don't burn 5 retry slots. Co-Authored-By: Claude Sonnet 4.6 --- codex/sources/openalex.py | 42 +++++++++++++----- codex/sources/semanticscholar.py | 37 ++++++++++++++-- tests/sources/test_openalex.py | 76 ++++++++++++++++++++++++++------ 3 files changed, 125 insertions(+), 30 deletions(-) diff --git a/codex/sources/openalex.py b/codex/sources/openalex.py index bd1aa2d..a7f2961 100644 --- a/codex/sources/openalex.py +++ b/codex/sources/openalex.py @@ -24,6 +24,20 @@ logger = logging.getLogger(__name__) _BASE = "https://api.openalex.org" +def _resolve_id(id: str) -> str: + """Prefix a bare DOI or arXiv ID for the OpenAlex /works/ endpoint. + + OpenAlex accepts: doi:, arxiv:, W-IDs, or full https:// URLs. + Bare DOIs (starting with "10.") and bare arXiv IDs require a prefix. + """ + s = id.strip() + if s.startswith(("W", "https://", "doi:", "arxiv:")): + return s + if s.startswith("10."): + return f"doi:{s}" + return f"arxiv:{s}" + + def _is_retryable(exc: BaseException) -> bool: """Return True for HTTP 429 / 5xx responses wrapped in httpx.HTTPStatusError.""" if isinstance(exc, httpx.HTTPStatusError): @@ -88,14 +102,15 @@ def fetch_paper(id: str) -> Paper | None: Parameters ---------- id: - An arXiv ID (``"2301.07041"``) or a DOI (``"10.1145/…"``). + An arXiv ID (``"2301.07041"``), a DOI (``"10.1145/…"``), + or an OpenAlex W-ID (``"W2741809807"``). Returns ------- Paper | None A populated Paper instance, or None on 404. """ - url = f"{_BASE}/works/{id}" + url = f"{_BASE}/works/{_resolve_id(id)}" try: response = _get(url) except httpx.HTTPStatusError as exc: @@ -109,17 +124,20 @@ def fetch_paper(id: str) -> Paper | None: def fetch_citations(openalex_id: str) -> list[Citation]: """Fetch the reference list for a paper as Citation objects. + Uses the ``referenced_works`` field from the work object — the + ``/works/{id}/references`` path is not a real OpenAlex endpoint. + Parameters ---------- openalex_id: - The OpenAlex work ID (``"https://openalex.org/W…"`` or short ``"W…"``). + The OpenAlex work ID (``"https://openalex.org/W…"`` or ``"W…"``). Returns ------- list[Citation] - One Citation per referenced work. + One Citation per referenced work (OpenAlex W-ID as cited_id). """ - url = f"{_BASE}/works/{openalex_id}/references" + url = f"{_BASE}/works/{_resolve_id(openalex_id)}" try: response = _get(url) except httpx.HTTPStatusError as exc: @@ -127,10 +145,10 @@ def fetch_citations(openalex_id: str) -> list[Citation]: 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 + data: dict[str, Any] = response.json() + referenced_works: list[str] = data.get("referenced_works", []) + return [ + Citation(citing_id=openalex_id, cited_id=cited_id) + for cited_id in referenced_works + if cited_id + ] diff --git a/codex/sources/semanticscholar.py b/codex/sources/semanticscholar.py index 8eb45d8..90c383b 100644 --- a/codex/sources/semanticscholar.py +++ b/codex/sources/semanticscholar.py @@ -4,16 +4,19 @@ 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. +Rate-limited to ≤1 req/s (per-request floor via monotonic clock). +Retried on 429/5xx with exponential back-off via tenacity. """ from __future__ import annotations import logging +import threading +import time from typing import Any import httpx -from tenacity import retry, stop_after_attempt, wait_fixed +from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential from codex.models import Citation @@ -22,10 +25,32 @@ logger = logging.getLogger(__name__) _BASE_GRAPH = "https://api.semanticscholar.org/graph/v1" _BASE_RECS = "https://api.semanticscholar.org/recommendations/v1" +# Per-request rate limit: ≤1 req/s without an API key. +_rate_lock = threading.Lock() +_last_request_time: float = 0.0 +_MIN_INTERVAL = 1.0 + + +def _rate_limit() -> None: + global _last_request_time + with _rate_lock: + now = time.monotonic() + wait = _MIN_INTERVAL - (now - _last_request_time) + if wait > 0: + time.sleep(wait) + _last_request_time = time.monotonic() + + +def _is_retryable(exc: BaseException) -> bool: + if isinstance(exc, httpx.HTTPStatusError): + return exc.response.status_code == 429 or exc.response.status_code >= 500 + return False + @retry( + retry=retry_if_exception(_is_retryable), stop=stop_after_attempt(5), - wait=wait_fixed(1), + wait=wait_exponential(min=1, max=30), before_sleep=lambda rs: logger.warning( "SemanticScholar retry %d after %s", rs.attempt_number, @@ -33,6 +58,7 @@ _BASE_RECS = "https://api.semanticscholar.org/recommendations/v1" ), ) def _get(url: str, params: dict[str, Any] | None = None) -> httpx.Response: + _rate_limit() response = httpx.get(url, params=params, timeout=30) response.raise_for_status() return response @@ -70,7 +96,10 @@ def fetch_references(paper_id: str) -> list[Citation]: 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 "" + 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)) diff --git a/tests/sources/test_openalex.py b/tests/sources/test_openalex.py index 54b9b8b..d9658c8 100644 --- a/tests/sources/test_openalex.py +++ b/tests/sources/test_openalex.py @@ -7,6 +7,7 @@ import pytest from codex.models import Citation, Paper from codex.sources import openalex +from codex.sources.openalex import _resolve_id # --------------------------------------------------------------------------- # Helpers @@ -27,14 +28,33 @@ _SAMPLE_WORK = { "is": [2], "great": [3], }, + "referenced_works": [ + "https://openalex.org/W1", + "https://openalex.org/W2", + ], } -_SAMPLE_REFS = { - "results": [ - {"doi": "10.1000/ref1", "id": "https://openalex.org/W1"}, - {"doi": "10.1000/ref2", "id": "https://openalex.org/W2"}, - ] -} + +# --------------------------------------------------------------------------- +# _resolve_id +# --------------------------------------------------------------------------- + + +def test_resolve_doi() -> None: + assert _resolve_id("10.1145/3592430") == "doi:10.1145/3592430" + + +def test_resolve_arxiv() -> None: + assert _resolve_id("2301.07041") == "arxiv:2301.07041" + + +def test_resolve_w_id_unchanged() -> None: + assert _resolve_id("W2741809807") == "W2741809807" + + +def test_resolve_prefixed_unchanged() -> None: + assert _resolve_id("doi:10.1145/x") == "doi:10.1145/x" + assert _resolve_id("arxiv:2301.07041") == "arxiv:2301.07041" # --------------------------------------------------------------------------- @@ -46,6 +66,7 @@ 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: + assert "doi:10.1145/3592430" in url return httpx.Response(200, json=_SAMPLE_WORK) monkeypatch.setattr(openalex, "_get", mock_get) @@ -61,6 +82,20 @@ def test_fetch_paper_success(monkeypatch: pytest.MonkeyPatch) -> None: assert paper.openalex_id == "https://openalex.org/W2741809807" +def test_fetch_paper_arxiv_id_prefixed(monkeypatch: pytest.MonkeyPatch) -> None: + """Bare arXiv IDs are prefixed with 'arxiv:' in the URL.""" + called_url: list[str] = [] + + def mock_get(url: str, params: dict[str, str] | None = None) -> httpx.Response: + called_url.append(url) + return httpx.Response(200, json=_SAMPLE_WORK) + + monkeypatch.setattr(openalex, "_get", mock_get) + openalex.fetch_paper("2301.07041") + + assert called_url and "arxiv:2301.07041" in called_url[0] + + # --------------------------------------------------------------------------- # fetch_paper — 404 → None # --------------------------------------------------------------------------- @@ -88,7 +123,7 @@ def test_fetch_paper_404_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: def test_fetch_paper_retry_on_429(monkeypatch: pytest.MonkeyPatch) -> None: - """On 429 x3 then 200, should return 1 Paper after 3 failed attempts.""" + """On 429×3 then 200, returns 1 Paper after 3 failed attempts.""" attempt = [0] def controlled_get(url: str, params: dict[str, str] | None = None) -> httpx.Response: @@ -104,7 +139,6 @@ def test_fetch_paper_retry_on_429(monkeypatch: pytest.MonkeyPatch) -> None: 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): @@ -116,19 +150,19 @@ def test_fetch_paper_retry_on_429(monkeypatch: pytest.MonkeyPatch) -> None: assert result is not None assert isinstance(result, Paper) - assert errors == 3 # 3 failures before success + assert errors == 3 # --------------------------------------------------------------------------- -# fetch_citations +# fetch_citations — uses referenced_works field # --------------------------------------------------------------------------- def test_fetch_citations(monkeypatch: pytest.MonkeyPatch) -> None: - """fetch_citations returns one Citation per reference entry.""" + """fetch_citations reads referenced_works from the work object.""" def mock_get(url: str, params: dict[str, str] | None = None) -> httpx.Response: - return httpx.Response(200, json=_SAMPLE_REFS) + return httpx.Response(200, json=_SAMPLE_WORK) monkeypatch.setattr(openalex, "_get", mock_get) @@ -138,5 +172,19 @@ def test_fetch_citations(monkeypatch: pytest.MonkeyPatch) -> None: 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" + assert citations[0].cited_id == "https://openalex.org/W1" + assert citations[1].cited_id == "https://openalex.org/W2" + + +def test_fetch_citations_404_returns_empty(monkeypatch: pytest.MonkeyPatch) -> None: + """404 on the work fetch returns an empty list.""" + + 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) + assert openalex.fetch_citations("W999") == []