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 <noreply@anthropic.com>
2026-06-04 23:49:35 +02:00
parent b7f19b6e2c
commit 52737abe3d
3 changed files with 125 additions and 30 deletions
--- 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:<doi>, arxiv:<id>, 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
+    ]
--- 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))
--- 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") == []