diff --git a/packages/markitdown/src/markitdown/converters/_pdf_converter.py b/packages/markitdown/src/markitdown/converters/_pdf_converter.py index ffbcbd990c..4a4c8dc36f 100644 --- a/packages/markitdown/src/markitdown/converters/_pdf_converter.py +++ b/packages/markitdown/src/markitdown/converters/_pdf_converter.py @@ -1,11 +1,12 @@ -import sys import io import re -from typing import BinaryIO, Any +import sys +from typing import Any, BinaryIO from .._base_converter import DocumentConverter, DocumentConverterResult +from .._exceptions import MISSING_DEPENDENCY_MESSAGE, MissingDependencyException from .._stream_info import StreamInfo -from .._exceptions import MissingDependencyException, MISSING_DEPENDENCY_MESSAGE +from ._llm_caption import llm_caption # Pattern for MasterFormat-style partial numbering (e.g., ".1", ".2", ".10") PARTIAL_NUMBERING_PATTERN = re.compile(r"^\.\d+$") @@ -492,6 +493,39 @@ def _extract_tables_from_words(page: Any) -> list[list[list[str]]]: return [table_rows] +def _extract_pdf_image_payload(image: dict[str, Any]) -> tuple[bytes, str, str] | None: + """Return image bytes and MIME metadata when the embedded format is recognized.""" + from pdfminer.pdfexceptions import PDFException + + image_stream = image.get("stream") + if image_stream is None: + return None + + try: + data = image_stream.get_data() + except PDFException: + return None + + if not isinstance(data, bytes) or not data: + return None + + signatures = ( + (b"\xff\xd8\xff", "image/jpeg", ".jpg"), + (b"\x89PNG\r\n\x1a\n", "image/png", ".png"), + (b"GIF87a", "image/gif", ".gif"), + (b"GIF89a", "image/gif", ".gif"), + (b"II*\x00", "image/tiff", ".tif"), + (b"MM\x00*", "image/tiff", ".tif"), + (b"RIFF", "image/webp", ".webp"), + ) + for signature, mimetype, extension in signatures: + if data.startswith(signature) and ( + mimetype != "image/webp" or data[8:12] == b"WEBP" + ): + return data, mimetype, extension + return None + + class PdfConverter(DocumentConverter): """ Converts PDFs to Markdown. @@ -539,6 +573,8 @@ def convert( # Read file stream into BytesIO for compatibility with pdfplumber pdf_bytes = io.BytesIO(file_stream.read()) + embedded_images: list[tuple[int, bytes, str, str]] = [] + try: # Single pass: check every page for form-style content. # Pages with tables/forms get rich extraction; plain-text @@ -563,6 +599,14 @@ def convert( if text and text.strip(): markdown_chunks.append(text.strip()) + for image in page.images: + payload = _extract_pdf_image_payload(image) + if payload is not None: + image_data, mimetype, extension = payload + embedded_images.append( + (page_idx + 1, image_data, mimetype, extension) + ) + page.close() # Free cached page data immediately # If no pages had form-style content, use pdfminer for @@ -583,6 +627,40 @@ def convert( pdf_bytes.seek(0) markdown = pdfminer.high_level.extract_text(pdf_bytes) + # Analyze embedded PDF images when an LLM client is configured. + llm_client = kwargs.get("llm_client") + llm_model = kwargs.get("llm_model") + + if llm_client is not None and llm_model and embedded_images: + image_prompt = kwargs.get("llm_prompt") or ( + "Analyze this image from a PDF. If it contains a chart or graph, " + "extract its data into a Markdown table, preserving labels, units, " + "and legible values. Do not invent unreadable values. Otherwise, " + "provide a concise description of the image." + ) + image_chunks: list[str] = [] + + for page_number, image_data, mimetype, extension in embedded_images: + caption = llm_caption( + io.BytesIO(image_data), + StreamInfo(mimetype=mimetype, extension=extension), + client=llm_client, + model=llm_model, + prompt=image_prompt, + ) + if caption and caption.strip(): + image_chunks.append( + f"### Image on page {page_number}\n\n{caption.strip()}" + ) + + if image_chunks: + image_markdown = "\n\n".join(image_chunks) + markdown = ( + f"{markdown.rstrip()}\n\n{image_markdown}" + if markdown and markdown.strip() + else image_markdown + ) + # Post-process to merge MasterFormat-style partial numbering with following text markdown = _merge_partial_numbering_lines(markdown) diff --git a/packages/markitdown/tests/test_pdf.py b/packages/markitdown/tests/test_pdf.py index c4f8c64a9c..1d27442c9e 100644 --- a/packages/markitdown/tests/test_pdf.py +++ b/packages/markitdown/tests/test_pdf.py @@ -5,14 +5,12 @@ from unittest.mock import patch import pytest - from markitdown import MarkItDown from markitdown.converters._pdf_converter import ( PARTIAL_NUMBERING_PATTERN, _merge_partial_numbering_lines, ) - # Table extraction TEST_FILES_DIR = os.path.join(os.path.dirname(__file__), "test_files") @@ -1384,5 +1382,58 @@ def test_markitdown_remote() -> None: assert test_string in result.text_content +def test_pdf_embedded_images_are_captioned_with_llm() -> None: + from unittest.mock import MagicMock + + pdf_path = os.path.join( + TEST_FILES_DIR, "REPAIR-2022-INV-001_multipage.pdf" + ) + client = MagicMock() + client.chat.completions.create.return_value.choices[0].message.content = ( + "Chart summary: quarterly values increased." + ) + prompt = "Describe this PDF image." + + result = MarkItDown( + llm_client=client, + llm_model="test-model", + llm_prompt=prompt, + ).convert(pdf_path) + + assert client.chat.completions.create.call_count == 2 + assert "### Image on page 1" in result.text_content + assert "### Image on page 2" in result.text_content + assert result.text_content.count( + "Chart summary: quarterly values increased." + ) == 2 + for call in client.chat.completions.create.call_args_list: + assert call.kwargs["model"] == "test-model" + assert call.kwargs["messages"][0]["content"][0]["text"] == prompt + + +def test_pdf_embedded_images_use_chart_extraction_prompt_by_default() -> None: + from unittest.mock import MagicMock + + pdf_path = os.path.join( + TEST_FILES_DIR, "REPAIR-2022-INV-001_multipage.pdf" + ) + client = MagicMock() + client.chat.completions.create.return_value.choices[0].message.content = ( + "Chart data extracted." + ) + + MarkItDown( + llm_client=client, + llm_model="test-model", + ).convert(pdf_path) + + assert client.chat.completions.create.call_count == 2 + for call in client.chat.completions.create.call_args_list: + prompt = call.kwargs["messages"][0]["content"][0]["text"] + assert "chart or graph" in prompt + assert "Markdown table" in prompt + assert "Do not invent unreadable values" in prompt + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__]))