kreuzberg 4.0.6__cp310-abi3-macosx_14_0_arm64.whl

kreuzberg/ocr/protocol.py

@@ -0,0 +1,150 @@
+"""Protocol for Python OCR backends compatible with Rust FFI bridge.
+
+This module defines the interface that Python OCR backends must implement
+to be registered with the Rust extraction core via the FFI bridge.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Protocol
+
+
+class OcrBackendProtocol(Protocol):
+    """Protocol for OCR backends compatible with Rust FFI bridge.
+
+    Python OCR backends implementing this protocol can be registered with the
+    Rust extraction core using `register_ocr_backend()`. Once registered, they
+    can be used by all extraction APIs (Python, Rust CLI, API server, MCP).
+
+    Required Methods:
+        name: Return backend name (e.g., 'easyocr', 'paddleocr')
+        supported_languages: Return list of supported language codes
+        process_image: Process image bytes and return extraction result
+
+    Optional Methods:
+        process_file: Process file from path (defaults to reading and calling process_image)
+        initialize: Called when backend is registered (e.g., load models)
+        shutdown: Called when backend is unregistered (e.g., cleanup resources)
+        version: Return backend version string (defaults to '1.0.0')
+
+    Example:
+        >>> class MyOcrBackend:
+        ...     def name(self) -> str:
+        ...         return "my-ocr"
+        ...
+        ...     def supported_languages(self) -> list[str]:
+        ...         return ["eng", "deu", "fra"]
+        ...
+        ...     def process_image(self, image_bytes: bytes, language: str) -> dict[str, Any]:
+        ...         # Process image and extract text
+        ...         return {
+        ...             "content": "extracted text",
+        ...             "metadata": {"confidence": 0.95, "width": 800, "height": 600},
+        ...             "tables": [],
+        ...         }
+        >>> from kreuzberg import register_ocr_backend
+        >>> backend = MyOcrBackend()
+        >>> register_ocr_backend(backend)
+
+    """
+
+    def name(self) -> str:
+        """Return backend name (e.g., 'easyocr', 'paddleocr').
+
+        The name must be unique across all registered backends and is used
+        to identify and select the backend for OCR operations.
+
+        Returns:
+            Backend name as lowercase string with no spaces
+
+        """
+        ...
+
+    def supported_languages(self) -> list[str]:
+        """Return list of supported language codes.
+
+        Language codes should use standard ISO 639 codes (e.g., 'eng', 'deu')
+        or library-specific codes (e.g., 'ch_sim', 'ch_tra' for Chinese).
+
+        Returns:
+            List of supported language code strings
+
+        """
+        ...
+
+    def process_image(self, image_bytes: bytes, language: str) -> dict[str, Any]:
+        r"""Process image bytes and return extraction result.
+
+        This is the core OCR method that takes raw image data and returns
+        extracted text with metadata. The method should handle all image
+        formats supported by the backend.
+
+        Args:
+            image_bytes: Raw image data (PNG, JPEG, TIFF, etc.)
+            language: Language code for OCR (must be in supported_languages())
+
+        Returns:
+            Dictionary with extraction result:
+            {
+                "content": "extracted text content",
+                "metadata": {
+                    "width": 800,           # Optional: image width
+                    "height": 600,          # Optional: image height
+                    "confidence": 0.95,     # Optional: OCR confidence score
+                    ...                     # Any other metadata
+                },
+                "tables": [                 # Optional: extracted tables (list of dicts)
+                    {
+                        "cells": [["row1col1", "row1col2"], ["row2col1", "row2col2"]],  # Required: 2D array of strings
+                        "markdown": "| Header |\\n| ------ |\\n| Cell   |",              # Required: markdown representation
+                        "page_number": 1                                                 # Required: 1-indexed page number
+                    }
+                ]
+            }
+
+        """
+        ...
+
+    def process_file(self, path: str, language: str) -> dict[str, Any]:
+        """Process file from path (optional, defaults to reading and calling process_image).
+
+        Backends can override this method if they have optimized file processing
+        that avoids loading the entire file into memory.
+
+        Args:
+            path: Path to image file
+            language: Language code for OCR
+
+        Returns:
+            Same format as process_image()
+
+        """
+        ...
+
+    def initialize(self) -> None:
+        """Initialize backend (optional, called when backend is registered).
+
+        Use this method to:
+        - Load ML models
+        - Initialize GPU/CUDA
+        - Download required data files
+        - Set up caching
+
+        This method is called exactly once when the backend is registered
+        via `register_ocr_backend()`.
+        """
+        ...
+
+    def shutdown(self) -> None:
+        """Shutdown backend (optional, called when backend is unregistered).
+
+        Use this method to:
+        - Cleanup temporary files
+        - Release GPU memory
+        - Close file handles
+        - Save cached data
+
+        This method is called when the backend is unregistered via
+        `unregister_ocr_backend()` or when the program exits.
+        """
+        ...

kreuzberg/postprocessors/__init__.py

@@ -0,0 +1,61 @@
+"""Custom PostProcessor support for Kreuzberg.
+
+This module provides the protocol interface for creating custom post-processors
+that can be registered with the Rust extraction pipeline.
+
+Example:
+    >>> from kreuzberg import PostProcessorProtocol, register_post_processor, ExtractionResult
+    >>>
+    >>> class MyCustomProcessor:
+    ...     '''Custom processor that adds custom metadata.'''
+    ...
+    ...     def name(self) -> str:
+    ...         return "my_custom_processor"
+    ...
+    ...     def process(self, result: ExtractionResult) -> ExtractionResult:
+    ...         # Add custom processing logic
+    ...         word_count = len(result.content.split())
+    ...         result.metadata["custom_word_count"] = word_count
+    ...         result.metadata["custom_tag"] = "processed"
+    ...         return result
+    ...
+    ...     def processing_stage(self) -> str:
+    ...         return "middle"  # or "early" or "late"
+    ...
+    ...     def initialize(self) -> None:
+    ...         # Optional: Initialize resources (e.g., load ML models)
+    ...         pass
+    ...
+    ...     def shutdown(self) -> None:
+    ...         # Optional: Release resources
+    ...         pass
+    >>>
+    >>> # Register the processor
+    >>> processor = MyCustomProcessor()
+    >>> register_post_processor(processor)
+    >>>
+    >>> # Now it will be called automatically during extraction
+    >>> from kreuzberg import extract_file_sync
+    >>> result = extract_file_sync("document.pdf")
+    >>> print(result.metadata.get("custom_word_count"))
+    >>> print(result.metadata.get("custom_tag"))
+
+Processing Stages:
+    - **early**: Runs first in the pipeline (e.g., language detection)
+    - **middle**: Runs in the middle (default, most processors)
+    - **late**: Runs last (e.g., final formatting, summaries)
+
+Thread Safety:
+    Processors are called from the Rust core which may use threading.
+    Ensure your processor is thread-safe or uses appropriate locking.
+
+Performance:
+    Keep processing fast - slow processors will impact extraction performance.
+    Consider lazy initialization for expensive resources (ML models, etc.).
+"""
+
+from __future__ import annotations
+
+from kreuzberg.postprocessors.protocol import PostProcessorProtocol
+
+__all__ = ["PostProcessorProtocol"]

kreuzberg/postprocessors/protocol.py

@@ -0,0 +1,83 @@
+"""PostProcessor protocol interface.
+
+This module defines the protocol that all Python postprocessors must implement
+to be registered with the Rust core via the FFI bridge.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal, Protocol
+
+if TYPE_CHECKING:
+    from kreuzberg._internal_bindings import ExtractionResult
+
+
+class PostProcessorProtocol(Protocol):
+    """Protocol for Python postprocessors.
+
+    All postprocessors must implement these methods to be compatible
+    with the Rust PostProcessor FFI bridge.
+    """
+
+    def name(self) -> str:
+        """Return the unique name of this postprocessor.
+
+        Returns:
+            str: Processor name (e.g., "entity_extraction", "keyword_extraction")
+
+        """
+        ...
+
+    def process(self, result: ExtractionResult) -> ExtractionResult:
+        """Process and enrich an extraction result.
+
+        Args:
+            result: ExtractionResult with extracted content, metadata, and tables
+
+        Returns:
+            ExtractionResult: Modified result with enriched metadata.
+                  New metadata keys are added, existing keys are preserved.
+
+        Note:
+            The processor should add its results to result.metadata and
+            return the modified ExtractionResult. Existing metadata keys will not be
+            overwritten by the FFI bridge.
+
+        Example:
+            >>> def process(self, result: ExtractionResult) -> ExtractionResult:
+            ...     text = result.content
+            ...     entities = extract_entities(text)
+            ...     result.metadata["entities"] = entities
+            ...     return result
+
+        """
+        ...
+
+    def processing_stage(self) -> Literal["early", "middle", "late"]:
+        """Return the processing stage for this processor.
+
+        Returns:
+            str: One of "early", "middle", or "late" (default: "middle")
+
+        Note:
+            Processing stages control the order in which processors are called:
+            - Early: Runs first (e.g., language detection)
+            - Middle: Runs in the middle (default, e.g., entity extraction)
+            - Late: Runs last (e.g., final formatting)
+
+        """
+        ...
+
+    def initialize(self) -> None:
+        """Initialize the processor (e.g., load ML models).
+
+        Called once when the processor is registered.
+        """
+        ...
+
+    def shutdown(self) -> None:
+        """Shutdown the processor and release resources.
+
+        Called when the processor is unregistered.
+        """
+        ...