kreuzberg 1.7.0__py3-none-any.whl → 2.0.1__py3-none-any.whl

kreuzberg/_pdf.py

@@ -0,0 +1,189 @@
+from __future__ import annotations
+
+from re import Pattern
+from re import compile as compile_regex
+from typing import TYPE_CHECKING, Final, cast
+
+import pypdfium2
+from anyio import Path as AsyncPath
+
+from kreuzberg import ExtractionResult
+from kreuzberg._mime_types import PLAIN_TEXT_MIME_TYPE
+from kreuzberg._string import normalize_spaces
+from kreuzberg._sync import run_sync
+from kreuzberg._tesseract import PSMMode, batch_process_images
+from kreuzberg.exceptions import ParsingError
+
+if TYPE_CHECKING:  # pragma: no cover
+    from pathlib import Path
+
+    from PIL.Image import Image
+
+
+# Pattern to detect common PDF text extraction corruption:
+# - Control and non-printable characters
+# - Unicode replacement and invalid characters
+# - Zero-width spaces and other invisible characters
+CORRUPTED_PATTERN: Final[Pattern[str]] = compile_regex(
+    r"[\x00-\x08\x0B-\x1F\x7F-\x9F]|\uFFFD|[\u200B-\u200F\u2028-\u202F]"
+)
+
+
+def _validate_extracted_text(text: str) -> bool:
+    """Check if text extracted from PDF is valid or corrupted.
+
+    This checks for common indicators of corrupted PDF text extraction:
+    1. Empty or whitespace-only text
+    2. Control characters and other non-printable characters
+    3. Unicode replacement characters
+    4. Zero-width spaces and other invisible characters
+
+    Args:
+        text: The extracted text to validate
+
+    Returns:
+        True if the text appears valid, False if it seems corrupted
+    """
+    # Check for empty or whitespace-only text
+    if not text or not text.strip():
+        return False
+
+    # Check for corruption indicators
+    return not bool(CORRUPTED_PATTERN.search(text))
+
+
+async def _convert_pdf_to_images(input_file: Path) -> list[Image]:
+    """Convert a PDF file to images.
+
+    Args:
+        input_file: The path to the PDF file.
+
+    Raises:
+        ParsingError: If the PDF file could not be converted to images.
+
+    Returns:
+        A list of Pillow Images.
+    """
+    document: pypdfium2.PdfDocument | None = None
+    try:
+        document = await run_sync(pypdfium2.PdfDocument, str(input_file))
+        return [page.render(scale=2.0).to_pil() for page in cast(pypdfium2.PdfDocument, document)]
+    except pypdfium2.PdfiumError as e:
+        raise ParsingError(
+            "Could not convert PDF to images", context={"file_path": str(input_file), "error": str(e)}
+        ) from e
+    finally:
+        if document:
+            await run_sync(document.close)
+
+
+async def _extract_pdf_text_with_ocr(
+    input_file: Path,
+    *,
+    language: str = "eng",
+    max_processes: int,
+    psm: PSMMode = PSMMode.AUTO,
+) -> ExtractionResult:
+    """Extract text from a scanned PDF file using pytesseract.
+
+    Args:
+        input_file: The path to the PDF file.
+        language: The language code for OCR. Defaults to "eng".
+        max_processes: Maximum number of concurrent processes. Defaults to CPU count / 2 (minimum 1).
+        psm: Page segmentation mode for Tesseract OCR. Defaults to PSMMode.AUTO.
+
+    Returns:
+        The extracted text.
+    """
+    images = await _convert_pdf_to_images(input_file)
+    ocr_results = await batch_process_images(images, max_processes=max_processes, psm=psm, language=language)
+    return ExtractionResult(
+        content="\n".join([v.content for v in ocr_results]), mime_type=PLAIN_TEXT_MIME_TYPE, metadata={}
+    )
+
+
+async def _extract_pdf_searchable_text(input_file: Path) -> str:
+    """Extract text from a searchable PDF file using pypdfium2.
+
+    Args:
+        input_file: The path to the PDF file.
+
+    Raises:
+        ParsingError: If the text could not be extracted from the PDF file.
+
+    Returns:
+        The extracted text.
+    """
+    document: pypdfium2.PdfDocument | None = None
+    try:
+        document = await run_sync(pypdfium2.PdfDocument, str(input_file))
+        text = "\n".join(page.get_textpage().get_text_bounded() for page in cast(pypdfium2.PdfDocument, document))
+        return normalize_spaces(text)
+    except pypdfium2.PdfiumError as e:
+        raise ParsingError(
+            "Could not extract text from PDF file", context={"file_path": str(input_file), "error": str(e)}
+        ) from e
+    finally:
+        if document:
+            await run_sync(document.close)
+
+
+async def extract_pdf_file(
+    input_file: Path,
+    *,
+    force_ocr: bool,
+    language: str = "eng",
+    max_processes: int,
+    psm: PSMMode = PSMMode.AUTO,
+) -> ExtractionResult:
+    """Extract text from a PDF file.
+
+    Args:
+        input_file: The path to the PDF file.
+        force_ocr: Whether to force OCR on PDF files that have a text layer.
+        language: The language code for OCR. Defaults to "eng".
+        max_processes: Maximum number of concurrent processes. Defaults to CPU count / 2 (minimum 1).
+        psm: Page segmentation mode for Tesseract OCR. Defaults to PSMMode.AUTO.
+
+    Returns:
+        The extracted text.
+    """
+    if (
+        not force_ocr
+        and (content := await _extract_pdf_searchable_text(input_file))
+        and _validate_extracted_text(content)
+    ):
+        return ExtractionResult(content=content, mime_type=PLAIN_TEXT_MIME_TYPE, metadata={})
+
+    return await _extract_pdf_text_with_ocr(input_file, max_processes=max_processes, language=language, psm=psm)
+
+
+async def extract_pdf_content(
+    content: bytes,
+    *,
+    force_ocr: bool,
+    language: str = "eng",
+    max_processes: int,
+    psm: PSMMode = PSMMode.AUTO,
+) -> ExtractionResult:
+    """Extract text from a PDF file content.
+
+    Args:
+        content: The PDF file content.
+        force_ocr: Whether to force OCR on PDF files that have a text layer.
+        language: The language code for OCR. Defaults to "eng".
+        max_processes: Maximum number of concurrent processes. Defaults to CPU count / 2 (minimum 1).
+        psm: Page segmentation mode for Tesseract OCR. Defaults to PSMMode.AUTO.
+
+    Returns:
+        The extracted text.
+    """
+    from kreuzberg._tmp import create_temp_file
+
+    file_path, unlink = await create_temp_file(".pdf")
+    await AsyncPath(file_path).write_bytes(content)
+    result = await extract_pdf_file(
+        file_path, force_ocr=force_ocr, max_processes=max_processes, psm=psm, language=language
+    )
+    await unlink()
+    return result

kreuzberg/_pptx.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import re
+from contextlib import suppress
+from html import escape
+from io import BytesIO
+from typing import TYPE_CHECKING
+
+import pptx
+from anyio import Path as AsyncPath
+from pptx.enum.shapes import MSO_SHAPE_TYPE
+
+from kreuzberg import ExtractionResult
+from kreuzberg._mime_types import MARKDOWN_MIME_TYPE
+from kreuzberg._string import normalize_spaces
+
+if TYPE_CHECKING:  # pragma: no cover
+    from pathlib import Path
+
+
+async def extract_pptx_file_content(file_path_or_contents: Path | bytes) -> ExtractionResult:
+    """Extract text from a PPTX file.
+
+    Notes:
+        This function is based on code vendored from `markitdown`, which has an MIT license as well.
+
+    Args:
+        file_path_or_contents: The path to the PPTX file or its contents as bytes.
+
+    Returns:
+        The extracted text content
+    """
+    md_content = ""
+    file_contents = (
+        file_path_or_contents
+        if isinstance(file_path_or_contents, bytes)
+        else await AsyncPath(file_path_or_contents).read_bytes()
+    )
+    presentation = pptx.Presentation(BytesIO(file_contents))
+
+    for index, slide in enumerate(presentation.slides):
+        md_content += f"\n\n<!-- Slide number: {index + 1} -->\n"
+
+        title = slide.shapes.title
+
+        for shape in slide.shapes:
+            if shape.shape_type == MSO_SHAPE_TYPE.PICTURE or (
+                shape.shape_type == MSO_SHAPE_TYPE.PLACEHOLDER and hasattr(shape, "image")
+            ):
+                alt_text = ""
+                with suppress(AttributeError):
+                    # access non-visual properties
+                    alt_text = shape._element._nvXxPr.cNvPr.attrib.get("descr", "")  # noqa: SLF001
+
+                filename = re.sub(r"\W", "", shape.name) + ".jpg"
+                md_content += f"\n![{alt_text if alt_text else shape.name}]({filename})\n"
+
+            elif shape.shape_type == MSO_SHAPE_TYPE.TABLE:
+                html_table = "<table>"
+                first_row = True
+
+                for row in shape.table.rows:
+                    html_table += "<tr>"
+
+                    for cell in row.cells:
+                        tag = "th" if first_row else "td"
+                        html_table += f"<{tag}>{escape(cell.text)}</{tag}>"
+
+                    html_table += "</tr>"
+                    first_row = False
+
+                html_table += "</table>"
+                md_content += "\n" + html_table + "\n"
+
+            elif shape.has_text_frame:
+                md_content += "# " + shape.text.lstrip() + "\n" if shape == title else shape.text + "\n"
+
+        md_content = md_content.strip()
+        if slide.has_notes_slide:
+            md_content += "\n\n### Notes:\n"
+            notes_frame = slide.notes_slide.notes_text_frame
+
+            if notes_frame is not None:  # pragma: no branch
+                md_content += notes_frame.text
+
+            md_content = md_content.strip()
+
+    return ExtractionResult(content=normalize_spaces(md_content), mime_type=MARKDOWN_MIME_TYPE, metadata={})

kreuzberg/_string.py

@@ -4,8 +4,6 @@ from contextlib import suppress
 
 from charset_normalizer import detect
 
-from kreuzberg.exceptions import ParsingError
-
 
 def safe_decode(byte_data: bytes, encoding: str | None = None) -> str:
     """Decode a byte string safely, removing invalid sequences.
@@ -14,22 +12,21 @@ def safe_decode(byte_data: bytes, encoding: str | None = None) -> str:
         byte_data: The byte string to decode.
         encoding: The encoding to use when decoding the byte string.
 
-    Raises:
-        ParsingError: If the byte string could not be decoded.
-
     Returns:
         The decoded string.
     """
     if not byte_data:
         return ""
 
-    encodings = [encoding, detect(byte_data).get("encoding", ""), "utf-8", "latin-1"]
+    # We try each encoding in order until one works
+    encodings = [encoding, detect(byte_data).get("encoding", ""), "utf-8"]
 
-    for enc in [e for e in encodings if e]:
+    for enc in [e for e in encodings if e]:  # pragma: no cover
         with suppress(UnicodeDecodeError):
             return byte_data.decode(enc)
 
-    raise ParsingError("Could not decode byte string. Please provide an encoding.")
+    # If all encodings fail, fall back to latin-1 which can handle any byte
+    return byte_data.decode("latin-1", errors="replace")
 
 
 def normalize_spaces(text: str) -> str:

kreuzberg/_sync.py

@@ -1,14 +1,19 @@
 from __future__ import annotations
 
+import sys
 from functools import partial
 from typing import TYPE_CHECKING, TypeVar, cast
 
 from anyio.to_thread import run_sync as any_io_run_sync
-from typing_extensions import ParamSpec
 
 if TYPE_CHECKING:  # pragma: no cover
     from collections.abc import Callable
 
+if sys.version_info >= (3, 10):
+    from typing import ParamSpec
+else:  # pragma: no cover
+    from typing_extensions import ParamSpec
+
 T = TypeVar("T")
 P = ParamSpec("P")