kreuzberg 3.14.0__py3-none-any.whl → 3.15.0__py3-none-any.whl

kreuzberg/__init__.py

@@ -4,9 +4,12 @@ from ._registry import ExtractorRegistry
 from ._types import (
     EasyOCRConfig,
     Entity,
+    ExtractedImage,
     ExtractionConfig,
     ExtractionResult,
     GMFTConfig,
+    ImageOCRConfig,
+    ImageOCRResult,
     LanguageDetectionConfig,
     Metadata,
     PaddleOCRConfig,
@@ -32,10 +35,13 @@ __version__ = version("kreuzberg")
 __all__ = [
     "EasyOCRConfig",
     "Entity",
+    "ExtractedImage",
     "ExtractionConfig",
     "ExtractionResult",
     "ExtractorRegistry",
     "GMFTConfig",
+    "ImageOCRConfig",
+    "ImageOCRResult",
     "KreuzbergError",
     "LanguageDetectionConfig",
     "Metadata",

kreuzberg/_api/_config_cache.py

@@ -0,0 +1,247 @@
+"""API Configuration Caching Module.
+
+This module provides LRU cached functions for API config operations to improve performance
+by avoiding repeated file system operations and object creation.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from functools import lru_cache
+from pathlib import Path
+from typing import Any
+
+from kreuzberg._config import discover_config
+from kreuzberg._types import (
+    EasyOCRConfig,
+    ExtractionConfig,
+    GMFTConfig,
+    HTMLToMarkdownConfig,
+    LanguageDetectionConfig,
+    PaddleOCRConfig,
+    SpacyEntityExtractionConfig,
+    TesseractConfig,
+)
+
+
+@lru_cache(maxsize=16)
+def _cached_discover_config(
+    search_path: str,
+    config_file_mtime: float,  # noqa: ARG001
+    config_file_size: int,  # noqa: ARG001
+) -> ExtractionConfig | None:
+    """Cache config discovery with file modification time validation."""
+    return discover_config(Path(search_path))
+
+
+def discover_config_cached(search_path: Path | str | None = None) -> ExtractionConfig | None:
+    """Cached version of discover_config with automatic invalidation.
+
+    This function caches the result of discover_config() and automatically invalidates
+    the cache when config files are modified.
+
+    Args:
+        search_path: Path to start searching for config files from
+
+    Returns:
+        ExtractionConfig if found, None otherwise
+    """
+    search_path = Path.cwd() if search_path is None else Path(search_path)
+
+    config_files = ["kreuzberg.toml", "pyproject.toml"]
+    for config_file_name in config_files:
+        config_path = search_path / config_file_name
+        if config_path.exists():
+            try:
+                stat = config_path.stat()
+                return _cached_discover_config(
+                    str(search_path),
+                    stat.st_mtime,
+                    stat.st_size,
+                )
+            except OSError:
+                return discover_config(search_path)
+
+    return _cached_discover_config(str(search_path), 0.0, 0)
+
+
+@lru_cache(maxsize=128)
+def _cached_create_ocr_config(
+    config_type: str,
+    config_json: str,
+) -> TesseractConfig | EasyOCRConfig | PaddleOCRConfig:
+    """Cache OCR config object creation."""
+    config_dict = json.loads(config_json)
+
+    if config_type == "tesseract":
+        return TesseractConfig(**config_dict)
+    if config_type == "easyocr":
+        return EasyOCRConfig(**config_dict)
+    if config_type == "paddleocr":
+        return PaddleOCRConfig(**config_dict)
+    msg = f"Unknown OCR config type: {config_type}"
+    raise ValueError(msg)
+
+
+@lru_cache(maxsize=64)
+def _cached_create_gmft_config(config_json: str) -> GMFTConfig:
+    """Cache GMFT config creation."""
+    return GMFTConfig(**json.loads(config_json))
+
+
+@lru_cache(maxsize=64)
+def _cached_create_language_detection_config(config_json: str) -> LanguageDetectionConfig:
+    """Cache language detection config creation."""
+    return LanguageDetectionConfig(**json.loads(config_json))
+
+
+@lru_cache(maxsize=64)
+def _cached_create_spacy_config(config_json: str) -> SpacyEntityExtractionConfig:
+    """Cache spaCy entity extraction config creation."""
+    return SpacyEntityExtractionConfig(**json.loads(config_json))
+
+
+@lru_cache(maxsize=64)
+def _cached_create_html_markdown_config(config_json: str) -> HTMLToMarkdownConfig:
+    """Cache HTML to Markdown config creation."""
+    return HTMLToMarkdownConfig(**json.loads(config_json))
+
+
+@lru_cache(maxsize=256)
+def _cached_parse_header_config(header_value: str) -> dict[str, Any]:
+    """Cache parsed header configurations."""
+    parsed_config: dict[str, Any] = json.loads(header_value)
+    return parsed_config
+
+
+def create_ocr_config_cached(
+    ocr_backend: str | None, config_dict: dict[str, Any]
+) -> TesseractConfig | EasyOCRConfig | PaddleOCRConfig:
+    """Cached version of OCR config creation.
+
+    Args:
+        ocr_backend: The OCR backend type
+        config_dict: Configuration dictionary
+
+    Returns:
+        Configured OCR config object
+    """
+    if not ocr_backend:
+        return TesseractConfig()
+
+    config_json = json.dumps(config_dict, sort_keys=True)
+    return _cached_create_ocr_config(ocr_backend, config_json)
+
+
+def create_gmft_config_cached(config_dict: dict[str, Any]) -> GMFTConfig:
+    """Cached version of GMFT config creation."""
+    config_json = json.dumps(config_dict, sort_keys=True)
+    return _cached_create_gmft_config(config_json)
+
+
+def create_language_detection_config_cached(config_dict: dict[str, Any]) -> LanguageDetectionConfig:
+    """Cached version of language detection config creation."""
+    config_json = json.dumps(config_dict, sort_keys=True)
+    return _cached_create_language_detection_config(config_json)
+
+
+def create_spacy_config_cached(config_dict: dict[str, Any]) -> SpacyEntityExtractionConfig:
+    """Cached version of spaCy config creation."""
+    config_json = json.dumps(config_dict, sort_keys=True)
+    return _cached_create_spacy_config(config_json)
+
+
+def create_html_markdown_config_cached(config_dict: dict[str, Any]) -> HTMLToMarkdownConfig:
+    """Cached version of HTML to Markdown config creation."""
+    config_json = json.dumps(config_dict, sort_keys=True)
+    return _cached_create_html_markdown_config(config_json)
+
+
+def parse_header_config_cached(header_value: str) -> dict[str, Any]:
+    """Cached version of header config parsing.
+
+    Args:
+        header_value: JSON string from X-Extraction-Config header
+
+    Returns:
+        Parsed configuration dictionary
+    """
+    return _cached_parse_header_config(header_value)
+
+
+@lru_cache(maxsize=512)
+def _cached_merge_configs(
+    static_config_hash: str,
+    query_params_hash: str,
+    header_config_hash: str,
+) -> ExtractionConfig:
+    """Cache the complete config merging process.
+
+    This is the ultimate optimization - cache the entire result of merge_configs()
+    based on content hashes of all inputs.
+    """
+    msg = "Not implemented yet - use individual component caching"
+    raise NotImplementedError(msg)
+
+
+def _hash_dict(data: dict[str, Any] | None) -> str:
+    """Create a hash string from a dictionary for cache keys."""
+    if data is None:
+        return "none"
+
+    json_str = json.dumps(data, sort_keys=True, default=str)
+    return hashlib.sha256(json_str.encode()).hexdigest()[:16]
+
+
+def get_cache_stats() -> dict[str, Any]:
+    """Get cache statistics for monitoring performance."""
+    return {
+        "discover_config": {
+            "hits": _cached_discover_config.cache_info().hits,
+            "misses": _cached_discover_config.cache_info().misses,
+            "size": _cached_discover_config.cache_info().currsize,
+            "max_size": _cached_discover_config.cache_info().maxsize,
+        },
+        "ocr_config": {
+            "hits": _cached_create_ocr_config.cache_info().hits,
+            "misses": _cached_create_ocr_config.cache_info().misses,
+            "size": _cached_create_ocr_config.cache_info().currsize,
+            "max_size": _cached_create_ocr_config.cache_info().maxsize,
+        },
+        "header_parsing": {
+            "hits": _cached_parse_header_config.cache_info().hits,
+            "misses": _cached_parse_header_config.cache_info().misses,
+            "size": _cached_parse_header_config.cache_info().currsize,
+            "max_size": _cached_parse_header_config.cache_info().maxsize,
+        },
+        "gmft_config": {
+            "hits": _cached_create_gmft_config.cache_info().hits,
+            "misses": _cached_create_gmft_config.cache_info().misses,
+            "size": _cached_create_gmft_config.cache_info().currsize,
+            "max_size": _cached_create_gmft_config.cache_info().maxsize,
+        },
+        "language_detection_config": {
+            "hits": _cached_create_language_detection_config.cache_info().hits,
+            "misses": _cached_create_language_detection_config.cache_info().misses,
+            "size": _cached_create_language_detection_config.cache_info().currsize,
+            "max_size": _cached_create_language_detection_config.cache_info().maxsize,
+        },
+        "spacy_config": {
+            "hits": _cached_create_spacy_config.cache_info().hits,
+            "misses": _cached_create_spacy_config.cache_info().misses,
+            "size": _cached_create_spacy_config.cache_info().currsize,
+            "max_size": _cached_create_spacy_config.cache_info().maxsize,
+        },
+    }
+
+
+def clear_all_caches() -> None:
+    """Clear all API configuration caches."""
+    _cached_discover_config.cache_clear()
+    _cached_create_ocr_config.cache_clear()
+    _cached_create_gmft_config.cache_clear()
+    _cached_create_language_detection_config.cache_clear()
+    _cached_create_spacy_config.cache_clear()
+    _cached_create_html_markdown_config.cache_clear()
+    _cached_parse_header_config.cache_clear()

kreuzberg/_api/main.py

@@ -1,17 +1,22 @@
 from __future__ import annotations
 
+import base64
+import io
 import traceback
-from functools import lru_cache
-from json import dumps, loads
+from json import dumps
 from typing import TYPE_CHECKING, Annotated, Any, Literal
 
 import msgspec
+import polars as pl
+from PIL import Image
 from typing_extensions import TypedDict
 
 from kreuzberg import (
     EasyOCRConfig,
+    ExtractedImage,
     ExtractionConfig,
     ExtractionResult,
+    ImageOCRResult,
     KreuzbergError,
     MissingDependencyError,
     PaddleOCRConfig,
@@ -20,12 +25,45 @@ from kreuzberg import (
     ValidationError,
     batch_extract_bytes,
 )
+from kreuzberg._api._config_cache import (
+    create_gmft_config_cached,
+    create_html_markdown_config_cached,
+    create_language_detection_config_cached,
+    create_ocr_config_cached,
+    create_spacy_config_cached,
+    discover_config_cached,
+    parse_header_config_cached,
+)
 from kreuzberg._config import discover_config
 
 if TYPE_CHECKING:
     from litestar.datastructures import UploadFile
 
 
+class ExtractedImageDict(TypedDict):
+    """TypedDict for extracted image JSON representation."""
+
+    data: str
+    format: str
+    filename: str | None
+    page_number: int | None
+    dimensions: tuple[int, int] | None
+    colorspace: str | None
+    bits_per_component: int | None
+    is_mask: bool
+    description: str | None
+
+
+class ImageOCRResultDict(TypedDict):
+    """TypedDict for image OCR result JSON representation."""
+
+    image: ExtractedImageDict
+    ocr_result: Any
+    confidence_score: float | None
+    processing_time: float | None
+    skipped_reason: str | None
+
+
 class HealthResponse(TypedDict):
     """Response model for health check endpoint."""
 
@@ -138,50 +176,63 @@ def _create_ocr_config(
     return config_dict
 
 
-@lru_cache(maxsize=128)
-def _merge_configs_cached(
+def _create_dimension_tuple(width: int | None, height: int | None) -> tuple[int, int] | None:
+    """Create a dimension tuple from width and height values.
+
+    Args:
+        width: Width value or None
+        height: Height value or None
+
+    Returns:
+        Tuple of (width, height) if both values are not None, otherwise None
+    """
+    if width is not None and height is not None:
+        return (width, height)
+    return None
+
+
+def merge_configs(
     static_config: ExtractionConfig | None,
-    query_params: tuple[tuple[str, Any], ...],
-    header_config: tuple[tuple[str, Any], ...] | None,
+    query_params: dict[str, Any],
+    header_config: dict[str, Any] | None,
 ) -> ExtractionConfig:
     base_config = static_config or ExtractionConfig()
     config_dict = base_config.to_dict()
 
-    query_dict = dict(query_params) if query_params else {}
-    for key, value in query_dict.items():
+    for key, value in query_params.items():
         if value is not None and key in config_dict:
             config_dict[key] = _convert_value_type(config_dict[key], value)
 
     if header_config:
-        header_dict = dict(header_config)
-        for key, value in header_dict.items():
+        for key, value in header_config.items():
             if key in config_dict:
                 config_dict[key] = value
 
     if "ocr_config" in config_dict and isinstance(config_dict["ocr_config"], dict):
         ocr_backend = config_dict.get("ocr_backend")
-        config_dict["ocr_config"] = _create_ocr_config(ocr_backend, config_dict["ocr_config"])
-
-    return ExtractionConfig(**config_dict)
+        config_dict["ocr_config"] = create_ocr_config_cached(ocr_backend, config_dict["ocr_config"])
 
+    if "gmft_config" in config_dict and isinstance(config_dict["gmft_config"], dict):
+        config_dict["gmft_config"] = create_gmft_config_cached(config_dict["gmft_config"])
 
-def _make_hashable(obj: Any) -> Any:
-    if isinstance(obj, dict):
-        return tuple(sorted((k, _make_hashable(v)) for k, v in obj.items()))
-    if isinstance(obj, list):
-        return tuple(_make_hashable(item) for item in obj)
-    return obj
+    if "language_detection_config" in config_dict and isinstance(config_dict["language_detection_config"], dict):
+        config_dict["language_detection_config"] = create_language_detection_config_cached(
+            config_dict["language_detection_config"]
+        )
 
+    if "spacy_entity_extraction_config" in config_dict and isinstance(
+        config_dict["spacy_entity_extraction_config"], dict
+    ):
+        config_dict["spacy_entity_extraction_config"] = create_spacy_config_cached(
+            config_dict["spacy_entity_extraction_config"]
+        )
 
-def merge_configs(
-    static_config: ExtractionConfig | None,
-    query_params: dict[str, Any],
-    header_config: dict[str, Any] | None,
-) -> ExtractionConfig:
-    query_tuple = tuple(sorted(query_params.items())) if query_params else ()
-    header_tuple = _make_hashable(header_config) if header_config else None
+    if "html_to_markdown_config" in config_dict and isinstance(config_dict["html_to_markdown_config"], dict):
+        config_dict["html_to_markdown_config"] = create_html_markdown_config_cached(
+            config_dict["html_to_markdown_config"]
+        )
 
-    return _merge_configs_cached(static_config, query_tuple, header_tuple)
+    return ExtractionConfig(**config_dict)
 
 
 @post("/extract", operation_id="ExtractFiles")
@@ -199,6 +250,13 @@ async def handle_files_upload(  # noqa: PLR0913
     ocr_backend: Literal["tesseract", "easyocr", "paddleocr"] | None = None,
     auto_detect_language: str | bool | None = None,
     pdf_password: str | None = None,
+    extract_images: str | bool | None = None,
+    ocr_extracted_images: str | bool | None = None,
+    image_ocr_backend: Literal["tesseract", "easyocr", "paddleocr"] | None = None,
+    image_ocr_min_width: int | None = None,
+    image_ocr_min_height: int | None = None,
+    image_ocr_max_width: int | None = None,
+    image_ocr_max_height: int | None = None,
 ) -> list[ExtractionResult]:
     """Extract text, metadata, and structured data from uploaded documents.
 
@@ -226,11 +284,30 @@ async def handle_files_upload(  # noqa: PLR0913
         ocr_backend: OCR engine to use (tesseract, easyocr, paddleocr)
         auto_detect_language: Enable automatic language detection
         pdf_password: Password for encrypted PDF files
+        extract_images: Enable image extraction for supported formats
+        ocr_extracted_images: Run OCR over extracted images
+        image_ocr_backend: Optional backend override for image OCR
+        image_ocr_min_width: Minimum image width for OCR eligibility
+        image_ocr_min_height: Minimum image height for OCR eligibility
+        image_ocr_max_width: Maximum image width for OCR eligibility
+        image_ocr_max_height: Maximum image height for OCR eligibility
 
     Returns:
         List of extraction results, one per uploaded file
+
+    Additional query parameters:
+        extract_images: Enable image extraction for supported formats
+        ocr_extracted_images: Run OCR over extracted images
+        image_ocr_backend: Optional backend override for image OCR
+        image_ocr_min_width: Minimum image width for OCR eligibility
+        image_ocr_min_height: Minimum image height for OCR eligibility
+        image_ocr_max_width: Maximum image width for OCR eligibility
+        image_ocr_max_height: Maximum image height for OCR eligibility
     """
-    static_config = discover_config()
+    static_config = discover_config_cached()
+
+    min_dims = _create_dimension_tuple(image_ocr_min_width, image_ocr_min_height)
+    max_dims = _create_dimension_tuple(image_ocr_max_width, image_ocr_max_height)
 
     query_params = {
         "chunk_content": chunk_content,
@@ -244,12 +321,17 @@ async def handle_files_upload(  # noqa: PLR0913
         "ocr_backend": ocr_backend,
         "auto_detect_language": auto_detect_language,
         "pdf_password": pdf_password,
+        "extract_images": extract_images,
+        "ocr_extracted_images": ocr_extracted_images,
+        "image_ocr_backend": image_ocr_backend,
+        "image_ocr_min_dimensions": min_dims,
+        "image_ocr_max_dimensions": max_dims,
     }
 
     header_config = None
     if config_header := request.headers.get("X-Extraction-Config"):
         try:
-            header_config = loads(config_header)
+            header_config = parse_header_config_cached(config_header)
         except Exception as e:
             raise ValidationError(f"Invalid JSON in X-Extraction-Config header: {e}", context={"error": str(e)}) from e
 
@@ -291,6 +373,42 @@ async def get_configuration() -> ConfigurationResponse:
     }
 
 
+def _polars_dataframe_encoder(obj: Any) -> Any:
+    return obj.to_dicts()
+
+
+def _pil_image_encoder(obj: Any) -> str:
+    buffer = io.BytesIO()
+    obj.save(buffer, format="PNG")
+    img_str = base64.b64encode(buffer.getvalue()).decode()
+    return f"data:image/png;base64,{img_str}"
+
+
+def _extracted_image_encoder(obj: ExtractedImage) -> ExtractedImageDict:
+    encoded_data = base64.b64encode(obj.data).decode()
+    return ExtractedImageDict(
+        data=f"data:image/{obj.format};base64,{encoded_data}",
+        format=obj.format,
+        filename=obj.filename,
+        page_number=obj.page_number,
+        dimensions=obj.dimensions,
+        colorspace=obj.colorspace,
+        bits_per_component=obj.bits_per_component,
+        is_mask=obj.is_mask,
+        description=obj.description,
+    )
+
+
+def _image_ocr_result_encoder(obj: ImageOCRResult) -> ImageOCRResultDict:
+    return ImageOCRResultDict(
+        image=_extracted_image_encoder(obj.image),
+        ocr_result=obj.ocr_result,
+        confidence_score=obj.confidence_score,
+        processing_time=obj.processing_time,
+        skipped_reason=obj.skipped_reason,
+    )
+
+
 openapi_config = OpenAPIConfig(
     title="Kreuzberg API",
     version="3.14.0",
@@ -307,6 +425,13 @@ openapi_config = OpenAPIConfig(
     create_examples=True,
 )
 
+type_encoders = {
+    pl.DataFrame: _polars_dataframe_encoder,
+    Image.Image: _pil_image_encoder,
+    ExtractedImage: _extracted_image_encoder,
+    ImageOCRResult: _image_ocr_result_encoder,
+}
+
 app = Litestar(
     route_handlers=[handle_files_upload, health_check, get_configuration],
     plugins=[OpenTelemetryPlugin(OpenTelemetryConfig())],
@@ -316,5 +441,6 @@ app = Litestar(
         KreuzbergError: exception_handler,
         Exception: general_exception_handler,
     },
-    request_max_body_size=1024 * 1024 * 1024,  # 1GB limit for large file uploads
+    type_encoders=type_encoders,
+    request_max_body_size=1024 * 1024 * 1024,
 )

kreuzberg/_chunker.py

@@ -20,14 +20,15 @@ def get_chunker(
     key = (max_characters, overlap_characters, mime_type)
     if key not in _chunkers:
         try:
-            if mime_type == MARKDOWN_MIME_TYPE:
-                from semantic_text_splitter import MarkdownSplitter  # noqa: PLC0415
+            match mime_type:
+                case x if x == MARKDOWN_MIME_TYPE:
+                    from semantic_text_splitter import MarkdownSplitter  # noqa: PLC0415
 
-                _chunkers[key] = MarkdownSplitter(max_characters, overlap_characters)
-            else:
-                from semantic_text_splitter import TextSplitter  # noqa: PLC0415
+                    _chunkers[key] = MarkdownSplitter(max_characters, overlap_characters)
+                case _:
+                    from semantic_text_splitter import TextSplitter  # noqa: PLC0415
 
-                _chunkers[key] = TextSplitter(max_characters, overlap_characters)
+                    _chunkers[key] = TextSplitter(max_characters, overlap_characters)
         except ImportError as e:  # pragma: no cover
             raise MissingDependencyError.create_for_package(
                 dependency_group="chunking", functionality="chunking", package_name="semantic-text-splitter"

kreuzberg/_constants.py

@@ -5,3 +5,5 @@ from typing import Final
 MINIMAL_SUPPORTED_PANDOC_VERSION: Final[int] = 2
 DEFAULT_MAX_CHARACTERS: Final[int] = 2000
 DEFAULT_MAX_OVERLAP: Final[int] = 100
+
+PDF_POINTS_PER_INCH: Final[float] = 72.0  # Standard PDF unit conversion

kreuzberg/_document_classification.py

@@ -65,12 +65,10 @@ def classify_document(result: ExtractionResult, config: ExtractionConfig) -> tup
         return None, None
 
     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
+    scores = {
+        doc_type: sum(1 for pattern in patterns if re.search(pattern, translated_text))
+        for doc_type, patterns in DOCUMENT_CLASSIFIERS.items()
+    }
 
     total_score = sum(scores.values())
     if total_score == 0:

kreuzberg/_entity_extraction.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import os
 import re
 from functools import lru_cache
+from itertools import chain
 from typing import TYPE_CHECKING, Any
 
 from kreuzberg._types import Entity, SpacyEntityExtractionConfig
@@ -21,11 +22,15 @@ def extract_entities(
 ) -> list[Entity]:
     entities: list[Entity] = []
     if custom_patterns:
-        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)
+        entities.extend(
+            chain.from_iterable(
+                (
+                    Entity(type=ent_type, text=match.group(), start=match.start(), end=match.end())
+                    for match in re.finditer(pattern, text)
+                )
+                for ent_type, pattern in custom_patterns
             )
+        )
 
     if spacy_config is None:
         spacy_config = SpacyEntityExtractionConfig()