diff --git a/CHANGELOG.md b/CHANGELOG.md index 04a6868..ed0a2cc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Changed + +- **MinerU engine upgraded to MinerU 2.x** — the adapter now targets the current + [`mineru`](https://github.com/opendatalab/MinerU) package (formerly `magic-pdf`) + via its supported `mineru.cli.common.do_parse` API. The dependency extra changed + from `magic-pdf[full]>=0.9` to `mineru[core]>=2.0`. `MinerUEngine` gains + `backend` (`pipeline` default, or `vlm`) and `parse_method` constructor options; + the public engine name (`mineru`), `process()` signature, and `EngineResult` + shape are unchanged. **Breaking for installs:** reinstall with + `pip install -U docfold[mineru]` to pull `mineru` 2.x. + ### Added - **MarkItDown engine adapter** — wraps Microsoft's [`markitdown`](https://github.com/microsoft/markitdown) pure-Python library that converts Office files, PDFs, HTML, images, CSV/JSON/XML, ePub, audio, and ZIP archives into LLM-friendly Markdown. Added to the `benchmark.py` harness alongside the other local engines. Install: `pip install docfold[markitdown]`. diff --git a/README.md b/README.md index c623586..463fedf 100644 --- a/README.md +++ b/README.md @@ -324,7 +324,7 @@ Docfold builds on and integrates with these excellent projects: | Project | Description | |---------|-------------| | [Docling](https://github.com/docling-project/docling) | IBM's document conversion toolkit — PDF, DOCX, PPTX, and more | -| [MinerU / PDF-Extract-Kit](https://github.com/opendatalab/MinerU) | End-to-end PDF structuring with layout analysis and formula recognition | +| [MinerU](https://github.com/opendatalab/MinerU) (2.x) | End-to-end PDF structuring with layout analysis and formula recognition (pipeline + VLM backends) | | [Marker](https://github.com/VikParuchuri/marker) | High-quality PDF to Markdown converter | | [PyMuPDF](https://github.com/pymupdf/PyMuPDF) | Fast PDF/XPS/EPUB processing library | | [PaddleOCR](https://github.com/PaddlePaddle/PaddleOCR) | Multilingual OCR toolkit (80+ languages) | diff --git a/docs/tasks/MINERU_2X_UPGRADE.md b/docs/tasks/MINERU_2X_UPGRADE.md new file mode 100644 index 0000000..0c4af23 --- /dev/null +++ b/docs/tasks/MINERU_2X_UPGRADE.md @@ -0,0 +1,89 @@ +--- +purpose: "Upgrade the MinerU engine adapter from legacy magic-pdf 0.9 to MinerU 2.x" +status: "IMPLEMENTED" +priority: "P1" +created: "2026-06-27" +--- + +# Feature: MinerU 2.x Upgrade + +## Problem +The `MinerUEngine` adapter is built on the **legacy** `magic-pdf` package (`>=0.9`) +and its old import surface (`magic_pdf.data.dataset.PymuDocDataset`, +`magic_pdf.operators.models.doc_analyze`, `pipe_txt_mode`/`pipe_ocr_mode`, …). + +Upstream [opendatalab/MinerU](https://github.com/opendatalab/MinerU) has been +renamed and rewritten as **MinerU 2.x**: + +- The PyPI package is now `mineru` (not `magic-pdf`). +- The Python import root is `mineru` (not `magic_pdf`). +- The legacy `PymuDocDataset` / `doc_analyze` / `pipe_*_mode` API is gone. + The supported programmatic entry point is `mineru.cli.common.do_parse`. +- New backends are available: `pipeline` (CPU-friendly), `vlm` (VLM engine). + +So our integration targets a dead API. Installing `docfold[mineru]` today pulls +an unmaintained version. This task updates the adapter to MinerU 2.x while +preserving the public docfold API (engine name `"mineru"`, `process()` signature, +`EngineResult` shape). + +## Proposed Solution +Rewrite `mineru_engine.py` around `mineru.cli.common.do_parse`: + +1. Lazy-import `do_parse` and `read_fn` from `mineru.cli.common`. +2. In `process()`, read the file via `read_fn`, run `do_parse` in a thread + executor into a temp `output_dir`, then read back the generated + `{name}.md` / `{name}_content_list.json`. +3. Output subdirectory depends on backend (mirrors upstream `do_parse`): + - `pipeline` → `output_dir/{name}/{parse_method}` (parse_method defaults `auto`) + - `vlm` → `output_dir/{name}/vlm` +4. Map kwargs: `lang`, `start_page`/`end_page` → `start_page_id`/`end_page_id`, + `backend`, `parse_method`. Disable bbox-drawing dumps we don't consume. +5. `is_available()` checks `import mineru`. +6. Constructor gains `backend` (default `"pipeline"`) and keeps + `config_path`/`gpu` for backward compatibility. + +Update the dependency extra to `mineru[core]>=2.0` and refresh docs/changelog. + +## Affected Files +- `src/docfold/engines/mineru_engine.py` - rewrite adapter for MinerU 2.x API +- `pyproject.toml` - `mineru` extra: `magic-pdf[full]>=0.9` → `mineru[core]>=2.0` +- `tests/engines/test_adapters.py` - update `TestMinerUEngine` to new API/mocks +- `README.md` - note MinerU 2.x; install hint unchanged (`docfold[mineru]`) +- `CHANGELOG.md` - record the breaking dependency upgrade + +## Test Plan + +### Unit / Functional Tests +- [ ] `test_name` / `test_supported_extensions` unchanged (`mineru`, `{pdf}`) +- [ ] `test_is_available_when_missing` patches `mineru` (not `magic_pdf`) +- [ ] `test_is_available_when_installed` patches `mineru` present → True +- [ ] `test_capabilities` unchanged +- [ ] `test_config_stored` includes new `backend` default `pipeline` +- [ ] `test_process_returns_engine_result` — mocks `do_parse`+`read_fn`, reads + generated `.md` from the pipeline output dir +- [ ] `test_process_json_output_format` — reads `_content_list.json` +- [ ] `test_process_with_page_range` — `start_page`/`end_page` forwarded as + `start_page_id`/`end_page_id` +- [ ] `test_process_vlm_backend` — backend=`vlm` reads from `vlm` subdir +- [ ] ABC conformance test still passes + +### Integration / E2E Tests +- [ ] E2E: real PDF through `docfold ... --engine mineru` (manual, slow, + downloads model weights) — verify markdown + JSON outputs + +### Test Commands +```bash +pytest tests/engines/test_adapters.py -k MinerU -v +pytest tests/ -m "not slow" +``` + +## Edge Cases +- MinerU writes multiple files; we only read `.md` and `_content_list.json`. +- `parse_method="auto"` is the directory name for pipeline (upstream does not + rewrite the subdir to the resolved txt/ocr method when `auto` is passed). +- Missing output file → raise a clear `RuntimeError`. + +## Out of Scope +- `hybrid` backend wiring (can be added later). +- Exposing bounding boxes (MinerU provides them in middle.json; not surfaced). +- Server/HTTP (`vlm-http-client`) backend. diff --git a/pyproject.toml b/pyproject.toml index fb25214..06cbd69 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,7 @@ docling = [ "docling>=2.0", ] mineru = [ - "magic-pdf[full]>=0.9", + "mineru[core]>=2.0", ] marker = [ "requests>=2.31", diff --git a/src/docfold/engines/mineru_engine.py b/src/docfold/engines/mineru_engine.py index 74d0617..16913eb 100644 --- a/src/docfold/engines/mineru_engine.py +++ b/src/docfold/engines/mineru_engine.py @@ -1,14 +1,19 @@ -"""MinerU / PDF-Extract-Kit engine adapter. +"""MinerU 2.x engine adapter. Install: ``pip install docfold[mineru]`` -Note: First run downloads model weights (~2-5 GB). +Built on `MinerU `_ (the ``mineru`` +package, formerly ``magic-pdf``). Uses the supported programmatic entry point +:func:`mineru.cli.common.do_parse`. + +Note: First run downloads model weights (~1-3 GB). License: AGPL-3.0 — see https://github.com/opendatalab/MinerU """ from __future__ import annotations import logging +import os import tempfile import time from typing import Any @@ -20,54 +25,39 @@ _SUPPORTED_EXTENSIONS = {"pdf"} # Lazy-loaded at first use; patchable in tests. -PymuDocDataset: Any = None -SupportedPdfParseMethod: Any = None -FileBasedDataWriter: Any = None -doc_analyze: Any = None -convert_pdf_bytes_to_bytes_by_pymupdf: Any = None +do_parse: Any = None +read_fn: Any = None def _ensure_imports() -> None: - """Import magic_pdf dependencies on first use.""" - global PymuDocDataset, SupportedPdfParseMethod, FileBasedDataWriter - global doc_analyze, convert_pdf_bytes_to_bytes_by_pymupdf - if PymuDocDataset is not None: + """Import the ``mineru`` programmatic API on first use.""" + global do_parse, read_fn + if do_parse is not None: return - from magic_pdf.config.enums import SupportedPdfParseMethod as _SPM # noqa: N814 - from magic_pdf.data.data_reader_writer import FileBasedDataWriter as _FBDW # noqa: N814 - from magic_pdf.data.dataset import PymuDocDataset as _PDD # noqa: N814 - try: - from magic_pdf.libs.pdf_utils import ( - convert_pdf_bytes_to_bytes_by_pymupdf as _convert, - ) - except (ImportError, ModuleNotFoundError): - from magic_pdf.tools.common import ( - convert_pdf_bytes_to_bytes_by_pymupdf as _convert, - ) + from mineru.cli.common import do_parse as _do_parse + from mineru.cli.common import read_fn as _read_fn - # doc_analyze location varies across magic-pdf versions. - try: - from magic_pdf.operators.models import doc_analyze as _da - except ImportError: - from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze as _da - - PymuDocDataset = _PDD - SupportedPdfParseMethod = _SPM - FileBasedDataWriter = _FBDW - doc_analyze = _da - convert_pdf_bytes_to_bytes_by_pymupdf = _convert + do_parse = _do_parse + read_fn = _read_fn class MinerUEngine(DocumentEngine): - """Adapter for MinerU (magic-pdf), the end-to-end PDF structuring tool - built on PDF-Extract-Kit. + """Adapter for MinerU 2.x, the end-to-end PDF structuring tool. See https://github.com/opendatalab/MinerU """ - def __init__(self, config_path: str | None = None, gpu: bool = False) -> None: + def __init__( + self, + config_path: str | None = None, + gpu: bool = False, + backend: str = "pipeline", + parse_method: str = "auto", + ) -> None: self._config_path = config_path self._gpu = gpu + self._backend = backend + self._parse_method = parse_method @property def name(self) -> str: @@ -85,7 +75,7 @@ def capabilities(self) -> EngineCapabilities: def is_available(self) -> bool: try: - import magic_pdf # noqa: F401 + import mineru # noqa: F401 return True except ImportError: return False @@ -121,69 +111,76 @@ def _run_mineru( ) -> tuple[str, dict]: _ensure_imports() - # PyTorch 2.6+ defaults weights_only=True which breaks loading - # doclayout_yolo model weights containing custom classes. - try: - import doclayout_yolo.nn.tasks as _tasks - import torch - _safe_classes = [ - cls for cls in vars(_tasks).values() - if isinstance(cls, type) - ] - # The YOLO checkpoint also requires dill._dill._load_type - try: - from dill._dill import _load_type - _safe_classes.append(_load_type) - except ImportError: - pass - if _safe_classes: - torch.serialization.add_safe_globals(_safe_classes) - except Exception: - pass - + backend = kwargs.get("backend", self._backend) + parse_method = kwargs.get("parse_method", self._parse_method) + lang = kwargs.get("lang") or "ch" start_page = kwargs.get("start_page") end_page = kwargs.get("end_page") - lang = kwargs.get("lang") - - with open(file_path, "rb") as f: - pdf_bytes = f.read() - - if start_page is not None or end_page is not None: - pdf_bytes = convert_pdf_bytes_to_bytes_by_pymupdf( - pdf_bytes, - start_page or 0, - end_page, + want_json = output_format == OutputFormat.JSON + + pdf_bytes = read_fn(file_path) + name = "document" + + with tempfile.TemporaryDirectory() as out_dir: + do_parse( + output_dir=out_dir, + pdf_file_names=[name], + pdf_bytes_list=[pdf_bytes], + p_lang_list=[lang], + backend=backend, + parse_method=parse_method, + start_page_id=start_page if start_page is not None else 0, + end_page_id=end_page, + f_dump_md=not want_json, + f_dump_content_list=want_json, + f_draw_layout_bbox=False, + f_draw_span_bbox=False, + f_dump_middle_json=False, + f_dump_model_output=False, + f_dump_orig_pdf=False, ) - ds = PymuDocDataset(pdf_bytes, lang=lang) - classify_result = ds.classify() - is_text_pdf = classify_result == SupportedPdfParseMethod.TXT - - with tempfile.TemporaryDirectory() as tmp_dir: - image_writer = FileBasedDataWriter(tmp_dir) - - infer_result = ds.apply( - doc_analyze, - ocr=not is_text_pdf, - lang=lang, - ) - - if is_text_pdf: - pipe_result = infer_result.pipe_txt_mode( - image_writer, debug_mode=False, lang=lang, - ) + md_dir = os.path.join(out_dir, name, self._output_subdir(backend, parse_method)) + if want_json: + target = os.path.join(md_dir, f"{name}_content_list.json") else: - pipe_result = infer_result.pipe_ocr_mode( - image_writer, debug_mode=False, lang=lang, + target = os.path.join(md_dir, f"{name}.md") + + if not os.path.exists(target): + raise RuntimeError( + f"MinerU did not produce expected output at {target!r}. " + f"Output dir contents: {self._list_dir(out_dir)}" ) - if output_format == OutputFormat.JSON: - content = pipe_result.get_content_list(tmp_dir) - else: - content = pipe_result.get_markdown(tmp_dir) + with open(target, encoding="utf-8") as f: + content = f.read() metadata = { - "method": "txt" if is_text_pdf else "ocr", + "backend": backend, + "parse_method": parse_method, + "lang": lang, } return content, metadata + + @staticmethod + def _output_subdir(backend: str, parse_method: str) -> str: + """Subdirectory ``do_parse`` writes into, per backend. + + Mirrors upstream ``mineru.cli.common.do_parse``: pipeline → the + ``parse_method`` name (e.g. ``auto``); vlm family → ``vlm``; + hybrid family → ``hybrid_``. + """ + if backend.startswith("vlm"): + return "vlm" + if backend.startswith("hybrid"): + return f"hybrid_{parse_method}" + return parse_method + + @staticmethod + def _list_dir(root: str) -> list[str]: + found: list[str] = [] + for dirpath, _dirs, files in os.walk(root): + for fn in files: + found.append(os.path.relpath(os.path.join(dirpath, fn), root)) + return found diff --git a/tests/engines/test_adapters.py b/tests/engines/test_adapters.py index 4d9b1a1..66d477e 100644 --- a/tests/engines/test_adapters.py +++ b/tests/engines/test_adapters.py @@ -55,7 +55,7 @@ def test_supported_extensions(self): def test_is_available_when_missing(self): from docfold.engines.mineru_engine import MinerUEngine e = MinerUEngine() - with patch.dict("sys.modules", {"magic_pdf": None}): + with patch.dict("sys.modules", {"mineru": None}): result = e.is_available() assert isinstance(result, bool) @@ -65,6 +65,17 @@ def test_config_stored(self): assert e._config_path == "/tmp/cfg.yaml" assert e._gpu is True + def test_default_backend_is_pipeline(self): + from docfold.engines.mineru_engine import MinerUEngine + e = MinerUEngine() + assert e._backend == "pipeline" + + def test_backend_config_stored(self): + from docfold.engines.mineru_engine import MinerUEngine + e = MinerUEngine(backend="vlm", parse_method="ocr") + assert e._backend == "vlm" + assert e._parse_method == "ocr" + def test_capabilities(self): from docfold.engines.mineru_engine import MinerUEngine e = MinerUEngine() @@ -76,180 +87,133 @@ def test_capabilities(self): assert caps.confidence is False def test_is_available_when_installed(self): - """When magic_pdf is importable, is_available returns True.""" + """When mineru is importable, is_available returns True.""" from unittest.mock import MagicMock, patch from docfold.engines.mineru_engine import MinerUEngine e = MinerUEngine() - with patch.dict("sys.modules", {"magic_pdf": MagicMock()}): + with patch.dict("sys.modules", {"mineru": MagicMock()}): result = e.is_available() assert result is True @pytest.mark.asyncio - async def test_process_returns_engine_result(self): - """MinerU engine processes a PDF and returns a valid EngineResult.""" + async def test_process_returns_engine_result(self, tmp_path): + """MinerU 2.x engine processes a PDF via do_parse and returns a result.""" import os - import tempfile - from unittest.mock import MagicMock, patch + from unittest.mock import patch from docfold.engines.base import EngineResult, OutputFormat from docfold.engines.mineru_engine import MinerUEngine - mock_pipe_result = MagicMock() - mock_pipe_result.get_markdown.return_value = "# Hello\n\nExtracted content" - - mock_infer_result = MagicMock() - mock_infer_result.pipe_txt_mode.return_value = mock_pipe_result - - mock_dataset = MagicMock() - mock_dataset.classify.return_value = "txt" - mock_dataset.apply.return_value = mock_infer_result - mock_dataset._lang = None - - mock_spm = MagicMock() - mock_spm.TXT = "txt" + def fake_do_parse(output_dir, pdf_file_names, *args, **kwargs): + # Mimic MinerU 2.x pipeline layout: output_dir///.md + name = pdf_file_names[0] + parse_method = kwargs.get("parse_method", "auto") + md_dir = os.path.join(output_dir, name, parse_method) + os.makedirs(md_dir, exist_ok=True) + with open(os.path.join(md_dir, f"{name}.md"), "w") as fh: + fh.write("# Hello\n\nExtracted content") with patch("docfold.engines.mineru_engine._ensure_imports"), \ - patch("docfold.engines.mineru_engine.PymuDocDataset", return_value=mock_dataset), \ - patch("docfold.engines.mineru_engine.SupportedPdfParseMethod", mock_spm), \ - patch("docfold.engines.mineru_engine.FileBasedDataWriter"), \ - patch("docfold.engines.mineru_engine.doc_analyze"): + patch("docfold.engines.mineru_engine.read_fn", return_value=b"%PDF-1.4"), \ + patch("docfold.engines.mineru_engine.do_parse", side_effect=fake_do_parse): e = MinerUEngine() - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f: - f.write(b"%PDF-1.4 minimal") - try: - result = await e.process(f.name, OutputFormat.MARKDOWN) - assert isinstance(result, EngineResult) - assert result.engine_name == "mineru" - assert result.content == "# Hello\n\nExtracted content" - assert result.format == OutputFormat.MARKDOWN - assert result.processing_time_ms >= 0 - finally: - os.unlink(f.name) + pdf = tmp_path / "doc.pdf" + pdf.write_bytes(b"%PDF-1.4 minimal") + result = await e.process(str(pdf), OutputFormat.MARKDOWN) + assert isinstance(result, EngineResult) + assert result.engine_name == "mineru" + assert result.content == "# Hello\n\nExtracted content" + assert result.format == OutputFormat.MARKDOWN + assert result.processing_time_ms >= 0 @pytest.mark.asyncio - async def test_process_ocr_mode_for_scanned_pdf(self): - """MinerU uses OCR mode when PDF is classified as scanned.""" + async def test_process_json_output_format(self, tmp_path): + """MinerU returns content_list JSON when output_format is JSON.""" import os - import tempfile - from unittest.mock import MagicMock, patch + from unittest.mock import patch from docfold.engines.base import OutputFormat from docfold.engines.mineru_engine import MinerUEngine - mock_pipe_result = MagicMock() - mock_pipe_result.get_markdown.return_value = "OCR content" - - mock_infer_result = MagicMock() - mock_infer_result.pipe_ocr_mode.return_value = mock_pipe_result - - mock_dataset = MagicMock() - mock_dataset.classify.return_value = "ocr" - mock_dataset.apply.return_value = mock_infer_result - mock_dataset._lang = None - - mock_spm = MagicMock() - mock_spm.TXT = "txt" + def fake_do_parse(output_dir, pdf_file_names, *args, **kwargs): + name = pdf_file_names[0] + parse_method = kwargs.get("parse_method", "auto") + md_dir = os.path.join(output_dir, name, parse_method) + os.makedirs(md_dir, exist_ok=True) + with open(os.path.join(md_dir, f"{name}_content_list.json"), "w") as fh: + fh.write('[{"type": "text", "text": "hello"}]') with patch("docfold.engines.mineru_engine._ensure_imports"), \ - patch("docfold.engines.mineru_engine.PymuDocDataset", return_value=mock_dataset), \ - patch("docfold.engines.mineru_engine.SupportedPdfParseMethod", mock_spm), \ - patch("docfold.engines.mineru_engine.FileBasedDataWriter"), \ - patch("docfold.engines.mineru_engine.doc_analyze"): + patch("docfold.engines.mineru_engine.read_fn", return_value=b"%PDF-1.4"), \ + patch("docfold.engines.mineru_engine.do_parse", side_effect=fake_do_parse): e = MinerUEngine() - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f: - f.write(b"%PDF-1.4 minimal") - try: - result = await e.process(f.name, OutputFormat.MARKDOWN) - assert result.content == "OCR content" - mock_infer_result.pipe_ocr_mode.assert_called_once() - finally: - os.unlink(f.name) + pdf = tmp_path / "doc.pdf" + pdf.write_bytes(b"%PDF-1.4 minimal") + result = await e.process(str(pdf), OutputFormat.JSON) + assert result.format == OutputFormat.JSON + assert "text" in result.content @pytest.mark.asyncio - async def test_process_json_output_format(self): - """MinerU returns content_list JSON when output_format is JSON.""" + async def test_process_forwards_page_range_and_lang(self, tmp_path): + """MinerU forwards start_page/end_page as start_page_id/end_page_id and lang.""" import os - import tempfile from unittest.mock import MagicMock, patch from docfold.engines.base import OutputFormat from docfold.engines.mineru_engine import MinerUEngine - mock_pipe_result = MagicMock() - mock_pipe_result.get_markdown.return_value = "md content" - mock_pipe_result.get_content_list.return_value = '[{"type": "text", "text": "hello"}]' - - mock_infer_result = MagicMock() - mock_infer_result.pipe_txt_mode.return_value = mock_pipe_result - - mock_dataset = MagicMock() - mock_dataset.classify.return_value = "txt" - mock_dataset.apply.return_value = mock_infer_result - mock_dataset._lang = None - - mock_spm = MagicMock() - mock_spm.TXT = "txt" + def fake_do_parse(output_dir, pdf_file_names, *args, **kwargs): + name = pdf_file_names[0] + parse_method = kwargs.get("parse_method", "auto") + md_dir = os.path.join(output_dir, name, parse_method) + os.makedirs(md_dir, exist_ok=True) + with open(os.path.join(md_dir, f"{name}.md"), "w") as fh: + fh.write("page content") + mock_do_parse = MagicMock(side_effect=fake_do_parse) with patch("docfold.engines.mineru_engine._ensure_imports"), \ - patch("docfold.engines.mineru_engine.PymuDocDataset", return_value=mock_dataset), \ - patch("docfold.engines.mineru_engine.SupportedPdfParseMethod", mock_spm), \ - patch("docfold.engines.mineru_engine.FileBasedDataWriter"), \ - patch("docfold.engines.mineru_engine.doc_analyze"): + patch("docfold.engines.mineru_engine.read_fn", return_value=b"%PDF-1.4"), \ + patch("docfold.engines.mineru_engine.do_parse", mock_do_parse): e = MinerUEngine() - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f: - f.write(b"%PDF-1.4 minimal") - try: - result = await e.process(f.name, OutputFormat.JSON) - assert result.format == OutputFormat.JSON - assert "text" in result.content - finally: - os.unlink(f.name) + pdf = tmp_path / "doc.pdf" + pdf.write_bytes(b"%PDF-1.4 minimal") + await e.process( + str(pdf), OutputFormat.MARKDOWN, + start_page=2, end_page=5, lang="ru", + ) + kwargs = mock_do_parse.call_args.kwargs + assert kwargs["start_page_id"] == 2 + assert kwargs["end_page_id"] == 5 + assert kwargs["p_lang_list"] == ["ru"] + assert kwargs["backend"] == "pipeline" @pytest.mark.asyncio - async def test_process_with_page_range(self): - """MinerU respects start_page and end_page kwargs.""" + async def test_process_vlm_backend_reads_vlm_subdir(self, tmp_path): + """With backend='vlm', output is read from the 'vlm' subdirectory.""" import os - import tempfile from unittest.mock import MagicMock, patch from docfold.engines.base import OutputFormat from docfold.engines.mineru_engine import MinerUEngine - mock_pipe_result = MagicMock() - mock_pipe_result.get_markdown.return_value = "page content" - - mock_infer_result = MagicMock() - mock_infer_result.pipe_txt_mode.return_value = mock_pipe_result - - mock_dataset = MagicMock() - mock_dataset.classify.return_value = "txt" - mock_dataset.apply.return_value = mock_infer_result - mock_dataset._lang = None - - mock_spm = MagicMock() - mock_spm.TXT = "txt" + def fake_do_parse(output_dir, pdf_file_names, *args, **kwargs): + name = pdf_file_names[0] + md_dir = os.path.join(output_dir, name, "vlm") + os.makedirs(md_dir, exist_ok=True) + with open(os.path.join(md_dir, f"{name}.md"), "w") as fh: + fh.write("vlm content") + mock_do_parse = MagicMock(side_effect=fake_do_parse) with patch("docfold.engines.mineru_engine._ensure_imports"), \ - patch("docfold.engines.mineru_engine.PymuDocDataset", return_value=mock_dataset), \ - patch("docfold.engines.mineru_engine.SupportedPdfParseMethod", mock_spm), \ - patch("docfold.engines.mineru_engine.FileBasedDataWriter"), \ - patch("docfold.engines.mineru_engine.doc_analyze"), \ - patch( - "docfold.engines.mineru_engine.convert_pdf_bytes_to_bytes_by_pymupdf", - ) as mock_convert: - mock_convert.return_value = b"%PDF-1.4 subset" - e = MinerUEngine() - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f: - f.write(b"%PDF-1.4 minimal") - try: - await e.process( - f.name, OutputFormat.MARKDOWN, - start_page=2, end_page=5, - ) - mock_convert.assert_called_once() - finally: - os.unlink(f.name) + patch("docfold.engines.mineru_engine.read_fn", return_value=b"%PDF-1.4"), \ + patch("docfold.engines.mineru_engine.do_parse", mock_do_parse): + e = MinerUEngine(backend="vlm") + pdf = tmp_path / "doc.pdf" + pdf.write_bytes(b"%PDF-1.4 minimal") + result = await e.process(str(pdf), OutputFormat.MARKDOWN) + assert result.content == "vlm content" + assert mock_do_parse.call_args.kwargs["backend"] == "vlm" class TestMarkerEngine: