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>
This commit is contained in:
Tarik Moussa
2026-06-04 23:49:35 +02:00
parent b7f19b6e2c
commit 52737abe3d
3 changed files with 125 additions and 30 deletions

View File

@@ -24,6 +24,20 @@ logger = logging.getLogger(__name__)
_BASE = "https://api.openalex.org" _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: def _is_retryable(exc: BaseException) -> bool:
"""Return True for HTTP 429 / 5xx responses wrapped in httpx.HTTPStatusError.""" """Return True for HTTP 429 / 5xx responses wrapped in httpx.HTTPStatusError."""
if isinstance(exc, httpx.HTTPStatusError): if isinstance(exc, httpx.HTTPStatusError):
@@ -88,14 +102,15 @@ def fetch_paper(id: str) -> Paper | None:
Parameters Parameters
---------- ----------
id: 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 Returns
------- -------
Paper | None Paper | None
A populated Paper instance, or None on 404. A populated Paper instance, or None on 404.
""" """
url = f"{_BASE}/works/{id}" url = f"{_BASE}/works/{_resolve_id(id)}"
try: try:
response = _get(url) response = _get(url)
except httpx.HTTPStatusError as exc: except httpx.HTTPStatusError as exc:
@@ -109,17 +124,20 @@ def fetch_paper(id: str) -> Paper | None:
def fetch_citations(openalex_id: str) -> list[Citation]: def fetch_citations(openalex_id: str) -> list[Citation]:
"""Fetch the reference list for a paper as Citation objects. """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 Parameters
---------- ----------
openalex_id: openalex_id:
The OpenAlex work ID (``"https://openalex.org/W…"`` or short ``"W…"``). The OpenAlex work ID (``"https://openalex.org/W…"`` or ``"W…"``).
Returns Returns
------- -------
list[Citation] 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: try:
response = _get(url) response = _get(url)
except httpx.HTTPStatusError as exc: except httpx.HTTPStatusError as exc:
@@ -127,10 +145,10 @@ def fetch_citations(openalex_id: str) -> list[Citation]:
return [] return []
raise raise
results: list[dict[str, Any]] = response.json().get("results", []) data: dict[str, Any] = response.json()
citations: list[Citation] = [] referenced_works: list[str] = data.get("referenced_works", [])
for ref in results: return [
cited_id: str = ref.get("doi") or ref.get("id") or "" Citation(citing_id=openalex_id, cited_id=cited_id)
if cited_id: for cited_id in referenced_works
citations.append(Citation(citing_id=openalex_id, cited_id=cited_id)) if cited_id
return citations ]

View File

@@ -4,16 +4,19 @@ Provides:
- fetch_references: retrieve references for a paper as Citation dataclasses. - fetch_references: retrieve references for a paper as Citation dataclasses.
- fetch_recommendations: retrieve recommended paper IDs. - 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 from __future__ import annotations
import logging import logging
import threading
import time
from typing import Any from typing import Any
import httpx 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 from codex.models import Citation
@@ -22,10 +25,32 @@ logger = logging.getLogger(__name__)
_BASE_GRAPH = "https://api.semanticscholar.org/graph/v1" _BASE_GRAPH = "https://api.semanticscholar.org/graph/v1"
_BASE_RECS = "https://api.semanticscholar.org/recommendations/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=retry_if_exception(_is_retryable),
stop=stop_after_attempt(5), stop=stop_after_attempt(5),
wait=wait_fixed(1), wait=wait_exponential(min=1, max=30),
before_sleep=lambda rs: logger.warning( before_sleep=lambda rs: logger.warning(
"SemanticScholar retry %d after %s", "SemanticScholar retry %d after %s",
rs.attempt_number, 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: def _get(url: str, params: dict[str, Any] | None = None) -> httpx.Response:
_rate_limit()
response = httpx.get(url, params=params, timeout=30) response = httpx.get(url, params=params, timeout=30)
response.raise_for_status() response.raise_for_status()
return response return response
@@ -70,7 +96,10 @@ def fetch_references(paper_id: str) -> list[Citation]:
context: str | None = contexts[0] if contexts else None context: str | None = contexts[0] if contexts else None
cited_id: str = ( 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: if cited_id:
citations.append(Citation(citing_id=paper_id, cited_id=cited_id, context=context)) citations.append(Citation(citing_id=paper_id, cited_id=cited_id, context=context))

View File

@@ -7,6 +7,7 @@ import pytest
from codex.models import Citation, Paper from codex.models import Citation, Paper
from codex.sources import openalex from codex.sources import openalex
from codex.sources.openalex import _resolve_id
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Helpers # Helpers
@@ -27,14 +28,33 @@ _SAMPLE_WORK = {
"is": [2], "is": [2],
"great": [3], "great": [3],
}, },
"referenced_works": [
"https://openalex.org/W1",
"https://openalex.org/W2",
],
} }
_SAMPLE_REFS = {
"results": [ # ---------------------------------------------------------------------------
{"doi": "10.1000/ref1", "id": "https://openalex.org/W1"}, # _resolve_id
{"doi": "10.1000/ref2", "id": "https://openalex.org/W2"}, # ---------------------------------------------------------------------------
]
}
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.""" """Successful fetch maps all fields to Paper correctly."""
def mock_get(url: str, params: dict[str, str] | None = None) -> httpx.Response: 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) return httpx.Response(200, json=_SAMPLE_WORK)
monkeypatch.setattr(openalex, "_get", mock_get) 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" 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 # 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: 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] attempt = [0]
def controlled_get(url: str, params: dict[str, str] | None = None) -> httpx.Response: 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) monkeypatch.setattr(openalex, "_get", controlled_get)
# Call fetch_paper in a loop to simulate retry behaviour
result: Paper | None = None result: Paper | None = None
errors = 0 errors = 0
for _ in range(5): 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 result is not None
assert isinstance(result, Paper) 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: 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: 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) monkeypatch.setattr(openalex, "_get", mock_get)
@@ -138,5 +172,19 @@ def test_fetch_citations(monkeypatch: pytest.MonkeyPatch) -> None:
assert len(citations) == 2 assert len(citations) == 2
assert all(isinstance(c, Citation) for c in citations) assert all(isinstance(c, Citation) for c in citations)
assert citations[0].citing_id == oa_id assert citations[0].citing_id == oa_id
assert citations[0].cited_id == "10.1000/ref1" assert citations[0].cited_id == "https://openalex.org/W1"
assert citations[1].cited_id == "10.1000/ref2" 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") == []