kreuzberg 3.6.2__py3-none-any.whl → 3.8.0__py3-none-any.whl

kreuzberg/_extractors/_spread_sheet.py

@@ -13,7 +13,7 @@ from python_calamine import CalamineWorkbook
 
 from kreuzberg._extractors._base import Extractor
 from kreuzberg._mime_types import MARKDOWN_MIME_TYPE, SPREADSHEET_MIME_TYPES
-from kreuzberg._types import ExtractionResult
+from kreuzberg._types import ExtractionResult, Metadata
 from kreuzberg._utils._string import normalize_spaces
 from kreuzberg._utils._sync import run_sync, run_taskgroup
 from kreuzberg._utils._tmp import create_temp_file
@@ -45,9 +45,14 @@ class SpreadSheetExtractor(Extractor):
             try:
                 results: list[str] = await run_taskgroup(*tasks)
 
-                return ExtractionResult(
-                    content="\n\n".join(results), mime_type=MARKDOWN_MIME_TYPE, metadata={}, chunks=[]
+                result = ExtractionResult(
+                    content="\n\n".join(results),
+                    mime_type=MARKDOWN_MIME_TYPE,
+                    metadata=self._extract_spreadsheet_metadata(workbook),
+                    chunks=[],
                 )
+
+                return self._apply_quality_processing(result)
             except ExceptionGroup as eg:
                 raise ParsingError(
                     "Failed to extract file data",
@@ -87,7 +92,14 @@ class SpreadSheetExtractor(Extractor):
                 sheet_text = self._convert_sheet_to_text_sync(workbook, sheet_name)
                 results.append(sheet_text)
 
-            return ExtractionResult(content="\n\n".join(results), mime_type=MARKDOWN_MIME_TYPE, metadata={}, chunks=[])
+            result = ExtractionResult(
+                content="\n\n".join(results),
+                mime_type=MARKDOWN_MIME_TYPE,
+                metadata=self._extract_spreadsheet_metadata(workbook),
+                chunks=[],
+            )
+
+            return self._apply_quality_processing(result)
         except Exception as e:
             raise ParsingError(
                 "Failed to extract file data",
@@ -181,3 +193,166 @@ class SpreadSheetExtractor(Extractor):
             result = "\n".join(markdown_lines)
 
         return f"## {sheet_name}\n\n{normalize_spaces(result)}"
+
+    def _enhance_sheet_with_table_data(self, workbook: CalamineWorkbook, sheet_name: str) -> str:
+        """Enhanced sheet processing with better table structure preservation."""
+        try:
+            # pandas is optional dependency
+            import pandas as pd
+
+            from kreuzberg._utils._table import enhance_table_markdown
+
+            sheet = workbook.get_sheet_by_name(sheet_name)
+            data = sheet.to_python()
+
+            if not data or not any(row for row in data):
+                return f"## {sheet_name}\n\n*Empty sheet*"
+
+            # Convert to DataFrame
+            df = pd.DataFrame(data)
+
+            # Clean up empty rows and columns
+            df = df.dropna(how="all").dropna(axis=1, how="all")
+
+            if df.empty:
+                return f"## {sheet_name}\n\n*No data*"
+
+            # Create a mock TableData for enhanced formatting
+            from PIL import Image
+
+            from kreuzberg._types import TableData
+
+            # Create a 1x1 transparent image as placeholder
+            placeholder_image = Image.new("RGBA", (1, 1), (0, 0, 0, 0))
+            mock_table: TableData = {"df": df, "text": "", "page_number": 0, "cropped_image": placeholder_image}
+
+            enhanced_markdown = enhance_table_markdown(mock_table)
+            return f"## {sheet_name}\n\n{enhanced_markdown}"
+
+        except (ImportError, AttributeError, ValueError):
+            # Fallback to original method if pandas/table enhancement fails
+            return self._convert_sheet_to_text_sync(workbook, sheet_name)
+
+    @staticmethod
+    def _extract_spreadsheet_metadata(workbook: CalamineWorkbook) -> Metadata:
+        """Extract metadata from spreadsheet using python-calamine.
+
+        Args:
+            workbook: CalamineWorkbook instance
+
+        Returns:
+            Metadata dict using existing metadata keys where possible
+        """
+        metadata: Metadata = {}
+
+        # Extract basic document properties
+        SpreadSheetExtractor._extract_document_properties(workbook, metadata)
+
+        # Add structural information
+        SpreadSheetExtractor._add_structure_info(workbook, metadata)
+
+        # Analyze content complexity
+        SpreadSheetExtractor._analyze_content_complexity(workbook, metadata)
+
+        return metadata
+
+    @staticmethod
+    def _extract_document_properties(workbook: CalamineWorkbook, metadata: Metadata) -> None:
+        """Extract basic document properties from workbook."""
+        with contextlib.suppress(AttributeError, Exception):
+            if not (hasattr(workbook, "metadata") and workbook.metadata):
+                return
+
+            props = workbook.metadata
+
+            # Basic properties mapping
+            property_mapping = {
+                "title": "title",
+                "author": "authors",  # Convert to list
+                "subject": "subject",
+                "comments": "comments",
+                "keywords": "keywords",  # Process separately
+                "category": "categories",  # Convert to list
+                "company": "organization",
+                "manager": "modified_by",
+            }
+
+            for prop_name, meta_key in property_mapping.items():
+                if hasattr(props, prop_name) and (value := getattr(props, prop_name)):
+                    if meta_key in ("authors", "categories"):
+                        metadata[meta_key] = [value]  # type: ignore[literal-required]
+                    elif meta_key == "keywords":
+                        keywords = [k.strip() for k in value.replace(";", ",").split(",") if k.strip()]
+                        if keywords:
+                            metadata[meta_key] = keywords  # type: ignore[literal-required]
+                    else:
+                        metadata[meta_key] = value  # type: ignore[literal-required]
+
+            # Handle dates separately
+            SpreadSheetExtractor._extract_date_properties(props, metadata)
+
+    @staticmethod
+    def _extract_date_properties(props: Any, metadata: Metadata) -> None:
+        """Extract and format date properties."""
+        date_mapping = {"created": "created_at", "modified": "modified_at"}
+
+        for prop_name, meta_key in date_mapping.items():
+            if hasattr(props, prop_name) and (date_value := getattr(props, prop_name)):
+                with contextlib.suppress(Exception):
+                    if hasattr(date_value, "isoformat"):
+                        metadata[meta_key] = date_value.isoformat()  # type: ignore[literal-required]
+                    else:
+                        metadata[meta_key] = str(date_value)  # type: ignore[literal-required]
+
+    @staticmethod
+    def _add_structure_info(workbook: CalamineWorkbook, metadata: Metadata) -> None:
+        """Add structural information about the spreadsheet."""
+        if not (hasattr(workbook, "sheet_names") and workbook.sheet_names):
+            return
+
+        sheet_count = len(workbook.sheet_names)
+        structure_info = f"Spreadsheet with {sheet_count} sheet{'s' if sheet_count != 1 else ''}"
+
+        # Don't list too many sheet names (magic number made constant)
+        max_sheet_names_to_list = 5
+        if sheet_count <= max_sheet_names_to_list:
+            structure_info += f": {', '.join(workbook.sheet_names)}"
+
+        metadata["description"] = structure_info
+
+    @staticmethod
+    def _analyze_content_complexity(workbook: CalamineWorkbook, metadata: Metadata) -> None:
+        """Analyze spreadsheet content for complexity indicators."""
+        with contextlib.suppress(Exception):
+            has_formulas = False
+            total_cells = 0
+
+            # Check only first few sheets for performance
+            max_sheets_to_check = 3
+            max_rows_to_check = 50
+
+            for sheet_name in workbook.sheet_names[:max_sheets_to_check]:
+                with contextlib.suppress(Exception):
+                    sheet = workbook.get_sheet_by_name(sheet_name)
+                    data = sheet.to_python()
+
+                    for row in data[:max_rows_to_check]:
+                        if not row:  # Skip empty rows
+                            continue
+
+                        total_cells += sum(1 for cell in row if cell is not None and str(cell).strip())
+
+                        # Check for formulas (simple heuristic)
+                        if any(isinstance(cell, str) and cell.startswith("=") for cell in row):
+                            has_formulas = True
+                            break
+
+            # Build summary
+            summary_parts = []
+            if total_cells > 0:
+                summary_parts.append(f"Contains {total_cells}+ data cells")
+            if has_formulas:
+                summary_parts.append("includes formulas")
+
+            if summary_parts and "summary" not in metadata:
+                metadata["summary"] = f"Spreadsheet that {', '.join(summary_parts)}."

kreuzberg/_extractors/_structured.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from anyio import Path as AsyncPath
+
+from kreuzberg._extractors._base import Extractor
+from kreuzberg._mime_types import JSON_MIME_TYPE, PLAIN_TEXT_MIME_TYPE, TOML_MIME_TYPE, YAML_MIME_TYPE
+from kreuzberg._types import ExtractionResult, normalize_metadata
+from kreuzberg._utils._string import normalize_spaces, safe_decode
+from kreuzberg._utils._sync import run_sync
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+class StructuredDataExtractor(Extractor):
+    SUPPORTED_MIME_TYPES: ClassVar[set[str]] = {
+        JSON_MIME_TYPE,
+        "text/json",
+        YAML_MIME_TYPE,
+        "text/yaml",
+        "text/x-yaml",
+        "application/yaml",
+        TOML_MIME_TYPE,
+        "text/toml",
+    }
+
+    async def extract_bytes_async(self, content: bytes) -> ExtractionResult:
+        return await run_sync(self.extract_bytes_sync, content)
+
+    async def extract_path_async(self, path: Path) -> ExtractionResult:
+        content = await AsyncPath(path).read_bytes()
+        return await self.extract_bytes_async(content)
+
+    def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
+        text_content = safe_decode(content)
+
+        try:
+            if self.mime_type in {JSON_MIME_TYPE, "text/json"}:
+                data = json.loads(text_content)
+            elif self.mime_type in {TOML_MIME_TYPE, "text/toml"}:
+                try:
+                    import tomllib  # type: ignore[import-not-found]
+                except ImportError:
+                    try:
+                        import tomli as tomllib  # type: ignore[import-not-found]
+                    except ImportError:
+                        return ExtractionResult(
+                            content=normalize_spaces(text_content),
+                            mime_type=PLAIN_TEXT_MIME_TYPE,
+                            metadata={"warning": "tomllib/tomli not available, returning raw text"},
+                            chunks=[],
+                        )
+                data = tomllib.loads(text_content)
+            else:
+                try:
+                    import yaml
+
+                    data = yaml.safe_load(text_content)
+                except ImportError:
+                    return ExtractionResult(
+                        content=normalize_spaces(text_content),
+                        mime_type=PLAIN_TEXT_MIME_TYPE,
+                        metadata={"warning": "PyYAML not available, returning raw text"},
+                        chunks=[],
+                    )
+
+            text_parts: list[str] = []
+            metadata: dict[str, Any] = {}
+
+            if isinstance(data, dict):
+                text_parts.extend(self._extract_from_dict(data, metadata))
+            elif isinstance(data, list):
+                text_parts.extend(self._extract_from_list(data, metadata))
+            else:
+                text_parts.append(str(data))
+
+            combined_text = "\n".join(text_parts) if text_parts else text_content
+
+            return ExtractionResult(
+                content=normalize_spaces(combined_text),
+                mime_type=PLAIN_TEXT_MIME_TYPE,
+                metadata=normalize_metadata(metadata),
+                chunks=[],
+            )
+
+        except (ValueError, TypeError, KeyError, AttributeError, UnicodeDecodeError) as e:
+            return ExtractionResult(
+                content=normalize_spaces(text_content),
+                mime_type=PLAIN_TEXT_MIME_TYPE,
+                metadata={"parse_error": str(e)},
+                chunks=[],
+            )
+
+    def extract_path_sync(self, path: Path) -> ExtractionResult:
+        content = path.read_bytes()
+        return self.extract_bytes_sync(content)
+
+    def _extract_from_dict(self, data: dict[str, Any], metadata: dict[str, Any], prefix: str = "") -> list[str]:
+        text_parts = []
+
+        for key, value in data.items():
+            full_key = f"{prefix}.{key}" if prefix else key
+
+            if isinstance(value, str) and value.strip():
+                text_parts.append(f"{full_key}: {value}")
+
+                if any(
+                    text_field in key.lower()
+                    for text_field in ["title", "name", "subject", "description", "content", "body", "text", "message"]
+                ):
+                    metadata[full_key] = value
+
+            elif isinstance(value, (int, float, bool)):
+                text_parts.append(f"{full_key}: {value}")
+
+            elif isinstance(value, dict):
+                text_parts.extend(self._extract_from_dict(value, metadata, full_key))
+
+            elif isinstance(value, list):
+                text_parts.extend(self._extract_from_list(value, metadata, full_key))
+
+            elif value is not None:
+                text_parts.append(f"{full_key}: {value!s}")
+
+        return text_parts
+
+    def _extract_from_list(self, data: list[Any], metadata: dict[str, Any], prefix: str = "") -> list[str]:
+        text_parts = []
+
+        for i, item in enumerate(data):
+            item_key = f"{prefix}[{i}]" if prefix else f"item_{i}"
+
+            if isinstance(item, str) and item.strip():
+                text_parts.append(f"{item_key}: {item}")
+
+            elif isinstance(item, dict):
+                text_parts.extend(self._extract_from_dict(item, metadata, item_key))
+
+            elif isinstance(item, list):
+                text_parts.extend(self._extract_from_list(item, metadata, item_key))
+
+            elif item is not None:
+                text_parts.append(f"{item_key}: {item!s}")
+
+        return text_parts

kreuzberg/_gmft.py

@@ -196,7 +196,7 @@ async def extract_tables(  # noqa: PLR0915
 
     try:
         if use_isolated_process:
-            from kreuzberg._multiprocessing.gmft_isolated import extract_tables_isolated_async
+            from kreuzberg._multiprocessing import extract_tables_isolated_async
 
             result = await extract_tables_isolated_async(file_path, config)
 
@@ -314,7 +314,7 @@ def extract_tables_sync(
         return cached_result  # type: ignore[no-any-return]
 
     if use_isolated_process:
-        from kreuzberg._multiprocessing.gmft_isolated import extract_tables_isolated
+        from kreuzberg._multiprocessing import extract_tables_isolated
 
         result = extract_tables_isolated(file_path, config)
 

kreuzberg/_mcp/__init__.py

@@ -0,0 +1,5 @@
+"""MCP server for Kreuzberg text extraction."""
+
+from .server import mcp
+
+__all__ = ["mcp"]

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/_multiprocessing/__init__.py

@@ -1,6 +1,5 @@
 """Multiprocessing utilities for kreuzberg."""
 
-from .process_manager import ProcessPoolManager
-from .tesseract_pool import TesseractProcessPool
+from .gmft_isolated import extract_tables_isolated, extract_tables_isolated_async
 
-__all__ = ["ProcessPoolManager", "TesseractProcessPool"]
+__all__ = ["extract_tables_isolated", "extract_tables_isolated_async"]