kreuzberg 3.8.1__py3-none-any.whl → 3.9.0__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/_chunker.py

@@ -2,9 +2,9 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from kreuzberg import MissingDependencyError
 from kreuzberg._constants import DEFAULT_MAX_CHARACTERS, DEFAULT_MAX_OVERLAP
 from kreuzberg._mime_types import MARKDOWN_MIME_TYPE
+from kreuzberg.exceptions import MissingDependencyError
 
 if TYPE_CHECKING:
     from semantic_text_splitter import MarkdownSplitter, TextSplitter
@@ -36,11 +36,11 @@ def get_chunker(
     if key not in _chunkers:
         try:
             if mime_type == MARKDOWN_MIME_TYPE:
-                from semantic_text_splitter import MarkdownSplitter
+                from semantic_text_splitter import MarkdownSplitter  # noqa: PLC0415
 
                 _chunkers[key] = MarkdownSplitter(max_characters, overlap_characters)
             else:
-                from semantic_text_splitter import TextSplitter
+                from semantic_text_splitter import TextSplitter  # noqa: PLC0415
 
                 _chunkers[key] = TextSplitter(max_characters, overlap_characters)
         except ImportError as e:

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  # noqa: PLC0415
+
+            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/_document_classification.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+import re
+from typing import TYPE_CHECKING
+
+from kreuzberg._ocr import get_ocr_backend
+from kreuzberg.exceptions import MissingDependencyError
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from kreuzberg._types import ExtractionConfig, ExtractionResult
+
+
+DOCUMENT_CLASSIFIERS = {
+    "invoice": [
+        r"invoice",
+        r"bill to",
+        r"invoice number",
+        r"total amount",
+        r"tax id",
+    ],
+    "receipt": [
+        r"receipt",
+        r"cash receipt",
+        r"payment",
+        r"subtotal",
+        r"total due",
+    ],
+    "contract": [
+        r"agreement",
+        r"contract",
+        r"party a",
+        r"party b",
+        r"terms and conditions",
+        r"signature",
+    ],
+    "report": [r"report", r"summary", r"analysis", r"findings", r"conclusion"],
+    "form": [r"form", r"fill out", r"signature", r"date", r"submit"],
+}
+
+
+def _get_translated_text(result: ExtractionResult) -> str:
+    """Translate extracted text to English using Google Translate API.
+
+    Args:
+        result: ExtractionResult containing the text to be translated
+
+    Returns:
+        str: The translated text in lowercase English
+
+    Raises:
+        MissingDependencyError: If the deep-translator package is not installed
+    """
+    try:
+        from deep_translator import GoogleTranslator  # noqa: PLC0415
+    except ImportError as e:
+        raise MissingDependencyError(
+            "The 'deep-translator' library is not installed. Please install it with: pip install 'kreuzberg[auto-classify-document-type]'"
+        ) from e
+
+    return str(GoogleTranslator(source="auto", target="en").translate(result.content).lower())
+
+
+def classify_document(result: ExtractionResult, config: ExtractionConfig) -> tuple[str | None, float | None]:
+    """Classifies the document type based on keywords and patterns.
+
+    Args:
+        result: The extraction result containing the content.
+        config: The extraction configuration.
+
+    Returns:
+        A tuple containing the detected document type and the confidence score,
+        or (None, None) if no type is detected with sufficient confidence.
+    """
+    translated_text = _get_translated_text(result)
+    scores = dict.fromkeys(DOCUMENT_CLASSIFIERS, 0)
+
+    for doc_type, patterns in DOCUMENT_CLASSIFIERS.items():
+        for pattern in patterns:
+            if re.search(pattern, translated_text):
+                scores[doc_type] += 1
+
+    total_score = sum(scores.values())
+    if total_score == 0:
+        return None, None
+
+    confidences = {doc_type: score / total_score for doc_type, score in scores.items()}
+
+    best_type, best_confidence = max(confidences.items(), key=lambda item: item[1])
+
+    if best_confidence >= config.document_type_confidence_threshold:
+        return best_type, best_confidence
+
+    return None, None
+
+
+def classify_document_from_layout(
+    result: ExtractionResult, config: ExtractionConfig
+) -> tuple[str | None, float | None]:
+    """Classifies the document type based on layout information from OCR.
+
+    Args:
+        result: The extraction result containing the layout data.
+        config: The extraction configuration.
+
+    Returns:
+        A tuple containing the detected document type and the confidence score,
+        or (None, None) if no type is detected with sufficient confidence.
+    """
+    translated_text = _get_translated_text(result)
+
+    if result.layout is None or result.layout.empty:
+        return None, None
+
+    layout_df = result.layout
+    if not all(col in layout_df.columns for col in ["text", "top", "height"]):
+        return None, None
+
+    layout_df["translated_text"] = translated_text
+
+    page_height = layout_df["top"].max() + layout_df["height"].max()
+    scores = dict.fromkeys(DOCUMENT_CLASSIFIERS, 0.0)
+
+    for doc_type, patterns in DOCUMENT_CLASSIFIERS.items():
+        for pattern in patterns:
+            found_words = layout_df[layout_df["translated_text"].str.contains(pattern, case=False, na=False)]
+            if not found_words.empty:
+                scores[doc_type] += 1.0
+                word_top = found_words.iloc[0]["top"]
+                if word_top < page_height * 0.3:
+                    scores[doc_type] += 0.5
+
+    total_score = sum(scores.values())
+    if total_score == 0:
+        return None, None
+
+    confidences = {doc_type: score / total_score for doc_type, score in scores.items()}
+
+    best_type, best_confidence = max(confidences.items(), key=lambda item: item[1])
+
+    if best_confidence >= config.document_type_confidence_threshold:
+        return best_type, best_confidence
+
+    return None, None
+
+
+def auto_detect_document_type(
+    result: ExtractionResult, config: ExtractionConfig, file_path: Path | None = None
+) -> ExtractionResult:
+    if config.document_classification_mode == "vision" and file_path:
+        layout_result = get_ocr_backend("tesseract").process_file_sync(file_path, **config.get_config_dict())
+        result.document_type, result.document_type_confidence = classify_document_from_layout(layout_result, config)
+    else:
+        result.document_type, result.document_type_confidence = classify_document(result, config)
+    return result

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)
@@ -138,7 +138,7 @@ def extract_entities(
         spacy_config = SpacyEntityExtractionConfig()
 
     try:
-        import spacy  # noqa: F401
+        import spacy  # noqa: F401, PLC0415
     except ImportError as e:
         raise MissingDependencyError.create_for_package(
             package_name="spacy",
@@ -179,7 +179,7 @@ def extract_entities(
 def _load_spacy_model(model_name: str, spacy_config: SpacyEntityExtractionConfig) -> Any:
     """Load a spaCy model with caching."""
     try:
-        import spacy
+        import spacy  # noqa: PLC0415
 
         if spacy_config.model_cache_dir:
             os.environ["SPACY_DATA"] = str(spacy_config.model_cache_dir)
@@ -223,7 +223,7 @@ def extract_keywords(
         MissingDependencyError: If `keybert` is not installed.
     """
     try:
-        from keybert import KeyBERT
+        from keybert import KeyBERT  # noqa: PLC0415
 
         kw_model = KeyBERT()
         keywords = kw_model.extract_keywords(text, top_n=keyword_count)