codex-py/codex/parsing/grobid.py

"""GROBID integration.

Extracts structured reference lists and full-text TEI XML from PDFs by
calling a self-hosted GROBID HTTP server.  No DB access, no embedding,
no network fetching of papers happens here.
"""

from __future__ import annotations

import xml.etree.ElementTree as ET

import httpx

from codex.config import Settings

_TEI_NS = "{http://www.tei-c.org/ns/1.0}"


def _text(element: ET.Element | None) -> str:
    """Return element text or empty string if element is None."""
    if element is None:
        return ""
    return (element.text or "").strip()


def extract_references(
    pdf_path: str,
    grobid_url: str | None = None,
    timeout: float = 60.0,
) -> list[dict[str, str]]:
    """Extract a structured reference list from a PDF via GROBID.

    Parameters
    ----------
    pdf_path:
        Path to the PDF file on disk.
    grobid_url:
        Base URL of the GROBID server.  Defaults to ``Settings().grobid_url``.
    timeout:
        HTTP client timeout in seconds.  Large PDFs (e.g. theses) can exceed the
        60s default, especially against a slower/emulated GROBID; raise it then.

    Returns
    -------
    list[dict[str, str]]
        One dict per reference.  All dicts contain the keys
        ``title``, ``authors``, ``year``, ``doi``, ``arxiv_id``
        (missing values are empty strings).
    """
    if grobid_url is None:
        grobid_url = Settings().grobid_url

    with open(pdf_path, "rb") as fh, httpx.Client(timeout=timeout) as client:
        response = client.post(
            f"{grobid_url}/api/processReferences",
            files={"input": (pdf_path, fh, "application/pdf")},
        )
    response.raise_for_status()

    root = ET.fromstring(response.text)
    results: list[dict[str, str]] = []

    for bib in root.iter(f"{_TEI_NS}biblStruct"):
        # Title
        title_el = bib.find(f".//{_TEI_NS}title[@level='a']")
        title = _text(title_el)

        # Authors: collect all persName elements
        authors_parts: list[str] = []
        for person in bib.iter(f"{_TEI_NS}persName"):
            forename_el = person.find(f"{_TEI_NS}forename")
            surname_el = person.find(f"{_TEI_NS}surname")
            forename = _text(forename_el)
            surname = _text(surname_el)
            full = " ".join(p for p in (forename, surname) if p)
            if full:
                authors_parts.append(full)
        authors = "; ".join(authors_parts)

        # Year
        date_el = bib.find(f".//{_TEI_NS}date[@type='published']")
        year = ""
        if date_el is not None:
            when = date_el.get("when", "")
            year = when[:4] if when else ""

        # DOI
        doi_el = bib.find(f".//{_TEI_NS}idno[@type='DOI']")
        doi = _text(doi_el)

        # arXiv ID — GROBID emits type="arXiv" (camel-case); match case-insensitively
        arxiv_id = ""
        for idno_el in bib.iter(f"{_TEI_NS}idno"):
            if idno_el.get("type", "").lower() == "arxiv":
                arxiv_id = (idno_el.text or "").strip()
                break

        results.append(
            {
                "title": title,
                "authors": authors,
                "year": year,
                "doi": doi,
                "arxiv_id": arxiv_id,
            }
        )

    return results


def extract_structure(
    pdf_path: str,
    grobid_url: str | None = None,
) -> str:
    """Extract full-text TEI XML from a PDF via GROBID.

    Parameters
    ----------
    pdf_path:
        Path to the PDF file on disk.
    grobid_url:
        Base URL of the GROBID server.  Defaults to ``Settings().grobid_url``.

    Returns
    -------
    str
        Raw TEI XML response text from GROBID.
    """
    if grobid_url is None:
        grobid_url = Settings().grobid_url

    with open(pdf_path, "rb") as fh, httpx.Client(timeout=120.0) as client:
        response = client.post(
            f"{grobid_url}/api/processFulltextDocument",
            files={"input": (pdf_path, fh, "application/pdf")},
        )
    response.raise_for_status()
    return response.text