Add LLM captioning for embedded PDF images

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)
 

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__]))