kreuzberg 3.3.0__py3-none-any.whl → 3.8.1__py3-none-any.whl

kreuzberg/_mcp/server.py

@@ -0,0 +1,227 @@
+"""Kreuzberg MCP server implementation."""
+
+from __future__ import annotations
+
+import base64
+from typing import Any
+
+from mcp.server import FastMCP
+from mcp.types import TextContent
+
+from kreuzberg._types import ExtractionConfig, OcrBackendType
+from kreuzberg.extraction import extract_bytes_sync, extract_file_sync
+
+# Create the MCP server
+mcp = FastMCP("Kreuzberg Text Extraction")
+
+
+@mcp.tool()
+def extract_document(  # noqa: PLR0913
+    file_path: str,
+    mime_type: str | None = None,
+    force_ocr: bool = False,
+    chunk_content: bool = False,
+    extract_tables: bool = False,
+    extract_entities: bool = False,
+    extract_keywords: bool = False,
+    ocr_backend: OcrBackendType = "tesseract",
+    max_chars: int = 1000,
+    max_overlap: int = 200,
+    keyword_count: int = 10,
+    auto_detect_language: bool = False,
+) -> dict[str, Any]:
+    """Extract text content from a document file.
+
+    Args:
+        file_path: Path to the document file
+        mime_type: MIME type of the document (auto-detected if not provided)
+        force_ocr: Force OCR even for text-based documents
+        chunk_content: Split content into chunks
+        extract_tables: Extract tables from the document
+        extract_entities: Extract named entities
+        extract_keywords: Extract keywords
+        ocr_backend: OCR backend to use (tesseract, easyocr, paddleocr)
+        max_chars: Maximum characters per chunk
+        max_overlap: Character overlap between chunks
+        keyword_count: Number of keywords to extract
+        auto_detect_language: Auto-detect document language
+
+    Returns:
+        Extracted content with metadata, tables, chunks, entities, and keywords
+    """
+    config = ExtractionConfig(
+        force_ocr=force_ocr,
+        chunk_content=chunk_content,
+        extract_tables=extract_tables,
+        extract_entities=extract_entities,
+        extract_keywords=extract_keywords,
+        ocr_backend=ocr_backend,
+        max_chars=max_chars,
+        max_overlap=max_overlap,
+        keyword_count=keyword_count,
+        auto_detect_language=auto_detect_language,
+    )
+
+    result = extract_file_sync(file_path, mime_type, config)
+    return result.to_dict()
+
+
+@mcp.tool()
+def extract_bytes(  # noqa: PLR0913
+    content_base64: str,
+    mime_type: str,
+    force_ocr: bool = False,
+    chunk_content: bool = False,
+    extract_tables: bool = False,
+    extract_entities: bool = False,
+    extract_keywords: bool = False,
+    ocr_backend: OcrBackendType = "tesseract",
+    max_chars: int = 1000,
+    max_overlap: int = 200,
+    keyword_count: int = 10,
+    auto_detect_language: bool = False,
+) -> dict[str, Any]:
+    """Extract text content from document bytes.
+
+    Args:
+        content_base64: Base64-encoded document content
+        mime_type: MIME type of the document
+        force_ocr: Force OCR even for text-based documents
+        chunk_content: Split content into chunks
+        extract_tables: Extract tables from the document
+        extract_entities: Extract named entities
+        extract_keywords: Extract keywords
+        ocr_backend: OCR backend to use (tesseract, easyocr, paddleocr)
+        max_chars: Maximum characters per chunk
+        max_overlap: Character overlap between chunks
+        keyword_count: Number of keywords to extract
+        auto_detect_language: Auto-detect document language
+
+    Returns:
+        Extracted content with metadata, tables, chunks, entities, and keywords
+    """
+    content_bytes = base64.b64decode(content_base64)
+
+    config = ExtractionConfig(
+        force_ocr=force_ocr,
+        chunk_content=chunk_content,
+        extract_tables=extract_tables,
+        extract_entities=extract_entities,
+        extract_keywords=extract_keywords,
+        ocr_backend=ocr_backend,
+        max_chars=max_chars,
+        max_overlap=max_overlap,
+        keyword_count=keyword_count,
+        auto_detect_language=auto_detect_language,
+    )
+
+    result = extract_bytes_sync(content_bytes, mime_type, config)
+    return result.to_dict()
+
+
+@mcp.tool()
+def extract_simple(
+    file_path: str,
+    mime_type: str | None = None,
+) -> str:
+    """Simple text extraction from a document file.
+
+    Args:
+        file_path: Path to the document file
+        mime_type: MIME type of the document (auto-detected if not provided)
+
+    Returns:
+        Extracted text content as a string
+    """
+    config = ExtractionConfig()
+    result = extract_file_sync(file_path, mime_type, config)
+    return result.content
+
+
+@mcp.resource("config://default")
+def get_default_config() -> str:
+    """Get the default extraction configuration."""
+    config = ExtractionConfig()
+    return str(config.__dict__)
+
+
+@mcp.resource("config://available-backends")
+def get_available_backends() -> str:
+    """Get available OCR backends."""
+    return "tesseract, easyocr, paddleocr"
+
+
+@mcp.resource("extractors://supported-formats")
+def get_supported_formats() -> str:
+    """Get supported document formats."""
+    return """
+    Supported formats:
+    - PDF documents
+    - Images (PNG, JPG, JPEG, TIFF, BMP, WEBP)
+    - Office documents (DOCX, PPTX, XLSX)
+    - HTML files
+    - Text files (TXT, CSV, TSV)
+    - And more...
+    """
+
+
+@mcp.prompt()
+def extract_and_summarize(file_path: str) -> list[TextContent]:
+    """Extract text from a document and provide a summary prompt.
+
+    Args:
+        file_path: Path to the document file
+
+    Returns:
+        Extracted content with summarization prompt
+    """
+    result = extract_file_sync(file_path, None, ExtractionConfig())
+
+    return [
+        TextContent(
+            type="text",
+            text=f"Document Content:\n{result.content}\n\nPlease provide a concise summary of this document.",
+        )
+    ]
+
+
+@mcp.prompt()
+def extract_structured(file_path: str) -> list[TextContent]:
+    """Extract text with structured analysis prompt.
+
+    Args:
+        file_path: Path to the document file
+
+    Returns:
+        Extracted content with structured analysis prompt
+    """
+    config = ExtractionConfig(
+        extract_entities=True,
+        extract_keywords=True,
+        extract_tables=True,
+    )
+    result = extract_file_sync(file_path, None, config)
+
+    content = f"Document Content:\n{result.content}\n\n"
+
+    if result.entities:
+        content += f"Entities: {[f'{e.text} ({e.type})' for e in result.entities]}\n\n"
+
+    if result.keywords:
+        content += f"Keywords: {[f'{kw[0]} ({kw[1]:.2f})' for kw in result.keywords]}\n\n"
+
+    if result.tables:
+        content += f"Tables found: {len(result.tables)}\n\n"
+
+    content += "Please analyze this document and provide structured insights."
+
+    return [TextContent(type="text", text=content)]
+
+
+def main() -> None:
+    """Main entry point for the MCP server."""
+    mcp.run()
+
+
+if __name__ == "__main__":
+    main()

kreuzberg/_mime_types.py

@@ -17,6 +17,12 @@ PLAIN_TEXT_MIME_TYPE: Final = "text/plain"
 POWER_POINT_MIME_TYPE: Final = "application/vnd.openxmlformats-officedocument.presentationml.presentation"
 DOCX_MIME_TYPE: Final = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
 
+EML_MIME_TYPE: Final = "message/rfc822"
+MSG_MIME_TYPE: Final = "application/vnd.ms-outlook"
+JSON_MIME_TYPE: Final = "application/json"
+YAML_MIME_TYPE: Final = "application/x-yaml"
+TOML_MIME_TYPE: Final = "application/toml"
+
 EXCEL_MIME_TYPE: Final = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
 EXCEL_BINARY_MIME_TYPE: Final = "application/vnd.ms-excel"
 EXCEL_MACRO_MIME_TYPE: Final = "application/vnd.ms-excel.sheet.macroEnabled.12"
@@ -127,6 +133,12 @@ EXT_TO_MIME_TYPE: Final[Mapping[str, str]] = {
     ".org": "text/x-org",
     ".epub": "application/epub+zip",
     ".rtf": "application/rtf",
+    ".eml": EML_MIME_TYPE,
+    ".msg": MSG_MIME_TYPE,
+    ".json": JSON_MIME_TYPE,
+    ".yaml": YAML_MIME_TYPE,
+    ".yml": YAML_MIME_TYPE,
+    ".toml": TOML_MIME_TYPE,
     ".odt": "application/vnd.oasis.opendocument.text",
     ".docx": DOCX_MIME_TYPE,
     ".bib": "application/x-bibtex",
@@ -139,7 +151,21 @@ SUPPORTED_MIME_TYPES: Final[set[str]] = (
     | IMAGE_MIME_TYPES
     | PANDOC_SUPPORTED_MIME_TYPES
     | SPREADSHEET_MIME_TYPES
-    | {PDF_MIME_TYPE, POWER_POINT_MIME_TYPE, HTML_MIME_TYPE}
+    | {
+        PDF_MIME_TYPE,
+        POWER_POINT_MIME_TYPE,
+        HTML_MIME_TYPE,
+        EML_MIME_TYPE,
+        MSG_MIME_TYPE,
+        JSON_MIME_TYPE,
+        YAML_MIME_TYPE,
+        TOML_MIME_TYPE,
+        "text/json",
+        "text/yaml",
+        "text/x-yaml",
+        "application/yaml",
+        "text/toml",
+    }
 )
 
 

kreuzberg/_ocr/__init__.py

@@ -4,9 +4,18 @@ from typing import Any
 from kreuzberg._ocr._base import OCRBackend
 from kreuzberg._ocr._easyocr import EasyOCRBackend
 from kreuzberg._ocr._paddleocr import PaddleBackend
-from kreuzberg._ocr._tesseract import TesseractBackend
+from kreuzberg._ocr._tesseract import TesseractBackend, TesseractProcessPool
 from kreuzberg._types import OcrBackendType
 
+__all__ = [
+    "EasyOCRBackend",
+    "OCRBackend",
+    "PaddleBackend",
+    "TesseractBackend",
+    "TesseractProcessPool",
+    "get_ocr_backend",
+]
+
 
 @lru_cache
 def get_ocr_backend(backend: OcrBackendType) -> OCRBackend[Any]:

kreuzberg/_ocr/_base.py

@@ -49,6 +49,65 @@ class OCRBackend(ABC, Generic[T]):
         """
         ...
 
+    @abstractmethod
+    def process_image_sync(self, image: Image, **kwargs: Unpack[T]) -> ExtractionResult:
+        """Synchronously process an image and extract its text and metadata.
+
+        Args:
+            image: An instance of PIL.Image representing the input image.
+            **kwargs: Any kwargs related to the given backend
+
+        Returns:
+            The extraction result object
+        """
+        ...
+
+    @abstractmethod
+    def process_file_sync(self, path: Path, **kwargs: Unpack[T]) -> ExtractionResult:
+        """Synchronously process a file and extract its text and metadata.
+
+        Args:
+            path: A Path object representing the file to be processed.
+            **kwargs: Any kwargs related to the given backend
+
+        Returns:
+            The extraction result object
+        """
+        ...
+
+    def process_batch_sync(self, paths: list[Path], **kwargs: Unpack[T]) -> list[ExtractionResult]:
+        """Synchronously process a batch of files and extract their text and metadata.
+
+        Default implementation processes files sequentially. Backends can override
+        for more efficient batch processing.
+
+        Args:
+            paths: List of Path objects representing files to be processed.
+            **kwargs: Any kwargs related to the given backend
+
+        Returns:
+            List of extraction result objects in the same order as input paths
+        """
+        return [self.process_file_sync(path, **kwargs) for path in paths]
+
+    async def process_batch(self, paths: list[Path], **kwargs: Unpack[T]) -> list[ExtractionResult]:
+        """Asynchronously process a batch of files and extract their text and metadata.
+
+        Default implementation processes files concurrently. Backends can override
+        for more efficient batch processing.
+
+        Args:
+            paths: List of Path objects representing files to be processed.
+            **kwargs: Any kwargs related to the given backend
+
+        Returns:
+            List of extraction result objects in the same order as input paths
+        """
+        from kreuzberg._utils._sync import run_taskgroup
+
+        tasks = [self.process_file(path, **kwargs) for path in paths]
+        return await run_taskgroup(*tasks)
+
     def __hash__(self) -> int:
         """Hash function for allowing caching."""
         return hash(type(self).__name__)

kreuzberg/_ocr/_easyocr.py

@@ -4,6 +4,7 @@ import warnings
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal
 
+import numpy as np
 from PIL import Image
 
 from kreuzberg._mime_types import PLAIN_TEXT_MIME_TYPE
@@ -319,7 +320,7 @@ class EasyOCRBackend(OCRBackend[EasyOCRConfig]):
         try:
             import torch
 
-            return torch.cuda.is_available()
+            return bool(torch.cuda.is_available())
         except ImportError:
             return False
 
@@ -440,3 +441,93 @@ class EasyOCRBackend(OCRBackend[EasyOCRConfig]):
             )
 
         return languages
+
+    def process_image_sync(self, image: Image.Image, **kwargs: Unpack[EasyOCRConfig]) -> ExtractionResult:
+        """Synchronously process an image and extract its text and metadata using EasyOCR.
+
+        Args:
+            image: An instance of PIL.Image representing the input image.
+            **kwargs: Configuration parameters for EasyOCR including language, detection thresholds, etc.
+
+        Returns:
+            ExtractionResult: The extraction result containing text content, mime type, and metadata.
+
+        Raises:
+            OCRError: If OCR processing fails.
+        """
+        self._init_easyocr_sync(**kwargs)
+
+        beam_width = kwargs.pop("beam_width")
+        kwargs.pop("language", None)
+        kwargs.pop("use_gpu", None)
+
+        try:
+            result = self._reader.readtext(
+                np.array(image),
+                beamWidth=beam_width,
+                **kwargs,
+            )
+
+            return self._process_easyocr_result(result, image)
+        except Exception as e:
+            raise OCRError(f"Failed to OCR using EasyOCR: {e}") from e
+
+    def process_file_sync(self, path: Path, **kwargs: Unpack[EasyOCRConfig]) -> ExtractionResult:
+        """Synchronously process a file and extract its text and metadata using EasyOCR.
+
+        Args:
+            path: A Path object representing the file to be processed.
+            **kwargs: Configuration parameters for EasyOCR including language, detection thresholds, etc.
+
+        Returns:
+            ExtractionResult: The extraction result containing text content, mime type, and metadata.
+
+        Raises:
+            OCRError: If file loading or OCR processing fails.
+        """
+        self._init_easyocr_sync(**kwargs)
+        try:
+            image = Image.open(path)
+            return self.process_image_sync(image, **kwargs)
+        except Exception as e:
+            raise OCRError(f"Failed to load or process image using EasyOCR: {e}") from e
+
+    @classmethod
+    def _init_easyocr_sync(cls, **kwargs: Unpack[EasyOCRConfig]) -> None:
+        """Synchronously initialize EasyOCR with the provided configuration.
+
+        Args:
+            **kwargs: Configuration parameters for EasyOCR including language, etc.
+
+        Raises:
+            MissingDependencyError: If EasyOCR is not installed.
+            OCRError: If initialization fails.
+        """
+        if cls._reader is not None:
+            return
+
+        try:
+            import easyocr
+        except ImportError as e:
+            raise MissingDependencyError.create_for_package(
+                dependency_group="easyocr", functionality="EasyOCR as an OCR backend", package_name="easyocr"
+            ) from e
+
+        languages = cls._validate_language_code(kwargs.pop("language", "en"))
+
+        device_info = cls._resolve_device_config(**kwargs)
+        use_gpu = device_info.device_type in ("cuda", "mps")
+
+        kwargs.setdefault("detector", True)
+        kwargs.setdefault("recognizer", True)
+        kwargs.setdefault("download_enabled", True)
+        kwargs.setdefault("recog_network", "standard")
+
+        try:
+            cls._reader = easyocr.Reader(
+                languages,
+                gpu=use_gpu,
+                verbose=False,
+            )
+        except Exception as e:
+            raise OCRError(f"Failed to initialize EasyOCR: {e}") from e

kreuzberg/_ocr/_paddleocr.py

@@ -4,8 +4,10 @@ import platform
 import warnings
 from dataclasses import dataclass
 from importlib.util import find_spec
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal
 
+import numpy as np
 from PIL import Image
 
 from kreuzberg._mime_types import PLAIN_TEXT_MIME_TYPE
@@ -364,3 +366,90 @@ class PaddleBackend(OCRBackend[PaddleOCRConfig]):
                 "supported_languages": ",".join(sorted(PADDLEOCR_SUPPORTED_LANGUAGE_CODES)),
             },
         )
+
+    def process_image_sync(self, image: Image.Image, **kwargs: Unpack[PaddleOCRConfig]) -> ExtractionResult:
+        """Synchronously process an image and extract its text and metadata using PaddleOCR.
+
+        Args:
+            image: An instance of PIL.Image representing the input image.
+            **kwargs: Configuration parameters for PaddleOCR including language, detection thresholds, etc.
+
+        Returns:
+            ExtractionResult: The extraction result containing text content, mime type, and metadata.
+
+        Raises:
+            OCRError: If OCR processing fails.
+        """
+        self._init_paddle_ocr_sync(**kwargs)
+
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        image_np = np.array(image)
+        try:
+            result = self._paddle_ocr.ocr(image_np, cls=kwargs.get("use_angle_cls", True))
+            return self._process_paddle_result(result, image)
+        except Exception as e:
+            raise OCRError(f"Failed to OCR using PaddleOCR: {e}") from e
+
+    def process_file_sync(self, path: Path, **kwargs: Unpack[PaddleOCRConfig]) -> ExtractionResult:
+        """Synchronously process a file and extract its text and metadata using PaddleOCR.
+
+        Args:
+            path: A Path object representing the file to be processed.
+            **kwargs: Configuration parameters for PaddleOCR including language, detection thresholds, etc.
+
+        Returns:
+            ExtractionResult: The extraction result containing text content, mime type, and metadata.
+
+        Raises:
+            OCRError: If file loading or OCR processing fails.
+        """
+        self._init_paddle_ocr_sync(**kwargs)
+        try:
+            image = Image.open(path)
+            return self.process_image_sync(image, **kwargs)
+        except Exception as e:
+            raise OCRError(f"Failed to load or process image using PaddleOCR: {e}") from e
+
+    @classmethod
+    def _init_paddle_ocr_sync(cls, **kwargs: Unpack[PaddleOCRConfig]) -> None:
+        """Synchronously initialize PaddleOCR with the provided configuration.
+
+        Args:
+            **kwargs: Configuration parameters for PaddleOCR including language, detection thresholds, etc.
+
+        Raises:
+            MissingDependencyError: If PaddleOCR is not installed.
+            OCRError: If initialization fails.
+        """
+        if cls._paddle_ocr is not None:
+            return
+
+        try:
+            from paddleocr import PaddleOCR
+        except ImportError as e:
+            raise MissingDependencyError.create_for_package(
+                dependency_group="paddleocr", functionality="PaddleOCR as an OCR backend", package_name="paddleocr"
+            ) from e
+
+        language = cls._validate_language_code(kwargs.pop("language", "en"))
+
+        device_info = cls._resolve_device_config(**kwargs)
+        use_gpu = device_info.device_type == "cuda"
+
+        has_gpu_package = bool(find_spec("paddlepaddle_gpu"))
+        kwargs.setdefault("use_angle_cls", True)
+        kwargs["use_gpu"] = use_gpu and has_gpu_package
+        kwargs.setdefault("enable_mkldnn", cls._is_mkldnn_supported() and not (use_gpu and has_gpu_package))
+        kwargs.setdefault("det_db_thresh", 0.3)
+        kwargs.setdefault("det_db_box_thresh", 0.5)
+        kwargs.setdefault("det_db_unclip_ratio", 1.6)
+
+        if device_info.device_type == "cuda" and kwargs.get("gpu_memory_limit"):
+            kwargs["gpu_mem"] = int(kwargs["gpu_memory_limit"] * 1024)
+
+        try:
+            cls._paddle_ocr = PaddleOCR(lang=language, show_log=False, **kwargs)
+        except Exception as e:
+            raise OCRError(f"Failed to initialize PaddleOCR: {e}") from e