docgun 0.1.0__py3-none-any.whl

docgun/__init__.py

@@ -0,0 +1,8 @@
+"""Public interface for the DocGun ingestion package."""
+
+from docgun.pipeline.graph import ingest
+from docgun.pipeline.state import PipelineState
+
+__all__ = ["PipelineState", "__version__", "ingest"]
+
+__version__ = "0.1.0"

docgun/agents/__init__.py

@@ -0,0 +1 @@
+"""Agent-supervised deterministic pipeline helpers."""

docgun/agents/confidence.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from docgun.domain.models import BoundingBox
+
+
+TextKind = Literal["heading", "paragraph", "list_item", "caption", "cell", "unknown"]
+
+
+def text_confidence(
+    text: str | None,
+    location: BoundingBox | None = None,
+    element_type: TextKind = "unknown",
+) -> float:
+    if not text:
+        return 0.0
+
+    printable_ratio = sum(char.isprintable() or char.isspace() for char in text) / len(text)
+    replacement_penalty = text.count("\ufffd") / len(text)
+    expected_length = {
+        "heading": 12,
+        "list_item": 20,
+        "caption": 20,
+        "cell": 6,
+        "paragraph": 80,
+        "unknown": 40,
+    }[element_type]
+    length_score = min(len(text.strip()) / expected_length, 1.0)
+    bbox_score = 1.0 if _valid_bbox(location) else 0.75
+    score = (0.55 * printable_ratio) + (0.20 * length_score) + (0.25 * bbox_score)
+    return round(max(0.0, min(1.0, score - replacement_penalty)), 4)
+
+
+def table_confidence(table: list[list[object]] | None) -> float:
+    if not table:
+        return 0.0
+    widths = [len(row) for row in table]
+    if not widths:
+        return 0.0
+    rectangular = 1.0 if len(set(widths)) == 1 else 0.4
+    non_empty = sum(1 for row in table for cell in row if cell not in (None, ""))
+    total = sum(widths) or 1
+    density = non_empty / total
+    header = table[0] if table else []
+    header_quality = sum(1 for cell in header if cell not in (None, "")) / (len(header) or 1)
+    return round(
+        max(0.0, min(1.0, (0.35 * rectangular) + (0.35 * density) + (0.30 * header_quality))),
+        4,
+    )
+
+
+def _valid_bbox(location: BoundingBox | None) -> bool:
+    if location is None:
+        return False
+    values = (location.x0, location.y0, location.x1, location.y1)
+    if any(value is None for value in values):
+        return False
+    return bool(location.x1 > location.x0 and location.y1 > location.y0)

docgun/agents/enrichment.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from docgun.domain.models import IngestedDocument
+
+
+def enrich_metadata(document: IngestedDocument) -> dict[str, object]:
+    text = "\n".join(element.text or "" for element in document.elements)
+    return {
+        "document_type": _guess_document_type(text),
+        "character_count": len(text),
+        "element_count": len(document.elements),
+    }
+
+
+def _guess_document_type(text: str) -> str:
+    lowered = text.lower()
+    if "invoice" in lowered:
+        return "invoice"
+    if "contract" in lowered or "agreement" in lowered:
+        return "contract"
+    if "statement" in lowered:
+        return "statement"
+    return "unknown"

docgun/agents/inspection.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+import hashlib
+import mimetypes
+from pathlib import Path
+
+from docgun.domain.errors import (
+    OptionalDependencyError,
+    PDFPageLimitExceeded,
+    WorkbookSheetLimitExceeded,
+    WorksheetCellLimitExceeded,
+)
+from docgun.domain.models import FileInspection, PageInspection
+from docgun.security.limits import IngestionLimits
+from docgun.security.validation import validate_file_signature, validate_file_size
+
+
+def sha256_file(path: Path, chunk_size: int = 1024 * 1024) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(chunk_size), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def inspect_file(
+    path: Path,
+    limits: IngestionLimits | None = None,
+    source_name: str | None = None,
+) -> FileInspection:
+    limits = limits or IngestionLimits()
+    validate_file_size(path, limits)
+    validate_file_signature(path)
+
+    suffix = path.suffix.lower()
+    mime_type = mimetypes.guess_type(path.name)[0] or "application/octet-stream"
+    warnings: list[str] = []
+    page_count: int | None = None
+    sheet_count: int | None = None
+    page_inspections: tuple[PageInspection, ...] = ()
+    likely_scanned = False
+    encrypted = False
+
+    if suffix == ".pdf":
+        pdf_info, encrypted = inspect_pdf(path)
+        page_inspections = tuple(pdf_info)
+        page_count = len(page_inspections)
+        if page_count > limits.max_pdf_pages:
+            raise PDFPageLimitExceeded(
+                f"{path.name} has {page_count} pages, exceeding limit {limits.max_pdf_pages}"
+            )
+        likely_scanned = bool(page_inspections) and all(
+            page.needs_vision for page in page_inspections
+        )
+    elif suffix in {".xlsx", ".xlsm"}:
+        sheet_count = inspect_workbook(path, limits)
+
+    return FileInspection(
+        mime_type=mime_type,
+        extension=suffix,
+        size_bytes=path.stat().st_size,
+        sha256=sha256_file(path),
+        source_name=source_name or path.name,
+        page_count=page_count,
+        sheet_count=sheet_count,
+        page_inspections=page_inspections,
+        likely_scanned=likely_scanned,
+        encrypted=encrypted,
+        warnings=tuple(warnings),
+    )
+
+
+def inspect_pdf(path: Path) -> tuple[list[PageInspection], bool]:
+    try:
+        import pymupdf
+    except ImportError as exc:
+        raise OptionalDependencyError("PyMuPDF is required for PDF inspection") from exc
+
+    results: list[PageInspection] = []
+    with pymupdf.open(path) as document:
+        encrypted = bool(getattr(document, "is_encrypted", False) or getattr(document, "needs_pass", False))
+        if encrypted:
+            return results, True
+        for index, page in enumerate(document):
+            text = page.get_text("text").strip()
+            words = page.get_text("words")
+            images = page.get_images(full=True)
+            usable_chars = sum(char.isprintable() for char in text)
+            needs_vision = usable_chars < 80 and len(words) < 15 and len(images) > 0
+            results.append(
+                PageInspection(
+                    page_number=index + 1,
+                    character_count=usable_chars,
+                    word_count=len(words),
+                    image_count=len(images),
+                    needs_vision=needs_vision,
+                )
+            )
+    return results, False
+
+
+def inspect_workbook(path: Path, limits: IngestionLimits) -> int:
+    try:
+        from openpyxl import load_workbook
+    except ImportError as exc:
+        raise OptionalDependencyError("openpyxl is required for workbook inspection") from exc
+
+    workbook = load_workbook(path, read_only=True, data_only=True)
+    try:
+        sheet_count = len(workbook.sheetnames)
+        if sheet_count > limits.max_workbook_sheets:
+            raise WorkbookSheetLimitExceeded(
+                f"{path.name} has {sheet_count} sheets, exceeding limit {limits.max_workbook_sheets}"
+            )
+        for sheet in workbook.worksheets:
+            cells = (sheet.max_row or 0) * (sheet.max_column or 0)
+            if cells > limits.max_sheet_cells:
+                raise WorksheetCellLimitExceeded(
+                    f"{sheet.title} has {cells} cells, exceeding limit {limits.max_sheet_cells}"
+                )
+        return sheet_count
+    finally:
+        workbook.close()

docgun/agents/quality.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from docgun.domain.models import DocumentElement, IngestedDocument
+
+
+@dataclass(frozen=True, slots=True)
+class QualityReport:
+    confidence: float
+    text_coverage: float
+    printable_ratio: float
+    structure_score: float
+    table_score: float
+    metadata_score: float
+    low_confidence_element_ids: tuple[str, ...]
+
+
+def score_document(document: IngestedDocument, threshold: float = 0.70) -> QualityReport:
+    if not document.elements:
+        return QualityReport(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, ())
+
+    text_coverage = _text_coverage(document.elements)
+    printable_ratio = _printable_ratio(document.elements)
+    structure_score = _structure_score(document.elements)
+    table_score = _table_score(document.elements)
+    metadata_score = 1.0 if document.metadata else 0.5
+    element_confidence = sum(element.confidence for element in document.elements) / len(
+        document.elements
+    )
+    confidence = (
+        0.30 * text_coverage
+        + 0.15 * printable_ratio
+        + 0.15 * structure_score
+        + 0.15 * table_score
+        + 0.15 * element_confidence
+        + 0.10 * metadata_score
+    )
+    low_confidence = tuple(
+        element.element_id for element in document.elements if element.confidence < threshold
+    )
+
+    return QualityReport(
+        confidence=round(confidence, 4),
+        text_coverage=text_coverage,
+        printable_ratio=printable_ratio,
+        structure_score=structure_score,
+        table_score=table_score,
+        metadata_score=metadata_score,
+        low_confidence_element_ids=low_confidence,
+    )
+
+
+def _text_coverage(elements: list[DocumentElement]) -> float:
+    with_text = sum(1 for element in elements if element.text or element.table)
+    return with_text / len(elements)
+
+
+def _printable_ratio(elements: list[DocumentElement]) -> float:
+    text = "\n".join(element.text or "" for element in elements)
+    if not text:
+        return 0.0
+    return sum(char.isprintable() or char.isspace() for char in text) / len(text)
+
+
+def _structure_score(elements: list[DocumentElement]) -> float:
+    structured = sum(1 for element in elements if element.hierarchy or element.location)
+    return max(0.5, structured / len(elements))
+
+
+def _table_score(elements: list[DocumentElement]) -> float:
+    tables = [element.table for element in elements if element.table]
+    if not tables:
+        return 1.0
+
+    rectangular = 0
+    for table in tables:
+        widths = {len(row) for row in table}
+        if len(widths) == 1:
+            rectangular += 1
+    return rectangular / len(tables)

docgun/agents/recovery.py

@@ -0,0 +1,206 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+from docgun.agents.quality import QualityReport, score_document
+from docgun.domain.models import DocumentElement, FileInspection, IngestedDocument
+from docgun.parsers.pdf_docling import DoclingParser
+from docgun.agents.vision import VisionPageRecovery
+
+
+@dataclass(frozen=True, slots=True)
+class RecoveryAction:
+    name: str
+    reason: str
+
+
+def choose_recovery_actions(parser: str, confidence: float) -> tuple[RecoveryAction, ...]:
+    if confidence >= 0.70:
+        return ()
+    if parser == "pymupdf":
+        return (
+            RecoveryAction("pymupdf-blocks", "Retry native extraction with block coordinates"),
+            RecoveryAction("docling", "Use layout-aware structured conversion"),
+            RecoveryAction("vision-page", "Render only failed pages for VLM fallback"),
+        )
+    return (RecoveryAction("alternate-parser", "Retry with the next deterministic parser"),)
+
+
+@dataclass(frozen=True, slots=True)
+class RecoveryResult:
+    document: IngestedDocument
+    quality: QualityReport
+    executed_actions: tuple[str, ...]
+    attempts: tuple[RecoveryAttempt, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class RecoveryAttempt:
+    action: str
+    status: Literal["succeeded", "failed", "skipped", "not_available"]
+    confidence_before: float
+    confidence_after: float | None = None
+    error_type: str | None = None
+    error_message: str | None = None
+
+
+class RecoveryEngine:
+    def __init__(self, vision_recovery: VisionPageRecovery | None = None):
+        self.vision_recovery = vision_recovery
+
+    def recover(
+        self,
+        path: Path,
+        inspection: FileInspection,
+        original: IngestedDocument,
+        quality: QualityReport,
+        parser: str,
+        threshold: float,
+    ) -> RecoveryResult:
+        actions = choose_recovery_actions(parser, quality.confidence)
+        attempts: list[RecoveryAttempt] = []
+        best = original
+        best_quality = quality
+
+        if parser == "pymupdf":
+            docling_result, docling_attempt = self._try_docling(path, inspection, quality.confidence)
+            attempts.append(docling_attempt)
+            if docling_result is not None:
+                docling_quality = score_document(docling_result, threshold=threshold)
+                attempts[-1] = RecoveryAttempt(
+                    action="docling",
+                    status="succeeded",
+                    confidence_before=quality.confidence,
+                    confidence_after=docling_quality.confidence,
+                )
+                if docling_quality.confidence > best_quality.confidence:
+                    best = docling_result
+                    best_quality = docling_quality
+
+            if best_quality.confidence < threshold or any(
+                page.needs_vision for page in inspection.page_inspections
+            ):
+                page_recovery, vision_attempt = self._recover_vision_pages(path, inspection, best, threshold)
+                if page_recovery is not best:
+                    best = page_recovery
+                    best_quality = score_document(best, threshold=threshold)
+                if vision_attempt is not None:
+                    attempts.append(vision_attempt)
+        elif actions:
+            attempts.append(
+                RecoveryAttempt(
+                    action=actions[0].name,
+                    status="skipped",
+                    confidence_before=quality.confidence,
+                    error_message="No alternate parser is configured for this document type",
+                )
+            )
+
+        executed = tuple(attempt.action for attempt in attempts if attempt.status in {"succeeded", "not_available"})
+        return RecoveryResult(best, best_quality, executed, tuple(attempts))
+
+    def _recover_vision_pages(self, path: Path, inspection: FileInspection, document: IngestedDocument, threshold: float) -> tuple[IngestedDocument, RecoveryAttempt | None]:
+        failed = [p.page_number for p in inspection.page_inspections if p.needs_vision and _page_confidence([
+            e for e in document.elements if e.location and e.location.page == p.page_number
+        ]) < threshold]
+        if not failed:
+            return document, None
+        before = score_document(document, threshold=threshold).confidence
+        if self.vision_recovery is None:
+            marked = self._mark_pages_for_vision(inspection, document)
+            return marked, RecoveryAttempt("vision-page", "not_available", before, score_document(marked, threshold=threshold).confidence, error_message="No vision backend is configured; page markers were emitted")
+        try:
+            recovered = self.vision_recovery.recover(path, failed)
+            recovered_pages = {e.location.page for e in recovered if e.location and e.location.page}
+            unresolved = set(failed) - recovered_pages
+            kept = [e for e in document.elements if not (e.location and e.location.page in recovered_pages)]
+            merged = IngestedDocument(document.document_id, document.source_name, document.mime_type, document.sha256, kept + recovered, {**document.metadata, "recovery": "vision-page"}, list(document.warnings))
+            if unresolved:
+                merged = self._mark_pages_for_vision(inspection, merged)
+            after = score_document(merged, threshold=threshold).confidence
+            status = "succeeded" if not unresolved else "failed"
+            return merged, RecoveryAttempt("vision-page", status, before, after, error_message=f"No usable result for pages {sorted(unresolved)}" if unresolved else None)
+        except Exception as exc:
+            marked = self._mark_pages_for_vision(inspection, document)
+            return marked, RecoveryAttempt("vision-page", "failed", before, score_document(marked, threshold=threshold).confidence, exc.__class__.__name__, str(exc))
+
+    def _try_docling(
+        self,
+        path: Path,
+        inspection: FileInspection,
+        confidence_before: float,
+    ) -> tuple[IngestedDocument | None, RecoveryAttempt]:
+        try:
+            return DoclingParser().parse(path, inspection), RecoveryAttempt(
+                action="docling",
+                status="succeeded",
+                confidence_before=confidence_before,
+            )
+        except Exception as exc:
+            status: Literal["failed", "not_available"] = (
+                "not_available" if exc.__class__.__name__ == "OptionalDependencyError" else "failed"
+            )
+            return None, RecoveryAttempt(
+                action="docling",
+                status=status,
+                confidence_before=confidence_before,
+                error_type=exc.__class__.__name__,
+                error_message=str(exc),
+            )
+
+    def _mark_pages_for_vision(
+        self,
+        inspection: FileInspection,
+        document: IngestedDocument,
+    ) -> IngestedDocument:
+        failed_pages = [
+            page for page in inspection.page_inspections if page.needs_vision
+        ]
+        if not failed_pages:
+            return document
+
+        elements = list(document.elements)
+        for page in failed_pages:
+            page_elements = [
+                element
+                for element in document.elements
+                if element.location and element.location.page == page.page_number
+            ]
+            native_confidence = _page_confidence(page_elements)
+            if native_confidence >= 0.70:
+                continue
+            elements.append(
+                DocumentElement(
+                    element_id=f"p{page.page_number}-vision-required",
+                    type="image",
+                    text=None,
+                    metadata={
+                        "recovery": "vision-page",
+                        "recovery_pending": True,
+                        "page": page.page_number,
+                        "reason": "native PDF text was below usable threshold",
+                        "native_page_confidence": native_confidence,
+                    },
+                    confidence=0.25,
+                )
+            )
+
+        warnings = list(document.warnings)
+        warnings.append("One or more pages require VLM/OCR recovery")
+        return IngestedDocument(
+            document_id=document.document_id,
+            source_name=document.source_name,
+            mime_type=document.mime_type,
+            sha256=document.sha256,
+            elements=elements,
+            metadata={**document.metadata, "recovery": "vision-page"},
+            warnings=warnings,
+        )
+
+
+def _page_confidence(elements: list[DocumentElement]) -> float:
+    if not elements:
+        return 0.0
+    return sum(element.confidence for element in elements) / len(elements)

docgun/agents/router.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from docgun.domain.errors import UnsupportedDocumentError
+from docgun.domain.models import FileInspection
+
+
+@dataclass(frozen=True, slots=True)
+class RouteDecision:
+    parser: str
+    reason: str
+    needs_recovery: bool = False
+
+
+def route_document(inspection: FileInspection) -> RouteDecision:
+    suffix = inspection.extension.lower()
+
+    if suffix == ".pdf":
+        return RouteDecision(
+            parser="pymupdf",
+            reason="PDF fast path uses native text extraction first",
+            needs_recovery=inspection.likely_scanned,
+        )
+    if suffix == ".docx":
+        return RouteDecision(parser="python-docx", reason="DOCX direct XML extraction")
+    if suffix in {".xlsx", ".xlsm"}:
+        return RouteDecision(parser="openpyxl", reason="Workbook structure is preserved")
+    if suffix == ".csv":
+        return RouteDecision(parser="pyarrow", reason="Fast deterministic CSV parsing")
+
+    raise UnsupportedDocumentError(f"Unsupported document extension: {suffix or '<none>'}")

docgun/agents/vision.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Protocol
+
+from docgun.agents.confidence import text_confidence
+from docgun.domain.errors import OptionalDependencyError
+from docgun.domain.models import BoundingBox, DocumentElement
+
+
+@dataclass(frozen=True, slots=True)
+class RenderedPage:
+    page_number: int
+    png: bytes
+    width: int
+    height: int
+
+
+class VisionBackend(Protocol):
+    def extract_batch(self, pages: Sequence[RenderedPage]) -> Sequence[Any]: ...
+
+
+class VisionPageRecovery:
+    """Render selected PDF pages and normalize a configured OCR/VLM response."""
+
+    def __init__(self, backend: VisionBackend | Callable[[Sequence[RenderedPage]], Sequence[Any]], dpi: int = 180):
+        self.backend = backend
+        self.dpi = dpi
+
+    def recover(self, pdf_path: Path, page_numbers: list[int]) -> list[DocumentElement]:
+        pages = self._render_pages(pdf_path, page_numbers)
+        extract = getattr(self.backend, "extract_batch", self.backend)
+        results = extract(pages)
+        if len(results) != len(pages):
+            raise ValueError("Vision backend returned a different number of page results")
+        elements: list[DocumentElement] = []
+        for page, result in zip(pages, results, strict=True):
+            elements.extend(self._normalize_page(page.page_number, result))
+        return elements
+
+    def _render_pages(self, path: Path, page_numbers: list[int]) -> list[RenderedPage]:
+        try:
+            import pymupdf
+        except ImportError as exc:
+            raise OptionalDependencyError("PyMuPDF is required for page rendering") from exc
+        scale = self.dpi / 72
+        rendered: list[RenderedPage] = []
+        with pymupdf.open(path) as document:
+            for page_number in page_numbers:
+                if page_number < 1 or page_number > len(document):
+                    raise ValueError(f"PDF page {page_number} is out of range")
+                pixmap = document[page_number - 1].get_pixmap(matrix=pymupdf.Matrix(scale, scale), alpha=False)
+                rendered.append(RenderedPage(page_number, pixmap.tobytes("png"), pixmap.width, pixmap.height))
+        return rendered
+
+    @staticmethod
+    def _normalize_page(page_number: int, result: Any) -> list[DocumentElement]:
+        items = result.get("elements", []) if isinstance(result, dict) else result
+        if isinstance(items, str):
+            items = [{"text": items}]
+        normalized: list[DocumentElement] = []
+        for index, item in enumerate(items or [], start=1):
+            if isinstance(item, str):
+                item = {"text": item}
+            text = item.get("text")
+            kind = item.get("type", "paragraph")
+            if kind not in {"heading", "paragraph", "table", "image", "formula"}:
+                kind = "paragraph"
+            bbox = item.get("bbox")
+            location = BoundingBox(page_number, *bbox) if bbox and len(bbox) == 4 else BoundingBox(page_number)
+            confidence = float(item.get("confidence", text_confidence(text, location, element_type="heading" if kind == "heading" else "paragraph")))
+            normalized.append(DocumentElement(
+                element_id=f"p{page_number}-vision-{index}", type=kind, text=text,
+                table=item.get("table"), location=location, confidence=confidence,
+                metadata={"parser": "vision", "recovery_source": "vision-page", "page": page_number},
+            ))
+        return normalized

docgun/api/__init__.py

@@ -0,0 +1 @@
+"""FastAPI routes."""