kreuzberg 3.8.0__py3-none-any.whl → 3.8.2__py3-none-any.whl

kreuzberg/_extractors/_pdf.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import contextlib
+import os
+import tempfile
 from multiprocessing import cpu_count
 from pathlib import Path
 from re import Pattern
@@ -15,8 +17,12 @@ from playa import parse
 from kreuzberg._extractors._base import Extractor
 from kreuzberg._mime_types import PDF_MIME_TYPE, PLAIN_TEXT_MIME_TYPE
 from kreuzberg._ocr import get_ocr_backend
+from kreuzberg._ocr._easyocr import EasyOCRConfig
+from kreuzberg._ocr._paddleocr import PaddleOCRConfig
+from kreuzberg._ocr._tesseract import TesseractConfig
 from kreuzberg._playa import extract_pdf_metadata, extract_pdf_metadata_sync
 from kreuzberg._types import ExtractionResult, OcrBackendType
+from kreuzberg._utils._errors import create_error_context, should_retry
 from kreuzberg._utils._pdf_lock import pypdfium_file_lock
 from kreuzberg._utils._string import normalize_spaces
 from kreuzberg._utils._sync import run_sync, run_taskgroup_batched
@@ -89,9 +95,6 @@ class PDFExtractor(Extractor):
 
     def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
         """Pure sync implementation of PDF extraction from bytes."""
-        import os
-        import tempfile
-
         fd, temp_path = tempfile.mkstemp(suffix=".pdf")
         try:
             with os.fdopen(fd, "wb") as f:
@@ -191,8 +194,6 @@ class PDFExtractor(Extractor):
         Returns:
             A list of Pillow Images.
         """
-        from kreuzberg._utils._errors import create_error_context, should_retry
-
         document: pypdfium2.PdfDocument | None = None
         last_error = None
 
@@ -247,9 +248,10 @@ class PDFExtractor(Extractor):
             *[backend.process_image(image, **self.config.get_config_dict()) for image in images],
             batch_size=cpu_count(),
         )
-        return ExtractionResult(
-            content="\n".join([v.content for v in ocr_results]), mime_type=PLAIN_TEXT_MIME_TYPE, metadata={}, chunks=[]
-        )
+        # Use list comprehension and join for efficient string building
+        content = "\n".join(result.content for result in ocr_results)
+
+        return ExtractionResult(content=content, mime_type=PLAIN_TEXT_MIME_TYPE, metadata={}, chunks=[])
 
     @staticmethod
     async def _extract_pdf_searchable_text(input_file: Path) -> str:
@@ -264,28 +266,28 @@ class PDFExtractor(Extractor):
         Returns:
             The extracted text.
         """
-        from kreuzberg._utils._errors import create_error_context
-
         document: pypdfium2.PdfDocument | None = None
         try:
             with pypdfium_file_lock(input_file):
                 document = await run_sync(pypdfium2.PdfDocument, str(input_file))
-                text_parts = []
+                pages_content = []
                 page_errors = []
 
                 for i, page in enumerate(cast("pypdfium2.PdfDocument", document)):
                     try:
                         text_page = page.get_textpage()
-                        text_parts.append(text_page.get_text_bounded())
+                        page_content = text_page.get_text_bounded()
+                        pages_content.append(page_content)
                     except Exception as e:  # noqa: PERF203, BLE001
                         page_errors.append({"page": i + 1, "error": str(e)})
-                        text_parts.append(f"[Error extracting page {i + 1}]")
+                        pages_content.append(f"[Error extracting page {i + 1}]")
 
-                text = "\n".join(text_parts)
+                text = "\n".join(pages_content)
+                has_content = bool(text.strip())
 
-                if page_errors and text_parts:
+                if page_errors and has_content:
                     return normalize_spaces(text)
-                if not text_parts:
+                if not has_content:
                     raise ParsingError(
                         "Could not extract any text from PDF",
                         context=create_error_context(
@@ -316,14 +318,14 @@ class PDFExtractor(Extractor):
         try:
             with pypdfium_file_lock(path):
                 pdf = pypdfium2.PdfDocument(str(path))
-                text_parts = []
+                pages_text = []
                 for page in pdf:
                     text_page = page.get_textpage()
                     text = text_page.get_text_bounded()
-                    text_parts.append(text)
+                    pages_text.append(text)
                     text_page.close()
                     page.close()
-                return "".join(text_parts)
+                return "\n".join(pages_text)
         except Exception as e:
             raise ParsingError(f"Failed to extract PDF text: {e}") from e
         finally:
@@ -345,9 +347,6 @@ class PDFExtractor(Extractor):
                     bitmap.close()
                     page.close()
 
-            import os
-            import tempfile
-
             image_paths = []
             temp_files = []
 
@@ -375,46 +374,29 @@ class PDFExtractor(Extractor):
 
     def _process_pdf_images_with_ocr(self, image_paths: list[str]) -> str:
         """Process PDF images with the configured OCR backend."""
-        if self.config.ocr_backend == "tesseract":
-            from kreuzberg._ocr._sync import process_batch_images_sync
-            from kreuzberg._ocr._tesseract import TesseractConfig
+        backend = get_ocr_backend(self.config.ocr_backend)
+        paths = [Path(p) for p in image_paths]
 
-            tesseract_config = (
+        if self.config.ocr_backend == "tesseract":
+            config = (
                 self.config.ocr_config if isinstance(self.config.ocr_config, TesseractConfig) else TesseractConfig()
             )
-            results = process_batch_images_sync([str(p) for p in image_paths], tesseract_config, backend="tesseract")
-            text_parts = [r.content for r in results]
-            return "\n\n".join(text_parts)
-
-        if self.config.ocr_backend == "paddleocr":
-            from kreuzberg._ocr._paddleocr import PaddleOCRConfig
-            from kreuzberg._ocr._sync import process_image_paddleocr_sync as paddle_process
-
+            results = backend.process_batch_sync(paths, **config.__dict__)
+        elif self.config.ocr_backend == "paddleocr":
             paddle_config = (
                 self.config.ocr_config if isinstance(self.config.ocr_config, PaddleOCRConfig) else PaddleOCRConfig()
             )
-
-            text_parts = []
-            for image_path in image_paths:
-                result = paddle_process(Path(image_path), paddle_config)
-                text_parts.append(result.content)
-            return "\n\n".join(text_parts)
-
-        if self.config.ocr_backend == "easyocr":
-            from kreuzberg._ocr._easyocr import EasyOCRConfig
-            from kreuzberg._ocr._sync import process_image_easyocr_sync as easy_process
-
+            results = backend.process_batch_sync(paths, **paddle_config.__dict__)
+        elif self.config.ocr_backend == "easyocr":
             easy_config = (
                 self.config.ocr_config if isinstance(self.config.ocr_config, EasyOCRConfig) else EasyOCRConfig()
             )
+            results = backend.process_batch_sync(paths, **easy_config.__dict__)
+        else:
+            raise NotImplementedError(f"Sync OCR not implemented for {self.config.ocr_backend}")
 
-            text_parts = []
-            for image_path in image_paths:
-                result = easy_process(Path(image_path), easy_config)
-                text_parts.append(result.content)
-            return "\n\n".join(text_parts)
-
-        raise NotImplementedError(f"Sync OCR not implemented for {self.config.ocr_backend}")
+        # Use list comprehension and join for efficient string building
+        return "\n\n".join(result.content for result in results)
 
     def _extract_with_playa_sync(self, path: Path, fallback_text: str) -> str:
         """Extract text using playa for better structure preservation."""
@@ -422,14 +404,14 @@ class PDFExtractor(Extractor):
             content = path.read_bytes()
             document = parse(content, max_workers=1)
 
-            text_parts = []
+            # Extract text while preserving structure
+            pages_text = []
             for page in document.pages:
-                # Extract text while preserving structure
                 page_text = page.extract_text()
                 if page_text and page_text.strip():
-                    text_parts.append(page_text)
+                    pages_text.append(page_text)
 
-            if text_parts:
-                return "\n\n".join(text_parts)
+            if pages_text:
+                return "\n\n".join(pages_text)
 
         return fallback_text

kreuzberg/_extractors/_spread_sheet.py

@@ -2,7 +2,9 @@ from __future__ import annotations
 
 import contextlib
 import csv
+import os
 import sys
+import tempfile
 from datetime import date, datetime, time, timedelta
 from io import StringIO
 from pathlib import Path
@@ -68,9 +70,6 @@ class SpreadSheetExtractor(Extractor):
 
     def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
         """Pure sync implementation of extract_bytes."""
-        import os
-        import tempfile
-
         fd, temp_path = tempfile.mkstemp(suffix=".xlsx")
 
         try:

kreuzberg/_extractors/_structured.py

@@ -14,6 +14,9 @@ from kreuzberg._utils._sync import run_sync
 if TYPE_CHECKING:
     from pathlib import Path
 
+# Define text field keywords as a set for O(1) membership testing
+_TEXT_FIELD_KEYWORDS = frozenset({"title", "name", "subject", "description", "content", "body", "text", "message"})
+
 
 class StructuredDataExtractor(Extractor):
     SUPPORTED_MIME_TYPES: ClassVar[set[str]] = {
@@ -70,12 +73,13 @@ class StructuredDataExtractor(Extractor):
             text_parts: list[str] = []
             metadata: dict[str, Any] = {}
 
+            # Use match statement for cleaner code and avoid multiple isinstance calls
             if isinstance(data, dict):
-                text_parts.extend(self._extract_from_dict(data, metadata))
+                text_parts = self._extract_from_dict(data, metadata)
             elif isinstance(data, list):
-                text_parts.extend(self._extract_from_list(data, metadata))
+                text_parts = self._extract_from_list(data, metadata)
             else:
-                text_parts.append(str(data))
+                text_parts = [str(data)]
 
             combined_text = "\n".join(text_parts) if text_parts else text_content
 
@@ -107,10 +111,9 @@ class StructuredDataExtractor(Extractor):
             if isinstance(value, str) and value.strip():
                 text_parts.append(f"{full_key}: {value}")
 
-                if any(
-                    text_field in key.lower()
-                    for text_field in ["title", "name", "subject", "description", "content", "body", "text", "message"]
-                ):
+                # Check if key contains any text field keywords efficiently
+                key_lower = key.lower()
+                if any(keyword in key_lower for keyword in _TEXT_FIELD_KEYWORDS):
                     metadata[full_key] = value
 
             elif isinstance(value, (int, float, bool)):

kreuzberg/_gmft.py

@@ -1,12 +1,20 @@
 from __future__ import annotations
 
+import io
+import multiprocessing as mp
 import os
+import queue
+import signal
+import traceback
 from dataclasses import dataclass, field
+from io import StringIO
 from typing import TYPE_CHECKING, Any, Literal
 
+import msgspec
+
 from kreuzberg._types import TableData
 from kreuzberg._utils._sync import run_sync
-from kreuzberg.exceptions import MissingDependencyError
+from kreuzberg.exceptions import MissingDependencyError, ParsingError
 
 if TYPE_CHECKING:
     from os import PathLike
@@ -15,7 +23,7 @@ if TYPE_CHECKING:
     from pandas import DataFrame
 
 
-@dataclass(unsafe_hash=True)
+@dataclass(unsafe_hash=True, slots=True)
 class GMFTConfig:
     """Configuration options for GMFT.
 
@@ -173,7 +181,7 @@ async def extract_tables(  # noqa: PLR0915
     cache_kwargs = {
         "file_info": str(sorted(file_info.items())),
         "extractor": "gmft",
-        "config": str(sorted(config.__dict__.items())),
+        "config": str(sorted(msgspec.to_builtins(config).items())),
     }
 
     table_cache = get_table_cache()
@@ -196,9 +204,7 @@ async def extract_tables(  # noqa: PLR0915
 
     try:
         if use_isolated_process:
-            from kreuzberg._multiprocessing import extract_tables_isolated_async
-
-            result = await extract_tables_isolated_async(file_path, config)
+            result = await _extract_tables_isolated_async(file_path, config)
 
             await table_cache.aset(result, **cache_kwargs)
 
@@ -305,7 +311,7 @@ def extract_tables_sync(
     cache_kwargs = {
         "file_info": str(sorted(file_info.items())),
         "extractor": "gmft",
-        "config": str(sorted(config.__dict__.items())),
+        "config": str(sorted(msgspec.to_builtins(config).items())),
     }
 
     table_cache = get_table_cache()
@@ -314,9 +320,7 @@ def extract_tables_sync(
         return cached_result  # type: ignore[no-any-return]
 
     if use_isolated_process:
-        from kreuzberg._multiprocessing import extract_tables_isolated
-
-        result = extract_tables_isolated(file_path, config)
+        result = _extract_tables_isolated(file_path, config)
 
         table_cache.set(result, **cache_kwargs)
 
@@ -378,3 +382,303 @@ def extract_tables_sync(
         raise MissingDependencyError.create_for_package(
             dependency_group="gmft", functionality="table extraction", package_name="gmft"
         ) from e
+
+
+def _extract_tables_in_process(
+    file_path: str | PathLike[str],
+    config_dict: dict[str, Any],
+    result_queue: queue.Queue[tuple[bool, Any]],
+) -> None:
+    """Extract tables in an isolated process to handle potential segfaults.
+
+    Args:
+        file_path: Path to the PDF file
+        config_dict: Serialized GMFTConfig as a dict
+        result_queue: Queue to put results or errors
+    """
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+    try:
+        from gmft.auto import AutoTableDetector, AutoTableFormatter  # type: ignore[attr-defined]
+        from gmft.detectors.tatr import TATRDetectorConfig  # type: ignore[attr-defined]
+        from gmft.formatters.tatr import TATRFormatConfig
+        from gmft.pdf_bindings.pdfium import PyPDFium2Document
+
+        config = GMFTConfig(**config_dict)
+
+        formatter = AutoTableFormatter(  # type: ignore[no-untyped-call]
+            config=TATRFormatConfig(
+                verbosity=config.verbosity,
+                formatter_base_threshold=config.formatter_base_threshold,
+                cell_required_confidence=config.cell_required_confidence,
+                remove_null_rows=config.remove_null_rows,
+                enable_multi_header=config.enable_multi_header,
+                semantic_spanning_cells=config.semantic_spanning_cells,
+                semantic_hierarchical_left_fill=config.semantic_hierarchical_left_fill,
+                large_table_if_n_rows_removed=config.large_table_if_n_rows_removed,
+                large_table_threshold=config.large_table_threshold,
+                large_table_row_overlap_threshold=config.large_table_row_overlap_threshold,
+                large_table_maximum_rows=config.large_table_maximum_rows,
+                force_large_table_assumption=config.force_large_table_assumption,
+            )
+        )
+        detector = AutoTableDetector(config=TATRDetectorConfig(detector_base_threshold=config.detector_base_threshold))  # type: ignore[no-untyped-call]
+
+        doc = PyPDFium2Document(str(file_path))
+        cropped_tables = []
+        dataframes = []
+
+        try:
+            for page in doc:
+                cropped_tables.extend(detector.extract(page))  # type: ignore[attr-defined]
+
+            for cropped_table in cropped_tables:
+                formatted_table = formatter.extract(cropped_table)  # type: ignore[attr-defined]
+                dataframes.append(formatted_table.df())
+
+            results = []
+            for data_frame, cropped_table in zip(dataframes, cropped_tables, strict=False):
+                img_bytes = io.BytesIO()
+                cropped_image = cropped_table.image()
+                cropped_image.save(img_bytes, format="PNG")
+                img_bytes.seek(0)
+
+                results.append(
+                    {
+                        "cropped_image_bytes": img_bytes.getvalue(),
+                        "page_number": cropped_table.page.page_number,
+                        "text": data_frame.to_markdown(),
+                        "df_csv": data_frame.to_csv(index=False),
+                    }
+                )
+
+            result_queue.put((True, results))
+
+        finally:
+            doc.close()  # type: ignore[no-untyped-call]
+
+    except Exception as e:  # noqa: BLE001
+        error_info = {"error": str(e), "type": type(e).__name__, "traceback": traceback.format_exc()}
+        result_queue.put((False, error_info))
+
+
+def _extract_tables_isolated(
+    file_path: str | PathLike[str],
+    config: GMFTConfig | None = None,
+    timeout: float = 300.0,
+) -> list[TableData]:
+    """Extract tables using an isolated process to handle segfaults.
+
+    Args:
+        file_path: Path to the PDF file
+        config: GMFT configuration
+        timeout: Maximum time to wait for extraction
+
+    Returns:
+        List of extracted tables
+
+    Raises:
+        RuntimeError: If extraction fails or times out
+    """
+    config = config or GMFTConfig()
+    config_dict = msgspec.to_builtins(config)
+
+    ctx = mp.get_context("spawn")
+    result_queue = ctx.Queue()
+
+    process = ctx.Process(
+        target=_extract_tables_in_process,
+        args=(str(file_path), config_dict, result_queue),
+    )
+
+    process.start()
+
+    try:
+        # Wait for result with timeout, checking for process death  # ~keep
+        import time
+
+        start_time = time.time()
+        while True:
+            try:
+                success, result = result_queue.get_nowait()
+                break
+            except queue.Empty:
+                if time.time() - start_time > timeout:
+                    raise
+
+                if not process.is_alive():
+                    # Process died without putting result  # ~keep
+                    if process.exitcode == -signal.SIGSEGV:
+                        raise ParsingError(
+                            "GMFT process crashed with segmentation fault",
+                            context={
+                                "file_path": str(file_path),
+                                "exit_code": process.exitcode,
+                            },
+                        ) from None
+                    raise ParsingError(
+                        f"GMFT process died unexpectedly with exit code {process.exitcode}",
+                        context={
+                            "file_path": str(file_path),
+                            "exit_code": process.exitcode,
+                        },
+                    ) from None
+
+                time.sleep(0.1)
+
+        if success:
+            tables = []
+            for table_dict in result:
+                from PIL import Image
+
+                img = Image.open(io.BytesIO(table_dict["cropped_image_bytes"]))
+                import pandas as pd
+
+                df = pd.read_csv(StringIO(table_dict["df_csv"]))
+
+                tables.append(
+                    TableData(
+                        cropped_image=img,
+                        page_number=table_dict["page_number"],
+                        text=table_dict["text"],
+                        df=df,
+                    )
+                )
+
+            return tables
+
+        error_info = result
+        raise ParsingError(
+            f"GMFT table extraction failed: {error_info['error']}",
+            context={
+                "file_path": str(file_path),
+                "error_type": error_info["type"],
+                "traceback": error_info["traceback"],
+            },
+        )
+
+    except queue.Empty as e:
+        raise ParsingError(
+            "GMFT table extraction timed out",
+            context={
+                "file_path": str(file_path),
+                "timeout": timeout,
+            },
+        ) from e
+    finally:
+        if process.is_alive():
+            process.terminate()
+            process.join(timeout=5)
+            if process.is_alive():
+                process.kill()
+                process.join()
+
+
+async def _extract_tables_isolated_async(
+    file_path: str | PathLike[str],
+    config: GMFTConfig | None = None,
+    timeout: float = 300.0,
+) -> list[TableData]:
+    """Async version of extract_tables_isolated using asyncio.
+
+    Args:
+        file_path: Path to the PDF file
+        config: GMFT configuration
+        timeout: Maximum time to wait for extraction
+
+    Returns:
+        List of extracted tables
+
+    Raises:
+        RuntimeError: If extraction fails or times out
+    """
+    import anyio
+
+    config = config or GMFTConfig()
+    config_dict = msgspec.to_builtins(config)
+
+    ctx = mp.get_context("spawn")
+    result_queue = ctx.Queue()
+
+    process = ctx.Process(
+        target=_extract_tables_in_process,
+        args=(str(file_path), config_dict, result_queue),
+    )
+
+    process.start()
+
+    try:
+
+        async def wait_for_result() -> tuple[bool, Any]:
+            while True:
+                try:
+                    return result_queue.get_nowait()  # type: ignore[no-any-return]
+                except queue.Empty:  # noqa: PERF203
+                    await anyio.sleep(0.1)
+                    if not process.is_alive():
+                        # Process died without putting result  # ~keep
+                        if process.exitcode == -signal.SIGSEGV:
+                            raise ParsingError(
+                                "GMFT process crashed with segmentation fault",
+                                context={
+                                    "file_path": str(file_path),
+                                    "exit_code": process.exitcode,
+                                },
+                            ) from None
+                        raise ParsingError(
+                            f"GMFT process died unexpectedly with exit code {process.exitcode}",
+                            context={
+                                "file_path": str(file_path),
+                                "exit_code": process.exitcode,
+                            },
+                        ) from None
+
+        with anyio.fail_after(timeout):
+            success, result = await wait_for_result()
+
+        if success:
+            tables = []
+            for table_dict in result:
+                from PIL import Image
+
+                img = Image.open(io.BytesIO(table_dict["cropped_image_bytes"]))
+                import pandas as pd
+
+                df = pd.read_csv(StringIO(table_dict["df_csv"]))
+
+                tables.append(
+                    TableData(
+                        cropped_image=img,
+                        page_number=table_dict["page_number"],
+                        text=table_dict["text"],
+                        df=df,
+                    )
+                )
+
+            return tables
+
+        error_info = result
+        raise ParsingError(
+            f"GMFT table extraction failed: {error_info['error']}",
+            context={
+                "file_path": str(file_path),
+                "error_type": error_info["type"],
+                "traceback": error_info["traceback"],
+            },
+        )
+
+    except TimeoutError as e:
+        raise ParsingError(
+            "GMFT table extraction timed out",
+            context={
+                "file_path": str(file_path),
+                "timeout": timeout,
+            },
+        ) from e
+    finally:
+        if process.is_alive():
+            process.terminate()
+            await anyio.to_thread.run_sync(lambda: process.join(timeout=5))
+            if process.is_alive():
+                process.kill()
+                await anyio.to_thread.run_sync(process.join)

kreuzberg/_language_detection.py

@@ -23,7 +23,7 @@ except ImportError:
 _CACHE_SIZE = 128
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, slots=True)
 class LanguageDetectionConfig:
     """Configuration for language detection.