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

kreuzberg/__init__.py

@@ -1,5 +1,6 @@
 from importlib.metadata import version
 
+from kreuzberg._config import discover_and_load_config, load_config_from_path, try_discover_config
 from kreuzberg._entity_extraction import SpacyEntityExtractionConfig
 from kreuzberg._gmft import GMFTConfig
 from kreuzberg._language_detection import LanguageDetectionConfig
@@ -48,8 +49,11 @@ __all__ = [
     "batch_extract_bytes_sync",
     "batch_extract_file",
     "batch_extract_file_sync",
+    "discover_and_load_config",
     "extract_bytes",
     "extract_bytes_sync",
     "extract_file",
     "extract_file_sync",
+    "load_config_from_path",
+    "try_discover_config",
 ]

kreuzberg/_api/main.py

@@ -3,7 +3,10 @@ from __future__ import annotations
 from json import dumps
 from typing import TYPE_CHECKING, Annotated, Any
 
+import msgspec
+
 from kreuzberg import (
+    ExtractionConfig,
     ExtractionResult,
     KreuzbergError,
     MissingDependencyError,
@@ -11,6 +14,7 @@ from kreuzberg import (
     ValidationError,
     batch_extract_bytes,
 )
+from kreuzberg._config import try_discover_config
 
 if TYPE_CHECKING:
     from litestar.datastructures import UploadFile
@@ -66,8 +70,12 @@ async def handle_files_upload(
     data: Annotated[list[UploadFile], Body(media_type=RequestEncodingType.MULTI_PART)],
 ) -> list[ExtractionResult]:
     """Extracts text content from an uploaded file."""
+    # Try to discover configuration from files
+    config = try_discover_config()
+
     return await batch_extract_bytes(
         [(await file.read(), file.content_type) for file in data],
+        config=config or ExtractionConfig(),
     )
 
 
@@ -77,8 +85,21 @@ async def health_check() -> dict[str, str]:
     return {"status": "ok"}
 
 
+@get("/config", operation_id="GetConfiguration")
+async def get_configuration() -> dict[str, Any]:
+    """Get the current configuration."""
+    config = try_discover_config()
+    if config is None:
+        return {"message": "No configuration file found", "config": None}
+
+    return {
+        "message": "Configuration loaded successfully",
+        "config": msgspec.to_builtins(config, order="deterministic"),
+    }
+
+
 app = Litestar(
-    route_handlers=[handle_files_upload, health_check],
+    route_handlers=[handle_files_upload, health_check, get_configuration],
     plugins=[OpenTelemetryPlugin(OpenTelemetryConfig())],
     logging_config=StructLoggingConfig(),
     exception_handlers={

kreuzberg/_config.py

@@ -0,0 +1,404 @@
+"""Configuration discovery and loading for Kreuzberg.
+
+This module provides configuration loading from both kreuzberg.toml and pyproject.toml files.
+Configuration is automatically discovered by searching up the directory tree from the current
+working directory.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib  # type: ignore[import-not-found]
+
+from kreuzberg._gmft import GMFTConfig
+from kreuzberg._ocr._easyocr import EasyOCRConfig
+from kreuzberg._ocr._paddleocr import PaddleOCRConfig
+from kreuzberg._ocr._tesseract import TesseractConfig
+from kreuzberg._types import ExtractionConfig, OcrBackendType
+from kreuzberg.exceptions import ValidationError
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+
+def load_config_from_file(config_path: Path) -> dict[str, Any]:
+    """Load configuration from a TOML file.
+
+    Args:
+        config_path: Path to the configuration file.
+
+    Returns:
+        Dictionary containing the loaded configuration.
+
+    Raises:
+        ValidationError: If the file cannot be read or parsed.
+    """
+    try:
+        with config_path.open("rb") as f:
+            data = tomllib.load(f)
+    except FileNotFoundError as e:
+        raise ValidationError(f"Configuration file not found: {config_path}") from e
+    except tomllib.TOMLDecodeError as e:
+        raise ValidationError(f"Invalid TOML in configuration file: {e}") from e
+
+    # Handle both kreuzberg.toml (root level) and pyproject.toml ([tool.kreuzberg])
+    if config_path.name == "kreuzberg.toml":
+        return data  # type: ignore[no-any-return]
+    return data.get("tool", {}).get("kreuzberg", {})  # type: ignore[no-any-return]
+
+
+def merge_configs(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
+    """Merge two configuration dictionaries recursively.
+
+    Args:
+        base: Base configuration dictionary.
+        override: Configuration dictionary to override base values.
+
+    Returns:
+        Merged configuration dictionary.
+    """
+    result = base.copy()
+    for key, value in override.items():
+        if isinstance(value, dict) and key in result and isinstance(result[key], dict):
+            result[key] = merge_configs(result[key], value)
+        else:
+            result[key] = value
+    return result
+
+
+def parse_ocr_backend_config(
+    config_dict: dict[str, Any], backend: OcrBackendType
+) -> TesseractConfig | EasyOCRConfig | PaddleOCRConfig | None:
+    """Parse OCR backend-specific configuration.
+
+    Args:
+        config_dict: Configuration dictionary.
+        backend: The OCR backend type.
+
+    Returns:
+        Backend-specific configuration object or None.
+    """
+    if backend not in config_dict:
+        return None
+
+    backend_config = config_dict[backend]
+    if not isinstance(backend_config, dict):
+        return None
+
+    if backend == "tesseract":
+        # Convert psm integer to PSMMode enum if needed
+        processed_config = backend_config.copy()
+        if "psm" in processed_config and isinstance(processed_config["psm"], int):
+            from kreuzberg._ocr._tesseract import PSMMode
+
+            processed_config["psm"] = PSMMode(processed_config["psm"])
+        return TesseractConfig(**processed_config)
+    if backend == "easyocr":
+        return EasyOCRConfig(**backend_config)
+    if backend == "paddleocr":
+        return PaddleOCRConfig(**backend_config)
+    return None
+
+
+def build_extraction_config_from_dict(config_dict: dict[str, Any]) -> ExtractionConfig:
+    """Build ExtractionConfig from a configuration dictionary.
+
+    Args:
+        config_dict: Configuration dictionary from TOML file.
+
+    Returns:
+        ExtractionConfig instance.
+    """
+    extraction_config: dict[str, Any] = {}
+
+    # Copy basic configuration fields using dictionary comprehension
+    basic_fields = {
+        "force_ocr",
+        "chunk_content",
+        "extract_tables",
+        "max_chars",
+        "max_overlap",
+        "ocr_backend",
+        "extract_entities",
+        "extract_keywords",
+        "auto_detect_language",
+        "enable_quality_processing",
+    }
+    extraction_config.update({field: config_dict[field] for field in basic_fields if field in config_dict})
+
+    # Handle OCR backend configuration
+    ocr_backend = extraction_config.get("ocr_backend")
+    if ocr_backend and ocr_backend != "none":
+        ocr_config = parse_ocr_backend_config(config_dict, ocr_backend)
+        if ocr_config:
+            extraction_config["ocr_config"] = ocr_config
+
+    # Handle GMFT configuration for table extraction
+    if extraction_config.get("extract_tables") and "gmft" in config_dict and isinstance(config_dict["gmft"], dict):
+        extraction_config["gmft_config"] = GMFTConfig(**config_dict["gmft"])
+
+    # Convert "none" to None for ocr_backend
+    if extraction_config.get("ocr_backend") == "none":
+        extraction_config["ocr_backend"] = None
+
+    return ExtractionConfig(**extraction_config)
+
+
+def find_config_file(start_path: Path | None = None) -> Path | None:
+    """Find configuration file by searching up the directory tree.
+
+    Searches for configuration files in the following order:
+    1. kreuzberg.toml
+    2. pyproject.toml (with [tool.kreuzberg] section)
+
+    Args:
+        start_path: Directory to start searching from. Defaults to current working directory.
+
+    Returns:
+        Path to the configuration file or None if not found.
+    """
+    current = start_path or Path.cwd()
+
+    while current != current.parent:
+        # First, look for kreuzberg.toml
+        kreuzberg_toml = current / "kreuzberg.toml"
+        if kreuzberg_toml.exists():
+            return kreuzberg_toml
+
+        # Then, look for pyproject.toml with [tool.kreuzberg] section
+        pyproject_toml = current / "pyproject.toml"
+        if pyproject_toml.exists():
+            try:
+                with pyproject_toml.open("rb") as f:
+                    data = tomllib.load(f)
+                if "tool" in data and "kreuzberg" in data["tool"]:
+                    return pyproject_toml
+            except Exception:  # noqa: BLE001
+                pass
+
+        current = current.parent
+    return None
+
+
+def load_default_config(start_path: Path | None = None) -> ExtractionConfig | None:
+    """Load the default configuration from discovered config file.
+
+    Args:
+        start_path: Directory to start searching from. Defaults to current working directory.
+
+    Returns:
+        ExtractionConfig instance or None if no configuration found.
+    """
+    config_path = find_config_file(start_path)
+    if not config_path:
+        return None
+
+    try:
+        config_dict = load_config_from_file(config_path)
+        if not config_dict:
+            return None
+        return build_extraction_config_from_dict(config_dict)
+    except Exception:  # noqa: BLE001
+        # Silently ignore configuration errors for default loading
+        return None
+
+
+def load_config_from_path(config_path: Path | str) -> ExtractionConfig:
+    """Load configuration from a specific file path.
+
+    Args:
+        config_path: Path to the configuration file.
+
+    Returns:
+        ExtractionConfig instance.
+
+    Raises:
+        ValidationError: If the file cannot be read, parsed, or is invalid.
+    """
+    path = Path(config_path)
+    config_dict = load_config_from_file(path)
+    return build_extraction_config_from_dict(config_dict)
+
+
+def discover_and_load_config(start_path: Path | str | None = None) -> ExtractionConfig:
+    """Load configuration by discovering config files in the directory tree.
+
+    Args:
+        start_path: Directory to start searching from. Defaults to current working directory.
+
+    Returns:
+        ExtractionConfig instance.
+
+    Raises:
+        ValidationError: If no configuration file is found or if the file is invalid.
+    """
+    search_path = Path(start_path) if start_path else None
+    config_path = find_config_file(search_path)
+
+    if not config_path:
+        raise ValidationError(
+            "No configuration file found. Searched for 'kreuzberg.toml' and 'pyproject.toml' with [tool.kreuzberg] section.",
+            context={"search_path": str(search_path or Path.cwd())},
+        )
+
+    config_dict = load_config_from_file(config_path)
+    if not config_dict:
+        raise ValidationError(
+            f"Configuration file found but contains no Kreuzberg configuration: {config_path}",
+            context={"config_path": str(config_path)},
+        )
+
+    return build_extraction_config_from_dict(config_dict)
+
+
+def try_discover_config(start_path: Path | str | None = None) -> ExtractionConfig | None:
+    """Try to discover and load configuration, returning None if not found.
+
+    Args:
+        start_path: Directory to start searching from. Defaults to current working directory.
+
+    Returns:
+        ExtractionConfig instance or None if no configuration found.
+    """
+    try:
+        return discover_and_load_config(start_path)
+    except ValidationError:
+        return None
+
+
+# Legacy functions for backward compatibility with CLI
+
+# Define common configuration fields to avoid repetition
+_CONFIG_FIELDS = [
+    "force_ocr",
+    "chunk_content",
+    "extract_tables",
+    "max_chars",
+    "max_overlap",
+    "ocr_backend",
+    "extract_entities",
+    "extract_keywords",
+    "auto_detect_language",
+    "enable_quality_processing",
+]
+
+
+def _merge_file_config(config_dict: dict[str, Any], file_config: dict[str, Any]) -> None:
+    """Merge file configuration into config dictionary."""
+    if not file_config:
+        return
+
+    for field in _CONFIG_FIELDS:
+        if field in file_config:
+            config_dict[field] = file_config[field]
+
+
+def _merge_cli_args(config_dict: dict[str, Any], cli_args: MutableMapping[str, Any]) -> None:
+    """Merge CLI arguments into config dictionary."""
+    for field in _CONFIG_FIELDS:
+        if field in cli_args and cli_args[field] is not None:
+            config_dict[field] = cli_args[field]
+
+
+def _build_ocr_config_from_cli(
+    ocr_backend: str, cli_args: MutableMapping[str, Any]
+) -> TesseractConfig | EasyOCRConfig | PaddleOCRConfig | None:
+    """Build OCR config from CLI arguments."""
+    config_key = f"{ocr_backend}_config"
+    if not cli_args.get(config_key):
+        return None
+
+    backend_args = cli_args[config_key]
+    if ocr_backend == "tesseract":
+        return TesseractConfig(**backend_args)
+    if ocr_backend == "easyocr":
+        return EasyOCRConfig(**backend_args)
+    if ocr_backend == "paddleocr":
+        return PaddleOCRConfig(**backend_args)
+    return None
+
+
+def _configure_ocr_backend(
+    config_dict: dict[str, Any],
+    file_config: dict[str, Any],
+    cli_args: MutableMapping[str, Any],
+) -> None:
+    """Configure OCR backend in config dictionary."""
+    ocr_backend = config_dict.get("ocr_backend")
+    if not ocr_backend or ocr_backend == "none":
+        return
+
+    # Try CLI config first, then file config
+    ocr_config = _build_ocr_config_from_cli(ocr_backend, cli_args)
+    if not ocr_config and file_config:
+        ocr_config = parse_ocr_backend_config(file_config, ocr_backend)
+
+    if ocr_config:
+        config_dict["ocr_config"] = ocr_config
+
+
+def _configure_gmft(
+    config_dict: dict[str, Any],
+    file_config: dict[str, Any],
+    cli_args: MutableMapping[str, Any],
+) -> None:
+    """Configure GMFT in config dictionary."""
+    if not config_dict.get("extract_tables"):
+        return
+
+    gmft_config = None
+    if cli_args.get("gmft_config"):
+        gmft_config = GMFTConfig(**cli_args["gmft_config"])
+    elif "gmft" in file_config and isinstance(file_config["gmft"], dict):
+        gmft_config = GMFTConfig(**file_config["gmft"])
+
+    if gmft_config:
+        config_dict["gmft_config"] = gmft_config
+
+
+def build_extraction_config(
+    file_config: dict[str, Any],
+    cli_args: MutableMapping[str, Any],
+) -> ExtractionConfig:
+    """Build ExtractionConfig from file config and CLI arguments.
+
+    Args:
+        file_config: Configuration loaded from file.
+        cli_args: CLI arguments.
+
+    Returns:
+        ExtractionConfig instance.
+    """
+    config_dict: dict[str, Any] = {}
+
+    # Merge configurations: file first, then CLI overrides
+    _merge_file_config(config_dict, file_config)
+    _merge_cli_args(config_dict, cli_args)
+
+    # Configure complex components
+    _configure_ocr_backend(config_dict, file_config, cli_args)
+    _configure_gmft(config_dict, file_config, cli_args)
+
+    # Convert "none" to None for ocr_backend
+    if config_dict.get("ocr_backend") == "none":
+        config_dict["ocr_backend"] = None
+
+    return ExtractionConfig(**config_dict)
+
+
+def find_default_config() -> Path | None:
+    """Find the default configuration file (pyproject.toml).
+
+    Returns:
+        Path to the configuration file or None if not found.
+
+    Note:
+        This function is deprecated. Use find_config_file() instead.
+    """
+    return find_config_file()

kreuzberg/_entity_extraction.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import re
 from dataclasses import dataclass
 from functools import lru_cache
@@ -13,7 +14,7 @@ if TYPE_CHECKING:
     from pathlib import Path
 
 
-@dataclass(unsafe_hash=True, frozen=True)
+@dataclass(unsafe_hash=True, frozen=True, slots=True)
 class SpacyEntityExtractionConfig:
     """Configuration for spaCy-based entity extraction."""
 
@@ -126,8 +127,8 @@ def extract_entities(
     """
     entities: list[Entity] = []
     if custom_patterns:
-        custom_patterns_dict = dict(custom_patterns)
-        for ent_type, pattern in custom_patterns_dict.items():
+        # Direct iteration over frozenset - no need to convert to dict
+        for ent_type, pattern in custom_patterns:
             entities.extend(
                 Entity(type=ent_type, text=match.group(), start=match.start(), end=match.end())
                 for match in re.finditer(pattern, text)
@@ -181,8 +182,6 @@ def _load_spacy_model(model_name: str, spacy_config: SpacyEntityExtractionConfig
         import spacy
 
         if spacy_config.model_cache_dir:
-            import os
-
             os.environ["SPACY_DATA"] = str(spacy_config.model_cache_dir)
 
         nlp = spacy.load(model_name)

kreuzberg/_extractors/_base.py

@@ -3,10 +3,12 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, ClassVar
 
+from kreuzberg._types import ExtractionResult, normalize_metadata
+from kreuzberg._utils._quality import calculate_quality_score, clean_extracted_text
+
 if TYPE_CHECKING:
     from pathlib import Path
 
-    from kreuzberg import ExtractionResult
     from kreuzberg._types import ExtractionConfig
 
 
@@ -104,8 +106,6 @@ class Extractor(ABC):
         if not self.config.enable_quality_processing:
             return result
 
-        from kreuzberg._utils._quality import calculate_quality_score, clean_extracted_text
-
         if not result.content:
             return result
 
@@ -120,8 +120,6 @@ class Extractor(ABC):
         enhanced_metadata["quality_score"] = quality_score
 
         # Return enhanced result
-        from kreuzberg._types import ExtractionResult, normalize_metadata
-
         return ExtractionResult(
             content=cleaned_content,
             mime_type=result.mime_type,

kreuzberg/_extractors/_image.py

@@ -11,13 +11,17 @@ from anyio import Path as AsyncPath
 from kreuzberg._extractors._base import Extractor
 from kreuzberg._mime_types import IMAGE_MIME_TYPES
 from kreuzberg._ocr import get_ocr_backend
-from kreuzberg._types import ExtractionResult
+from kreuzberg._ocr._easyocr import EasyOCRConfig
+from kreuzberg._ocr._paddleocr import PaddleOCRConfig
+from kreuzberg._ocr._tesseract import TesseractConfig
 from kreuzberg._utils._tmp import create_temp_file
 from kreuzberg.exceptions import ValidationError
 
 if TYPE_CHECKING:  # pragma: no cover
     from collections.abc import Mapping
 
+    from kreuzberg._types import ExtractionResult
+
 
 class ImageExtractor(Extractor):
     SUPPORTED_MIME_TYPES: ClassVar[set[str]] = IMAGE_MIME_TYPES
@@ -78,44 +82,26 @@ class ImageExtractor(Extractor):
         if self.config.ocr_backend is None:
             raise ValidationError("ocr_backend is None, cannot perform OCR")
 
-        if self.config.ocr_backend == "tesseract":
-            from kreuzberg._ocr._sync import process_batch_images_sync
-            from kreuzberg._ocr._tesseract import TesseractConfig
-
-            if isinstance(self.config.ocr_config, TesseractConfig):
-                config = self.config.ocr_config
-            else:
-                config = TesseractConfig()
-
-            results = process_batch_images_sync([str(path)], config, backend="tesseract")
-            if results:
-                result = results[0]
-                return self._apply_quality_processing(result)
-            return ExtractionResult(content="", mime_type="text/plain", metadata={}, chunks=[])
-
-        if self.config.ocr_backend == "paddleocr":
-            from kreuzberg._ocr._paddleocr import PaddleOCRConfig
-            from kreuzberg._ocr._sync import process_image_paddleocr_sync as paddle_process
+        backend = get_ocr_backend(self.config.ocr_backend)
 
+        if self.config.ocr_backend == "tesseract":
+            config = (
+                self.config.ocr_config if isinstance(self.config.ocr_config, TesseractConfig) else TesseractConfig()
+            )
+            result = backend.process_file_sync(path, **config.__dict__)
+        elif self.config.ocr_backend == "paddleocr":
             paddle_config = (
                 self.config.ocr_config if isinstance(self.config.ocr_config, PaddleOCRConfig) else PaddleOCRConfig()
             )
-
-            result = paddle_process(path, paddle_config)
-            return self._apply_quality_processing(result)
-
-        if self.config.ocr_backend == "easyocr":
-            from kreuzberg._ocr._easyocr import EasyOCRConfig
-            from kreuzberg._ocr._sync import process_image_easyocr_sync as easy_process
-
+            result = backend.process_file_sync(path, **paddle_config.__dict__)
+        elif self.config.ocr_backend == "easyocr":
             easy_config = (
                 self.config.ocr_config if isinstance(self.config.ocr_config, EasyOCRConfig) else EasyOCRConfig()
             )
-
-            result = easy_process(path, easy_config)
-            return self._apply_quality_processing(result)
-
-        raise NotImplementedError(f"Sync OCR not implemented for {self.config.ocr_backend}")
+            result = backend.process_file_sync(path, **easy_config.__dict__)
+        else:
+            raise NotImplementedError(f"Sync OCR not implemented for {self.config.ocr_backend}")
+        return self._apply_quality_processing(result)
 
     def _get_extension_from_mime_type(self, mime_type: str) -> str:
         if mime_type in self.IMAGE_MIME_TYPE_EXT_MAP:

kreuzberg/_extractors/_pandoc.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
 import contextlib
+import os
 import re
+import subprocess
 import sys
+import tempfile
 from json import JSONDecodeError, loads
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast
@@ -203,10 +206,6 @@ class PandocExtractor(Extractor):
         Returns:
             ExtractionResult with the extracted text and metadata.
         """
-        import os
-        import tempfile
-        from pathlib import Path
-
         extension = self._get_pandoc_type_from_mime_type(self.mime_type)
         fd, temp_path = tempfile.mkstemp(suffix=f".{extension}")
 
@@ -579,8 +578,6 @@ class PandocExtractor(Extractor):
 
     def _validate_pandoc_version_sync(self) -> None:
         """Synchronous version of _validate_pandoc_version."""
-        import subprocess
-
         try:
             if self._checked_version:
                 return
@@ -625,10 +622,6 @@ class PandocExtractor(Extractor):
 
     def _extract_metadata_sync(self, path: Path) -> Metadata:
         """Synchronous version of _handle_extract_metadata."""
-        import os
-        import subprocess
-        import tempfile
-
         pandoc_type = self._get_pandoc_type_from_mime_type(self.mime_type)
         fd, metadata_file = tempfile.mkstemp(suffix=".json")
         os.close(fd)
@@ -663,10 +656,6 @@ class PandocExtractor(Extractor):
 
     def _extract_file_sync(self, path: Path) -> str:
         """Synchronous version of _handle_extract_file."""
-        import os
-        import subprocess
-        import tempfile
-
         pandoc_type = self._get_pandoc_type_from_mime_type(self.mime_type)
         fd, output_path = tempfile.mkstemp(suffix=".md")
         os.close(fd)