langparse 0.1.0__py3-none-any.whl

langparse/engines/pdf/deepdoc_engine.py

@@ -0,0 +1,164 @@
+import threading
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Any
+
+from langparse.core.engine import PageResult
+from langparse.engines.pdf.simple import BasePDFEngine
+from langparse.logging import get_logger
+from langparse.progress import ProgressCallback, ProgressReporter
+from langparse.types import ParsedDocumentResult
+
+logger = get_logger(__name__)
+
+
+class DeepDocEngine(BasePDFEngine):
+    """
+    Adapter for the ported DeepDoc (RAGFlow) OCR + layout + table-structure
+    pipeline. CPU-only ONNX inference, no separate runtime service --
+    unlike MinerU, it runs in-process (see langparse/engines/pdf/deepdoc/).
+    """
+
+    def __init__(
+        self,
+        device: str = "cpu",
+        model_dir: str | None = None,
+        download_dir: str | None = None,
+        model_policy: str = "download_if_missing",
+        parser: Any = None,
+        **kwargs: Any,
+    ):
+        if device != "cpu":
+            raise ValueError(
+                f"DeepDocEngine only supports device='cpu' in this version, got: {device!r}"
+            )
+        self.device = device
+        self.model_dir = model_dir
+        self.download_dir = download_dir
+        self.model_policy = model_policy
+        self._parser = parser
+        # Batch runs share one engine (and thus one RAGFlowPdfParser) across
+        # worker threads. RAGFlowPdfParser is deeply per-parse stateful --
+        # __images__ resets self.boxes/self.page_images/etc. at the top of
+        # every call, and the downstream chain mutates self.boxes in place --
+        # so two threads parsing concurrently would corrupt each other's
+        # state. The lock also prevents two threads from racing to build (and
+        # load ~100MB of models for) separate parsers. Same precedent as
+        # SimplePDFEngine._ocr_lock: correctness over throughput on an
+        # already-slow path.
+        self._parser_lock = threading.Lock()
+
+    def _build_parser(self):
+        try:
+            from langparse.engines.pdf.deepdoc.model_loader import ensure_deepdoc_models
+            from langparse.engines.pdf.deepdoc.pdf_parser import RAGFlowPdfParser
+        except ImportError as exc:
+            raise ImportError(
+                'DeepDoc engine needs extra dependencies. Install them with `pip install "langparse[deepdoc]"`.'
+            ) from exc
+
+        resolved_model_dir = ensure_deepdoc_models(
+            model_dir=self.model_dir,
+            download_dir=self.download_dir,
+            model_policy=self.model_policy,
+        )
+        return RAGFlowPdfParser(model_dir=resolved_model_dir)
+
+    def _classify_ocr_pages(self, file_path: Path) -> dict[int, bool]:
+        """Per-page {page_number: bool} saying whether a page's text should be
+        credited to OCR, using the same needs_ocr() heuristic simple/ocr.py
+        already uses -- independent of deepdoc's own (unconditional) internal
+        OCR pass, so the two engines report comparable metadata for comparable
+        input. 1-indexed to match render_pages()'s page_number convention;
+        pdfplumber.pages itself is 0-indexed.
+
+        This runs after the real (and potentially minutes-long) deepdoc parse
+        has already completed, purely to annotate advisory metadata. Any
+        failure here (missing pdfplumber, malformed/encrypted/oddly-structured
+        PDF, etc.) is logged and swallowed rather than raised, so it can never
+        discard an otherwise-successful parse. An empty dict is already the
+        documented contract: render_pages() defaults every absent page to
+        ocr_applied=False / ocr_text_chars=0.
+        """
+        try:
+            try:
+                import pdfplumber
+            except ImportError as exc:
+                raise ImportError(
+                    'DeepDoc engine needs extra dependencies. Install them with `pip install "langparse[deepdoc]"`.'
+                ) from exc
+            from langparse.engines.pdf.ocr import needs_ocr
+
+            with pdfplumber.open(file_path) as pdf:
+                return {
+                    page_number: needs_ocr(page)
+                    for page_number, page in enumerate(pdf.pages, start=1)
+                }
+        except Exception as exc:
+            logger.warning("Skipping OCR-applied classification for %s: %s", file_path, exc)
+            return {}
+
+    def process_document(
+        self,
+        file_path: Path,
+        *,
+        progress_callback: ProgressCallback | None = None,
+        **kwargs: Any,
+    ) -> ParsedDocumentResult:
+        reporter = ProgressReporter(str(file_path), progress_callback)
+        reporter.emit("preparing", message="Preparing DeepDoc parser")
+        try:
+            from langparse.engines.pdf.deepdoc.rendering import render_pages
+        except ImportError as exc:
+            raise ImportError(
+                'DeepDoc engine needs extra dependencies. Install them with `pip install "langparse[deepdoc]"`.'
+            ) from exc
+
+        with self._parser_lock:
+            if self._parser is None:
+                self._parser = self._build_parser()
+            if progress_callback is None:
+                boxes = self._parser.parse_into_bboxes(str(file_path))
+            else:
+                last_percent = 0.0
+
+                def on_progress(progress, message=""):
+                    nonlocal last_percent
+                    last_percent = max(last_percent, min(100.0, max(0.0, progress * 100)))
+                    reporter.emit("parsing", percent=last_percent, message=message)
+
+                boxes = self._parser.parse_into_bboxes(str(file_path), callback=on_progress)
+
+        reporter.emit("rendering")
+        ocr_pages = self._classify_ocr_pages(file_path)
+        pages = render_pages(boxes, ocr_pages=ocr_pages)
+        return ParsedDocumentResult(
+            source=str(file_path),
+            filename=Path(file_path).name,
+            engine="deepdoc",
+            pages=pages,
+            markdown_content="\n\n".join(page.markdown_content for page in pages),
+            metadata={
+                "device": self.device,
+                "model_dir": self.model_dir,
+                "ocr_applied": any(page.metadata["ocr_applied"] for page in pages),
+                "ocr_text_chars": sum(page.metadata["ocr_text_chars"] for page in pages),
+            },
+        )
+
+    def process(
+        self, file_path: Path, *, progress_callback: ProgressCallback | None = None, **kwargs
+    ) -> Iterator[PageResult]:
+        if progress_callback is not None:
+            kwargs["progress_callback"] = progress_callback
+        parsed = self.process_document(file_path, **kwargs)
+        for page in parsed.pages:
+            yield PageResult(
+                page_number=page.page_number,
+                markdown_content=page.markdown_content,
+                plain_text=page.plain_text,
+                elements=page.elements,
+                tables=page.tables,
+                images=page.images,
+                metadata=page.metadata,
+            )

langparse/engines/pdf/mineru.py

@@ -0,0 +1,259 @@
+from collections.abc import Iterator, Mapping
+from pathlib import Path
+from typing import Any
+
+from langparse.core.engine import PageResult
+from langparse.engines.pdf.mineru_client import MinerUClient
+from langparse.engines.pdf.mineru_service import MinerUServiceManager
+from langparse.engines.pdf.simple import BasePDFEngine
+from langparse.progress import ProgressCallback, ProgressReporter
+from langparse.types import ParsedDocumentResult, ParsedElement, ParsedPageResult
+
+_SENSITIVE_OPTION_KEYS = frozenset(
+    {
+        "api_key",
+        "openai_api_key",
+        "model",
+        "base_url",
+        "openai_base_url",
+        "openai_model",
+        "workbook_disambiguation",
+        "disambiguation",
+        "token",
+        "secret",
+        "apikey",
+        "auth",
+        "authorization",
+        "password",
+        "bearer",
+        "credentials",
+    }
+)
+_SENSITIVE_OPTION_SUFFIXES = ("_key", "_secret", "_token", "_password", "_auth")
+
+
+def _sanitize_extra_options(options: Mapping[str, Any] | None) -> dict[str, Any]:
+    if options is None:
+        return {}
+    return {
+        key: value
+        for key, value in options.items()
+        if key.lower() not in _SENSITIVE_OPTION_KEYS
+        and not key.lower().endswith(_SENSITIVE_OPTION_SUFFIXES)
+    }
+
+
+class MinerUEngine(BasePDFEngine):
+    """
+    Adapter for MinerU (Magic-PDF).
+    High precision parsing for complex documents (papers, textbooks).
+    """
+
+    def __init__(
+        self,
+        device: str = "auto",
+        model_dir: str | None = None,
+        download_dir: str | None = None,
+        enable_ocr: bool = True,
+        api_url: str | None = None,
+        api_host: str = "127.0.0.1",
+        api_port: int = 8000,
+        api_command: str = "mineru-api",
+        api_start_timeout: float = 30.0,
+        request_timeout: float = 300.0,
+        backend: str | None = None,
+        server_url: str | None = None,
+        model_policy: str = "download_if_missing",
+        model_source: str | None = None,
+        auto_install_runtime: bool = False,
+        runtime_package: str = "mineru>=3.4,<4",
+        extra_options: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ):
+        self.device = device
+        self.model_dir = model_dir
+        self.download_dir = download_dir
+        self.enable_ocr = enable_ocr
+        self.api_url = api_url
+        self.api_host = api_host
+        self.api_port = api_port
+        self.api_command = api_command
+        self.api_start_timeout = api_start_timeout
+        self.request_timeout = request_timeout
+        self.backend = backend
+        self.server_url = server_url
+        self.model_policy = model_policy
+        self.model_source = model_source
+        self.auto_install_runtime = auto_install_runtime
+        self.runtime_package = runtime_package
+        self.extra_options = {
+            **_sanitize_extra_options(extra_options),
+            **_sanitize_extra_options(kwargs),
+        }
+        if self.backend is not None:
+            self.extra_options["backend"] = self.backend
+        if self.server_url is not None:
+            self.extra_options["server_url"] = self.server_url
+
+    def _cuda_available(self) -> bool:
+        try:
+            import torch
+        except ImportError:
+            return False
+
+        return bool(torch.cuda.is_available())
+
+    def _resolve_device(self, device: str | None = None) -> str:
+        requested_device = self.device if device is None else device
+
+        if requested_device == "auto":
+            return "cuda" if self._cuda_available() else "cpu"
+
+        if requested_device == "cuda" and not self._cuda_available():
+            raise RuntimeError("CUDA was requested for MinerU but is not available.")
+
+        return requested_device
+
+    def _ensure_runtime(self) -> None:
+        # The real runtime path is mineru-api. Dependency validation happens when
+        # starting or connecting to that service.
+        return None
+
+    def _build_runtime_config(self, **kwargs: Any) -> dict[str, Any]:
+        requested_device = kwargs.get("device", self.device)
+        runtime_extra_options = _sanitize_extra_options(kwargs.get("extra_options"))
+        return {
+            "device": self._resolve_device(requested_device),
+            "model_dir": kwargs.get("model_dir", self.model_dir),
+            "download_dir": kwargs.get("download_dir", self.download_dir),
+            "enable_ocr": kwargs.get("enable_ocr", self.enable_ocr),
+            "extra_options": {**self.extra_options, **runtime_extra_options},
+        }
+
+    def _build_service_config(self) -> dict[str, Any]:
+        return {
+            "api_url": self.api_url,
+            "host": self.api_host,
+            "port": self.api_port,
+            "command": self.api_command,
+            "start_timeout": self.api_start_timeout,
+            "request_timeout": self.request_timeout,
+            "model_dir": self.model_dir,
+            "download_dir": self.download_dir,
+            "model_policy": self.model_policy,
+            "model_source": self.model_source,
+            "auto_install_runtime": self.auto_install_runtime,
+            "runtime_package": self.runtime_package,
+        }
+
+    def _create_client(self, base_url: str) -> MinerUClient:
+        return MinerUClient(base_url, timeout=self.request_timeout)
+
+    def _create_service_manager(self) -> MinerUServiceManager:
+        return MinerUServiceManager(**self._build_service_config())
+
+    def _run_mineru(
+        self,
+        file_path: Path,
+        runtime_config: dict[str, Any],
+        progress_callback: ProgressCallback | None = None,
+    ) -> list[dict[str, Any]]:
+        reporter = ProgressReporter(str(file_path), progress_callback)
+        manager = self._create_service_manager()
+        with manager.running_service() as base_url:
+            client = self._create_client(base_url)
+            reporter.emit("parsing", message="Waiting for MinerU response")
+            return client.parse_file(file_path, runtime_config)
+
+    def process_document(
+        self,
+        file_path: Path,
+        *,
+        progress_callback: ProgressCallback | None = None,
+        **kwargs: Any,
+    ) -> ParsedDocumentResult:
+        reporter = ProgressReporter(str(file_path), progress_callback)
+        reporter.emit("preparing", message="Preparing MinerU service")
+        self._ensure_runtime()
+        runtime_config = self._build_runtime_config(**kwargs)
+        if progress_callback is None:
+            raw_pages = self._run_mineru(file_path, runtime_config)
+        else:
+            raw_pages = self._run_mineru(file_path, runtime_config, progress_callback)
+        reporter.emit("rendering")
+        pages = [
+            ParsedPageResult(
+                page_number=item["page_number"],
+                markdown_content=item.get("markdown", ""),
+                plain_text=item.get("plain_text", ""),
+                elements=[
+                    element
+                    if isinstance(element, ParsedElement)
+                    else ParsedElement(
+                        kind=element.get("kind", "text"),
+                        text=element.get("text", ""),
+                        bbox=element.get("bbox"),
+                        metadata=element.get("metadata", {}),
+                    )
+                    for element in item.get("elements", [])
+                ],
+                tables=item.get("tables", []),
+                images=item.get("images", []),
+                metadata={
+                    "engine_name": "mineru",
+                    "device": runtime_config["device"],
+                    "engine_specific": item.get("engine_specific", {}),
+                },
+            )
+            for item in raw_pages
+        ]
+        engine_specific_items = [page.metadata.get("engine_specific", {}) for page in pages]
+        quality_metadata = {
+            "ocr_applied": any(bool(item.get("ocr_applied")) for item in engine_specific_items),
+            "ocr_text_chars": sum(
+                int(item.get("ocr_text_chars", 0) or 0) for item in engine_specific_items
+            ),
+            "multi_column_detected": any(
+                bool(item.get("multi_column_detected")) for item in engine_specific_items
+            ),
+            "reading_order_warnings": sum(
+                int(item.get("reading_order_warnings", 0) or 0) for item in engine_specific_items
+            ),
+            "header_footer_removed_count": sum(
+                int(item.get("header_footer_removed_count", 0) or 0)
+                for item in engine_specific_items
+            ),
+        }
+        return ParsedDocumentResult(
+            source=str(file_path),
+            filename=file_path.name,
+            engine="mineru",
+            pages=pages,
+            markdown_content="\n".join(page.markdown_content for page in pages),
+            metadata={
+                "device": runtime_config["device"],
+                "model_dir": runtime_config["model_dir"],
+                "download_dir": runtime_config["download_dir"],
+                "enable_ocr": runtime_config["enable_ocr"],
+                "model_policy": self.model_policy,
+                "model_source": self.model_source,
+                **quality_metadata,
+            },
+        )
+
+    def process(
+        self, file_path: Path, *, progress_callback: ProgressCallback | None = None, **kwargs
+    ) -> Iterator[PageResult]:
+        if progress_callback is not None:
+            kwargs["progress_callback"] = progress_callback
+        parsed = self.process_document(file_path, **kwargs)
+        for page in parsed.pages:
+            yield PageResult(
+                page_number=page.page_number,
+                markdown_content=page.markdown_content,
+                plain_text=page.plain_text,
+                elements=page.elements,
+                tables=page.tables,
+                images=page.images,
+                metadata=page.metadata,
+            )