feat: add Mathpix PDF extraction

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.
+"""

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 ``<out_dir>/media/`` and
+        referenced from the markdown as ``./media/<name>``. 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

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"))