From 112ecab35699ae4c69d45b3b8132ef70bd3d12f4 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Mon, 31 Aug 2026 12:25:05 +0100 Subject: [PATCH] feat: add Mathpix PDF extraction in2lambda/wizard/mathpix.py: pdf_to_markdown() uploads a PDF to the Mathpix OCR API, polls for the rendered markdown, downloads any remote figures into /media/, and repoints the markdown at ./media/ so the Markdown filter's image resolution finds them. - Credentials from $MATHPIX_APP_ID / $MATHPIX_API_KEY; a missing pair raises a clear RuntimeError. - Only needs `requests` (already a core dep), so the module imports without the llm extra. - Ported and cleaned up from conversion2025/converter.py on Summer2025: print/exit calls become exceptions, the PIL round-trip is dropped (bytes are streamed straight to disk), poll interval/count are parameters. Tests mock all HTTP. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_017VXb8aZqgFBjoeuuddjW6r --- in2lambda/wizard/__init__.py | 6 +++ in2lambda/wizard/mathpix.py | 102 +++++++++++++++++++++++++++++++++++ tests/test_mathpix.py | 83 ++++++++++++++++++++++++++++ 3 files changed, 191 insertions(+) create mode 100644 in2lambda/wizard/__init__.py create mode 100644 in2lambda/wizard/mathpix.py create mode 100644 tests/test_mathpix.py diff --git a/in2lambda/wizard/__init__.py b/in2lambda/wizard/__init__.py new file mode 100644 index 0000000..10bdd25 --- /dev/null +++ b/in2lambda/wizard/__init__.py @@ -0,0 +1,6 @@ +"""Turn unstructured documents into the ``#``/``##`` markdown in2lambda understands. + +The pieces here (Mathpix OCR, LLM extraction) are driven by the +``in2lambda wizard`` command and need the optional ``llm`` extra plus API +credentials. +""" diff --git a/in2lambda/wizard/mathpix.py b/in2lambda/wizard/mathpix.py new file mode 100644 index 0000000..22e7199 --- /dev/null +++ b/in2lambda/wizard/mathpix.py @@ -0,0 +1,102 @@ +"""Convert a PDF into markdown with the Mathpix OCR API. + +Needs ``MATHPIX_APP_ID`` and ``MATHPIX_API_KEY`` in the environment (a ``.env`` +file is honoured by the wizard). Figures referenced by the returned markdown are +downloaded next to it so the ``Markdown`` filter can pick them up. +""" + +import os +import re +import time +from pathlib import Path + +import requests + +MATHPIX_PDF_ENDPOINT = "https://api.mathpix.com/v3/pdf" + +# Matches ``![alt](https://...)`` image references in Mathpix markdown. +_REMOTE_IMAGE = re.compile(r"!\[.*?\]\((https?://[^)]+)\)") + + +def _headers() -> dict: + """Return the Mathpix auth headers, or raise if credentials are missing.""" + app_id = os.getenv("MATHPIX_APP_ID") + app_key = os.getenv("MATHPIX_API_KEY") + if not app_id or not app_key: + raise RuntimeError( + "MATHPIX_APP_ID and MATHPIX_API_KEY must be set to convert PDFs " + "(see https://mathpix.com/ocr)." + ) + return {"app_id": app_id, "app_key": app_key} + + +def pdf_to_markdown( + pdf_path: str, + out_dir: str, + poll_interval: float = 5.0, + max_polls: int = 60, +) -> Path: + """Convert ``pdf_path`` to markdown, writing it and its figures under ``out_dir``. + + Args: + pdf_path: Path to the source PDF. + out_dir: Directory to write ``.md`` and a ``media/`` folder into. + poll_interval: Seconds to wait between Mathpix "is it ready yet" polls. + max_polls: How many times to poll before giving up. + + Returns: + The path to the written markdown file. Figures are saved in + ``/media/`` and referenced from the markdown as + ``./media/``. + + Raises: + RuntimeError: if credentials are missing or Mathpix does not finish in time. + """ + headers = _headers() + out = Path(out_dir) + (out / "media").mkdir(parents=True, exist_ok=True) + + with open(pdf_path, "rb") as pdf: + response = requests.post( + MATHPIX_PDF_ENDPOINT, headers=headers, files={"file": pdf} + ) + response.raise_for_status() + pdf_id = response.json()["pdf_id"] + + markdown = _poll_for_markdown(pdf_id, headers, poll_interval, max_polls) + markdown = _localise_figures(markdown, out) + + md_path = out / f"{Path(pdf_path).stem}.md" + md_path.write_text(markdown, encoding="utf-8") + return md_path + + +def _poll_for_markdown( + pdf_id: str, headers: dict, poll_interval: float, max_polls: int +) -> str: + """Poll Mathpix until the ``.md`` render of ``pdf_id`` is ready.""" + url = f"{MATHPIX_PDF_ENDPOINT}/{pdf_id}.md" + for _ in range(max_polls): + response = requests.get(url, headers=headers) + if response.status_code == 200: + return response.text + time.sleep(poll_interval) + raise RuntimeError(f"Mathpix did not finish converting {pdf_id} in time.") + + +def _localise_figures(markdown: str, out_dir: Path) -> str: + """Download remote figures into ``out_dir/media`` and repoint the markdown at them.""" + markdown = markdown.replace("![]", "![pictureTag]") + + for idx, url in enumerate(dict.fromkeys(_REMOTE_IMAGE.findall(markdown))): + basename = os.path.basename(url).split("?")[0] or f"figure_{idx}.png" + local_name = f"{idx}_{basename}" + + image = requests.get(url) + if image.status_code != 200: + continue + + (out_dir / "media" / local_name).write_bytes(image.content) + markdown = markdown.replace(url, f"./media/{local_name}") + + return markdown diff --git a/tests/test_mathpix.py b/tests/test_mathpix.py new file mode 100644 index 0000000..1f16529 --- /dev/null +++ b/tests/test_mathpix.py @@ -0,0 +1,83 @@ +"""Tests for the Mathpix PDF -> markdown helper. All HTTP is mocked.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from in2lambda.wizard.mathpix import pdf_to_markdown + + +@pytest.fixture(autouse=True) +def _mathpix_creds(monkeypatch): + monkeypatch.setenv("MATHPIX_APP_ID", "test-id") + monkeypatch.setenv("MATHPIX_API_KEY", "test-key") + + +def _pdf(tmp_path): + pdf = tmp_path / "paper.pdf" + pdf.write_bytes(b"%PDF-1.4 fake") + return pdf + + +def test_pdf_to_markdown_writes_md_and_localises_figures(tmp_path): + pdf = _pdf(tmp_path) + out_dir = tmp_path / "out" + + post = MagicMock(status_code=200) + post.json.return_value = {"pdf_id": "abc123"} + md = MagicMock( + status_code=200, + text="# Heading\n\n![](https://cdn.mathpix.com/x/fig.png?width=8) done\n", + ) + image = MagicMock(status_code=200, content=b"PNGBYTES") + + with patch("in2lambda.wizard.mathpix.requests") as req: + req.post.return_value = post + req.get.side_effect = [md, image] + md_path = pdf_to_markdown(str(pdf), str(out_dir), poll_interval=0.0) + + assert md_path == out_dir / "paper.md" + text = md_path.read_text() + assert "![pictureTag](./media/0_fig.png)" in text + assert (out_dir / "media" / "0_fig.png").read_bytes() == b"PNGBYTES" + + +def test_pdf_to_markdown_polls_until_ready(tmp_path): + pdf = _pdf(tmp_path) + + post = MagicMock(status_code=200) + post.json.return_value = {"pdf_id": "abc123"} + not_ready = MagicMock(status_code=202) + ready = MagicMock(status_code=200, text="# Only text, no figures\n") + + with patch("in2lambda.wizard.mathpix.requests") as req: + req.post.return_value = post + req.get.side_effect = [not_ready, not_ready, ready] + md_path = pdf_to_markdown( + str(pdf), str(tmp_path / "out"), poll_interval=0.0, max_polls=5 + ) + + assert md_path.read_text().startswith("# Only text") + + +def test_pdf_to_markdown_times_out(tmp_path): + pdf = _pdf(tmp_path) + + post = MagicMock(status_code=200) + post.json.return_value = {"pdf_id": "abc123"} + + with patch("in2lambda.wizard.mathpix.requests") as req: + req.post.return_value = post + req.get.return_value = MagicMock(status_code=202) + with pytest.raises(RuntimeError, match="did not finish"): + pdf_to_markdown( + str(pdf), str(tmp_path / "out"), poll_interval=0.0, max_polls=3 + ) + + +def test_missing_credentials_raise(tmp_path, monkeypatch): + monkeypatch.delenv("MATHPIX_APP_ID", raising=False) + monkeypatch.delenv("MATHPIX_API_KEY", raising=False) + + with pytest.raises(RuntimeError, match="MATHPIX_APP_ID"): + pdf_to_markdown(str(_pdf(tmp_path)), str(tmp_path / "out"))