codex-py/codex/sources/openalex.py

"""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 _resolve_id(id: str) -> str:
    """Resolve a bare DOI or arXiv ID for the OpenAlex /works/ endpoint.

    OpenAlex resolves DOIs, W-IDs and full https:// URLs in the /works/ path,
    but NOT the ``arxiv:`` namespace (that path 404s). arXiv papers are looked
    up via their arXiv DOI ``10.48550/arXiv.<id>`` instead — accepted by
    OpenAlex for both modern (``2301.07041``) and legacy (``math/0603097``) IDs.
    """
    s = id.strip()
    if s.lower().startswith("arxiv:"):
        s = s[len("arxiv:") :]
    if s.startswith(("W", "https://", "doi:")):
        return s
    if s.startswith("10."):
        return f"doi:{s}"
    return f"doi:10.48550/arXiv.{s}"


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"``), 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/{_resolve_id(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.

    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 ``"W…"``).

    Returns
    -------
    list[Citation]
        One Citation per referenced work (OpenAlex W-ID as cited_id).
    """
    url = f"{_BASE}/works/{_resolve_id(openalex_id)}"
    try:
        response = _get(url)
    except httpx.HTTPStatusError as exc:
        if exc.response.status_code == 404:
            return []
        raise

    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
    ]