feat(sources): OpenAlex, SemanticScholar, arXiv API clients
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Tarik Moussa
100
codex/sources/arxiv.py
Normal file
100
codex/sources/arxiv.py
Normal 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
136
codex/sources/openalex.py
Normal 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
|
||||
106
codex/sources/semanticscholar.py
Normal file
106
codex/sources/semanticscholar.py
Normal 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")]
|
||||
Reference in New Issue
Block a user