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..3b656e0 --- /dev/null +++ b/in2lambda/wizard/mathpix.py @@ -0,0 +1,136 @@ +"""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. + +The PDF is uploaded to Mathpix, a third-party OCR service, for processing. +Instructors converting student work should be told their PDFs leave the +local machine. Mathpix also offers an opt-out from using submitted data to +improve its models; see https://mathpix.com/privacy for how to enable it. +""" + +import os +import re +import time +import warnings +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, + timeout: float = 30.0, +) -> str: + """Convert ``pdf_path`` to markdown, downloading its figures under ``out_dir``. + + Args: + pdf_path: Path to the source PDF. + out_dir: Directory to write a ``media/`` folder of figures into. + poll_interval: Seconds to wait between Mathpix "is it ready yet" polls. + max_polls: How many times to poll before giving up. + timeout: Seconds to wait for each individual HTTP request. + + Returns: + The converted markdown, with figures saved in ``/media/`` and + referenced from the markdown as ``./media/``. The caller is + responsible for writing the markdown out wherever it belongs. + + Raises: + RuntimeError: if credentials are missing, Mathpix rejects the PDF or + fails to convert it, or the conversion 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}, + timeout=timeout, + ) + response.raise_for_status() + data = response.json() + if "error" in data: + raise RuntimeError(f"Mathpix rejected the PDF: {data['error']}") + pdf_id = data["pdf_id"] + + markdown = _poll_for_markdown(pdf_id, headers, poll_interval, max_polls, timeout) + return _localise_figures(markdown, out, timeout) + + +def _poll_for_markdown( + pdf_id: str, + headers: dict, + poll_interval: float, + max_polls: int, + timeout: float, +) -> str: + """Poll Mathpix until ``pdf_id`` finishes converting, then return its markdown.""" + status_url = f"{MATHPIX_PDF_ENDPOINT}/{pdf_id}" + for _ in range(max_polls): + response = requests.get(status_url, headers=headers, timeout=timeout) + response.raise_for_status() + data = response.json() + status = data.get("status") + if status == "completed": + break + if status == "error": + raise RuntimeError( + f"Mathpix failed to convert {pdf_id}: {data.get('error', 'unknown error')}" + ) + time.sleep(poll_interval) + else: + raise RuntimeError(f"Mathpix did not finish converting {pdf_id} in time.") + + md_response = requests.get( + f"{MATHPIX_PDF_ENDPOINT}/{pdf_id}.md", headers=headers, timeout=timeout + ) + md_response.raise_for_status() + return md_response.text + + +def _localise_figures(markdown: str, out_dir: Path, timeout: float) -> 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, timeout=timeout) + if image.status_code != 200: + warnings.warn( + f"Mathpix figure download failed for {url} " + f"(status {image.status_code}); markdown will reference a " + f"missing file: ./media/{local_name}" + ) + 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..d6a9468 --- /dev/null +++ b/tests/test_mathpix.py @@ -0,0 +1,137 @@ +"""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 _post(pdf_id="abc123", error=None): + post = MagicMock(status_code=200) + post.json.return_value = {"error": error} if error else {"pdf_id": pdf_id} + return post + + +def _status(status, error=None): + body = {"status": status} + if error: + body["error"] = error + resp = MagicMock(status_code=200) + resp.json.return_value = body + return resp + + +def test_pdf_to_markdown_returns_markdown_and_localises_figures(tmp_path): + pdf = _pdf(tmp_path) + out_dir = tmp_path / "out" + + completed = _status("completed") + 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 = [completed, md, image] + markdown = pdf_to_markdown(str(pdf), str(out_dir), poll_interval=0.0) + + assert "![pictureTag](./media/0_fig.png)" in markdown + assert not (out_dir / "paper.md").exists() + 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) + + processing = _status("processing") + completed = _status("completed") + md = 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 = [processing, processing, completed, md] + markdown = pdf_to_markdown( + str(pdf), str(tmp_path / "out"), poll_interval=0.0, max_polls=5 + ) + + assert markdown.startswith("# Only text") + + +def test_pdf_to_markdown_times_out(tmp_path): + pdf = _pdf(tmp_path) + + with patch("in2lambda.wizard.mathpix.requests") as req: + req.post.return_value = _post() + req.get.return_value = _status("processing") + 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_pdf_to_markdown_raises_on_rejected_upload(tmp_path): + pdf = _pdf(tmp_path) + + with patch("in2lambda.wizard.mathpix.requests") as req: + req.post.return_value = _post(error="Invalid file type") + with pytest.raises(RuntimeError, match="Mathpix rejected the PDF"): + pdf_to_markdown(str(pdf), str(tmp_path / "out")) + + +def test_pdf_to_markdown_raises_immediately_on_conversion_error(tmp_path): + pdf = _pdf(tmp_path) + + with patch("in2lambda.wizard.mathpix.requests") as req: + req.post.return_value = _post() + req.get.return_value = _status("error", error="conversion failed") + with pytest.raises(RuntimeError, match="conversion failed"): + pdf_to_markdown( + str(pdf), str(tmp_path / "out"), poll_interval=0.0, max_polls=60 + ) + + # Only the single status poll should have happened, not all 60. + assert req.get.call_count == 1 + + +def test_pdf_to_markdown_warns_on_failed_figure_download(tmp_path): + pdf = _pdf(tmp_path) + out_dir = tmp_path / "out" + + completed = _status("completed") + md = MagicMock( + status_code=200, + text="![](https://cdn.mathpix.com/x/fig.png) done\n", + ) + image = MagicMock(status_code=404, content=b"") + + with patch("in2lambda.wizard.mathpix.requests") as req: + req.post.return_value = _post() + req.get.side_effect = [completed, md, image] + with pytest.warns(UserWarning, match="figure download failed"): + markdown = pdf_to_markdown(str(pdf), str(out_dir), poll_interval=0.0) + + assert "https://cdn.mathpix.com/x/fig.png" in markdown + assert not (out_dir / "media" / "0_fig.png").exists() + + +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"))