From c6a428d335a26edad47ca51e4d40531461f540aa Mon Sep 17 00:00:00 2001 From: Tarik Moussa Date: Thu, 4 Jun 2026 23:30:58 +0200 Subject: [PATCH] feat(parsing): LaTeX, Nougat, GROBID parsers with chunking Co-Authored-By: Claude Sonnet 4.6 --- codex/config.py | 8 +- codex/parsing/grobid.py | 131 +++++++++++++++++++++++++++++ codex/parsing/nougat.py | 54 ++++++++++++ codex/parsing/tex.py | 159 +++++++++++++++++++++++++++++++++++ tests/parsing/__init__.py | 0 tests/parsing/test_grobid.py | 139 ++++++++++++++++++++++++++++++ tests/parsing/test_nougat.py | 70 +++++++++++++++ tests/parsing/test_tex.py | 112 ++++++++++++++++++++++++ 8 files changed, 672 insertions(+), 1 deletion(-) create mode 100644 codex/parsing/grobid.py create mode 100644 codex/parsing/nougat.py create mode 100644 codex/parsing/tex.py create mode 100644 tests/parsing/__init__.py create mode 100644 tests/parsing/test_grobid.py create mode 100644 tests/parsing/test_nougat.py create mode 100644 tests/parsing/test_tex.py diff --git a/codex/config.py b/codex/config.py index 2e02ccc..a6f8bd7 100644 --- a/codex/config.py +++ b/codex/config.py @@ -9,7 +9,7 @@ from __future__ import annotations from functools import lru_cache -from pydantic import Field +from pydantic import AliasChoices, Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -42,6 +42,12 @@ class Settings(BaseSettings): description="Base URL of the GROBID HTTP API (containerised).", ) + nougat_url: str = Field( + default="http://localhost:8080", + validation_alias=AliasChoices("NOUGAT_URL", "nougat_url"), + description="Base URL of the Nougat OCR HTTP API (containerised).", + ) + ollama_base_url: str = Field( default="http://localhost:11434", description="Base URL of the local Ollama endpoint (optional Q&A layer).", diff --git a/codex/parsing/grobid.py b/codex/parsing/grobid.py new file mode 100644 index 0000000..03bb186 --- /dev/null +++ b/codex/parsing/grobid.py @@ -0,0 +1,131 @@ +"""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, +) -> 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``. + + 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=60.0) 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 + arxiv_el = bib.find(f".//{_TEI_NS}idno[@type='arxiv']") + arxiv_id = _text(arxiv_el) + + 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 diff --git a/codex/parsing/nougat.py b/codex/parsing/nougat.py new file mode 100644 index 0000000..48021a2 --- /dev/null +++ b/codex/parsing/nougat.py @@ -0,0 +1,54 @@ +"""Nougat OCR integration. + +Converts PDF files to Mathpix Markdown (mmd) by calling a self-hosted +Nougat HTTP server. No network fetching of papers happens here — the +caller is expected to pass a local PDF path. +""" + +from __future__ import annotations + +import httpx +import tenacity + +from codex.config import Settings + + +def pdf_to_markdown(pdf_path: str, nougat_url: str | None = None) -> str: + """Convert a local PDF to Mathpix Markdown via the Nougat HTTP API. + + Parameters + ---------- + pdf_path: + Absolute (or relative) path to the PDF file on disk. + nougat_url: + Base URL of the Nougat server. Defaults to ``Settings().nougat_url``. + + Returns + ------- + str + The raw ``.mmd`` text returned by the server. + + Raises + ------ + httpx.HTTPStatusError + If the server returns a non-2xx status code. + """ + if nougat_url is None: + nougat_url = Settings().nougat_url + + @tenacity.retry( + retry=tenacity.retry_if_exception_type(httpx.ConnectError), + stop=tenacity.stop_after_attempt(3), # 1 original + 2 retries + wait=tenacity.wait_fixed(0), + reraise=True, + ) + def _post() -> str: + with open(pdf_path, "rb") as fh, httpx.Client(timeout=120.0) as client: + response = client.post( + f"{nougat_url}/predict", + files={"file": (pdf_path, fh, "application/pdf")}, + ) + response.raise_for_status() + return response.text + + return _post() diff --git a/codex/parsing/tex.py b/codex/parsing/tex.py new file mode 100644 index 0000000..43514d6 --- /dev/null +++ b/codex/parsing/tex.py @@ -0,0 +1,159 @@ +"""LaTeX parsing utilities. + +Provides helpers to extract sections, chunk text, and convert LaTeX to plain +readable prose by stripping markup that is not useful for NLP/embedding. +""" + +from __future__ import annotations + +import re + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + +_SECTION_RE = re.compile( + r"\\(?:sub)*section\*?\s*\{([^}]*)\}", + re.DOTALL, +) + +# Patterns for LaTeX noise removal (applied in order). +_COMMENT_RE = re.compile(r"%[^\n]*") +_CITE_RE = re.compile(r"\\cite\*?\{[^}]*\}") +_LABEL_RE = re.compile(r"\\label\{[^}]*\}") +_REF_RE = re.compile(r"\\ref\{[^}]*\}") +# Display math: $$...$$ (before single $ to avoid greedy mismatch) +_DISPLAY_DOLLAR_RE = re.compile(r"\$\$.*?\$\$", re.DOTALL) +# Inline math: $...$ +_INLINE_MATH_RE = re.compile(r"\$[^$\n]*?\$", re.DOTALL) +# \begin{equation}...\end{equation} +_ENV_EQUATION_RE = re.compile( + r"\\begin\{equation\*?\}.*?\\end\{equation\*?\}", + re.DOTALL, +) +# \begin{align}...\end{align} (covers align, align*, aligned, …) +_ENV_ALIGN_RE = re.compile( + r"\\begin\{align[^}]*\}.*?\\end\{align[^}]*\}", + re.DOTALL, +) +# Collapse excess whitespace +_WHITESPACE_RE = re.compile(r"[ \t]+") +_BLANK_LINES_RE = re.compile(r"\n{3,}") + + +def _clean_latex(text: str) -> str: + """Strip LaTeX markup and return readable prose.""" + text = _COMMENT_RE.sub("", text) + text = _ENV_EQUATION_RE.sub("", text) + text = _ENV_ALIGN_RE.sub("", text) + text = _DISPLAY_DOLLAR_RE.sub("", text) + text = _INLINE_MATH_RE.sub("", text) + text = _CITE_RE.sub("", text) + text = _LABEL_RE.sub("", text) + text = _REF_RE.sub("", text) + text = _WHITESPACE_RE.sub(" ", text) + text = _BLANK_LINES_RE.sub("\n\n", text) + return text.strip() + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def extract_sections(latex: str) -> list[tuple[str, str]]: + """Split *latex* on \\section / \\subsection boundaries. + + Returns a list of ``(title, cleaned_text)`` tuples, one per section. + Text between the document start and the first section command is + discarded (preamble / abstract handling is out of scope). + """ + # Find all section positions + matches = list(_SECTION_RE.finditer(latex)) + if not matches: + return [] + + sections: list[tuple[str, str]] = [] + for idx, match in enumerate(matches): + title = match.group(1).strip() + body_start = match.end() + body_end = matches[idx + 1].start() if idx + 1 < len(matches) else len(latex) + body = latex[body_start:body_end] + cleaned = _clean_latex(body) + sections.append((title, cleaned)) + + return sections + + +def chunk_text(text: str, size: int = 512, overlap: int = 64) -> list[str]: + """Split *text* into overlapping word-based chunks. + + Parameters + ---------- + text: + Plain text to chunk (not raw LaTeX). + size: + Target number of words per chunk. + overlap: + Number of words to carry over into the next chunk. + + Each chunk is at most ``size + overlap`` words. Boundaries are snapped + to the nearest ``". "`` (sentence end) within ±20 words of the nominal + boundary when possible. + """ + words = text.split() + if not words: + return [] + + snap_window = 20 + chunks: list[str] = [] + start = 0 + + while start < len(words): + end = min(start + size, len(words)) + + # Try to snap *end* to a sentence boundary within ±snap_window words. + if end < len(words): + best = end + # Build a small search window + lo = max(start + 1, end - snap_window) + hi = min(len(words), end + snap_window + 1) + # Prefer the closest sentence-end ". " to *end* + for offset in range(0, snap_window + 1): + for candidate in (end - offset, end + offset): + if ( + lo <= candidate < hi + and candidate > start + and words[candidate - 1].endswith(".") + ): + # Sentence boundary: word at (candidate-1) ends with ".". + best = candidate + break + if best != end: + break + end = best + + chunk_words = words[start:end] + chunks.append(" ".join(chunk_words)) + + if end >= len(words): + break + + # Next chunk starts *overlap* words before *end*. + start = max(start + 1, end - overlap) + + return chunks + + +def latex_to_text(latex: str) -> str: + """Convert a LaTeX document to plain text. + + Extracts all sections, cleans each one, and joins them with a blank line. + If no section commands are found, the whole document is cleaned and + returned as a single block. + """ + sections = extract_sections(latex) + if sections: + return "\n\n".join(body for _, body in sections) + # Fallback: no section structure — clean the whole string. + return _clean_latex(latex) diff --git a/tests/parsing/__init__.py b/tests/parsing/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/parsing/test_grobid.py b/tests/parsing/test_grobid.py new file mode 100644 index 0000000..eaa7dfd --- /dev/null +++ b/tests/parsing/test_grobid.py @@ -0,0 +1,139 @@ +"""Tests for codex.parsing.grobid.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +from codex.parsing.grobid import extract_references + +# Minimal valid TEI XML with 2 biblStruct elements +_TEI_XML = """\ + + + + +
+ + + + Deep Learning for NLP + + + Jane + Doe + + + + + + + + + 10.1234/nlp.2021 + 2101.00001 + + + + Attention Is All You Need + + + Ashish + Vaswani + + + + + Noam + Shazeer + + + + + + + + + 10.5678/attention + + +
+ + + +""" + + +class _FakeResponse: + text = _TEI_XML + status_code = 200 + + def raise_for_status(self) -> None: + pass + + +class TestExtractReferences: + def test_returns_two_dicts(self, tmp_path: pytest.TempdirFactory) -> None: + pdf = tmp_path / "paper.pdf" # type: ignore[operator] + pdf.write_bytes(b"%PDF fake") + + with patch("httpx.Client") as mock_client_cls: + mock_client = MagicMock() + mock_client_cls.return_value.__enter__.return_value = mock_client + mock_client.post.return_value = _FakeResponse() + + refs = extract_references(str(pdf), grobid_url="http://localhost:8070") + + assert len(refs) == 2 + + def test_all_keys_present(self, tmp_path: pytest.TempdirFactory) -> None: + pdf = tmp_path / "paper.pdf" # type: ignore[operator] + pdf.write_bytes(b"%PDF fake") + + with patch("httpx.Client") as mock_client_cls: + mock_client = MagicMock() + mock_client_cls.return_value.__enter__.return_value = mock_client + mock_client.post.return_value = _FakeResponse() + + refs = extract_references(str(pdf), grobid_url="http://localhost:8070") + + required_keys = {"title", "authors", "year", "doi", "arxiv_id"} + for ref in refs: + assert required_keys == set(ref.keys()), f"Missing keys in {ref}" + + def test_first_ref_values(self, tmp_path: pytest.TempdirFactory) -> None: + pdf = tmp_path / "paper.pdf" # type: ignore[operator] + pdf.write_bytes(b"%PDF fake") + + with patch("httpx.Client") as mock_client_cls: + mock_client = MagicMock() + mock_client_cls.return_value.__enter__.return_value = mock_client + mock_client.post.return_value = _FakeResponse() + + refs = extract_references(str(pdf), grobid_url="http://localhost:8070") + + first = refs[0] + assert first["title"] == "Deep Learning for NLP" + assert "Jane Doe" in first["authors"] + assert first["year"] == "2021" + assert first["doi"] == "10.1234/nlp.2021" + assert first["arxiv_id"] == "2101.00001" + + def test_second_ref_missing_arxiv_is_empty_string( + self, tmp_path: pytest.TempdirFactory + ) -> None: + pdf = tmp_path / "paper.pdf" # type: ignore[operator] + pdf.write_bytes(b"%PDF fake") + + with patch("httpx.Client") as mock_client_cls: + mock_client = MagicMock() + mock_client_cls.return_value.__enter__.return_value = mock_client + mock_client.post.return_value = _FakeResponse() + + refs = extract_references(str(pdf), grobid_url="http://localhost:8070") + + second = refs[1] + assert second["arxiv_id"] == "" + assert second["title"] == "Attention Is All You Need" + assert second["year"] == "2017" diff --git a/tests/parsing/test_nougat.py b/tests/parsing/test_nougat.py new file mode 100644 index 0000000..ead894c --- /dev/null +++ b/tests/parsing/test_nougat.py @@ -0,0 +1,70 @@ +"""Tests for codex.parsing.nougat.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +from codex.parsing.nougat import pdf_to_markdown + + +class _FakeResponse: + """Minimal stand-in for httpx.Response.""" + + def __init__(self, text: str, status_code: int = 200) -> None: + self.text = text + self.status_code = status_code + + def raise_for_status(self) -> None: + if self.status_code >= 400: + raise httpx.HTTPStatusError( + "error", + request=MagicMock(), + response=MagicMock(status_code=self.status_code), + ) + + +class TestPdfToMarkdown: + def test_success_returns_response_text(self, tmp_path: pytest.TempdirFactory) -> None: + pdf = tmp_path / "paper.pdf" # type: ignore[operator] + pdf.write_bytes(b"%PDF fake") + + expected = "# Title\nContent" + fake_response = _FakeResponse(expected) + + with patch("httpx.Client") as mock_client_cls: + mock_client = MagicMock() + mock_client_cls.return_value.__enter__.return_value = mock_client + mock_client.post.return_value = fake_response + + result = pdf_to_markdown(str(pdf), nougat_url="http://localhost:8080") + + assert result == expected + mock_client.post.assert_called_once() + call_kwargs = mock_client.post.call_args + assert "/predict" in call_kwargs[0][0] + + def test_connect_error_retried_exactly_twice(self, tmp_path: pytest.TempdirFactory) -> None: + """ConnectError must trigger 2 retries (3 total attempts).""" + pdf = tmp_path / "paper.pdf" # type: ignore[operator] + pdf.write_bytes(b"%PDF fake") + + call_count = 0 + + def _raising_post(*args: object, **kwargs: object) -> _FakeResponse: + nonlocal call_count + call_count += 1 + raise httpx.ConnectError("refused") + + with patch("httpx.Client") as mock_client_cls: + mock_client = MagicMock() + mock_client_cls.return_value.__enter__.return_value = mock_client + mock_client.post.side_effect = _raising_post + + with pytest.raises(httpx.ConnectError): + pdf_to_markdown(str(pdf), nougat_url="http://localhost:8080") + + # tenacity: stop_after_attempt(3) → 1 original + 2 retries = 3 calls + assert call_count == 3 diff --git a/tests/parsing/test_tex.py b/tests/parsing/test_tex.py new file mode 100644 index 0000000..002cc23 --- /dev/null +++ b/tests/parsing/test_tex.py @@ -0,0 +1,112 @@ +"""Tests for codex.parsing.tex.""" + +from __future__ import annotations + +import pytest + +from codex.parsing.tex import chunk_text, extract_sections, latex_to_text + + +class TestExtractSections: + def test_two_sections_returned(self) -> None: + latex = r"\section{Intro}" + "\nHello \\cite{foo} world\n" + latex += r"\section{Method}" + "\nResult" + sections = extract_sections(latex) + assert len(sections) == 2 + + def test_section_titles(self) -> None: + latex = r"\section{Intro}" + "\nHello\n" + r"\section{Method}" + "\nResult" + sections = extract_sections(latex) + assert sections[0][0] == "Intro" + assert sections[1][0] == "Method" + + def test_cite_removed_from_text(self) -> None: + latex = r"\section{Intro}" + "\nHello \\cite{foo} world\n" + latex += r"\section{Method}" + "\nResult" + sections = extract_sections(latex) + body = sections[0][1] + assert r"\cite" not in body + assert "foo" not in body + + def test_no_sections_returns_empty(self) -> None: + assert extract_sections("No sections here at all.") == [] + + def test_subsection_recognised(self) -> None: + latex = r"\subsection{Background}" + "\nSome text." + sections = extract_sections(latex) + assert len(sections) == 1 + assert sections[0][0] == "Background" + + def test_comments_stripped(self) -> None: + latex = r"\section{A}" + "\nHello % this is a comment\nWorld" + sections = extract_sections(latex) + assert "%" not in sections[0][1] + assert "comment" not in sections[0][1] + + def test_math_stripped(self) -> None: + latex = r"\section{A}" + "\nLet $x = 1$ and $$y = 2$$ be numbers." + sections = extract_sections(latex) + body = sections[0][1] + assert "$" not in body + + +class TestChunkText: + @pytest.fixture() + def long_text(self) -> str: + # Build a 600-word string with real sentences + sentence = "The quick brown fox jumps over the lazy dog near the river. " + words_per_sentence = len(sentence.split()) + repeats = (600 // words_per_sentence) + 1 + return (sentence * repeats).strip() + + def test_at_least_two_chunks(self, long_text: str) -> None: + chunks = chunk_text(long_text, size=512, overlap=64) + assert len(chunks) >= 2 + + def test_each_chunk_within_size_bound(self, long_text: str) -> None: + size, overlap = 512, 64 + chunks = chunk_text(long_text, size=size, overlap=overlap) + for chunk in chunks: + word_count = len(chunk.split()) + assert word_count <= size + overlap, ( + f"Chunk has {word_count} words, expected <= {size + overlap}" + ) + + def test_empty_text_returns_empty_list(self) -> None: + assert chunk_text("") == [] + + def test_short_text_returns_single_chunk(self) -> None: + chunks = chunk_text("Hello world.", size=512, overlap=64) + assert len(chunks) == 1 + + def test_chunks_cover_all_content(self) -> None: + # First chunk must start with first word, last chunk must end with last word + text = " ".join(f"word{i}" for i in range(600)) + chunks = chunk_text(text, size=200, overlap=50) + assert chunks[0].startswith("word0") + assert chunks[-1].endswith("word599") + + +class TestLatexToText: + def test_no_percent_in_output(self) -> None: + latex = r"\section{A}" + "\nSome % comment\ntext." + result = latex_to_text(latex) + assert "%" not in result + + def test_no_cite_in_output(self) -> None: + latex = r"\section{A}" + "\nHello \\cite{bar} world." + result = latex_to_text(latex) + assert r"\cite" not in result + assert "bar" not in result + + def test_returns_string(self) -> None: + latex = r"\section{A}" + "\nSimple text." + result = latex_to_text(latex) + assert isinstance(result, str) + assert len(result) > 0 + + def test_no_section_fallback(self) -> None: + latex = r"Some \cite{foo} text with % comment" + result = latex_to_text(latex) + assert r"\cite" not in result + assert "%" not in result