feat(sources): OpenAlex, SemanticScholar, arXiv API clients

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Tarik Moussa
2026-06-04 23:31:15 +02:00
parent b52a6a7412
commit b7f19b6e2c
8 changed files with 750 additions and 0 deletions

100
codex/sources/arxiv.py Normal file
View File

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

136
codex/sources/openalex.py Normal file
View File

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

View File

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

0
tests/__init__.py Normal file
View File

View File

140
tests/sources/test_arxiv.py Normal file
View File

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

View File

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

View File

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