multixtract 0.1.1__py3-none-any.whl

multixtract/extractors/registry.py

@@ -0,0 +1,72 @@
+"""Extractor registry — maps file extensions to DocumentExtractor factories.
+
+The registry lets the pipeline pick the right extractor automatically based on
+file type, and lets users register their own extractors for new formats::
+
+    from multixtract.extractors import register_extractor
+    from multixtract.interfaces import DocumentExtractor
+
+    class MarkdownExtractor:
+        extensions = (".md",)
+        def extract(self, path): ...
+
+    register_extractor(MarkdownExtractor())
+"""
+from __future__ import annotations
+
+import os
+from typing import Dict, Iterable
+
+from ..interfaces import DocumentExtractor
+
+
+class ExtractorRegistry:
+    """Dispatches a file path to the DocumentExtractor for its extension."""
+
+    def __init__(self) -> None:
+        self._by_ext: Dict[str, DocumentExtractor] = {}
+
+    def register(self, extractor: DocumentExtractor, extensions: Iterable[str] | None = None) -> DocumentExtractor:
+        """Register *extractor* for the given *extensions* (defaults to the
+        extractor's own ``extensions`` attribute). Later registrations for the
+        same extension win, so users can override built-ins."""
+        exts = extensions if extensions is not None else getattr(extractor, "extensions", ())
+        for ext in exts:
+            self._by_ext[_norm(ext)] = extractor
+        return extractor
+
+    def get(self, path: str) -> DocumentExtractor:
+        """Return the extractor for *path*'s extension, or raise."""
+        ext = _norm(os.path.splitext(path)[1])
+        extractor = self._by_ext.get(ext)
+        if extractor is None:
+            raise ValueError(
+                f"No extractor registered for '{ext or path}'. "
+                f"Supported: {sorted(self._by_ext)}. "
+                f"Register one with register_extractor()."
+            )
+        return extractor
+
+    @property
+    def supported_extensions(self):
+        return sorted(self._by_ext)
+
+
+def _norm(ext: str) -> str:
+    ext = ext.lower().strip()
+    return ext if ext.startswith(".") or ext == "" else f".{ext}"
+
+
+# Process-wide default registry. Built-in extractors register themselves on
+# import of multixtract.extractors.
+default_registry = ExtractorRegistry()
+
+
+def register_extractor(extractor: DocumentExtractor, extensions: Iterable[str] | None = None) -> DocumentExtractor:
+    """Register *extractor* on the process-wide default registry."""
+    return default_registry.register(extractor, extensions)
+
+
+def get_extractor(path: str) -> DocumentExtractor:
+    """Look up the default-registry extractor for *path*."""
+    return default_registry.get(path)

multixtract/extractors/rtf.py

@@ -0,0 +1,54 @@
+"""RTF extractor (.rtf).
+
+Requires striprtf (optional extra [rtf]).
+"""
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, Dict, List, Optional, Tuple
+
+log = logging.getLogger("multixtract")
+
+
+class RtfExtractor:
+    """DocumentExtractor for RTF files."""
+
+    extensions: Tuple[str, ...] = (".rtf",)
+
+    def extract(
+        self,
+        path: str,
+        image_filter: Optional[Any] = None,
+    ) -> Tuple[Dict[str, Any], List[Dict[str, Any]]]:
+        base_name = os.path.splitext(os.path.basename(path))[0]
+        empty = {"_base_name": base_name, "metadata": {}, "pgs": []}
+        try:
+            from striprtf.striprtf import rtf_to_text
+        except ImportError as exc:
+            raise ImportError(
+                "RTF support requires striprtf: pip install 'multixtract[rtf]'"
+            ) from exc
+        try:
+            with open(path, encoding="utf-8", errors="replace") as fh:
+                raw = fh.read()
+
+            txt = rtf_to_text(raw).strip()
+
+            document = {
+                "_base_name": base_name,
+                "metadata": {
+                    "page_count": 1,
+                    "format": "rtf",
+                    "char_count": len(txt),
+                },
+                "pgs": [
+                    {"pg_num": 1, "txt": txt, "tables": [], "imgs": []}
+                ],
+            }
+            return document, []
+        except ImportError:
+            raise
+        except Exception:
+            log.warning("RtfExtractor failed for %s", path, exc_info=True)
+            return empty, []

multixtract/extractors/text.py

@@ -0,0 +1,50 @@
+"""Plain-text extractor (.txt / .log / .conf / .ini / .md).
+
+The entire file becomes a single page. A dedicated MarkdownExtractor
+registered later will override the .md extension with richer handling.
+"""
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, Dict, List, Optional, Tuple
+
+log = logging.getLogger("multixtract")
+
+
+class TextExtractor:
+    """DocumentExtractor for plain-text files."""
+
+    extensions: Tuple[str, ...] = (".txt", ".log", ".conf", ".ini", ".md")
+
+    def extract(
+        self,
+        path: str,
+        image_filter: Optional[Any] = None,
+    ) -> Tuple[Dict[str, Any], List[Dict[str, Any]]]:
+        base_name = os.path.splitext(os.path.basename(path))[0]
+        fmt = os.path.splitext(path)[1].lstrip(".").lower() or "txt"
+        empty = {"_base_name": base_name, "metadata": {}, "pgs": []}
+        try:
+            try:
+                with open(path, encoding="utf-8") as fh:
+                    text = fh.read()
+            except UnicodeDecodeError:
+                with open(path, encoding="latin-1", errors="replace") as fh:
+                    text = fh.read()
+
+            document = {
+                "_base_name": base_name,
+                "metadata": {
+                    "page_count": 1,
+                    "format": fmt,
+                    "char_count": len(text),
+                },
+                "pgs": [
+                    {"pg_num": 1, "txt": text, "tables": [], "imgs": []}
+                ],
+            }
+            return document, []
+        except Exception:
+            log.warning("TextExtractor failed for %s", path, exc_info=True)
+            return empty, []

multixtract/filters.py

@@ -0,0 +1,181 @@
+"""Image filtering pipeline (vendor-neutral).
+
+Decides which extracted images are worth analyzing. Depends only on Pillow
+and ImageHash — no cloud SDKs.
+
+Filter order (correctness + performance):
+  1. Absolute minimum  — max(w,h) < 20 -> instant reject (no PIL decode)
+  2. PIL decode        — single decode shared by later stages
+  3. pHash compute     — shared by logo detection
+  4. Reference-logo    — highest priority so small logos are tagged, not kept
+  5. Dimension filter  — two-threshold (max & min side)
+  6. Content quality   — solid_color, tiny_icon
+
+Cross-page dedup is handled by the extractor via xref tracking, not here.
+"""
+from __future__ import annotations
+
+import io
+import os
+from collections import defaultdict
+from typing import Any, Dict, List, Optional, Tuple
+
+import imagehash
+from PIL import Image
+
+_ABSOLUTE_MIN_DIM = 20
+
+
+class ImageFilterPipeline:
+    """Stateful image filtering and preparation.
+
+    One instance per process; call :meth:`reset` at the start of each document
+    to clear per-document filter statistics.
+    """
+
+    SOLID_RANGE_MAX = 35
+    ICON_MAX_DIM = 200
+    ICON_MAX_COLORS = 8
+    SAMPLE_SIZE = 64
+    COLOR_SAMPLE_MAX = 256
+    LOGO_PHASH_THRESHOLD = 60
+    LOGO_ASPECT_RANGE = (0.2, 5.0)
+
+    def __init__(
+        self,
+        min_image_size: int = 100,
+        min_image_size_minor: int = 75,
+        reference_img_dir: str = "",
+    ) -> None:
+        self.min_image_size = min_image_size
+        self.min_image_size_minor = min_image_size_minor
+        self._reference_hashes: List[Tuple[imagehash.ImageHash, str]] = []
+        if reference_img_dir:
+            self._load_reference_images(reference_img_dir)
+        self.reset()
+
+    # ---- reference logos -------------------------------------------------
+
+    def _load_reference_images(self, ref_dir: str) -> None:
+        if not os.path.isdir(ref_dir):
+            return
+        exts = {".png", ".jpg", ".jpeg", ".bmp", ".gif", ".tiff"}
+        for fname in sorted(os.listdir(ref_dir)):
+            if os.path.splitext(fname)[1].lower() not in exts:
+                continue
+            try:
+                img = Image.open(os.path.join(ref_dir, fname)).convert("RGB")
+                try:
+                    self._reference_hashes.append((imagehash.phash(img, hash_size=16), fname))
+                finally:
+                    img.close()
+            except Exception:
+                pass
+
+    def _is_reference_logo(self, phash, width: int, height: int) -> Tuple[bool, str]:
+        if not self._reference_hashes:
+            return False, ""
+        aspect = width / max(height, 1)
+        aspect_min, aspect_max = self.LOGO_ASPECT_RANGE
+        if not (aspect_min <= aspect <= aspect_max):
+            return False, ""
+        best_dist, best_ref = min(
+            ((phash - rh, rn) for rh, rn in self._reference_hashes),
+            key=lambda x: x[0],
+        )
+        return (best_dist <= self.LOGO_PHASH_THRESHOLD, best_ref)
+
+    # ---- per-document state ---------------------------------------------
+
+    def reset(self) -> None:
+        """Reset per-document filter statistics."""
+        self._filter_stats: Dict[str, int] = defaultdict(int)
+
+    @property
+    def filter_stats(self) -> Dict[str, int]:
+        return dict(self._filter_stats)
+
+    def note_duplicate(self) -> None:
+        """Record a cross-page duplicate (called by the extractor)."""
+        self._filter_stats["duplicate"] += 1
+
+    # ---- content quality -------------------------------------------------
+
+    def _is_low_value(self, small, small_nearest) -> Tuple[bool, str]:
+        if max(channel_max - channel_min for channel_min, channel_max in small.getextrema()) < self.SOLID_RANGE_MAX:
+            return True, "solid_color"
+        if small_nearest is not None:
+            try:
+                ic = small_nearest.getcolors(maxcolors=self.COLOR_SAMPLE_MAX)
+                if ic and len(ic) <= self.ICON_MAX_COLORS:
+                    return True, "tiny_icon"
+            except Exception:
+                pass
+        return False, ""
+
+    # ---- main entry point ------------------------------------------------
+
+    def prepare_image(
+        self,
+        image_bytes: bytes,
+        ext: str,
+        width: int,
+        height: int,
+        image_id: str,
+        page_number: int,
+        img_idx: int,
+    ) -> Optional[Dict[str, Any]]:
+        """Apply all filters to one image. Returns a metadata dict for images
+        that pass, or ``None`` for images that are filtered out. The returned
+        dict carries ``image_bytes`` so the caller can run vision on it.
+        """
+        if max(width, height) < _ABSOLUTE_MIN_DIM:
+            self._filter_stats["dimension"] += 1
+            return None
+
+        try:
+            img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        except Exception:
+            self._filter_stats["decode_error"] += 1
+            return None
+
+        try:
+            phash = imagehash.phash(img, hash_size=16)
+
+            is_logo, _ = self._is_reference_logo(phash, width, height)
+            if is_logo:
+                self._filter_stats["ref_logo"] += 1
+                return None
+
+            if max(width, height) < self.min_image_size or min(width, height) < self.min_image_size_minor:
+                self._filter_stats["dimension"] += 1
+                return None
+
+            sample_size = self.SAMPLE_SIZE
+            small = img.resize((sample_size, sample_size), Image.Resampling.NEAREST)
+            # Reuse `small` for the icon colour check — same dimensions, no second decode.
+            small_nearest = small if max(width, height) < self.ICON_MAX_DIM else None
+            try:
+                skip, reason = self._is_low_value(small, small_nearest)
+            except Exception:
+                skip, reason = False, ""
+            finally:
+                small.close()
+                # small_nearest is either None or the same object as small; don't double-close.
+            if skip:
+                self._filter_stats[reason] += 1
+                return None
+        finally:
+            img.close()
+
+        self._filter_stats["kept"] += 1
+        return {
+            "image_id": image_id,
+            "page_number": page_number,
+            "img_idx": img_idx,
+            "image_bytes": image_bytes,
+            "ext": ext,
+            "width": width,
+            "height": height,
+            "img_path": f"pg{page_number}_img{img_idx}.{ext}",
+        }

multixtract/interfaces.py

@@ -0,0 +1,151 @@
+"""Vendor-neutral provider interfaces.
+
+The core pipeline depends ONLY on these contracts — never on a concrete
+vendor SDK. Swap OpenAI for a local model, Azure Blob for S3 or local disk,
+by providing a class that satisfies the relevant Protocol.
+
+Protocols use structural typing: any object with the right methods works,
+no explicit subclassing required.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Protocol, Tuple, runtime_checkable
+
+
+@dataclass
+class VisionResult:
+    """Structured output of a vision model describing a single image."""
+
+    caption: str = ""
+    ocr_text: str = ""
+    description: str = ""
+
+    def best_text(self) -> str:
+        """Text most suitable for embedding (description, else caption)."""
+        return self.description or self.caption or ""
+
+
+@runtime_checkable
+class VisionModel(Protocol):
+    """Turns image bytes into a caption + OCR text + description.
+
+    Implement this to plug in any vision-capable model (GPT-4o, Claude,
+    Gemini, a local Llama 3.2 Vision, etc.).
+    """
+
+    def analyze(
+        self,
+        image_bytes: bytes,
+        ext: str = "png",
+        width: int = 0,
+        height: int = 0,
+    ) -> VisionResult:
+        """Return a :class:`VisionResult` for one image. Must not raise on
+        model errors — return an empty ``VisionResult`` instead."""
+        ...
+
+
+@runtime_checkable
+class Embedder(Protocol):
+    """Turns text into fixed-length numeric vectors.
+
+    Implement this to plug in any embedding model (OpenAI, Cohere,
+    sentence-transformers, etc.).
+    """
+
+    #: Dimensionality of the vectors this embedder produces.
+    dim: int
+
+    def embed(self, texts: List[str]) -> List[Optional[List[float]]]:
+        """Embed a batch of texts. Returns one vector per input (same order);
+        ``None`` for empty or failed items. Implementations are encouraged to
+        batch internally to minimise API calls."""
+        ...
+
+
+@runtime_checkable
+class BlobStore(Protocol):
+    """Persists bytes and JSON to some backing store.
+
+    Implement this to plug in any storage target (local disk, Azure Blob,
+    S3, GCS, etc.). ``path`` is always a relative, forward-slash path.
+    """
+
+    def put_bytes(self, path: str, data: bytes, content_type: str = "") -> str:
+        """Store raw bytes at *path*. Returns a locator (URL or absolute path)."""
+        ...
+
+    def put_json(self, path: str, obj: object, compact: bool = False) -> str:
+        """Serialize *obj* to JSON and store at *path*. When *compact* is True,
+        use separator-only JSON (~30-40% smaller). Returns a locator."""
+        ...
+
+    def exists(self, path: str) -> bool:
+        """Return True if an object already exists at *path* (resume support)."""
+        ...
+
+
+@runtime_checkable
+class DocumentExtractor(Protocol):
+    """Reads a document file into the normalized document structure.
+
+    Implement this to support a new file format (PDF, Word, PowerPoint, Excel,
+    HTML, ...). The downstream pipeline (filtering, vision, chunking, embedding,
+    storage) is format-agnostic and consumes the structure returned here.
+
+    The returned document is a dict::
+
+        {
+            "metadata": {...},
+            "_base_name": "<file stem>",
+            "pgs": [                      # a "page" = page / slide / sheet / section
+                {"pg_num": 1, "kind": "page", "txt": "...",
+                 "tables": [[[...]]], "imgs": []},
+                ...
+            ],
+        }
+
+    and ``prepared_images`` is the list of images that passed filtering, each a
+    dict carrying ``image_bytes`` for downstream vision analysis.
+    """
+
+    #: Lower-case file extensions this extractor handles, e.g. ``(".pdf",)``.
+    extensions: Tuple[str, ...]
+
+    def extract(
+        self,
+        path: str,
+        image_filter: Optional[Any] = None,
+    ) -> Tuple[Dict[str, Any], List[Dict[str, Any]]]:
+        """Return ``(document, prepared_images)`` for the file at *path*.
+
+        *image_filter* is an optional :class:`ImageFilterPipeline`; when omitted
+        the extractor should create a default one. Text-only extractors may
+        ignore it.
+        """
+        ...
+
+
+@dataclass
+class PipelineConfig:
+    """Tunable knobs shared across the pipeline."""
+
+    # Image filtering
+    min_image_size: int = 100
+    min_image_size_minor: int = 75
+    reference_img_dir: str = ""
+
+    # Vision / embedding
+    vision_workers: int = 6
+    embed_text_limit: int = 8000
+
+    # Chunking
+    chunk_target_tokens: int = 500
+    chunk_overlap_tokens: int = 50
+
+    # Storage layout (relative sub-folders under the store root)
+    images_subdir: str = "extracted_images"
+    doc_json_subdir: str = "jsons"
+    image_json_subdir: str = "image_jsons"
+    chunks_subdir: str = "chunks"