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

kreuzberg/_api/main.py

@@ -6,6 +6,7 @@ from json import dumps, loads
 from typing import TYPE_CHECKING, Annotated, Any, Literal
 
 import msgspec
+from typing_extensions import TypedDict
 
 from kreuzberg import (
     EasyOCRConfig,
@@ -24,11 +25,28 @@ from kreuzberg._config import discover_config
 if TYPE_CHECKING:
     from litestar.datastructures import UploadFile
 
+
+class HealthResponse(TypedDict):
+    """Response model for health check endpoint."""
+
+    status: str
+
+
+class ConfigurationResponse(TypedDict):
+    """Response model for configuration endpoint."""
+
+    message: str
+    config: dict[str, Any] | None
+
+
 try:
     from litestar import Litestar, Request, Response, get, post
     from litestar.contrib.opentelemetry import OpenTelemetryConfig, OpenTelemetryPlugin
     from litestar.enums import RequestEncodingType
     from litestar.logging import StructLoggingConfig
+    from litestar.openapi.config import OpenAPIConfig
+    from litestar.openapi.spec.contact import Contact
+    from litestar.openapi.spec.license import License
     from litestar.params import Body
     from litestar.status_codes import (
         HTTP_400_BAD_REQUEST,
@@ -71,7 +89,6 @@ def exception_handler(request: Request[Any, Any, Any], exception: KreuzbergError
 
 
 def general_exception_handler(request: Request[Any, Any, Any], exception: Exception) -> Response[Any]:
-    """Temporary handler to catch ALL exceptions for debugging."""
     error_type = type(exception).__name__
     error_message = str(exception)
     traceback_str = traceback.format_exc()
@@ -127,7 +144,6 @@ def _merge_configs_cached(
     query_params: tuple[tuple[str, Any], ...],
     header_config: tuple[tuple[str, Any], ...] | None,
 ) -> ExtractionConfig:
-    """Cached implementation of merge_configs with hashable parameters."""
     base_config = static_config or ExtractionConfig()
     config_dict = base_config.to_dict()
 
@@ -150,7 +166,6 @@ def _merge_configs_cached(
 
 
 def _make_hashable(obj: Any) -> Any:
-    """Convert nested dicts/lists to hashable tuples."""
     if isinstance(obj, dict):
         return tuple(sorted((k, _make_hashable(v)) for k, v in obj.items()))
     if isinstance(obj, list):
@@ -163,7 +178,6 @@ def merge_configs(
     query_params: dict[str, Any],
     header_config: dict[str, Any] | None,
 ) -> ExtractionConfig:
-    """Merge configurations with precedence: header > query > static > default."""
     query_tuple = tuple(sorted(query_params.items())) if query_params else ()
     header_tuple = _make_hashable(header_config) if header_config else None
 
@@ -186,14 +200,35 @@ async def handle_files_upload(  # noqa: PLR0913
     auto_detect_language: str | bool | None = None,
     pdf_password: str | None = None,
 ) -> list[ExtractionResult]:
-    """Extracts text content from uploaded files with optional runtime configuration.
-
-    Configuration can be provided via:
-    1. Query parameters for common settings
-    2. X-Extraction-Config header for complex nested configurations (JSON format)
-    3. Static configuration file (kreuzberg.toml or pyproject.toml)
-
-    Precedence: Header config > Query params > Static config > Defaults
+    """Extract text, metadata, and structured data from uploaded documents.
+
+    This endpoint processes multiple file uploads and extracts comprehensive information including:
+    - Text content with metadata
+    - Tables (if enabled)
+    - Named entities (if enabled)
+    - Keywords (if enabled)
+    - Language detection (if enabled)
+
+    Supports various file formats including PDF, Office documents, images, and more.
+    Maximum file size: 1GB per file.
+
+    Args:
+        request: The HTTP request object
+        data: List of files to process (multipart form data)
+        chunk_content: Enable text chunking for large documents
+        max_chars: Maximum characters per chunk (default: 1000)
+        max_overlap: Character overlap between chunks (default: 200)
+        extract_tables: Extract tables from documents
+        extract_entities: Extract named entities from text
+        extract_keywords: Extract keywords from text
+        keyword_count: Number of keywords to extract (default: 10)
+        force_ocr: Force OCR processing even for text-based documents
+        ocr_backend: OCR engine to use (tesseract, easyocr, paddleocr)
+        auto_detect_language: Enable automatic language detection
+        pdf_password: Password for encrypted PDF files
+
+    Returns:
+        List of extraction results, one per uploaded file
     """
     static_config = discover_config()
 
@@ -227,14 +262,25 @@ async def handle_files_upload(  # noqa: PLR0913
 
 
 @get("/health", operation_id="HealthCheck")
-async def health_check() -> dict[str, str]:
-    """A simple health check endpoint."""
+async def health_check() -> HealthResponse:
+    """Check the health status of the API.
+
+    Returns:
+        Simple status response indicating the API is operational
+    """
     return {"status": "ok"}
 
 
 @get("/config", operation_id="GetConfiguration")
-async def get_configuration() -> dict[str, Any]:
-    """Get the current configuration."""
+async def get_configuration() -> ConfigurationResponse:
+    """Get the current extraction configuration.
+
+    Returns the loaded configuration from kreuzberg.toml file if available,
+    or indicates that no configuration file was found.
+
+    Returns:
+        Configuration data with status message
+    """
     config = discover_config()
     if config is None:
         return {"message": "No configuration file found", "config": None}
@@ -245,12 +291,30 @@ async def get_configuration() -> dict[str, Any]:
     }
 
 
+openapi_config = OpenAPIConfig(
+    title="Kreuzberg API",
+    version="3.14.0",
+    description="Document intelligence framework API for extracting text, metadata, and structured data from diverse file formats",
+    contact=Contact(
+        name="Kreuzberg",
+        url="https://github.com/Goldziher/kreuzberg",
+    ),
+    license=License(
+        name="MIT",
+        identifier="MIT",
+    ),
+    use_handler_docstrings=True,
+    create_examples=True,
+)
+
 app = Litestar(
     route_handlers=[handle_files_upload, health_check, get_configuration],
     plugins=[OpenTelemetryPlugin(OpenTelemetryConfig())],
-    logging_config=StructLoggingConfig(),  # Use default config
+    logging_config=StructLoggingConfig(),
+    openapi_config=openapi_config,
     exception_handlers={
         KreuzbergError: exception_handler,
-        Exception: general_exception_handler,  # Catch all exceptions for debugging
+        Exception: general_exception_handler,
     },
+    request_max_body_size=1024 * 1024 * 1024,  # 1GB limit for large file uploads
 )

kreuzberg/_config.py

@@ -162,7 +162,6 @@ def load_config_from_file(config_path: Path) -> dict[str, Any]:
     if config_path.name == "pyproject.toml":
         return data.get("tool", {}).get("kreuzberg", {})  # type: ignore[no-any-return]
 
-    # For any other TOML file, check if it has [tool.kreuzberg] section
     if "tool" in data and "kreuzberg" in data["tool"]:
         return data["tool"]["kreuzberg"]  # type: ignore[no-any-return]
 

kreuzberg/_extractors/_image.py

@@ -7,10 +7,13 @@ from pathlib import Path
 from typing import TYPE_CHECKING, ClassVar
 
 from anyio import Path as AsyncPath
+from PIL import Image
 
 from kreuzberg._extractors._base import Extractor
 from kreuzberg._mime_types import IMAGE_MIME_TYPES
 from kreuzberg._ocr import get_ocr_backend
+from kreuzberg._utils._image_preprocessing import normalize_image_dpi
+from kreuzberg._utils._sync import run_sync
 from kreuzberg._utils._tmp import create_temp_file
 from kreuzberg.exceptions import ValidationError
 
@@ -57,7 +60,15 @@ class ImageExtractor(Extractor):
         if self.config.ocr_backend is None:
             raise ValidationError("ocr_backend is None, cannot perform OCR")
 
-        result = await get_ocr_backend(self.config.ocr_backend).process_file(path, **self.config.get_config_dict())
+        image = await run_sync(Image.open, str(path))
+        normalized_image, preprocessing_metadata = normalize_image_dpi(image, self.config)
+
+        backend = get_ocr_backend(self.config.ocr_backend)
+        result = await backend.process_image(normalized_image, **self.config.get_config_dict())
+
+        if preprocessing_metadata:
+            result.metadata["image_preprocessing"] = preprocessing_metadata
+
         return self._apply_quality_processing(result)
 
     def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
@@ -77,8 +88,15 @@ class ImageExtractor(Extractor):
         if self.config.ocr_backend is None:
             raise ValidationError("ocr_backend is None, cannot perform OCR")
 
+        image = Image.open(str(path))
+        normalized_image, preprocessing_metadata = normalize_image_dpi(image, self.config)
+
         backend = get_ocr_backend(self.config.ocr_backend)
-        result = backend.process_file_sync(path, **self.config.get_config_dict())
+        result = backend.process_image_sync(normalized_image, **self.config.get_config_dict())
+
+        if preprocessing_metadata:
+            result.metadata["image_preprocessing"] = preprocessing_metadata
+
         return self._apply_quality_processing(result)
 
     def _get_extension_from_mime_type(self, mime_type: str) -> str:

kreuzberg/_extractors/_pdf.py

@@ -21,6 +21,7 @@ from kreuzberg._ocr import get_ocr_backend
 from kreuzberg._playa import extract_pdf_metadata, extract_pdf_metadata_sync
 from kreuzberg._types import EasyOCRConfig, ExtractionResult, Metadata, OcrBackendType, PaddleOCRConfig, TesseractConfig
 from kreuzberg._utils._errors import create_error_context, should_retry
+from kreuzberg._utils._image_preprocessing import calculate_optimal_dpi
 from kreuzberg._utils._pdf_lock import pypdfium_file_lock
 from kreuzberg._utils._string import normalize_spaces
 from kreuzberg._utils._sync import run_sync, run_taskgroup_batched
@@ -170,7 +171,26 @@ class PDFExtractor(Extractor):
             try:
                 with pypdfium_file_lock(input_file):
                     document = await run_sync(pypdfium2.PdfDocument, str(input_file))
-                    return [page.render(scale=200 / 72).to_pil() for page in cast("pypdfium2.PdfDocument", document)]
+                    images = []
+                    for page in cast("pypdfium2.PdfDocument", document):
+                        width, height = page.get_size()
+
+                        if self.config.auto_adjust_dpi:
+                            optimal_dpi = calculate_optimal_dpi(
+                                page_width=width,
+                                page_height=height,
+                                target_dpi=self.config.target_dpi,
+                                max_dimension=self.config.max_image_dimension,
+                                min_dpi=self.config.min_dpi,
+                                max_dpi=self.config.max_dpi,
+                            )
+                        else:
+                            optimal_dpi = self.config.target_dpi
+
+                        scale = optimal_dpi / 72.0
+
+                        images.append(page.render(scale=scale).to_pil())
+                    return images
             except pypdfium2.PdfiumError as e:  # noqa: PERF203
                 last_error = e
                 if not should_retry(e, attempt + 1):

kreuzberg/_extractors/_spread_sheet.py

@@ -35,7 +35,6 @@ class SpreadSheetExtractor(Extractor):
     SUPPORTED_MIME_TYPES = SPREADSHEET_MIME_TYPES
 
     def _get_file_extension(self) -> str:
-        """Get the appropriate file extension based on MIME type."""
         mime_to_ext = {
             "application/vnd.ms-excel": ".xls",
             "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx",

kreuzberg/_gmft.py

@@ -9,11 +9,11 @@ import time
 import traceback
 from io import StringIO
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 import anyio
 import msgspec
-import pandas as pd
+import polars as pl
 from PIL import Image
 
 from kreuzberg._types import GMFTConfig, TableData
@@ -25,7 +25,63 @@ if TYPE_CHECKING:
     from os import PathLike
 
     from gmft.detectors.base import CroppedTable
-    from pandas import DataFrame
+
+
+def _pandas_to_polars(pandas_df: Any) -> pl.DataFrame:
+    if pandas_df is None:
+        return pl.DataFrame()
+
+    try:
+        return pl.from_pandas(pandas_df)
+    except (TypeError, ValueError, AttributeError):
+        if hasattr(pandas_df, "columns") and hasattr(pandas_df.columns, "duplicated"):
+            mask = ~pandas_df.columns.duplicated()
+            pandas_df = pandas_df.loc[:, mask]
+            return pl.from_pandas(pandas_df)
+        return pl.DataFrame()
+
+
+def _dataframe_to_markdown(df: Any) -> str:
+    if df is None:
+        return ""
+
+    if isinstance(df, pl.DataFrame):
+        if df.is_empty():
+            return ""
+        return str(df)
+
+    if hasattr(df, "to_markdown"):
+        return cast("str", df.to_markdown())
+
+    return str(df)
+
+
+def _dataframe_to_csv(df: Any) -> str:
+    if df is None:
+        return ""
+
+    if isinstance(df, pl.DataFrame):
+        if df.is_empty():
+            return ""
+        return df.write_csv()
+
+    if hasattr(df, "to_csv"):
+        return cast("str", df.to_csv(index=False))
+
+    return ""
+
+
+def _is_dataframe_empty(df: Any) -> bool:
+    if df is None:
+        return True
+
+    if isinstance(df, pl.DataFrame):
+        return df.is_empty()
+
+    if hasattr(df, "empty"):
+        return cast("bool", df.empty)
+
+    return True
 
 
 async def extract_tables(
@@ -111,7 +167,7 @@ async def extract_tables(
             )
             doc = await run_sync(PyPDFium2Document, str(file_path))
             cropped_tables: list[CroppedTable] = []
-            dataframes: list[DataFrame] = []
+            dataframes: list[pl.DataFrame] = []
             try:
                 for page in doc:
                     cropped_tables.extend(await run_sync(detector.extract, page))
@@ -124,8 +180,8 @@ async def extract_tables(
                     TableData(
                         cropped_image=cropped_table.image(),
                         page_number=cropped_table.page.page_number,
-                        text=data_frame.to_markdown(),
-                        df=data_frame,
+                        text=_dataframe_to_markdown(data_frame),
+                        df=_pandas_to_polars(data_frame),
                     )
                     for data_frame, cropped_table in zip(dataframes, cropped_tables, strict=False)
                 ]
@@ -225,7 +281,7 @@ def extract_tables_sync(
                 TableData(
                     cropped_image=cropped_table.image(),
                     page_number=cropped_table.page.page_number,
-                    text=data_frame.to_markdown(),
+                    text=_dataframe_to_markdown(data_frame),
                     df=data_frame,
                 )
                 for data_frame, cropped_table in zip(dataframes, cropped_tables, strict=False)
@@ -295,26 +351,16 @@ def _extract_tables_in_process(
                 cropped_image.save(img_bytes, format="PNG")
                 img_bytes.seek(0)
 
-                if data_frame.empty:
-                    results.append(
-                        {
-                            "cropped_image_bytes": img_bytes.getvalue(),
-                            "page_number": cropped_table.page.page_number,
-                            "text": data_frame.to_markdown(),
-                            "df_columns": data_frame.columns.tolist(),
-                            "df_csv": None,
-                        }
-                    )
-                else:
-                    results.append(
-                        {
-                            "cropped_image_bytes": img_bytes.getvalue(),
-                            "page_number": cropped_table.page.page_number,
-                            "text": data_frame.to_markdown(),
-                            "df_columns": None,
-                            "df_csv": data_frame.to_csv(index=False),
-                        }
-                    )
+                csv_data = _dataframe_to_csv(data_frame) if not _is_dataframe_empty(data_frame) else ""
+                results.append(
+                    {
+                        "cropped_image_bytes": img_bytes.getvalue(),
+                        "page_number": cropped_table.page.page_number,
+                        "text": _dataframe_to_markdown(data_frame),
+                        "df_columns": data_frame.columns,
+                        "df_csv": csv_data if csv_data else None,
+                    }
+                )
 
             result_queue.put((True, results))
 
@@ -381,10 +427,10 @@ def _extract_tables_isolated(
             for table_dict in result:
                 img = Image.open(io.BytesIO(table_dict["cropped_image_bytes"]))
 
-                if table_dict["df_csv"] is None:
-                    df = pd.DataFrame(columns=table_dict["df_columns"])
+                if table_dict["df_csv"] is None or table_dict["df_csv"] == "":
+                    df = pl.DataFrame()
                 else:
-                    df = pd.read_csv(StringIO(table_dict["df_csv"]))
+                    df = pl.read_csv(StringIO(table_dict["df_csv"]), truncate_ragged_lines=True)
 
                 tables.append(
                     TableData(
@@ -468,10 +514,10 @@ async def _extract_tables_isolated_async(
             for table_dict in result:
                 img = Image.open(io.BytesIO(table_dict["cropped_image_bytes"]))
 
-                if table_dict["df_csv"] is None:
-                    df = pd.DataFrame(columns=table_dict["df_columns"])
+                if table_dict["df_csv"] is None or table_dict["df_csv"] == "":
+                    df = pl.DataFrame()
                 else:
-                    df = pd.read_csv(StringIO(table_dict["df_csv"]))
+                    df = pl.read_csv(StringIO(table_dict["df_csv"]), truncate_ragged_lines=True)
 
                 tables.append(
                     TableData(

kreuzberg/_mcp/server.py

@@ -16,14 +16,6 @@ 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.
-    """
     base_config = discover_config()
 
     if base_config is None:
@@ -64,25 +56,6 @@ def extract_document(  # noqa: PLR0913
     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 = _create_config_with_overrides(
         force_ocr=force_ocr,
         chunk_content=chunk_content,
@@ -115,25 +88,6 @@ def extract_bytes(  # noqa: PLR0913
     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 = _create_config_with_overrides(
@@ -158,15 +112,6 @@ 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 = _create_config_with_overrides()
     result = extract_file_sync(file_path, mime_type, config)
     return result.content
@@ -174,14 +119,12 @@ def extract_simple(
 
 @mcp.resource("config://default")
 def get_default_config() -> str:
-    """Get the default extraction configuration."""
     config = ExtractionConfig()
     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 = discover_config()
     if config is None:
         return "No configuration file found"
@@ -190,13 +133,11 @@ def get_discovered_config() -> str:
 
 @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
@@ -210,14 +151,6 @@ def get_supported_formats() -> str:
 
 @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, _create_config_with_overrides())
 
     return [
@@ -230,14 +163,6 @@ def extract_and_summarize(file_path: str) -> list[TextContent]:
 
 @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 = _create_config_with_overrides(
         extract_entities=True,
         extract_keywords=True,
@@ -262,7 +187,6 @@ def extract_structured(file_path: str) -> list[TextContent]:
 
 
 def main() -> None:  # pragma: no cover
-    """Main entry point for the MCP server."""
     mcp.run()
 
 

kreuzberg/_ocr/_base.py

@@ -5,6 +5,7 @@ from typing import Generic, TypeVar
 from PIL.Image import Image
 
 from kreuzberg._types import ExtractionResult
+from kreuzberg._utils._sync import run_taskgroup
 
 try:  # pragma: no cover
     from typing import Unpack  # type: ignore[attr-defined]
@@ -32,8 +33,6 @@ class OCRBackend(ABC, Generic[T]):
         return [self.process_file_sync(path, **kwargs) for path in paths]  # pragma: no cover
 
     async def process_batch(self, paths: list[Path], **kwargs: Unpack[T]) -> list[ExtractionResult]:
-        from kreuzberg._utils._sync import run_taskgroup  # noqa: PLC0415
-
         tasks = [self.process_file(path, **kwargs) for path in paths]
         return await run_taskgroup(*tasks)  # pragma: no cover