feat: merge F-03 parsing layer (LaTeX, Nougat, GROBID)
Review-Gate: APPROVE (Opus, cold, 2 passes) 23 tests, ruff+mypy clean. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -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).",
|
||||
|
||||
134
codex/parsing/grobid.py
Normal file
134
codex/parsing/grobid.py
Normal file
@@ -0,0 +1,134 @@
|
||||
"""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 — 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
|
||||
54
codex/parsing/nougat.py
Normal file
54
codex/parsing/nougat.py
Normal file
@@ -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()
|
||||
159
codex/parsing/tex.py
Normal file
159
codex/parsing/tex.py
Normal file
@@ -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)
|
||||
17
conftest.py
Normal file
17
conftest.py
Normal file
@@ -0,0 +1,17 @@
|
||||
"""Root conftest — ensures the project root is on sys.path.
|
||||
|
||||
Required when the virtual-environment Python is symlinked to an external
|
||||
interpreter (e.g. Anaconda) that does not process the venv's .pth files
|
||||
automatically. Adding the project root here makes ``codex`` importable
|
||||
regardless of how pytest is invoked (``pytest``, ``python -m pytest``, etc.).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Insert the project root so `import codex` always resolves.
|
||||
_project_root = str(Path(__file__).parent)
|
||||
if _project_root not in sys.path:
|
||||
sys.path.insert(0, _project_root)
|
||||
0
tests/parsing/__init__.py
Normal file
0
tests/parsing/__init__.py
Normal file
161
tests/parsing/test_grobid.py
Normal file
161
tests/parsing/test_grobid.py
Normal file
@@ -0,0 +1,161 @@
|
||||
"""Tests for codex.parsing.grobid."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from codex.parsing.grobid import extract_references, extract_structure
|
||||
|
||||
# Real GROBID TEI output uses type="arXiv" (camel-case), not "arxiv".
|
||||
_TEI_XML = """\
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<TEI xmlns="http://www.tei-c.org/ns/1.0">
|
||||
<text>
|
||||
<back>
|
||||
<div type="references">
|
||||
<listBibl>
|
||||
<biblStruct>
|
||||
<analytic>
|
||||
<title level="a">Deep Learning for NLP</title>
|
||||
<author>
|
||||
<persName>
|
||||
<forename>Jane</forename>
|
||||
<surname>Doe</surname>
|
||||
</persName>
|
||||
</author>
|
||||
</analytic>
|
||||
<monogr>
|
||||
<imprint>
|
||||
<date type="published" when="2021-06"/>
|
||||
</imprint>
|
||||
</monogr>
|
||||
<idno type="DOI">10.1234/nlp.2021</idno>
|
||||
<idno type="arXiv">2101.00001</idno>
|
||||
</biblStruct>
|
||||
<biblStruct>
|
||||
<analytic>
|
||||
<title level="a">Attention Is All You Need</title>
|
||||
<author>
|
||||
<persName>
|
||||
<forename>Ashish</forename>
|
||||
<surname>Vaswani</surname>
|
||||
</persName>
|
||||
</author>
|
||||
<author>
|
||||
<persName>
|
||||
<forename>Noam</forename>
|
||||
<surname>Shazeer</surname>
|
||||
</persName>
|
||||
</author>
|
||||
</analytic>
|
||||
<monogr>
|
||||
<imprint>
|
||||
<date type="published" when="2017"/>
|
||||
</imprint>
|
||||
</monogr>
|
||||
<idno type="DOI">10.5678/attention</idno>
|
||||
</biblStruct>
|
||||
</listBibl>
|
||||
</div>
|
||||
</back>
|
||||
</text>
|
||||
</TEI>
|
||||
"""
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
text = _TEI_XML
|
||||
status_code = 200
|
||||
|
||||
def raise_for_status(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _mock_post(xml: str = _TEI_XML) -> MagicMock:
|
||||
fake = MagicMock()
|
||||
fake.text = xml
|
||||
fake.status_code = 200
|
||||
fake.raise_for_status = MagicMock()
|
||||
return fake
|
||||
|
||||
|
||||
class TestExtractReferences:
|
||||
def test_returns_two_dicts(self, tmp_path: Path) -> None:
|
||||
pdf = tmp_path / "paper.pdf"
|
||||
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 = _mock_post()
|
||||
|
||||
refs = extract_references(str(pdf), grobid_url="http://localhost:8070")
|
||||
|
||||
assert len(refs) == 2
|
||||
|
||||
def test_all_keys_present(self, tmp_path: Path) -> None:
|
||||
pdf = tmp_path / "paper.pdf"
|
||||
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 = _mock_post()
|
||||
|
||||
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: Path) -> None:
|
||||
pdf = tmp_path / "paper.pdf"
|
||||
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 = _mock_post()
|
||||
|
||||
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: Path) -> None:
|
||||
pdf = tmp_path / "paper.pdf"
|
||||
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 = _mock_post()
|
||||
|
||||
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"
|
||||
|
||||
|
||||
class TestExtractStructure:
|
||||
def test_returns_xml_string(self, tmp_path: Path) -> None:
|
||||
pdf = tmp_path / "paper.pdf"
|
||||
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 = _mock_post("<TEI>raw xml</TEI>")
|
||||
|
||||
result = extract_structure(str(pdf), grobid_url="http://localhost:8070")
|
||||
|
||||
assert result == "<TEI>raw xml</TEI>"
|
||||
mock_client.post.assert_called_once()
|
||||
assert "processFulltextDocument" in mock_client.post.call_args[0][0]
|
||||
70
tests/parsing/test_nougat.py
Normal file
70
tests/parsing/test_nougat.py
Normal file
@@ -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
|
||||
112
tests/parsing/test_tex.py
Normal file
112
tests/parsing/test_tex.py
Normal file
@@ -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
|
||||
Reference in New Issue
Block a user