kreuzberg 3.14.1__py3-none-any.whl → 3.16.0__py3-none-any.whl

kreuzberg/_extractors/_base.py

@@ -1,16 +1,41 @@
 from __future__ import annotations
 
+import io
+import logging
+import time
+import zlib
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, ClassVar
+from dataclasses import asdict
+from multiprocessing import cpu_count
+from typing import TYPE_CHECKING, Any, ClassVar
 
-from kreuzberg._types import ExtractionResult, normalize_metadata
+from PIL import Image
+
+from kreuzberg._ocr import get_ocr_backend
+from kreuzberg._types import (
+    EasyOCRConfig,
+    ExtractedImage,
+    ExtractionResult,
+    ImageOCRResult,
+    PaddleOCRConfig,
+    TesseractConfig,
+    normalize_metadata,
+)
 from kreuzberg._utils._quality import calculate_quality_score, clean_extracted_text
+from kreuzberg._utils._sync import run_taskgroup_batched
 
 if TYPE_CHECKING:
     from pathlib import Path
 
     from kreuzberg._types import ExtractionConfig
 
+MAX_TOTAL_IMAGE_SIZE_MB = 100
+MAX_SINGLE_IMAGE_SIZE_MB = 50
+MAX_TOTAL_IMAGE_SIZE = MAX_TOTAL_IMAGE_SIZE_MB * 1024 * 1024
+MAX_SINGLE_IMAGE_SIZE = MAX_SINGLE_IMAGE_SIZE_MB * 1024 * 1024
+
+logger = logging.getLogger(__name__)
+
 
 class Extractor(ABC):
     __slots__ = ("config", "mime_type")
@@ -52,11 +77,252 @@ class Extractor(ABC):
 
         enhanced_metadata = (dict(result.metadata) if result.metadata else {}) | {"quality_score": quality_score}
 
+        deduplicated_images = self._deduplicate_images(result.images) if result.images else []
+
         return ExtractionResult(
             content=cleaned_content,
             mime_type=result.mime_type,
             metadata=normalize_metadata(enhanced_metadata),
+            tables=result.tables,
             chunks=result.chunks,
+            images=deduplicated_images,
+            image_ocr_results=result.image_ocr_results,
+            entities=result.entities,
+            keywords=result.keywords,
             detected_languages=result.detected_languages,
-            tables=result.tables,
+            document_type=result.document_type,
+            document_type_confidence=result.document_type_confidence,
+            layout=result.layout,
+        )
+
+    def _check_image_memory_limits(self, images: list[ExtractedImage]) -> list[ExtractedImage]:
+        """Filter images based on memory safety limits."""
+        if not images:
+            return []
+
+        images_with_sizes = [(img, len(img.data)) for img in images]
+
+        valid_images = []
+        for img, size in images_with_sizes:
+            if size <= MAX_SINGLE_IMAGE_SIZE:
+                valid_images.append((img, size))
+            else:
+                logger.warning(
+                    "Skipping image %s: size %d MB exceeds limit of %d MB",
+                    img.filename or "unknown",
+                    size // (1024 * 1024),
+                    MAX_SINGLE_IMAGE_SIZE_MB,
+                )
+
+        total_size = sum(size for _, size in valid_images)
+
+        if total_size <= MAX_TOTAL_IMAGE_SIZE:
+            return [img for img, _ in valid_images]
+
+        logger.warning(
+            "Total image size %d MB exceeds limit of %d MB, selecting subset",
+            total_size // (1024 * 1024),
+            MAX_TOTAL_IMAGE_SIZE_MB,
         )
+
+        sorted_images = sorted(valid_images, key=lambda x: x[1])
+        selected = []
+        current_size = 0
+
+        for img, img_size in sorted_images:
+            if current_size + img_size <= MAX_TOTAL_IMAGE_SIZE:
+                selected.append(img)
+                current_size += img_size
+            else:
+                logger.debug("Skipping image %s: would exceed total memory limit", img.filename or "unknown")
+
+        return selected
+
+    _SMALL_IMAGE_THRESHOLD = 1024
+    _HASH_SAMPLE_SIZE = 512
+
+    def _compute_image_hash(self, img: ExtractedImage) -> int:
+        """Compute hash for image deduplication using progressive hashing.
+
+        For small images (<1KB), hash the entire content.
+        For larger images, use size + first/last bytes for quick comparison.
+
+        Args:
+            img: Image to hash
+
+        Returns:
+            Hash value for deduplication
+        """
+        data_len = len(img.data)
+
+        if data_len < self._SMALL_IMAGE_THRESHOLD:
+            return zlib.crc32(img.data) & 0xFFFFFFFF
+
+        hash_components = [
+            str(data_len).encode(),
+            img.data[: self._HASH_SAMPLE_SIZE],
+            img.data[-self._HASH_SAMPLE_SIZE :],
+            img.format.encode() if img.format else b"",
+        ]
+
+        combined = b"".join(hash_components)
+        return zlib.crc32(combined) & 0xFFFFFFFF
+
+    def _deduplicate_images(self, images: list[ExtractedImage]) -> list[ExtractedImage]:
+        if not self.config.deduplicate_images or not images:
+            return images
+
+        seen_hashes = set()
+        unique_images = []
+
+        for img in images:
+            img_hash = self._compute_image_hash(img)
+            if img_hash not in seen_hashes:
+                seen_hashes.add(img_hash)
+                unique_images.append(img)
+            else:
+                logger.debug("Filtered duplicate image: %s", img.filename)
+
+        if len(unique_images) < len(images):
+            logger.info("Deduplicated %d images to %d unique", len(images), len(unique_images))
+
+        return unique_images
+
+    def _prepare_ocr_config(self, backend_name: str) -> dict[str, Any]:
+        """Prepare OCR configuration for the specified backend.
+
+        Args:
+            backend_name: Name of the OCR backend
+
+        Returns:
+            Configuration dictionary for the backend
+        """
+        default_config: TesseractConfig | EasyOCRConfig | PaddleOCRConfig
+        config_class: type[TesseractConfig | EasyOCRConfig | PaddleOCRConfig]
+
+        if backend_name == "tesseract":
+            default_config = TesseractConfig()
+            config_class = TesseractConfig
+        elif backend_name == "easyocr":
+            default_config = EasyOCRConfig()
+            config_class = EasyOCRConfig
+        elif backend_name == "paddleocr":
+            default_config = PaddleOCRConfig()
+            config_class = PaddleOCRConfig
+        else:
+            raise ValueError(f"Unknown OCR backend: {backend_name}")
+
+        cfg: dict[str, Any] = asdict(default_config)
+
+        if self.config.ocr_config and isinstance(self.config.ocr_config, config_class):
+            user_cfg: dict[str, Any] = asdict(self.config.ocr_config)
+            cfg.update(user_cfg)
+
+        cfg["use_cache"] = self.config.use_cache
+        return cfg
+
+    def _validate_image_for_ocr(self, img: ExtractedImage) -> str | None:
+        """Validate if an image is suitable for OCR processing.
+
+        Args:
+            img: Image to validate
+
+        Returns:
+            Reason for skipping if invalid, None if valid
+        """
+        fmt = img.format.lower()
+        if fmt not in self.config.image_ocr_formats:
+            return f"Unsupported format: {img.format}"
+
+        if img.dimensions is not None:
+            w, h = img.dimensions
+            min_w, min_h = self.config.image_ocr_min_dimensions
+            max_w, max_h = self.config.image_ocr_max_dimensions
+
+            if w < min_w or h < min_h:
+                return f"Too small: {w}x{h}"
+            if w > max_w or h > max_h:
+                return f"Too large: {w}x{h}"
+
+        return None
+
+    async def _ocr_single_image(self, target: ExtractedImage, backend: Any, cfg: dict[str, Any]) -> ImageOCRResult:
+        """Process a single image with OCR.
+
+        Args:
+            target: Image to process
+            backend: OCR backend instance
+            cfg: Configuration for the backend
+
+        Returns:
+            OCR result for the image
+        """
+        try:
+            start = time.time()
+            pil_img = Image.open(io.BytesIO(target.data))
+            ocr_res = await backend.process_image(pil_img, **cfg)
+            duration = time.time() - start
+            return ImageOCRResult(
+                image=target,
+                ocr_result=ocr_res,
+                confidence_score=None,
+                processing_time=duration,
+            )
+        except (OSError, ValueError) as e:  # pragma: no cover
+            return ImageOCRResult(
+                image=target,
+                ocr_result=ExtractionResult(content="", mime_type="text/plain", metadata={}),
+                skipped_reason=f"OCR failed: {type(e).__name__}: {e}",
+            )
+        except (RuntimeError, TypeError) as e:  # pragma: no cover
+            return ImageOCRResult(
+                image=target,
+                ocr_result=ExtractionResult(content="", mime_type="text/plain", metadata={}),
+                skipped_reason=f"Backend error: {type(e).__name__}: {e}",
+            )
+
+    async def _process_images_with_ocr(
+        self, images: tuple[ExtractedImage, ...] | list[ExtractedImage]
+    ) -> list[ImageOCRResult]:
+        """Process multiple images with OCR.
+
+        Args:
+            images: Tuple or list of images to process
+
+        Returns:
+            List of OCR results
+        """
+        if not images or not self.config.ocr_extracted_images:
+            return []
+
+        images_list = list(self._deduplicate_images(list(images)))
+        images_list = self._check_image_memory_limits(images_list)
+
+        backend_name = self.config.image_ocr_backend or self.config.ocr_backend
+        if backend_name is None:
+            return []
+
+        cfg = self._prepare_ocr_config(backend_name)
+        backend = get_ocr_backend(backend_name)
+
+        results: list[ImageOCRResult] = []
+        tasks = []
+
+        for img in images_list:
+            skip_reason = self._validate_image_for_ocr(img)
+            if skip_reason:
+                results.append(
+                    ImageOCRResult(
+                        image=img,
+                        ocr_result=ExtractionResult(content="", mime_type="text/plain", metadata={}),
+                        skipped_reason=skip_reason,
+                    )
+                )
+            else:
+                tasks.append(self._ocr_single_image(img, backend, cfg))
+
+        if tasks:
+            batch_size = max(1, min(len(tasks), cpu_count()))
+            results.extend(await run_taskgroup_batched(*tasks, batch_size=batch_size))
+
+        return results

kreuzberg/_extractors/_email.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import base64
 import re
 from html import unescape
 from typing import TYPE_CHECKING, Any, ClassVar
@@ -8,9 +9,8 @@ from anyio import Path as AsyncPath
 
 from kreuzberg._extractors._base import Extractor
 from kreuzberg._mime_types import EML_MIME_TYPE, PLAIN_TEXT_MIME_TYPE
-from kreuzberg._types import ExtractionResult, normalize_metadata
-from kreuzberg._utils._string import normalize_spaces
-from kreuzberg._utils._sync import run_sync
+from kreuzberg._types import ExtractedImage, ExtractionResult, ImageOCRResult, normalize_metadata
+from kreuzberg._utils._sync import run_maybe_async, run_sync
 from kreuzberg.exceptions import MissingDependencyError
 
 if TYPE_CHECKING:
@@ -27,6 +27,8 @@ except ImportError:  # pragma: no cover
     html2text = None
 
 _HTML_TAG_PATTERN = re.compile(r"<[^>]+>")
+_UNICODE_QUOTES_PATTERN = re.compile(r"[\u201c\u201d]")
+_UNICODE_SINGLE_QUOTES_PATTERN = re.compile(r"[\u2018\u2019]")
 
 
 class EmailExtractor(Extractor):
@@ -84,24 +86,25 @@ class EmailExtractor(Extractor):
             text_parts.append(f"BCC: {bcc_formatted}")
 
     def _format_email_field(self, field: Any) -> str:
-        if isinstance(field, list):
-            emails = []
-            for item in field:
-                if isinstance(item, dict):
-                    email = item.get("email", "")
-                    if email:
-                        emails.append(email)
-                else:
-                    emails.append(str(item))
-            return ", ".join(emails)
-        if isinstance(field, dict):
-            return str(field.get("email", ""))
-        return str(field)
+        match field:
+            case list():
+                emails = []
+                for item in field:
+                    if isinstance(item, dict):
+                        if email := item.get("email", ""):
+                            emails.append(str(email))
+                    else:
+                        emails.append(str(item))
+                return ", ".join(emails)
+            case dict():
+                return str(field.get("email", ""))
+            case _:
+                return str(field)
 
     def _extract_email_body(self, parsed_email: dict[str, Any], text_parts: list[str]) -> None:
         text_content = parsed_email.get("text")
         if text_content:
-            text_parts.append(f"\n{text_content}")
+            text_parts.append(str(text_content))
             return
 
         html_content = parsed_email.get("html")
@@ -111,20 +114,80 @@ class EmailExtractor(Extractor):
                 h.ignore_links = True
                 h.ignore_images = True
                 converted_text = h.handle(html_content)
-                text_parts.append(f"\n{converted_text}")
+                text_parts.append(converted_text)
             else:
-                clean_html = _HTML_TAG_PATTERN.sub("", html_content)
+                cleaned = re.sub(r"<script[^>]*>.*?</script>", "", html_content, flags=re.IGNORECASE | re.DOTALL)
+                cleaned = re.sub(r"<style[^>]*>.*?</style>", "", cleaned, flags=re.IGNORECASE | re.DOTALL)
+                clean_html = _HTML_TAG_PATTERN.sub("", cleaned)
                 clean_html = unescape(clean_html)
-                text_parts.append(f"\n{clean_html}")
+                clean_html = _UNICODE_QUOTES_PATTERN.sub('"', clean_html)
+                clean_html = _UNICODE_SINGLE_QUOTES_PATTERN.sub("'", clean_html)
+                text_parts.append(clean_html)
 
     def _extract_email_attachments(
         self, parsed_email: dict[str, Any], text_parts: list[str], metadata: dict[str, Any]
     ) -> None:
-        if parsed_email.get("attachments"):
-            attachment_names = [att.get("name", "unknown") for att in parsed_email["attachments"]]
-            metadata["attachments"] = attachment_names
-            if attachment_names:
-                text_parts.append(f"\nAttachments: {', '.join(attachment_names)}")
+        attachments = parsed_email.get("attachments")
+        if not isinstance(attachments, list):
+            return
+        names: list[str] = []
+        for att in attachments:
+            name_val: str = "unknown"
+            if isinstance(att, dict):
+                n = att.get("name") or att.get("filename")
+                if isinstance(n, str) and n:
+                    name_val = n
+            names.append(name_val)
+        if names:
+            metadata["attachments"] = names
+            text_parts.append("Attachments: " + ", ".join(names))
+
+    def _extract_images_from_attachments(self, parsed_email: dict[str, Any]) -> list[ExtractedImage]:
+        images: list[ExtractedImage] = []
+        attachments = parsed_email.get("attachments") or []
+        if not isinstance(attachments, list):
+            return []
+
+        for idx, att in enumerate(attachments, start=1):
+            if not isinstance(att, dict):
+                continue
+
+            mime = att.get("mime") or att.get("content_type") or att.get("type")
+            if not isinstance(mime, str) or not mime.startswith("image/"):
+                continue
+
+            name = att.get("name") or att.get("filename")
+            name = name if isinstance(name, str) else None
+            data = att.get("data") or att.get("content") or att.get("payload")
+            raw: bytes | None = None
+            if isinstance(data, (bytes, bytearray)):
+                raw = bytes(data)
+            elif isinstance(data, str):
+                try:
+                    raw = base64.b64decode(data)
+                except Exception:  # noqa: BLE001
+                    raw = data.encode()
+
+            if raw is None:
+                continue
+
+            fmt = mime.split("/", 1)[1].lower()
+            if name and "." in name:
+                ext = name.rsplit(".", 1)[-1].lower()
+                if ext:
+                    fmt = ext
+
+            filename = name or f"attachment_image_{idx}.{fmt}"
+            images.append(
+                ExtractedImage(
+                    data=raw,
+                    format=fmt,
+                    filename=filename,
+                    page_number=None,
+                )
+            )
+
+        return images
 
     def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
         if mailparse is None:
@@ -142,13 +205,24 @@ class EmailExtractor(Extractor):
 
             combined_text = "\n".join(text_parts)
 
-            return ExtractionResult(
-                content=normalize_spaces(combined_text),
+            result = ExtractionResult(
+                content=combined_text,
                 mime_type=PLAIN_TEXT_MIME_TYPE,
                 metadata=normalize_metadata(metadata),
                 chunks=[],
             )
 
+            if self.config.extract_images:
+                images = self._extract_images_from_attachments(parsed_email)
+                result.images = images
+                if self.config.ocr_extracted_images and result.images:
+                    image_ocr_results: list[ImageOCRResult] = run_maybe_async(
+                        self._process_images_with_ocr, result.images
+                    )
+                    result.image_ocr_results = image_ocr_results
+
+            return result
+
         except Exception as e:
             msg = f"Failed to parse email content: {e}"
             raise RuntimeError(msg) from e

kreuzberg/_extractors/_html.py

@@ -1,29 +1,44 @@
 from __future__ import annotations
 
+import base64
+import binascii
+import io
+import logging
 from typing import TYPE_CHECKING, ClassVar
 
 import html_to_markdown
 from anyio import Path as AsyncPath
+from bs4 import BeautifulSoup
+from PIL import Image
 
-from kreuzberg._extractors._base import Extractor
+from kreuzberg._extractors._base import MAX_SINGLE_IMAGE_SIZE, Extractor
 from kreuzberg._mime_types import HTML_MIME_TYPE, MARKDOWN_MIME_TYPE
-from kreuzberg._types import ExtractionResult, HTMLToMarkdownConfig
+from kreuzberg._types import ExtractedImage, ExtractionResult, HTMLToMarkdownConfig
+from kreuzberg._utils._html_streaming import should_use_streaming
 from kreuzberg._utils._string import safe_decode
-from kreuzberg._utils._sync import run_sync
+from kreuzberg._utils._sync import run_maybe_async, run_sync
 
 if TYPE_CHECKING:
     from pathlib import Path
 
+logger = logging.getLogger(__name__)
+
 
 class HTMLExtractor(Extractor):
     SUPPORTED_MIME_TYPES: ClassVar[set[str]] = {HTML_MIME_TYPE}
 
     async def extract_bytes_async(self, content: bytes) -> ExtractionResult:
-        return await run_sync(self.extract_bytes_sync, content)
+        result = await run_sync(self.extract_bytes_sync, content)
+        if self.config.extract_images and self.config.ocr_extracted_images and result.images:
+            result.image_ocr_results = await self._process_images_with_ocr(result.images)
+        return result
 
     async def extract_path_async(self, path: Path) -> ExtractionResult:
         content = await AsyncPath(path).read_bytes()
-        return await run_sync(self.extract_bytes_sync, content)
+        result = await run_sync(self.extract_bytes_sync, content)
+        if self.config.extract_images and self.config.ocr_extracted_images and result.images:
+            result.image_ocr_results = await self._process_images_with_ocr(result.images)
+        return result
 
     def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
         config = self.config.html_to_markdown_config if self.config else None
@@ -32,12 +47,102 @@ class HTMLExtractor(Extractor):
 
         config_dict = config.to_dict()
 
-        result = html_to_markdown.convert_to_markdown(safe_decode(content), **config_dict)
+        html_content = safe_decode(content)
+
+        use_streaming, chunk_size = should_use_streaming(len(content))
+        config_dict["stream_processing"] = use_streaming
+        config_dict["chunk_size"] = chunk_size
+
+        result = html_to_markdown.convert_to_markdown(html_content, **config_dict)
 
-        extraction_result = ExtractionResult(content=result, mime_type=MARKDOWN_MIME_TYPE, metadata={}, chunks=[])
+        extraction_result = ExtractionResult(content=result, mime_type=MARKDOWN_MIME_TYPE, metadata={})
+
+        if self.config.extract_images:
+            extraction_result.images = self._extract_images_from_html(html_content)
+            if self.config.ocr_extracted_images and extraction_result.images:
+                extraction_result.image_ocr_results = run_maybe_async(
+                    self._process_images_with_ocr, extraction_result.images
+                )
 
         return self._apply_quality_processing(extraction_result)
 
     def extract_path_sync(self, path: Path) -> ExtractionResult:
         content = path.read_bytes()
         return self.extract_bytes_sync(content)
+
+    def _extract_images_from_html(self, html_content: str) -> list[ExtractedImage]:
+        images: list[ExtractedImage] = []
+        soup = BeautifulSoup(html_content, "xml")
+
+        for img in soup.find_all("img"):
+            src_val = img.get("src")  # type: ignore[union-attr]
+            if isinstance(src_val, str) and src_val.startswith("data:image/"):
+                try:
+                    header, data = src_val.split(",", 1)
+                    mime_type = header.split(";")[0].split(":")[1]
+                    format_name = mime_type.split("/")[1]
+
+                    if not data or len(data) < 4:
+                        logger.debug("Skipping empty or too small base64 data")
+                        continue
+
+                    if len(data) > 67 * 1024 * 1024:
+                        logger.warning("Skipping base64 image larger than 67MB")
+                        continue
+
+                    image_data = base64.b64decode(data)
+
+                    if len(image_data) > MAX_SINGLE_IMAGE_SIZE:
+                        logger.warning(
+                            "Skipping decoded image larger than %dMB", MAX_SINGLE_IMAGE_SIZE // (1024 * 1024)
+                        )
+                        continue
+
+                    dimensions = None
+                    try:
+                        with Image.open(io.BytesIO(image_data)) as pil_img:
+                            dimensions = pil_img.size
+                    except (OSError, ValueError) as e:
+                        logger.debug("Could not determine image dimensions for %s: %s", format_name, e)
+
+                    alt_val = img.get("alt")  # type: ignore[union-attr]
+                    desc = alt_val if isinstance(alt_val, str) else None
+                    images.append(
+                        ExtractedImage(
+                            data=image_data,
+                            format=format_name,
+                            filename=f"embedded_image_{len(images) + 1}.{format_name}",
+                            description=desc,
+                            dimensions=dimensions,
+                        )
+                    )
+                except (ValueError, binascii.Error) as e:
+                    logger.warning("Failed to extract base64 image: %s", e)
+
+        def extract_svg_safe(svg_element: object) -> ExtractedImage | None:
+            try:
+                svg_content = str(svg_element).encode("utf-8")
+
+                def _get_attr_safe(obj: object, attr: str) -> str | None:
+                    get_method = getattr(obj, "get", None)
+                    if callable(get_method):
+                        result = get_method(attr)
+                        return result if isinstance(result, str) else None
+                    return None
+
+                title_or_aria = _get_attr_safe(svg_element, "title") or _get_attr_safe(svg_element, "aria-label")
+                desc_svg = title_or_aria if isinstance(title_or_aria, str) else None
+                return ExtractedImage(
+                    data=svg_content,
+                    format="svg",
+                    filename=f"inline_svg_{len(images) + 1}.svg",
+                    description=desc_svg,
+                )
+            except (UnicodeEncodeError, AttributeError) as e:
+                logger.warning("Failed to extract SVG: %s", e)
+                return None
+
+        svg_images = [extract_svg_safe(svg) for svg in soup.find_all("svg")]
+        images.extend(img for img in svg_images if img is not None)
+
+        return images