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

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