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

@@ -191,7 +191,7 @@ def validate_mime_type(
         return _validate_explicit_mime_type(mime_type)
 
     if file_path:
-        from kreuzberg._utils._cache import get_mime_cache
+        from kreuzberg._utils._cache import get_mime_cache  # noqa: PLC0415
 
         path = Path(file_path)
 

kreuzberg/_ocr/_base.py

@@ -103,7 +103,7 @@ class OCRBackend(ABC, Generic[T]):
         Returns:
             List of extraction result objects in the same order as input paths
         """
-        from kreuzberg._utils._sync import run_taskgroup
+        from kreuzberg._utils._sync import run_taskgroup  # noqa: PLC0415
 
         tasks = [self.process_file(path, **kwargs) for path in paths]
         return await run_taskgroup(*tasks)

kreuzberg/_ocr/_easyocr.py

@@ -111,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."""
 
@@ -180,7 +180,7 @@ class EasyOCRBackend(OCRBackend[EasyOCRConfig]):
         Raises:
             OCRError: If OCR processing fails.
         """
-        import numpy as np
+        import numpy as np  # noqa: PLC0415
 
         await self._init_easyocr(**kwargs)
 
@@ -318,7 +318,7 @@ class EasyOCRBackend(OCRBackend[EasyOCRConfig]):
             bool: True if GPU support is available.
         """
         try:
-            import torch
+            import torch  # noqa: PLC0415
 
             return bool(torch.cuda.is_available())
         except ImportError:
@@ -339,7 +339,7 @@ class EasyOCRBackend(OCRBackend[EasyOCRConfig]):
             return
 
         try:
-            import easyocr
+            import easyocr  # noqa: PLC0415
         except ImportError as e:
             raise MissingDependencyError.create_for_package(
                 dependency_group="easyocr", functionality="EasyOCR as an OCR backend", package_name="easyocr"
@@ -507,7 +507,7 @@ class EasyOCRBackend(OCRBackend[EasyOCRConfig]):
             return
 
         try:
-            import easyocr
+            import easyocr  # noqa: PLC0415
         except ImportError as e:
             raise MissingDependencyError.create_for_package(
                 dependency_group="easyocr", functionality="EasyOCR as an OCR backend", package_name="easyocr"

kreuzberg/_ocr/_paddleocr.py

@@ -31,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.
 
@@ -124,7 +124,7 @@ class PaddleBackend(OCRBackend[PaddleOCRConfig]):
         Raises:
             OCRError: If OCR processing fails.
         """
-        import numpy as np
+        import numpy as np  # noqa: PLC0415
 
         await self._init_paddle_ocr(**kwargs)
 
@@ -260,7 +260,7 @@ class PaddleBackend(OCRBackend[PaddleOCRConfig]):
             return
 
         try:
-            from paddleocr import PaddleOCR
+            from paddleocr import PaddleOCR  # noqa: PLC0415
         except ImportError as e:
             raise MissingDependencyError.create_for_package(
                 dependency_group="paddleocr", functionality="PaddleOCR as an OCR backend", package_name="paddleocr"
@@ -427,7 +427,7 @@ class PaddleBackend(OCRBackend[PaddleOCRConfig]):
             return
 
         try:
-            from paddleocr import PaddleOCR
+            from paddleocr import PaddleOCR  # noqa: PLC0415
         except ImportError as e:
             raise MissingDependencyError.create_for_package(
                 dependency_group="paddleocr", functionality="PaddleOCR as an OCR backend", package_name="paddleocr"

kreuzberg/_ocr/_tesseract.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import hashlib
+import io
 import os
 import re
 import subprocess
@@ -11,8 +12,10 @@ from enum import Enum
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, Final
 
+import anyio
 from anyio import Path as AsyncPath
 from anyio import run_process
+from PIL import Image
 from typing_extensions import Self
 
 from kreuzberg._mime_types import PLAIN_TEXT_MIME_TYPE
@@ -24,7 +27,7 @@ from kreuzberg._utils._tmp import create_temp_file
 from kreuzberg.exceptions import MissingDependencyError, OCRError, ValidationError
 
 if TYPE_CHECKING:
-    from PIL.Image import Image
+    from PIL.Image import Image as PILImage
 
 try:  # pragma: no cover
     from typing import Unpack  # type: ignore[attr-defined]
@@ -192,7 +195,7 @@ class PSMMode(Enum):
     """Treat the image as a single character."""
 
 
-@dataclass(unsafe_hash=True, frozen=True)
+@dataclass(unsafe_hash=True, frozen=True, slots=True)
 class TesseractConfig:
     """Configuration options for Tesseract OCR engine."""
 
@@ -232,12 +235,10 @@ class TesseractBackend(OCRBackend[TesseractConfig]):
 
     async def process_image(
         self,
-        image: Image,
+        image: PILImage,
         **kwargs: Unpack[TesseractConfig],
     ) -> ExtractionResult:
-        import io
-
-        from kreuzberg._utils._cache import get_ocr_cache
+        from kreuzberg._utils._cache import get_ocr_cache  # noqa: PLC0415
 
         image_buffer = io.BytesIO()
         await run_sync(image.save, image_buffer, format="PNG")
@@ -255,8 +256,6 @@ class TesseractBackend(OCRBackend[TesseractConfig]):
             return cached_result
 
         if ocr_cache.is_processing(**cache_kwargs):
-            import anyio
-
             event = ocr_cache.mark_processing(**cache_kwargs)
             await anyio.to_thread.run_sync(event.wait)
 
@@ -287,7 +286,7 @@ class TesseractBackend(OCRBackend[TesseractConfig]):
         path: Path,
         **kwargs: Unpack[TesseractConfig],
     ) -> ExtractionResult:
-        from kreuzberg._utils._cache import get_ocr_cache
+        from kreuzberg._utils._cache import get_ocr_cache  # noqa: PLC0415
 
         try:
             stat = path.stat()
@@ -315,8 +314,6 @@ class TesseractBackend(OCRBackend[TesseractConfig]):
             return cached_result
 
         if ocr_cache.is_processing(**cache_kwargs):
-            import anyio
-
             event = ocr_cache.mark_processing(**cache_kwargs)
             await anyio.to_thread.run_sync(event.wait)
 
@@ -412,7 +409,7 @@ class TesseractBackend(OCRBackend[TesseractConfig]):
 
     def process_image_sync(
         self,
-        image: Image,
+        image: PILImage,
         **kwargs: Unpack[TesseractConfig],
     ) -> ExtractionResult:
         """Synchronously process an image and extract its text and metadata.
@@ -424,9 +421,7 @@ class TesseractBackend(OCRBackend[TesseractConfig]):
         Returns:
             The extraction result object
         """
-        import io
-
-        from kreuzberg._utils._cache import get_ocr_cache
+        from kreuzberg._utils._cache import get_ocr_cache  # noqa: PLC0415
 
         image_buffer = io.BytesIO()
         image.save(image_buffer, format="PNG")
@@ -485,7 +480,7 @@ class TesseractBackend(OCRBackend[TesseractConfig]):
         Returns:
             The extraction result object
         """
-        from kreuzberg._utils._cache import get_ocr_cache
+        from kreuzberg._utils._cache import get_ocr_cache  # noqa: PLC0415
 
         file_info = self._get_file_info(path)
 
@@ -774,10 +769,6 @@ def _process_image_bytes_with_tesseract(
         OCR result as dictionary.
     """
     try:
-        import io
-
-        from PIL import Image
-
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_image:
             with Image.open(io.BytesIO(image_bytes)) as image:
                 image.save(tmp_image.name, format="PNG")
@@ -815,7 +806,7 @@ class TesseractProcessPool:
             max_processes: Maximum number of processes.
             memory_limit_gb: Memory limit in GB.
         """
-        from kreuzberg._utils._process_pool import ProcessPoolManager
+        from kreuzberg._utils._process_pool import ProcessPoolManager  # noqa: PLC0415
 
         self.config = config or TesseractConfig()
         self.process_manager = ProcessPoolManager(

kreuzberg/_playa.py

@@ -114,9 +114,8 @@ def _extract_keyword_metadata(pdf_info: dict[str, Any], result: Metadata) -> Non
     if keywords := pdf_info.get("keywords"):
         if isinstance(keywords, (str, bytes)):
             kw_str = decode_text(keywords)
-            kw_list = [k.strip() for k in kw_str.split(",")]
-            kw_list = [k.strip() for k in " ".join(kw_list).split(";")]
-            result["keywords"] = [k for k in kw_list if k]
+            # Combine multiple operations into a single comprehension
+            result["keywords"] = [k.strip() for part in kw_str.replace(";", ",").split(",") if (k := part.strip())]
         elif isinstance(keywords, list):
             result["keywords"] = [decode_text(k) for k in keywords]
 

kreuzberg/_types.py

@@ -5,7 +5,14 @@ from collections.abc import Awaitable, Callable
 from dataclasses import asdict, dataclass, field
 from typing import TYPE_CHECKING, Any, Literal, TypedDict
 
+import msgspec
+
 from kreuzberg._constants import DEFAULT_MAX_CHARACTERS, DEFAULT_MAX_OVERLAP
+from kreuzberg._utils._table import (
+    export_table_to_csv,
+    export_table_to_tsv,
+    extract_table_structure_info,
+)
 from kreuzberg.exceptions import ValidationError
 
 if sys.version_info < (3, 11):  # pragma: no cover
@@ -191,7 +198,7 @@ def normalize_metadata(data: dict[str, Any] | None) -> Metadata:
     return normalized
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, slots=True)
 class Entity:
     """Represents an extracted entity with type, text, and position."""
 
@@ -205,7 +212,7 @@ class Entity:
     """End character offset in the content"""
 
 
-@dataclass
+@dataclass(slots=True)
 class ExtractionResult:
     """The result of a file extraction."""
 
@@ -225,10 +232,36 @@ class ExtractionResult:
     """Extracted keywords and their scores, if keyword extraction is enabled."""
     detected_languages: list[str] | None = None
     """Languages detected in the extracted content, if language detection is enabled."""
+    document_type: str | None = None
+    """Detected document type, if document type detection is enabled."""
+    document_type_confidence: float | None = None
+    """Confidence of the detected document type."""
+    layout: DataFrame | None = field(default=None, repr=False, hash=False)
+    """Internal layout data from OCR, not for public use."""
+
+    def to_dict(self, include_none: bool = False) -> dict[str, Any]:
+        """Converts the ExtractionResult to a dictionary.
+
+        Args:
+            include_none: If True, include fields with None values.
+                         If False (default), exclude None values.
+
+        Returns:
+            Dictionary representation of the ExtractionResult.
+        """
+        # Use msgspec.to_builtins for efficient conversion
+        # The builtin_types parameter allows DataFrames to pass through
+        result = msgspec.to_builtins(
+            self,
+            builtin_types=(type(None),),  # Allow None to pass through
+            order="deterministic",  # Ensure consistent output
+        )
+
+        if include_none:
+            return result  # type: ignore[no-any-return]
 
-    def to_dict(self) -> dict[str, Any]:
-        """Converts the ExtractionResult to a dictionary."""
-        return asdict(self)
+        # Remove None values to match expected behavior
+        return {k: v for k, v in result.items() if v is not None}
 
     def export_tables_to_csv(self) -> list[str]:
         """Export all tables to CSV format.
@@ -239,8 +272,6 @@ class ExtractionResult:
         if not self.tables:
             return []
 
-        from kreuzberg._utils._table import export_table_to_csv
-
         return [export_table_to_csv(table) for table in self.tables]
 
     def export_tables_to_tsv(self) -> list[str]:
@@ -252,8 +283,6 @@ class ExtractionResult:
         if not self.tables:
             return []
 
-        from kreuzberg._utils._table import export_table_to_tsv
-
         return [export_table_to_tsv(table) for table in self.tables]
 
     def get_table_summaries(self) -> list[dict[str, Any]]:
@@ -265,8 +294,6 @@ class ExtractionResult:
         if not self.tables:
             return []
 
-        from kreuzberg._utils._table import extract_table_structure_info
-
         return [extract_table_structure_info(table) for table in self.tables]
 
 
@@ -274,7 +301,7 @@ PostProcessingHook = Callable[[ExtractionResult], ExtractionResult | Awaitable[E
 ValidationHook = Callable[[ExtractionResult], None | Awaitable[None]]
 
 
-@dataclass(unsafe_hash=True)
+@dataclass(unsafe_hash=True, slots=True)
 class ExtractionConfig:
     """Represents configuration settings for an extraction process.
 
@@ -322,6 +349,12 @@ class ExtractionConfig:
     """Configuration for language detection. If None, uses default settings."""
     spacy_entity_extraction_config: SpacyEntityExtractionConfig | None = None
     """Configuration for spaCy entity extraction. If None, uses default settings."""
+    auto_detect_document_type: bool = False
+    """Whether to automatically detect the document type."""
+    document_type_confidence_threshold: float = 0.7
+    """Confidence threshold for document type detection."""
+    document_classification_mode: Literal["text", "vision"] = "text"
+    """The mode to use for document classification."""
     enable_quality_processing: bool = True
     """Whether to apply quality post-processing to improve extraction results."""
 
@@ -332,9 +365,9 @@ class ExtractionConfig:
             object.__setattr__(self, "post_processing_hooks", tuple(self.post_processing_hooks))
         if self.validators is not None and isinstance(self.validators, list):
             object.__setattr__(self, "validators", tuple(self.validators))
-        from kreuzberg._ocr._easyocr import EasyOCRConfig
-        from kreuzberg._ocr._paddleocr import PaddleOCRConfig
-        from kreuzberg._ocr._tesseract import TesseractConfig
+        from kreuzberg._ocr._easyocr import EasyOCRConfig  # noqa: PLC0415
+        from kreuzberg._ocr._paddleocr import PaddleOCRConfig  # noqa: PLC0415
+        from kreuzberg._ocr._tesseract import TesseractConfig  # noqa: PLC0415
 
         if self.ocr_backend is None and self.ocr_config is not None:
             raise ValidationError("'ocr_backend' is None but 'ocr_config' is provided")
@@ -355,18 +388,23 @@ class ExtractionConfig:
         Returns:
             A dict of the OCR configuration or an empty dict if no backend is provided.
         """
-        if self.ocr_backend is not None:
-            if self.ocr_config is not None:
-                return asdict(self.ocr_config)
-            if self.ocr_backend == "tesseract":
-                from kreuzberg._ocr._tesseract import TesseractConfig
+        if self.ocr_backend is None:
+            return {}
 
-                return asdict(TesseractConfig())
-            if self.ocr_backend == "easyocr":
-                from kreuzberg._ocr._easyocr import EasyOCRConfig
+        if self.ocr_config is not None:
+            # Use asdict for OCR configs to preserve enum objects correctly
+            return asdict(self.ocr_config)
 
-                return asdict(EasyOCRConfig())
-            from kreuzberg._ocr._paddleocr import PaddleOCRConfig
+        # Lazy load and cache default configs instead of creating new instances
+        if self.ocr_backend == "tesseract":
+            from kreuzberg._ocr._tesseract import TesseractConfig  # noqa: PLC0415
 
-            return asdict(PaddleOCRConfig())
-        return {}
+            return asdict(TesseractConfig())
+        if self.ocr_backend == "easyocr":
+            from kreuzberg._ocr._easyocr import EasyOCRConfig  # noqa: PLC0415
+
+            return asdict(EasyOCRConfig())
+        # paddleocr
+        from kreuzberg._ocr._paddleocr import PaddleOCRConfig  # noqa: PLC0415
+
+        return asdict(PaddleOCRConfig())

kreuzberg/_utils/_cache.py

@@ -7,6 +7,7 @@ import os
 import threading
 import time
 from contextlib import suppress
+from io import StringIO
 from pathlib import Path
 from typing import Any, Generic, TypeVar
 
@@ -64,11 +65,10 @@ class KreuzbergCache(Generic[T]):
         Returns:
             Unique cache key string
         """
-        # Use more efficient string building for cache key
         if not kwargs:
             return "empty"
 
-        # Build key string efficiently
+        # Build cache key using list + join (faster than StringIO)
         parts = []
         for key in sorted(kwargs):
             value = kwargs[key]
@@ -81,6 +81,7 @@ class KreuzbergCache(Generic[T]):
                 parts.append(f"{key}={type(value).__name__}:{value!s}")
 
         cache_str = "&".join(parts)
+        # SHA256 is secure and fast enough for cache keys
         return hashlib.sha256(cache_str.encode()).hexdigest()[:16]
 
     def _get_cache_path(self, cache_key: str) -> Path:
@@ -107,15 +108,14 @@ class KreuzbergCache(Generic[T]):
             serialized_data = []
             for item in result:
                 if isinstance(item, dict) and "df" in item:
-                    # Create a copy and serialize the DataFrame as CSV
-                    item_copy = item.copy()
+                    # Build new dict without unnecessary copy
+                    serialized_item = {k: v for k, v in item.items() if k != "df"}
                     if hasattr(item["df"], "to_csv"):
-                        item_copy["df_csv"] = item["df"].to_csv(index=False)
+                        serialized_item["df_csv"] = item["df"].to_csv(index=False)
                     else:
                         # Fallback for non-DataFrame objects
-                        item_copy["df_csv"] = str(item["df"])
-                    del item_copy["df"]
-                    serialized_data.append(item_copy)
+                        serialized_item["df_csv"] = str(item["df"])
+                    serialized_data.append(serialized_item)
                 else:
                     serialized_data.append(item)
             return {"type": "TableDataList", "data": serialized_data, "cached_at": time.time()}
@@ -127,18 +127,15 @@ class KreuzbergCache(Generic[T]):
         data = cached_data["data"]
 
         if cached_data.get("type") == "TableDataList" and isinstance(data, list):
+            import pandas as pd  # noqa: PLC0415
+
             deserialized_data = []
             for item in data:
                 if isinstance(item, dict) and "df_csv" in item:
-                    # Restore the DataFrame from CSV
-                    item_copy = item.copy()
-                    from io import StringIO
-
-                    import pandas as pd
-
-                    item_copy["df"] = pd.read_csv(StringIO(item["df_csv"]))
-                    del item_copy["df_csv"]
-                    deserialized_data.append(item_copy)
+                    # Build new dict without unnecessary copy
+                    deserialized_item = {k: v for k, v in item.items() if k != "df_csv"}
+                    deserialized_item["df"] = pd.read_csv(StringIO(item["df_csv"]))
+                    deserialized_data.append(deserialized_item)
                 else:
                     deserialized_data.append(item)
             return deserialized_data  # type: ignore[return-value]