fix(parsers): harden MinerU cloud response handling, timeout, md selection; drop redundant image rewrite (#77)

Author: KylinMountain

openkb/parsers/mineru.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
 import io
+import logging
 import os
-import re
 import time
 import zipfile
 from pathlib import Path
 from typing import Any
 
 from openkb.parsers.base import ParseResult, Parser
 
+logger = logging.getLogger(__name__)
+
 _SUPPORTED = {".pdf", ".docx", ".pptx", ".xlsx", ".xls", ".html", ".htm"}
 _CLOUD_BASE = "https://mineru.net/api/v4"
 
@@ -29,24 +31,32 @@ def _result_from_zip(zip_bytes: bytes) -> ParseResult:
     images: dict[str, bytes] = {}
     markdown = ""
     with zipfile.ZipFile(io.BytesIO(zip_bytes)) as zf:
-        md_names = [n for n in zf.namelist() if n.lower().endswith(".md")]
+        names = zf.namelist()
+        md_names = sorted(n for n in names if n.lower().endswith(".md"))
         if md_names:
-            chosen = next((n for n in md_names if n.endswith("full.md")), md_names[0])
+            chosen = next((n for n in md_names if Path(n).name == "full.md"), md_names[0])
             markdown = zf.read(chosen).decode("utf-8", errors="replace")
-        for name in zf.namelist():
+        for name in names:
             if name.lower().endswith((".png", ".jpg", ".jpeg", ".gif", ".webp")):
-                images[Path(name).name] = zf.read(name)
-    # Markdown references images as 'images/<file>'; localize_images matches on
-    # the bare filename, so rewrite 'images/fig.png' -> 'fig.png'.
-    for fname in images:
-        # Rewrite only `![alt](images/<fname>)` links (anchored on markdown image
-        # syntax) to the bare filename, for localize_images to canonicalize. A
-        # replacement function avoids regex-escape injection from arbitrary names.
-        pattern = re.compile(r"(!\[[^\]]*\]\()" + re.escape("images/" + fname) + r"(\))")
-        markdown = pattern.sub(lambda m, f=fname: m.group(1) + f + m.group(2), markdown)
+                base = Path(name).name
+                if base in images:
+                    logger.warning(
+                        "MinerU result has multiple images named %r in different "
+                        "folders; keeping the last. Earlier one may be lost.", base
+                    )
+                images[base] = zf.read(name)
     return ParseResult(markdown=markdown, images=images)
 
 
+def _mineru_body(resp):
+    """Return the 'data' dict from a MinerU v4 JSON response, raising on API errors."""
+    body = resp.json()
+    code = body.get("code")
+    if code not in (0, None):
+        raise RuntimeError(f"MinerU API error (code={code}): {body.get('msg')}")
+    return body.get("data") or {}
+
+
 class MineruParser(Parser):
     """MinerU via HTTP — self-hosted server or hosted cloud API."""
 
@@ -58,7 +68,8 @@ def __init__(self, opts: dict[str, Any] | None = None):
         self.base_url = self.opts.get("base_url")
         pi = self.opts.get("poll_interval", 3)
         self.poll_interval = pi if isinstance(pi, (int, float)) and pi > 0 else 3
-        self.timeout = self.opts.get("timeout", 600)
+        t = self.opts.get("timeout", 600)
+        self.timeout = t if isinstance(t, (int, float)) and t > 0 else 600
 
     def supports(self, suffix: str) -> bool:
         return suffix.lower() in _SUPPORTED
@@ -92,37 +103,43 @@ def _parse_cloud(self, src: Path) -> ParseResult:
             )
         httpx = _httpx()
         headers = {"Authorization": f"Bearer {api_key}"}
-        with httpx.Client(timeout=self.timeout) as client:
+        with httpx.Client(timeout=min(self.timeout, 120)) as client:
             r = client.post(
                 f"{_CLOUD_BASE}/file-urls/batch",
                 headers=headers,
                 json={"files": [{"name": src.name, "is_ocr": True}]},
             )
             r.raise_for_status()
-            data = r.json()["data"]
-            batch_id = data["batch_id"]
-            upload_url = data["file_urls"][0]
+            data = _mineru_body(r)
+            batch_id = data.get("batch_id")
+            file_urls = data.get("file_urls") or []
+            if not batch_id or not file_urls:
+                raise RuntimeError(f"MinerU returned no upload URL: {data}")
+            upload_url = file_urls[0]
             client.put(upload_url, content=src.read_bytes()).raise_for_status()
-            elapsed = 0
+            deadline = time.monotonic() + self.timeout
             zip_url = None
-            while elapsed < self.timeout:
+            while time.monotonic() < deadline:
                 pr = client.get(
                     f"{_CLOUD_BASE}/extract-results/batch/{batch_id}", headers=headers
                 )
                 pr.raise_for_status()
-                results = pr.json()["data"]["extract_result"]
+                data = _mineru_body(pr)
+                results = data.get("extract_result") or []
                 if not results:
                     time.sleep(self.poll_interval)
-                    elapsed += self.poll_interval
                     continue
                 state = results[0].get("state")
                 if state == "done":
-                    zip_url = results[0]["full_zip_url"]
+                    zip_url = results[0].get("full_zip_url")
+                    if not zip_url:
+                        raise RuntimeError(
+                            f"MinerU reported done but no full_zip_url: {results[0]}"
+                        )
                     break
                 if state == "failed":
                     raise RuntimeError(f"MinerU extraction failed: {results[0]}")
                 time.sleep(self.poll_interval)
-                elapsed += self.poll_interval
             if zip_url is None:
                 raise RuntimeError("MinerU extraction timed out.")
             zr = client.get(zip_url)

tests/test_parsers_mineru.py

@@ -63,9 +63,12 @@ def test_self_hosted_parses_zip(monkeypatch, tmp_path):
     assert isinstance(result, ParseResult)
     assert "Mineru" in result.markdown
     assert result.images["fig.png"] == b"PNGBYTES"
-    # the images/ prefix should be rewritten to the bare filename for localize_images
-    assert "images/fig.png" not in result.markdown
-    assert "![p](fig.png)" in result.markdown
+    # _result_from_zip no longer rewrites links; the raw 'images/fig.png' survives
+    assert "images/fig.png" in result.markdown
+    # localize_images (which now rewrites by basename) canonicalizes it
+    from openkb.images import localize_images
+    md2 = localize_images(result.markdown, result.images, "d", tmp_path / "imgs")
+    assert "sources/images/d/fig.png" in md2
 
 
 def test_cloud_flow_polls_then_downloads(monkeypatch, tmp_path):
@@ -74,7 +77,7 @@ def test_cloud_flow_polls_then_downloads(monkeypatch, tmp_path):
 
     buf = io.BytesIO()
     with zipfile.ZipFile(buf, "w") as zf:
-        zf.writestr("full.md", "# Cloud\n\n![p](images/fig.png)")
+        zf.writestr("full.md", "# Cloud")
         zf.writestr("images/fig.png", b"ZBYTES")
     zip_bytes = buf.getvalue()
 
@@ -124,8 +127,6 @@ def _get(url, *a, **k):
     assert isinstance(result, ParseResult)
     assert "Cloud" in result.markdown
     assert result.images["fig.png"] == b"ZBYTES"
-    assert "images/fig.png" not in result.markdown
-    assert "![p](fig.png)" in result.markdown
     # drove the full poll loop: running once, then done
     assert _get.calls == 2
 
@@ -137,19 +138,19 @@ def test_poll_interval_zero_is_clamped_to_positive():
     assert MineruParser({"poll_interval": 2}).poll_interval == 2
 
 
-def test_image_prefix_rewrite_is_anchored(tmp_path):
-    import io, sys, types, zipfile
-    from unittest.mock import MagicMock
-    # markdown has a real image link AND an unrelated 'images/fig.png' substring in prose
+def test_result_from_zip_does_not_rewrite_links(tmp_path):
+    import io, zipfile
+    # The images/ -> bare rewrite moved OUT of _result_from_zip into
+    # localize_images; _result_from_zip must leave the markdown link text intact.
     buf = io.BytesIO()
     with zipfile.ZipFile(buf, "w") as zf:
         zf.writestr("full.md", "See path other_images/fig.png in text.\n\n![p](images/fig.png)")
         zf.writestr("images/fig.png", b"PNG")
     from openkb.parsers.mineru import _result_from_zip
     result = _result_from_zip(buf.getvalue())
-    assert "![p](fig.png)" in result.markdown          # link rewritten
+    assert "![p](images/fig.png)" in result.markdown   # link text unchanged
     assert "other_images/fig.png" in result.markdown    # unrelated prose untouched
-    assert result.images["fig.png"] == b"PNG"
+    assert result.images["fig.png"] == b"PNG"           # images keyed by basename
 
 
 def test_cloud_empty_extract_result_then_done(monkeypatch, tmp_path):
@@ -186,3 +187,73 @@ def _get(url, *a, **k):
     src = tmp_path / "d.pdf"; src.write_bytes(b"%PDF")
     result = MineruParser({"mode": "cloud", "poll_interval": 1}).parse(src)
     assert "Ok" in result.markdown   # survived the empty-list poll without crashing
+
+
+def test_timeout_invalid_is_clamped():
+    from openkb.parsers.mineru import MineruParser
+    assert MineruParser({"timeout": 0}).timeout == 600
+    assert MineruParser({"timeout": "x"}).timeout == 600
+    assert MineruParser({"timeout": 30}).timeout == 30
+
+
+def test_cloud_api_error_envelope_raises(monkeypatch, tmp_path):
+    import sys, types
+    from unittest.mock import MagicMock
+    monkeypatch.setenv("MINERU_API_KEY", "key")
+    r = MagicMock(); r.raise_for_status = MagicMock()
+    r.json.return_value = {"code": -10001, "msg": "token expired", "data": None}
+    client = MagicMock()
+    client.__enter__ = MagicMock(return_value=client); client.__exit__ = MagicMock(return_value=False)
+    client.post.return_value = r
+    httpx_mod = types.ModuleType("httpx"); httpx_mod.Client = MagicMock(return_value=client)
+    monkeypatch.setitem(sys.modules, "httpx", httpx_mod)
+    from openkb.parsers.mineru import MineruParser
+    src = tmp_path / "d.pdf"; src.write_bytes(b"%PDF")
+    import pytest
+    with pytest.raises(RuntimeError) as exc:
+        MineruParser({"mode": "cloud"}).parse(src)
+    assert "token expired" in str(exc.value) or "-10001" in str(exc.value)
+
+
+def test_cloud_empty_file_urls_raises(monkeypatch, tmp_path):
+    import sys, types
+    from unittest.mock import MagicMock
+    monkeypatch.setenv("MINERU_API_KEY", "key")
+    r = MagicMock(); r.raise_for_status = MagicMock()
+    r.json.return_value = {"code": 0, "data": {"batch_id": "b1", "file_urls": []}}
+    client = MagicMock()
+    client.__enter__ = MagicMock(return_value=client); client.__exit__ = MagicMock(return_value=False)
+    client.post.return_value = r
+    httpx_mod = types.ModuleType("httpx"); httpx_mod.Client = MagicMock(return_value=client)
+    monkeypatch.setitem(sys.modules, "httpx", httpx_mod)
+    from openkb.parsers.mineru import MineruParser
+    src = tmp_path / "d.pdf"; src.write_bytes(b"%PDF")
+    import pytest
+    with pytest.raises(RuntimeError) as exc:
+        MineruParser({"mode": "cloud"}).parse(src)
+    assert "upload URL" in str(exc.value)
+
+
+def test_full_md_basename_preferred_over_endswith(tmp_path):
+    import io, zipfile
+    from openkb.parsers.mineru import _result_from_zip
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, "w") as zf:
+        zf.writestr("careful.md", "# WRONG")     # ends with 'full.md' but isn't it
+        zf.writestr("full.md", "# RIGHT")
+    result = _result_from_zip(buf.getvalue())
+    assert "RIGHT" in result.markdown
+    assert "WRONG" not in result.markdown
+
+
+def test_image_basename_collision_warns(tmp_path, caplog):
+    import io, zipfile, logging as _logging
+    from openkb.parsers.mineru import _result_from_zip
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, "w") as zf:
+        zf.writestr("full.md", "# x")
+        zf.writestr("images/fig.png", b"A")
+        zf.writestr("sub/fig.png", b"B")
+    with caplog.at_level(_logging.WARNING):
+        result = _result_from_zip(buf.getvalue())
+    assert any("fig.png" in r.message for r in caplog.records)