docfold 0.3.0__py3-none-any.whl

docfold/engines/google_docai_engine.py

@@ -0,0 +1,215 @@
+"""Google Document AI engine adapter — cloud document understanding.
+
+Install: ``pip install docfold[google-docai]``
+
+Requires Google Cloud credentials and a Document AI processor.
+Set ``GOOGLE_APPLICATION_CREDENTIALS`` environment variable for auth,
+and configure processor via constructor or environment variables:
+``GOOGLE_DOCAI_PROJECT_ID``, ``GOOGLE_DOCAI_LOCATION``, ``GOOGLE_DOCAI_PROCESSOR_ID``.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import time
+from typing import Any
+
+from docfold.engines.base import DocumentEngine, EngineCapabilities, EngineResult, OutputFormat
+
+logger = logging.getLogger(__name__)
+
+_SUPPORTED_EXTENSIONS = {"pdf", "png", "jpg", "jpeg", "tiff", "tif", "gif", "bmp", "webp"}
+
+_MIME_MAP = {
+    "pdf": "application/pdf",
+    "png": "image/png",
+    "jpg": "image/jpeg",
+    "jpeg": "image/jpeg",
+    "tiff": "image/tiff",
+    "tif": "image/tiff",
+    "gif": "image/gif",
+    "bmp": "image/bmp",
+    "webp": "image/webp",
+}
+
+
+class GoogleDocAIEngine(DocumentEngine):
+    """Adapter for Google Document AI.
+
+    Processes documents using a configured Document AI processor.
+    Supports OCR, layout analysis, table extraction, and more.
+
+    See https://cloud.google.com/document-ai
+    """
+
+    def __init__(
+        self,
+        project_id: str | None = None,
+        location: str | None = None,
+        processor_id: str | None = None,
+    ) -> None:
+        self._project_id = project_id or os.getenv("GOOGLE_DOCAI_PROJECT_ID")
+        self._location = location or os.getenv("GOOGLE_DOCAI_LOCATION", "us")
+        self._processor_id = processor_id or os.getenv("GOOGLE_DOCAI_PROCESSOR_ID")
+
+    @property
+    def name(self) -> str:
+        return "google_docai"
+
+    @property
+    def supported_extensions(self) -> set[str]:
+        return _SUPPORTED_EXTENSIONS
+
+    @property
+    def capabilities(self) -> EngineCapabilities:
+        return EngineCapabilities(
+            bounding_boxes=True, confidence=True, table_structure=True,
+            heading_detection=True, reading_order=True,
+        )
+
+    def is_available(self) -> bool:
+        try:
+            from google.cloud import documentai  # noqa: F401
+
+            return bool(self._project_id and self._processor_id)
+        except ImportError:
+            return False
+
+    async def process(
+        self,
+        file_path: str,
+        output_format: OutputFormat = OutputFormat.MARKDOWN,
+        **kwargs: Any,
+    ) -> EngineResult:
+        import asyncio
+
+        start = time.perf_counter()
+
+        loop = asyncio.get_running_loop()
+        content, metadata, boxes, conf, tables = await loop.run_in_executor(
+            None, self._process_document, file_path, output_format
+        )
+
+        elapsed_ms = int((time.perf_counter() - start) * 1000)
+
+        return EngineResult(
+            content=content,
+            format=output_format,
+            engine_name=self.name,
+            processing_time_ms=elapsed_ms,
+            metadata=metadata,
+            bounding_boxes=boxes,
+            confidence=conf,
+            tables=tables,
+        )
+
+    def _process_document(
+        self,
+        file_path: str,
+        output_format: OutputFormat,
+    ) -> tuple[str, dict, list[dict], float | None, list[dict] | None]:
+        from google.cloud import documentai
+
+        client = documentai.DocumentProcessorServiceClient()
+
+        processor_name = client.processor_path(
+            self._project_id, self._location, self._processor_id
+        )
+
+        ext = os.path.splitext(file_path)[1].lstrip(".").lower()
+        mime_type = _MIME_MAP.get(ext, "application/octet-stream")
+
+        with open(file_path, "rb") as f:
+            raw_document = documentai.RawDocument(content=f.read(), mime_type=mime_type)
+
+        request = documentai.ProcessRequest(name=processor_name, raw_document=raw_document)
+        result = client.process_document(request=request)
+        document = result.document
+
+        # Extract text
+        full_text = document.text or ""
+
+        # Extract bounding boxes and confidence
+        bounding_boxes: list[dict[str, Any]] = []
+        confidences: list[float] = []
+
+        for page in document.pages:
+            page_num = page.page_number
+
+            for paragraph in page.paragraphs:
+                text_segment = self._get_text_segment(paragraph.layout, document.text)
+                conf = paragraph.layout.confidence
+                if conf:
+                    confidences.append(conf)
+
+                vertices = self._get_vertices(paragraph.layout)
+                if vertices:
+                    bounding_boxes.append({
+                        "type": "paragraph",
+                        "text": text_segment,
+                        "vertices": vertices,
+                        "page": page_num,
+                        "confidence": conf,
+                    })
+
+        avg_conf = sum(confidences) / len(confidences) if confidences else None
+
+        # Extract tables
+        tables: list[dict[str, Any]] = []
+        for page in document.pages:
+            for table in page.tables:
+                table_data = self._extract_table(table, document.text)
+                if table_data:
+                    tables.append(table_data)
+
+        # Format output
+        if output_format == OutputFormat.JSON:
+            import json
+            data = {"text": full_text, "page_count": len(document.pages)}
+            content = json.dumps(data, ensure_ascii=False)
+        elif output_format == OutputFormat.HTML:
+            content = f"<html><body><pre>{full_text}</pre></body></html>"
+        else:
+            content = full_text
+
+        metadata = {
+            "page_count": len(document.pages),
+            "processor_id": self._processor_id,
+            "mime_type": mime_type,
+        }
+
+        return content, metadata, bounding_boxes, avg_conf, tables or None
+
+    def _get_text_segment(self, layout: Any, full_text: str) -> str:
+        """Extract text from a layout's text_anchor."""
+        segments = layout.text_anchor.text_segments if layout.text_anchor else []
+        parts = []
+        for segment in segments:
+            start = int(segment.start_index) if segment.start_index else 0
+            end = int(segment.end_index) if segment.end_index else 0
+            parts.append(full_text[start:end])
+        return "".join(parts).strip()
+
+    def _get_vertices(self, layout: Any) -> list[dict[str, float]] | None:
+        """Extract normalized vertices from layout bounding poly."""
+        bp = layout.bounding_poly
+        if not bp or not bp.normalized_vertices:
+            return None
+        return [{"x": v.x, "y": v.y} for v in bp.normalized_vertices]
+
+    def _extract_table(self, table: Any, full_text: str) -> dict[str, Any] | None:
+        """Extract table rows from Document AI table object."""
+        rows_data = []
+
+        for row in table.header_rows:
+            cells = [self._get_text_segment(cell.layout, full_text) for cell in row.cells]
+            rows_data.append({"type": "header", "cells": cells})
+
+        for row in table.body_rows:
+            cells = [self._get_text_segment(cell.layout, full_text) for cell in row.cells]
+            rows_data.append({"type": "body", "cells": cells})
+
+        if not rows_data:
+            return None
+        return {"rows": rows_data}

docfold/engines/llamaparse_engine.py

@@ -0,0 +1,107 @@
+"""LlamaParse engine adapter — LLM-powered document parsing by LlamaIndex.
+
+Install: ``pip install docfold[llamaparse]``
+
+Requires an API key: https://cloud.llamaindex.ai/
+Set ``LLAMA_CLOUD_API_KEY`` environment variable.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import time
+from typing import Any
+
+from docfold.engines.base import DocumentEngine, EngineCapabilities, EngineResult, OutputFormat
+
+logger = logging.getLogger(__name__)
+
+_SUPPORTED_EXTENSIONS = {
+    "pdf", "docx", "doc", "pptx", "ppt", "xlsx", "xls",
+    "html", "htm", "png", "jpg", "jpeg", "csv", "epub",
+}
+
+
+class LlamaParseEngine(DocumentEngine):
+    """Adapter for LlamaParse (LlamaIndex Cloud).
+
+    LLM-powered parsing with excellent table and layout understanding.
+    Free tier: 1000 pages/day.
+
+    See https://docs.llamaindex.ai/en/stable/llama_cloud/llama_parse/
+    """
+
+    def __init__(self, api_key: str | None = None, result_type: str = "markdown") -> None:
+        self._api_key = api_key or os.getenv("LLAMA_CLOUD_API_KEY")
+        self._result_type = result_type
+
+    @property
+    def name(self) -> str:
+        return "llamaparse"
+
+    @property
+    def supported_extensions(self) -> set[str]:
+        return _SUPPORTED_EXTENSIONS
+
+    @property
+    def capabilities(self) -> EngineCapabilities:
+        return EngineCapabilities(table_structure=True, heading_detection=True)
+
+    def is_available(self) -> bool:
+        try:
+            import llama_parse  # noqa: F401
+
+            return bool(self._api_key)
+        except ImportError:
+            return False
+
+    async def process(
+        self,
+        file_path: str,
+        output_format: OutputFormat = OutputFormat.MARKDOWN,
+        **kwargs: Any,
+    ) -> EngineResult:
+        start = time.perf_counter()
+
+        content, metadata = await self._parse(file_path, output_format)
+
+        elapsed_ms = int((time.perf_counter() - start) * 1000)
+
+        return EngineResult(
+            content=content,
+            format=output_format,
+            engine_name=self.name,
+            processing_time_ms=elapsed_ms,
+            metadata=metadata,
+        )
+
+    async def _parse(
+        self, file_path: str, output_format: OutputFormat
+    ) -> tuple[str, dict]:
+        from llama_parse import LlamaParse
+
+        fmt_map = {
+            OutputFormat.MARKDOWN: "markdown",
+            OutputFormat.HTML: "html",
+            OutputFormat.JSON: "markdown",
+            OutputFormat.TEXT: "text",
+        }
+        result_type = fmt_map[output_format]
+
+        parser = LlamaParse(api_key=self._api_key, result_type=result_type)
+        documents = await parser.aload_data(file_path)
+
+        content = "\n\n".join(doc.text for doc in documents)
+
+        if output_format == OutputFormat.JSON:
+            import json
+
+            data = [{"page": i + 1, "text": doc.text} for i, doc in enumerate(documents)]
+            content = json.dumps(data, ensure_ascii=False)
+
+        metadata = {
+            "result_type": result_type,
+            "document_count": len(documents),
+        }
+        return content, metadata

docfold/engines/marker_engine.py

@@ -0,0 +1,146 @@
+"""Marker API (Datalab) engine adapter.
+
+Install: ``pip install docfold[marker]``
+
+Requires a Datalab API key: https://www.datalab.to/
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+from docfold.engines.base import DocumentEngine, EngineCapabilities, EngineResult, OutputFormat
+
+logger = logging.getLogger(__name__)
+
+_SUPPORTED_EXTENSIONS = {
+    "pdf", "docx", "doc", "pptx", "ppt", "xlsx", "xls",
+    "odt", "odp", "ods", "html", "epub",
+    "png", "jpg", "jpeg", "webp", "gif", "tiff",
+}
+
+_API_BASE = "https://www.datalab.to/api/v1/marker"
+_DEFAULT_POLL_INTERVAL = 2
+_DEFAULT_MAX_POLLS = 300
+
+
+class MarkerEngine(DocumentEngine):
+    """Adapter for the Marker API (Datalab SaaS).
+
+    See https://documentation.datalab.to/
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        use_llm: bool = False,
+        force_ocr: bool = False,
+    ) -> None:
+        self._api_key = api_key or os.getenv("MARKER_API_KEY") or os.getenv("DATALAB_API_KEY")
+        self._use_llm = use_llm
+        self._force_ocr = force_ocr
+
+    @property
+    def name(self) -> str:
+        return "marker"
+
+    @property
+    def supported_extensions(self) -> set[str]:
+        return _SUPPORTED_EXTENSIONS
+
+    @property
+    def capabilities(self) -> EngineCapabilities:
+        return EngineCapabilities(
+            bounding_boxes=True, images=True, table_structure=True,
+            heading_detection=True,
+        )
+
+    def is_available(self) -> bool:
+        try:
+            import requests  # noqa: F401
+            return bool(self._api_key)
+        except ImportError:
+            return False
+
+    async def process(
+        self,
+        file_path: str,
+        output_format: OutputFormat = OutputFormat.MARKDOWN,
+        **kwargs: Any,
+    ) -> EngineResult:
+        import asyncio
+
+        start = time.perf_counter()
+
+        loop = asyncio.get_running_loop()
+        content, images, meta = await loop.run_in_executor(
+            None, self._call_marker, file_path, output_format
+        )
+
+        elapsed_ms = int((time.perf_counter() - start) * 1000)
+
+        return EngineResult(
+            content=content,
+            format=output_format,
+            engine_name=self.name,
+            images=images,
+            pages=meta.get("page_count"),
+            processing_time_ms=elapsed_ms,
+            metadata=meta,
+        )
+
+    def _call_marker(
+        self,
+        file_path: str,
+        output_format: OutputFormat,
+    ) -> tuple[str, dict | None, dict]:
+        import requests
+
+        fmt_map = {
+            OutputFormat.MARKDOWN: "markdown",
+            OutputFormat.HTML: "html",
+            OutputFormat.JSON: "json",
+            OutputFormat.TEXT: "markdown",  # Marker doesn't have plain text; use markdown
+        }
+        marker_fmt = fmt_map[output_format]
+
+        headers = {"X-Api-Key": self._api_key}
+
+        with open(file_path, "rb") as f:
+            form_data = {
+                "file": (Path(file_path).name, f, "application/octet-stream"),
+                "output_format": (None, marker_fmt),
+                "use_llm": (None, str(self._use_llm)),
+                "force_ocr": (None, str(self._force_ocr)),
+                "paginate": (None, "False"),
+                "strip_existing_ocr": (None, "False"),
+                "disable_image_extraction": (None, "False"),
+            }
+            resp = requests.post(_API_BASE, files=form_data, headers=headers, timeout=30)
+            resp.raise_for_status()
+            data = resp.json()
+
+        check_url = data["request_check_url"]
+
+        for _ in range(_DEFAULT_MAX_POLLS):
+            time.sleep(_DEFAULT_POLL_INTERVAL)
+            resp = requests.get(check_url, headers=headers, timeout=30)
+            result = resp.json()
+
+            if result.get("status") == "complete":
+                content = result.get(marker_fmt, "")
+                images = result.get("images")
+                meta = {
+                    "page_count": result.get("page_count"),
+                    "marker_output_format": marker_fmt,
+                }
+                return content, images, meta
+
+            if result.get("status") == "failed":
+                raise RuntimeError(f"Marker API failed: {result.get('error')}")
+
+        raise TimeoutError("Marker API did not complete within the polling window.")

docfold/engines/mineru_engine.py

@@ -0,0 +1,102 @@
+"""MinerU / PDF-Extract-Kit engine adapter.
+
+Install: ``pip install docfold[mineru]``
+
+Note: First run downloads model weights (~2-5 GB).
+License: AGPL-3.0 — see https://github.com/opendatalab/MinerU
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Any
+
+from docfold.engines.base import DocumentEngine, EngineCapabilities, EngineResult, OutputFormat
+
+logger = logging.getLogger(__name__)
+
+_SUPPORTED_EXTENSIONS = {"pdf"}
+
+
+class MinerUEngine(DocumentEngine):
+    """Adapter for MinerU (magic-pdf), the end-to-end PDF structuring tool
+    built on PDF-Extract-Kit.
+
+    See https://github.com/opendatalab/MinerU
+    """
+
+    def __init__(self, config_path: str | None = None, gpu: bool = False) -> None:
+        self._config_path = config_path
+        self._gpu = gpu
+
+    @property
+    def name(self) -> str:
+        return "mineru"
+
+    @property
+    def supported_extensions(self) -> set[str]:
+        return _SUPPORTED_EXTENSIONS
+
+    @property
+    def capabilities(self) -> EngineCapabilities:
+        return EngineCapabilities(
+            table_structure=True, heading_detection=True, reading_order=True,
+        )
+
+    def is_available(self) -> bool:
+        try:
+            import magic_pdf  # noqa: F401
+            return True
+        except ImportError:
+            return False
+
+    async def process(
+        self,
+        file_path: str,
+        output_format: OutputFormat = OutputFormat.MARKDOWN,
+        **kwargs: Any,
+    ) -> EngineResult:
+        import asyncio
+
+        start = time.perf_counter()
+
+        loop = asyncio.get_running_loop()
+        content, metadata = await loop.run_in_executor(
+            None, self._run_mineru, file_path, output_format
+        )
+
+        elapsed_ms = int((time.perf_counter() - start) * 1000)
+
+        return EngineResult(
+            content=content,
+            format=output_format,
+            engine_name=self.name,
+            processing_time_ms=elapsed_ms,
+            metadata=metadata,
+        )
+
+    def _run_mineru(self, file_path: str, output_format: OutputFormat) -> tuple[str, dict]:
+        """Synchronous MinerU processing.
+
+        This is a placeholder implementation. The actual integration will
+        depend on MinerU's Python API which may change across versions.
+        Adapt the import paths and function calls to the installed version.
+        """
+        # TODO: Replace with actual MinerU API calls once version is pinned.
+        # The general pattern is:
+        #
+        #   from magic_pdf.pipe.UNIPipe import UNIPipe
+        #   from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
+        #
+        #   reader = DiskReaderWriter(parent_dir)
+        #   pipe = UNIPipe(pdf_bytes, model_list, reader)
+        #   pipe.pipe_classify()
+        #   pipe.pipe_analyze()
+        #   pipe.pipe_parse()
+        #   md_content = pipe.pipe_mk_markdown(...)
+
+        raise NotImplementedError(
+            "MinerU adapter requires magic-pdf to be installed and configured. "
+            "Install with: pip install docfold[mineru]"
+        )

docfold/engines/mistral_ocr_engine.py

@@ -0,0 +1,128 @@
+"""Mistral OCR engine adapter — Vision LLM-powered document understanding.
+
+Install: ``pip install docfold[mistral-ocr]``
+
+Requires an API key: https://console.mistral.ai/
+Set ``MISTRAL_API_KEY`` environment variable.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import time
+from typing import Any
+
+from docfold.engines.base import DocumentEngine, EngineCapabilities, EngineResult, OutputFormat
+
+logger = logging.getLogger(__name__)
+
+_SUPPORTED_EXTENSIONS = {"pdf", "png", "jpg", "jpeg", "tiff", "tif", "webp", "bmp"}
+
+
+class MistralOCREngine(DocumentEngine):
+    """Adapter for Mistral's OCR API.
+
+    Uses Mistral's document understanding capabilities via the
+    ``mistral.ocr.process`` endpoint for high-quality structured
+    extraction from PDFs and images.
+
+    See https://docs.mistral.ai/capabilities/document/
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        model: str = "mistral-ocr-latest",
+    ) -> None:
+        self._api_key = api_key or os.getenv("MISTRAL_API_KEY")
+        self._model = model
+
+    @property
+    def name(self) -> str:
+        return "mistral_ocr"
+
+    @property
+    def supported_extensions(self) -> set[str]:
+        return _SUPPORTED_EXTENSIONS
+
+    @property
+    def capabilities(self) -> EngineCapabilities:
+        return EngineCapabilities(table_structure=True, heading_detection=True)
+
+    def is_available(self) -> bool:
+        try:
+            import mistralai  # noqa: F401
+
+            return bool(self._api_key)
+        except ImportError:
+            return False
+
+    async def process(
+        self,
+        file_path: str,
+        output_format: OutputFormat = OutputFormat.MARKDOWN,
+        **kwargs: Any,
+    ) -> EngineResult:
+        import asyncio
+
+        start = time.perf_counter()
+
+        loop = asyncio.get_running_loop()
+        content, metadata = await loop.run_in_executor(
+            None, self._call_ocr, file_path, output_format
+        )
+
+        elapsed_ms = int((time.perf_counter() - start) * 1000)
+
+        return EngineResult(
+            content=content,
+            format=output_format,
+            engine_name=self.name,
+            processing_time_ms=elapsed_ms,
+            metadata=metadata,
+        )
+
+    def _call_ocr(self, file_path: str, output_format: OutputFormat) -> tuple[str, dict]:
+        from mistralai import Mistral
+
+        client = Mistral(api_key=self._api_key)
+
+        # Upload file and process with OCR
+        with open(file_path, "rb") as f:
+            file_data = {"file_name": os.path.basename(file_path), "content": f}
+            uploaded = client.files.upload(file=file_data)
+
+        ocr_response = client.ocr.process(
+            model=self._model,
+            document={"type": "file_id", "file_id": uploaded.id},
+        )
+
+        # Combine pages into single output
+        pages_md = []
+        for page in ocr_response.pages:
+            pages_md.append(page.markdown)
+
+        content = "\n\n".join(pages_md)
+
+        if output_format == OutputFormat.JSON:
+            import json
+
+            data = [
+                {"page": i + 1, "text": page.markdown}
+                for i, page in enumerate(ocr_response.pages)
+            ]
+            content = json.dumps(data, ensure_ascii=False)
+        elif output_format == OutputFormat.HTML:
+            html_parts = [
+                f"<div class='page' data-page='{i + 1}'><p>{page.markdown}</p></div>"
+                for i, page in enumerate(ocr_response.pages)
+            ]
+            content = "<html><body>" + "\n".join(html_parts) + "</body></html>"
+
+        metadata = {
+            "model": self._model,
+            "page_count": len(ocr_response.pages),
+            "file_id": uploaded.id,
+        }
+        return content, metadata