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

kreuzberg/_mcp/server.py

@@ -3,11 +3,14 @@
 from __future__ import annotations
 
 import base64
+import json
 from typing import Any
 
+import msgspec
 from mcp.server import FastMCP
 from mcp.types import TextContent
 
+from kreuzberg._config import try_discover_config
 from kreuzberg._types import ExtractionConfig, OcrBackendType
 from kreuzberg.extraction import extract_bytes_sync, extract_file_sync
 
@@ -15,6 +18,44 @@ from kreuzberg.extraction import extract_bytes_sync, extract_file_sync
 mcp = FastMCP("Kreuzberg Text Extraction")
 
 
+def _create_config_with_overrides(**kwargs: Any) -> ExtractionConfig:
+    """Create ExtractionConfig with discovered config as base and tool parameters as overrides.
+
+    Args:
+        **kwargs: Tool parameters to override defaults/discovered config.
+
+    Returns:
+        ExtractionConfig instance.
+    """
+    # Try to discover configuration from files
+    base_config = try_discover_config()
+
+    if base_config is None:
+        # No config file found, use defaults
+        return ExtractionConfig(**kwargs)
+
+    # Merge discovered config with tool parameters (tool params take precedence)
+    config_dict: dict[str, Any] = {
+        "force_ocr": base_config.force_ocr,
+        "chunk_content": base_config.chunk_content,
+        "extract_tables": base_config.extract_tables,
+        "extract_entities": base_config.extract_entities,
+        "extract_keywords": base_config.extract_keywords,
+        "ocr_backend": base_config.ocr_backend,
+        "max_chars": base_config.max_chars,
+        "max_overlap": base_config.max_overlap,
+        "keyword_count": base_config.keyword_count,
+        "auto_detect_language": base_config.auto_detect_language,
+        "ocr_config": base_config.ocr_config,
+        "gmft_config": base_config.gmft_config,
+    }
+
+    # Override with provided parameters
+    config_dict.update(kwargs)
+
+    return ExtractionConfig(**config_dict)
+
+
 @mcp.tool()
 def extract_document(  # noqa: PLR0913
     file_path: str,
@@ -49,7 +90,7 @@ def extract_document(  # noqa: PLR0913
     Returns:
         Extracted content with metadata, tables, chunks, entities, and keywords
     """
-    config = ExtractionConfig(
+    config = _create_config_with_overrides(
         force_ocr=force_ocr,
         chunk_content=chunk_content,
         extract_tables=extract_tables,
@@ -63,7 +104,7 @@ def extract_document(  # noqa: PLR0913
     )
 
     result = extract_file_sync(file_path, mime_type, config)
-    return result.to_dict()
+    return result.to_dict(include_none=True)
 
 
 @mcp.tool()
@@ -102,7 +143,7 @@ def extract_bytes(  # noqa: PLR0913
     """
     content_bytes = base64.b64decode(content_base64)
 
-    config = ExtractionConfig(
+    config = _create_config_with_overrides(
         force_ocr=force_ocr,
         chunk_content=chunk_content,
         extract_tables=extract_tables,
@@ -116,7 +157,7 @@ def extract_bytes(  # noqa: PLR0913
     )
 
     result = extract_bytes_sync(content_bytes, mime_type, config)
-    return result.to_dict()
+    return result.to_dict(include_none=True)
 
 
 @mcp.tool()
@@ -133,7 +174,7 @@ def extract_simple(
     Returns:
         Extracted text content as a string
     """
-    config = ExtractionConfig()
+    config = _create_config_with_overrides()
     result = extract_file_sync(file_path, mime_type, config)
     return result.content
 
@@ -142,7 +183,16 @@ def extract_simple(
 def get_default_config() -> str:
     """Get the default extraction configuration."""
     config = ExtractionConfig()
-    return str(config.__dict__)
+    return json.dumps(msgspec.to_builtins(config, order="deterministic"), indent=2)
+
+
+@mcp.resource("config://discovered")
+def get_discovered_config() -> str:
+    """Get the discovered configuration from config files."""
+    config = try_discover_config()
+    if config is None:
+        return "No configuration file found"
+    return json.dumps(msgspec.to_builtins(config, order="deterministic"), indent=2)
 
 
 @mcp.resource("config://available-backends")
@@ -175,7 +225,7 @@ def extract_and_summarize(file_path: str) -> list[TextContent]:
     Returns:
         Extracted content with summarization prompt
     """
-    result = extract_file_sync(file_path, None, ExtractionConfig())
+    result = extract_file_sync(file_path, None, _create_config_with_overrides())
 
     return [
         TextContent(
@@ -195,7 +245,7 @@ def extract_structured(file_path: str) -> list[TextContent]:
     Returns:
         Extracted content with structured analysis prompt
     """
-    config = ExtractionConfig(
+    config = _create_config_with_overrides(
         extract_entities=True,
         extract_keywords=True,
         extract_tables=True,

kreuzberg/_ocr/__init__.py

@@ -4,19 +4,7 @@ from typing import Any
 from kreuzberg._ocr._base import OCRBackend
 from kreuzberg._ocr._easyocr import EasyOCRBackend
 from kreuzberg._ocr._paddleocr import PaddleBackend
-from kreuzberg._ocr._pool import TesseractProcessPool
-from kreuzberg._ocr._sync import (
-    process_batch_images_process_pool,
-    process_batch_images_sync,
-    process_batch_images_threaded,
-    process_image_bytes_easyocr_sync,
-    process_image_bytes_paddleocr_sync,
-    process_image_bytes_tesseract_sync,
-    process_image_easyocr_sync,
-    process_image_paddleocr_sync,
-    process_image_tesseract_sync,
-)
-from kreuzberg._ocr._tesseract import TesseractBackend
+from kreuzberg._ocr._tesseract import TesseractBackend, TesseractProcessPool
 from kreuzberg._types import OcrBackendType
 
 __all__ = [
@@ -26,15 +14,6 @@ __all__ = [
     "TesseractBackend",
     "TesseractProcessPool",
     "get_ocr_backend",
-    "process_batch_images_process_pool",
-    "process_batch_images_sync",
-    "process_batch_images_threaded",
-    "process_image_bytes_easyocr_sync",
-    "process_image_bytes_paddleocr_sync",
-    "process_image_bytes_tesseract_sync",
-    "process_image_easyocr_sync",
-    "process_image_paddleocr_sync",
-    "process_image_tesseract_sync",
 ]
 
 

kreuzberg/_ocr/_base.py

@@ -49,6 +49,65 @@ class OCRBackend(ABC, Generic[T]):
         """
         ...
 
+    @abstractmethod
+    def process_image_sync(self, image: Image, **kwargs: Unpack[T]) -> ExtractionResult:
+        """Synchronously process an image and extract its text and metadata.
+
+        Args:
+            image: An instance of PIL.Image representing the input image.
+            **kwargs: Any kwargs related to the given backend
+
+        Returns:
+            The extraction result object
+        """
+        ...
+
+    @abstractmethod
+    def process_file_sync(self, path: Path, **kwargs: Unpack[T]) -> ExtractionResult:
+        """Synchronously process a file and extract its text and metadata.
+
+        Args:
+            path: A Path object representing the file to be processed.
+            **kwargs: Any kwargs related to the given backend
+
+        Returns:
+            The extraction result object
+        """
+        ...
+
+    def process_batch_sync(self, paths: list[Path], **kwargs: Unpack[T]) -> list[ExtractionResult]:
+        """Synchronously process a batch of files and extract their text and metadata.
+
+        Default implementation processes files sequentially. Backends can override
+        for more efficient batch processing.
+
+        Args:
+            paths: List of Path objects representing files to be processed.
+            **kwargs: Any kwargs related to the given backend
+
+        Returns:
+            List of extraction result objects in the same order as input paths
+        """
+        return [self.process_file_sync(path, **kwargs) for path in paths]
+
+    async def process_batch(self, paths: list[Path], **kwargs: Unpack[T]) -> list[ExtractionResult]:
+        """Asynchronously process a batch of files and extract their text and metadata.
+
+        Default implementation processes files concurrently. Backends can override
+        for more efficient batch processing.
+
+        Args:
+            paths: List of Path objects representing files to be processed.
+            **kwargs: Any kwargs related to the given backend
+
+        Returns:
+            List of extraction result objects in the same order as input paths
+        """
+        from kreuzberg._utils._sync import run_taskgroup
+
+        tasks = [self.process_file(path, **kwargs) for path in paths]
+        return await run_taskgroup(*tasks)
+
     def __hash__(self) -> int:
         """Hash function for allowing caching."""
         return hash(type(self).__name__)

kreuzberg/_ocr/_easyocr.py

@@ -4,6 +4,7 @@ import warnings
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal
 
+import numpy as np
 from PIL import Image
 
 from kreuzberg._mime_types import PLAIN_TEXT_MIME_TYPE
@@ -110,7 +111,7 @@ EASYOCR_SUPPORTED_LANGUAGE_CODES: Final[set[str]] = {
 }
 
 
-@dataclass(unsafe_hash=True, frozen=True)
+@dataclass(unsafe_hash=True, frozen=True, slots=True)
 class EasyOCRConfig:
     """Configuration options for EasyOCR."""
 
@@ -440,3 +441,93 @@ class EasyOCRBackend(OCRBackend[EasyOCRConfig]):
             )
 
         return languages
+
+    def process_image_sync(self, image: Image.Image, **kwargs: Unpack[EasyOCRConfig]) -> ExtractionResult:
+        """Synchronously process an image and extract its text and metadata using EasyOCR.
+
+        Args:
+            image: An instance of PIL.Image representing the input image.
+            **kwargs: Configuration parameters for EasyOCR including language, detection thresholds, etc.
+
+        Returns:
+            ExtractionResult: The extraction result containing text content, mime type, and metadata.
+
+        Raises:
+            OCRError: If OCR processing fails.
+        """
+        self._init_easyocr_sync(**kwargs)
+
+        beam_width = kwargs.pop("beam_width")
+        kwargs.pop("language", None)
+        kwargs.pop("use_gpu", None)
+
+        try:
+            result = self._reader.readtext(
+                np.array(image),
+                beamWidth=beam_width,
+                **kwargs,
+            )
+
+            return self._process_easyocr_result(result, image)
+        except Exception as e:
+            raise OCRError(f"Failed to OCR using EasyOCR: {e}") from e
+
+    def process_file_sync(self, path: Path, **kwargs: Unpack[EasyOCRConfig]) -> ExtractionResult:
+        """Synchronously process a file and extract its text and metadata using EasyOCR.
+
+        Args:
+            path: A Path object representing the file to be processed.
+            **kwargs: Configuration parameters for EasyOCR including language, detection thresholds, etc.
+
+        Returns:
+            ExtractionResult: The extraction result containing text content, mime type, and metadata.
+
+        Raises:
+            OCRError: If file loading or OCR processing fails.
+        """
+        self._init_easyocr_sync(**kwargs)
+        try:
+            image = Image.open(path)
+            return self.process_image_sync(image, **kwargs)
+        except Exception as e:
+            raise OCRError(f"Failed to load or process image using EasyOCR: {e}") from e
+
+    @classmethod
+    def _init_easyocr_sync(cls, **kwargs: Unpack[EasyOCRConfig]) -> None:
+        """Synchronously initialize EasyOCR with the provided configuration.
+
+        Args:
+            **kwargs: Configuration parameters for EasyOCR including language, etc.
+
+        Raises:
+            MissingDependencyError: If EasyOCR is not installed.
+            OCRError: If initialization fails.
+        """
+        if cls._reader is not None:
+            return
+
+        try:
+            import easyocr
+        except ImportError as e:
+            raise MissingDependencyError.create_for_package(
+                dependency_group="easyocr", functionality="EasyOCR as an OCR backend", package_name="easyocr"
+            ) from e
+
+        languages = cls._validate_language_code(kwargs.pop("language", "en"))
+
+        device_info = cls._resolve_device_config(**kwargs)
+        use_gpu = device_info.device_type in ("cuda", "mps")
+
+        kwargs.setdefault("detector", True)
+        kwargs.setdefault("recognizer", True)
+        kwargs.setdefault("download_enabled", True)
+        kwargs.setdefault("recog_network", "standard")
+
+        try:
+            cls._reader = easyocr.Reader(
+                languages,
+                gpu=use_gpu,
+                verbose=False,
+            )
+        except Exception as e:
+            raise OCRError(f"Failed to initialize EasyOCR: {e}") from e

kreuzberg/_ocr/_paddleocr.py

@@ -4,8 +4,10 @@ import platform
 import warnings
 from dataclasses import dataclass
 from importlib.util import find_spec
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal
 
+import numpy as np
 from PIL import Image
 
 from kreuzberg._mime_types import PLAIN_TEXT_MIME_TYPE
@@ -29,7 +31,7 @@ except ImportError:  # pragma: no cover
 PADDLEOCR_SUPPORTED_LANGUAGE_CODES: Final[set[str]] = {"ch", "en", "french", "german", "japan", "korean"}
 
 
-@dataclass(unsafe_hash=True, frozen=True)
+@dataclass(unsafe_hash=True, frozen=True, slots=True)
 class PaddleOCRConfig:
     """Configuration options for PaddleOCR.
 
@@ -364,3 +366,90 @@ class PaddleBackend(OCRBackend[PaddleOCRConfig]):
                 "supported_languages": ",".join(sorted(PADDLEOCR_SUPPORTED_LANGUAGE_CODES)),
             },
         )
+
+    def process_image_sync(self, image: Image.Image, **kwargs: Unpack[PaddleOCRConfig]) -> ExtractionResult:
+        """Synchronously process an image and extract its text and metadata using PaddleOCR.
+
+        Args:
+            image: An instance of PIL.Image representing the input image.
+            **kwargs: Configuration parameters for PaddleOCR including language, detection thresholds, etc.
+
+        Returns:
+            ExtractionResult: The extraction result containing text content, mime type, and metadata.
+
+        Raises:
+            OCRError: If OCR processing fails.
+        """
+        self._init_paddle_ocr_sync(**kwargs)
+
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        image_np = np.array(image)
+        try:
+            result = self._paddle_ocr.ocr(image_np, cls=kwargs.get("use_angle_cls", True))
+            return self._process_paddle_result(result, image)
+        except Exception as e:
+            raise OCRError(f"Failed to OCR using PaddleOCR: {e}") from e
+
+    def process_file_sync(self, path: Path, **kwargs: Unpack[PaddleOCRConfig]) -> ExtractionResult:
+        """Synchronously process a file and extract its text and metadata using PaddleOCR.
+
+        Args:
+            path: A Path object representing the file to be processed.
+            **kwargs: Configuration parameters for PaddleOCR including language, detection thresholds, etc.
+
+        Returns:
+            ExtractionResult: The extraction result containing text content, mime type, and metadata.
+
+        Raises:
+            OCRError: If file loading or OCR processing fails.
+        """
+        self._init_paddle_ocr_sync(**kwargs)
+        try:
+            image = Image.open(path)
+            return self.process_image_sync(image, **kwargs)
+        except Exception as e:
+            raise OCRError(f"Failed to load or process image using PaddleOCR: {e}") from e
+
+    @classmethod
+    def _init_paddle_ocr_sync(cls, **kwargs: Unpack[PaddleOCRConfig]) -> None:
+        """Synchronously initialize PaddleOCR with the provided configuration.
+
+        Args:
+            **kwargs: Configuration parameters for PaddleOCR including language, detection thresholds, etc.
+
+        Raises:
+            MissingDependencyError: If PaddleOCR is not installed.
+            OCRError: If initialization fails.
+        """
+        if cls._paddle_ocr is not None:
+            return
+
+        try:
+            from paddleocr import PaddleOCR
+        except ImportError as e:
+            raise MissingDependencyError.create_for_package(
+                dependency_group="paddleocr", functionality="PaddleOCR as an OCR backend", package_name="paddleocr"
+            ) from e
+
+        language = cls._validate_language_code(kwargs.pop("language", "en"))
+
+        device_info = cls._resolve_device_config(**kwargs)
+        use_gpu = device_info.device_type == "cuda"
+
+        has_gpu_package = bool(find_spec("paddlepaddle_gpu"))
+        kwargs.setdefault("use_angle_cls", True)
+        kwargs["use_gpu"] = use_gpu and has_gpu_package
+        kwargs.setdefault("enable_mkldnn", cls._is_mkldnn_supported() and not (use_gpu and has_gpu_package))
+        kwargs.setdefault("det_db_thresh", 0.3)
+        kwargs.setdefault("det_db_box_thresh", 0.5)
+        kwargs.setdefault("det_db_unclip_ratio", 1.6)
+
+        if device_info.device_type == "cuda" and kwargs.get("gpu_memory_limit"):
+            kwargs["gpu_mem"] = int(kwargs["gpu_memory_limit"] * 1024)
+
+        try:
+            cls._paddle_ocr = PaddleOCR(lang=language, show_log=False, **kwargs)
+        except Exception as e:
+            raise OCRError(f"Failed to initialize PaddleOCR: {e}") from e