kreuzberg 3.8.1__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

@@ -14,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."""
 
@@ -127,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)

kreuzberg/_extractors/_pdf.py

@@ -248,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:
@@ -269,22 +270,24 @@ class PDFExtractor(Extractor):
         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(
@@ -315,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:
@@ -392,8 +395,8 @@ class PDFExtractor(Extractor):
         else:
             raise NotImplementedError(f"Sync OCR not implemented for {self.config.ocr_backend}")
 
-        text_parts = [r.content for r in results]
-        return "\n\n".join(text_parts)
+        # 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."""
@@ -401,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)):