pdf-extractor-ai 0.1.0__tar.gz

pdf_extractor_ai-0.1.0/PKG-INFO

@@ -0,0 +1,218 @@
+Metadata-Version: 2.3
+Name: pdf-extractor-ai
+Version: 0.1.0
+Summary: Add your description here
+Author: Oleh Lokshyn
+Author-email: Oleh Lokshyn <olokshyn@gmail.com>
+Requires-Dist: docling>=2.130.0
+Requires-Dist: pillow>=12.3.0
+Requires-Dist: pypdfium2>=5.13.0
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+
+# pdf-extractor-ai
+
+Extract a PDF from bytes into one [Docling](https://github.com/docling-project/docling) document, optionally replace figures with summaries from a vision model you supply, then serialize that document to markdown.
+
+The package does not download URLs and does not call a vision API. You pass PDF bytes, and you pass the model when you want summaries.
+
+```text
+PDF bytes  →  PdfExtractor.extract  →  DoclingDocument
+                                              ↓
+                                    summarize_images (optional)
+                                              ↓
+                                         to_markdown
+```
+
+Reading order is newspaper order: physical page, then left-to-right column, then top to bottom. Headings become the tree. Text, lists, tables, and pictures that follow a heading are its children. A section that continues onto later pages is inserted only after those pages have been read.
+
+## Install
+
+This project uses [uv](https://docs.astral.sh/uv/). From the repository:
+
+```bash
+uv sync
+```
+
+Run scripts with `uv run` so they use that environment:
+
+```bash
+uv run python your_script.py
+```
+
+## Extract a PDF
+
+`extract` returns a `DoclingDocument`. It does not call a vision model and it does not return markdown. `do_ocr` defaults to `False`. Set it to `True` for scanned pages. Tables stay on either way.
+
+```python
+from pdf_extractor_ai import PdfExtractor, document_title, to_markdown
+
+pdf_bytes = open("report.pdf", "rb").read()
+
+with PdfExtractor(num_workers=2) as extractor:
+    document = extractor.extract(pdf_bytes, name="report.pdf")
+
+print(document_title(document, fallback="report.pdf"))
+print(to_markdown(document))
+```
+
+`num_workers` threads share one queue. Each thread builds its own Docling converter on first use and keeps a private converter for each `do_ocr` value. Use the extractor as a context manager, or call `shutdown()` when you are done.
+
+`extract_async` is the same job, awaited from an event loop:
+
+```python
+document = await extractor.extract_async(pdf_bytes, do_ocr=False, max_pages=20)
+```
+
+`max_pages` stops after that many pages. An empty or unreadable PDF raises `ValueError`.
+
+## Markdown
+
+`to_markdown` walks the document in tree order. Page banners stay open across a page and close when the next page starts, or at the end. A picture that was not summarized becomes `<!-- image -->`. Text that sits inside a picture box is left out. A summary block is written with its line breaks intact.
+
+A section heading of level 1 is `##`, because a document title uses `#` and each section level adds one more hash.
+
+For a two-page report whose first heading is “Crop prospects” and whose second page starts “Regional outlook”, `to_markdown` returns:
+
+```text
+========== page 1 start ==========
+
+## Crop prospects
+
+Global cereal production is expected to rise.
+
+========== page 1 end ==========
+
+========== page 2 start ==========
+
+### Regional outlook
+
+West Africa faces a rainfall deficit.
+
+========== page 2 end ==========
+```
+
+`document_title` reads that tree. It prefers a page-1 title over a page-1 section heading, then the longest heading at that rank, then the first shallowest heading anywhere, then the fallback you pass. For the document above it returns `Crop prospects`.
+
+`split_pages` returns the text inside each banner:
+
+```python
+from pdf_extractor_ai import split_pages
+
+pages = split_pages(markdown)
+pages[1]
+# '## Crop prospects\n\nGlobal cereal production is expected to rise.'
+pages[2]
+# '### Regional outlook\n\nWest Africa faces a rainfall deficit.'
+```
+
+`is_low_text` is true when the text, after banners and image placeholders are removed, is shorter than `min_chars` (default 200). The sample above is short, so `is_low_text(markdown)` is `True` and `is_low_text(markdown, min_chars=40)` is `False`.
+
+`layout_from_markdown` maps those banners and headings onto line spans:
+
+```python
+from pdf_extractor_ai.markdown import layout_from_markdown
+
+layout = layout_from_markdown(markdown, document)
+layout.document_title
+# 'Crop prospects'
+layout.sections[1].heading_path
+# ('Crop prospects', 'Regional outlook')
+layout.sections[1].page_start, layout.sections[1].page_end
+# (2, 2)
+```
+
+## Summarize figures
+
+`summarize_images` runs after `extract` and before `to_markdown`. Group pictures by their parent heading and call your model once per section that has at least one picture. Sections with no pictures are skipped. The default concurrency is 3. The input document is left unchanged, and the function returns a new one.
+
+Your model matches this protocol. `visual_id` values look like `s0-v0`. Return a map of those ids to summary text.
+
+```python
+import asyncio
+
+from pdf_extractor_ai import summarize_images, to_markdown
+from pdf_extractor_ai.image_summarizer import SectionVisual, VisualLanguageModel
+
+
+class GeminiCharts:
+    async def describe_section(
+        self,
+        *,
+        section_title: str,
+        section_text: str,
+        visuals: list[SectionVisual],
+    ) -> dict[str, str]:
+        # section_text is the heading plus the other text in that section.
+        # visuals[i].png_bytes is a PNG crop. Send those to your vision API.
+        return {
+            visual.visual_id: "Cereal production rises through 2026."
+            for visual in visuals
+        }
+
+
+async def main() -> None:
+    summarized = await summarize_images(
+        document,
+        GeminiCharts(),
+        model_name="gemini/gemini-3.7-flash",
+    )
+    print(to_markdown(summarized))
+
+
+asyncio.run(main())
+```
+
+A successful summary replaces that picture with a text block. A missing id, an empty summary, or a raised error leaves the picture in place, so `to_markdown` still emits `<!-- image -->` for it.
+
+Before a summary, a figure under the heading “Charts” looks like this:
+
+```text
+========== page 1 start ==========
+
+## Charts
+
+Chart context
+
+<!-- image -->
+
+========== page 1 end ==========
+```
+
+After the call above, the same figure is the summary block. `model_name`, page, and bounding box come from the picture, and the last lines are the text your model returned:
+
+```text
+========== page 1 start ==========
+
+## Charts
+
+Chart context
+
+[Summarized Image: gemini/gemini-3.7-flash]
+page: 1
+bounding box: l=40.0, t=120.0, r=280.0, b=360.0 (TOPLEFT)
+Summary:
+Cereal production rises through 2026.
+
+========== page 1 end ==========
+```
+
+Pictures whose shorter side is under 80 pixels are omitted during extraction.
+
+## Development
+
+```bash
+uv sync
+uv run ruff format
+uv run ruff check
+uv run mypy
+uv run pytest
+```
+
+`tests/test_pdf_visual_integration.py` calls Gemini and is skipped unless `GEMINI_API_KEY` is in the environment. Pass the repo env file when you want that test to run:
+
+```bash
+uv run --env-file .env pytest
+```
+
+The rest of the suite does not load Docling weights.

pdf_extractor_ai-0.1.0/README.md

@@ -0,0 +1,206 @@
+# pdf-extractor-ai
+
+Extract a PDF from bytes into one [Docling](https://github.com/docling-project/docling) document, optionally replace figures with summaries from a vision model you supply, then serialize that document to markdown.
+
+The package does not download URLs and does not call a vision API. You pass PDF bytes, and you pass the model when you want summaries.
+
+```text
+PDF bytes  →  PdfExtractor.extract  →  DoclingDocument
+                                              ↓
+                                    summarize_images (optional)
+                                              ↓
+                                         to_markdown
+```
+
+Reading order is newspaper order: physical page, then left-to-right column, then top to bottom. Headings become the tree. Text, lists, tables, and pictures that follow a heading are its children. A section that continues onto later pages is inserted only after those pages have been read.
+
+## Install
+
+This project uses [uv](https://docs.astral.sh/uv/). From the repository:
+
+```bash
+uv sync
+```
+
+Run scripts with `uv run` so they use that environment:
+
+```bash
+uv run python your_script.py
+```
+
+## Extract a PDF
+
+`extract` returns a `DoclingDocument`. It does not call a vision model and it does not return markdown. `do_ocr` defaults to `False`. Set it to `True` for scanned pages. Tables stay on either way.
+
+```python
+from pdf_extractor_ai import PdfExtractor, document_title, to_markdown
+
+pdf_bytes = open("report.pdf", "rb").read()
+
+with PdfExtractor(num_workers=2) as extractor:
+    document = extractor.extract(pdf_bytes, name="report.pdf")
+
+print(document_title(document, fallback="report.pdf"))
+print(to_markdown(document))
+```
+
+`num_workers` threads share one queue. Each thread builds its own Docling converter on first use and keeps a private converter for each `do_ocr` value. Use the extractor as a context manager, or call `shutdown()` when you are done.
+
+`extract_async` is the same job, awaited from an event loop:
+
+```python
+document = await extractor.extract_async(pdf_bytes, do_ocr=False, max_pages=20)
+```
+
+`max_pages` stops after that many pages. An empty or unreadable PDF raises `ValueError`.
+
+## Markdown
+
+`to_markdown` walks the document in tree order. Page banners stay open across a page and close when the next page starts, or at the end. A picture that was not summarized becomes `<!-- image -->`. Text that sits inside a picture box is left out. A summary block is written with its line breaks intact.
+
+A section heading of level 1 is `##`, because a document title uses `#` and each section level adds one more hash.
+
+For a two-page report whose first heading is “Crop prospects” and whose second page starts “Regional outlook”, `to_markdown` returns:
+
+```text
+========== page 1 start ==========
+
+## Crop prospects
+
+Global cereal production is expected to rise.
+
+========== page 1 end ==========
+
+========== page 2 start ==========
+
+### Regional outlook
+
+West Africa faces a rainfall deficit.
+
+========== page 2 end ==========
+```
+
+`document_title` reads that tree. It prefers a page-1 title over a page-1 section heading, then the longest heading at that rank, then the first shallowest heading anywhere, then the fallback you pass. For the document above it returns `Crop prospects`.
+
+`split_pages` returns the text inside each banner:
+
+```python
+from pdf_extractor_ai import split_pages
+
+pages = split_pages(markdown)
+pages[1]
+# '## Crop prospects\n\nGlobal cereal production is expected to rise.'
+pages[2]
+# '### Regional outlook\n\nWest Africa faces a rainfall deficit.'
+```
+
+`is_low_text` is true when the text, after banners and image placeholders are removed, is shorter than `min_chars` (default 200). The sample above is short, so `is_low_text(markdown)` is `True` and `is_low_text(markdown, min_chars=40)` is `False`.
+
+`layout_from_markdown` maps those banners and headings onto line spans:
+
+```python
+from pdf_extractor_ai.markdown import layout_from_markdown
+
+layout = layout_from_markdown(markdown, document)
+layout.document_title
+# 'Crop prospects'
+layout.sections[1].heading_path
+# ('Crop prospects', 'Regional outlook')
+layout.sections[1].page_start, layout.sections[1].page_end
+# (2, 2)
+```
+
+## Summarize figures
+
+`summarize_images` runs after `extract` and before `to_markdown`. Group pictures by their parent heading and call your model once per section that has at least one picture. Sections with no pictures are skipped. The default concurrency is 3. The input document is left unchanged, and the function returns a new one.
+
+Your model matches this protocol. `visual_id` values look like `s0-v0`. Return a map of those ids to summary text.
+
+```python
+import asyncio
+
+from pdf_extractor_ai import summarize_images, to_markdown
+from pdf_extractor_ai.image_summarizer import SectionVisual, VisualLanguageModel
+
+
+class GeminiCharts:
+    async def describe_section(
+        self,
+        *,
+        section_title: str,
+        section_text: str,
+        visuals: list[SectionVisual],
+    ) -> dict[str, str]:
+        # section_text is the heading plus the other text in that section.
+        # visuals[i].png_bytes is a PNG crop. Send those to your vision API.
+        return {
+            visual.visual_id: "Cereal production rises through 2026."
+            for visual in visuals
+        }
+
+
+async def main() -> None:
+    summarized = await summarize_images(
+        document,
+        GeminiCharts(),
+        model_name="gemini/gemini-3.7-flash",
+    )
+    print(to_markdown(summarized))
+
+
+asyncio.run(main())
+```
+
+A successful summary replaces that picture with a text block. A missing id, an empty summary, or a raised error leaves the picture in place, so `to_markdown` still emits `<!-- image -->` for it.
+
+Before a summary, a figure under the heading “Charts” looks like this:
+
+```text
+========== page 1 start ==========
+
+## Charts
+
+Chart context
+
+<!-- image -->
+
+========== page 1 end ==========
+```
+
+After the call above, the same figure is the summary block. `model_name`, page, and bounding box come from the picture, and the last lines are the text your model returned:
+
+```text
+========== page 1 start ==========
+
+## Charts
+
+Chart context
+
+[Summarized Image: gemini/gemini-3.7-flash]
+page: 1
+bounding box: l=40.0, t=120.0, r=280.0, b=360.0 (TOPLEFT)
+Summary:
+Cereal production rises through 2026.
+
+========== page 1 end ==========
+```
+
+Pictures whose shorter side is under 80 pixels are omitted during extraction.
+
+## Development
+
+```bash
+uv sync
+uv run ruff format
+uv run ruff check
+uv run mypy
+uv run pytest
+```
+
+`tests/test_pdf_visual_integration.py` calls Gemini and is skipped unless `GEMINI_API_KEY` is in the environment. Pass the repo env file when you want that test to run:
+
+```bash
+uv run --env-file .env pytest
+```
+
+The rest of the suite does not load Docling weights.

pdf_extractor_ai-0.1.0/pyproject.toml

@@ -0,0 +1,54 @@
+[project]
+name = "pdf-extractor-ai"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+authors = [
+    { name = "Oleh Lokshyn", email = "olokshyn@gmail.com" }
+]
+requires-python = ">=3.13"
+dependencies = [
+    "docling>=2.130.0",
+    "pillow>=12.3.0",
+    "pypdfium2>=5.13.0",
+]
+
+[build-system]
+requires = ["uv_build>=0.11.0,<0.12.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "ipykernel>=7.3.0",
+    "mypy>=2.3.1",
+    "pytest>=9.1.1",
+    "pytest-asyncio>=1.4.0",
+    "ruff>=0.16.4",
+]
+
+[tool.ruff]
+extend-exclude = ["notebooks", "src/pdf_extractor_ai/old"]
+
+[tool.ruff.lint]
+extend-select = ["I", "ANN"]
+ignore = ["ANN401"]  # Allow explicitly chosen Any as an escape hatch
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = ["."]
+python_files = ["*.py"]
+asyncio_mode = "auto"
+markers = [
+    "integration: tests that call external services",
+]
+
+[tool.mypy]
+strict = true
+files = ["src/pdf_extractor_ai", "tests"]
+mypy_path = "src"
+explicit_package_bases = true
+exclude = ["^notebooks/", "^src/pdf_extractor_ai/old/"]
+
+[[tool.mypy.overrides]]
+module = ["pypdfium2", "pypdfium2.*"]
+ignore_missing_imports = true

pdf_extractor_ai-0.1.0/src/pdf_extractor_ai/__init__.py

@@ -0,0 +1,19 @@
+"""Bytes-only PDF extraction that returns a reordered Docling document."""
+
+from pdf_extractor_ai.extractor import PdfExtractor
+from pdf_extractor_ai.image_summarizer import summarize_images
+from pdf_extractor_ai.markdown import (
+    document_title,
+    is_low_text,
+    split_pages,
+    to_markdown,
+)
+
+__all__ = [
+    "PdfExtractor",
+    "document_title",
+    "is_low_text",
+    "split_pages",
+    "summarize_images",
+    "to_markdown",
+]

pdf_extractor_ai-0.1.0/src/pdf_extractor_ai/extractor.py

@@ -0,0 +1,193 @@
+"""Page-by-page PDF extraction on a pool of private Docling converters."""
+
+from __future__ import annotations
+
+import asyncio
+import gc
+import io
+import logging
+import queue
+import threading
+from collections.abc import Callable
+from concurrent.futures import Future
+from dataclasses import dataclass
+from typing import Any, Self
+
+from docling_core.types.doc.document import DoclingDocument
+
+from pdf_extractor_ai.reflow import SectionAssembler
+
+logger = logging.getLogger(__name__)
+
+
+def build_converter(do_ocr: bool) -> Any:
+    """Build a Docling converter. Call this inside the worker that will use it."""
+    from docling.datamodel.accelerator_options import (
+        AcceleratorDevice,
+        AcceleratorOptions,
+    )
+    from docling.datamodel.base_models import InputFormat
+    from docling.datamodel.pipeline_options import PdfPipelineOptions
+    from docling.document_converter import DocumentConverter, PdfFormatOption
+
+    pipeline_options = PdfPipelineOptions(
+        do_ocr=do_ocr,
+        do_table_structure=True,
+        generate_picture_images=True,
+        images_scale=2.0,
+        accelerator_options=AcceleratorOptions(device=AcceleratorDevice.CPU),
+    )
+    return DocumentConverter(
+        format_options={
+            InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options),
+        }
+    )
+
+
+@dataclass(slots=True)
+class _Job:
+    pdf_bytes: bytes
+    do_ocr: bool
+    max_pages: int | None
+    name: str
+    future: Future[DoclingDocument]
+
+
+class PdfExtractor:
+    """Extract PDFs from a shared queue. Each worker owns its Docling converters."""
+
+    def __init__(
+        self,
+        num_workers: int = 1,
+        *,
+        converter_factory: Callable[[bool], Any] | None = None,
+    ) -> None:
+        if num_workers < 1:
+            raise ValueError("num_workers must be at least 1")
+        self._factory = converter_factory or build_converter
+        self._queue: queue.Queue[_Job | None] = queue.Queue()
+        self._shutdown_lock = threading.Lock()
+        self._stopped = False
+        self._threads = [
+            threading.Thread(target=self._run, name=f"pdf-extract-{index}", daemon=True)
+            for index in range(num_workers)
+        ]
+        for thread in self._threads:
+            thread.start()
+
+    def extract(
+        self,
+        pdf_bytes: bytes,
+        *,
+        do_ocr: bool = False,
+        max_pages: int | None = None,
+        name: str = "document.pdf",
+    ) -> DoclingDocument:
+        return self._submit(
+            pdf_bytes,
+            do_ocr=do_ocr,
+            max_pages=max_pages,
+            name=name,
+        ).result()
+
+    async def extract_async(
+        self,
+        pdf_bytes: bytes,
+        *,
+        do_ocr: bool = False,
+        max_pages: int | None = None,
+        name: str = "document.pdf",
+    ) -> DoclingDocument:
+        future = self._submit(
+            pdf_bytes,
+            do_ocr=do_ocr,
+            max_pages=max_pages,
+            name=name,
+        )
+        return await asyncio.wrap_future(future)
+
+    def shutdown(self) -> None:
+        with self._shutdown_lock:
+            if self._stopped:
+                return
+            self._stopped = True
+            for _thread in self._threads:
+                self._queue.put(None)
+        for thread in self._threads:
+            thread.join(timeout=30)
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, *_args: object) -> None:
+        self.shutdown()
+
+    def _submit(
+        self,
+        pdf_bytes: bytes,
+        *,
+        do_ocr: bool,
+        max_pages: int | None,
+        name: str,
+    ) -> Future[DoclingDocument]:
+        if self._stopped:
+            raise RuntimeError("PdfExtractor is shut down")
+        future: Future[DoclingDocument] = Future()
+        self._queue.put(
+            _Job(
+                pdf_bytes=pdf_bytes,
+                do_ocr=do_ocr,
+                max_pages=max_pages,
+                name=name,
+                future=future,
+            )
+        )
+        return future
+
+    def _run(self) -> None:
+        converters: dict[bool, Any] = {}
+        while True:
+            job = self._queue.get()
+            if job is None:
+                return
+            try:
+                converter = converters.get(job.do_ocr)
+                if converter is None:
+                    converter = self._factory(job.do_ocr)
+                    converters[job.do_ocr] = converter
+                job.future.set_result(_convert_pdf(converter, job))
+            except Exception as exc:
+                logger.exception("PDF extract failed for %s", job.name)
+                if not job.future.done():
+                    job.future.set_exception(exc)
+
+
+def _convert_pdf(converter: Any, job: _Job) -> DoclingDocument:
+    import pypdfium2 as pdfium
+    from docling_core.types.io import DocumentStream
+
+    try:
+        pdf = pdfium.PdfDocument(job.pdf_bytes)
+    except Exception as exc:
+        if type(exc).__name__ != "PdfiumError":
+            raise
+        raise ValueError(f"PDF has no pages: {job.name}") from exc
+    try:
+        page_count = len(pdf)
+        if page_count < 1:
+            raise ValueError(f"PDF has no pages: {job.name}")
+        limit = page_count if job.max_pages is None else min(page_count, job.max_pages)
+        if limit < 1:
+            raise ValueError(f"PDF has no pages: {job.name}")
+        assembler = SectionAssembler(name=job.name)
+        for page_no in range(1, limit + 1):
+            source = DocumentStream(name=job.name, stream=io.BytesIO(job.pdf_bytes))
+            result = converter.convert(source, page_range=(page_no, page_no))
+            try:
+                assembler.add_page(result.document, pdf=pdf)
+            finally:
+                del result
+                gc.collect()
+        return assembler.finish()
+    finally:
+        pdf.close()