132 lines
5 KiB
Python
132 lines
5 KiB
Python
import asyncio
|
|
import io
|
|
import tempfile
|
|
import threading
|
|
from collections.abc import Generator
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING
|
|
|
|
import pypdfium2 as pdfium
|
|
|
|
from haiku.rag.client.exceptions import UnsupportedSourceError
|
|
from haiku.rag.telemetry import logfire
|
|
|
|
# pypdfium2 wraps libpdfium, which has global C state and is not thread-safe.
|
|
# Multiple workers calling iter_pdf_slices concurrently race on that state
|
|
# and corrupt it — first error surfaces as e.g. "Failed to import pages",
|
|
# and after that every subsequent PDF load fails with "Failed to load
|
|
# document" until the process restarts. Hold this lock around every pdfium
|
|
# call so only one worker's slicing runs at a time. Crucially we release it
|
|
# between slices so other workers' slicing can interleave; the heavy work
|
|
# (docling convert) happens between yields with the lock free.
|
|
_PDFIUM_LOCK = threading.Lock()
|
|
|
|
if TYPE_CHECKING:
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
|
|
from haiku.rag.converters.base import DocumentConverter
|
|
|
|
_SENTINEL: object = object()
|
|
|
|
|
|
def iter_pdf_slices(
|
|
path: Path, slice_size: int
|
|
) -> Generator[tuple[int, int, bytes], None, None]:
|
|
"""Yield ``(start_page, end_page_inclusive, pdf_bytes)`` for each
|
|
``slice_size``-page slice of ``path``. Page numbers are 1-based to match
|
|
docling's ``prov.page_no`` convention. Each yielded byte string is a
|
|
standalone PDF that docling can convert.
|
|
|
|
The pdfium lock is held only around each pdfium call (open, slice
|
|
extraction, close) — never across a ``yield``. The source ``PdfDocument``
|
|
handle stays open across yields; per-document handles coexist safely as
|
|
long as no two pdfium calls execute concurrently.
|
|
"""
|
|
if slice_size <= 0:
|
|
raise ValueError(f"slice_size must be >= 1, got {slice_size}")
|
|
with _PDFIUM_LOCK:
|
|
try:
|
|
src = pdfium.PdfDocument(str(path))
|
|
except pdfium.PdfiumError as exc:
|
|
raise UnsupportedSourceError(
|
|
f"pypdfium2 cannot open PDF {path}: {exc}"
|
|
) from exc
|
|
total = len(src)
|
|
try:
|
|
for start in range(0, total, slice_size):
|
|
end = min(start + slice_size, total)
|
|
with _PDFIUM_LOCK:
|
|
dst = pdfium.PdfDocument.new()
|
|
try:
|
|
dst.import_pages(src, list(range(start, end)))
|
|
buf = io.BytesIO()
|
|
dst.save(buf)
|
|
slice_bytes = buf.getvalue()
|
|
finally:
|
|
dst.close()
|
|
yield (start + 1, end, slice_bytes)
|
|
finally:
|
|
with _PDFIUM_LOCK:
|
|
src.close()
|
|
|
|
|
|
async def convert_pdf_with_splitting(
|
|
converter: "DocumentConverter",
|
|
path: Path,
|
|
source_uri: str | None,
|
|
slice_size: int,
|
|
) -> "DoclingDocument":
|
|
"""Split a PDF, convert each slice through ``converter``, return the
|
|
merged ``DoclingDocument``.
|
|
|
|
Slices are produced lazily — only one slice's bytes live in memory at a
|
|
time, so peak working set stays bounded regardless of page count. A
|
|
single slice failure aborts the whole job (raised as ``ValueError`` so
|
|
the ingester pipeline classifies it as ``TransientError`` and the queue
|
|
retries the entire document). Per-slice retry would risk interleaved
|
|
partial state with subsequent runs and is not worth the complexity.
|
|
"""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
|
|
def _next(it):
|
|
return next(it, _SENTINEL)
|
|
|
|
it = iter_pdf_slices(path, slice_size)
|
|
converted: list[DoclingDocument] = []
|
|
try:
|
|
while True:
|
|
slice_item = await asyncio.to_thread(_next, it)
|
|
if slice_item is _SENTINEL:
|
|
break
|
|
start, end, pdf_bytes = slice_item
|
|
with tempfile.NamedTemporaryFile(
|
|
mode="wb", suffix=".pdf", delete=False
|
|
) as tmp:
|
|
tmp.write(pdf_bytes)
|
|
tmp.flush()
|
|
tmp_path = Path(tmp.name)
|
|
try:
|
|
with logfire.span(
|
|
"document.convert_slice",
|
|
uri=source_uri,
|
|
start_page=start,
|
|
end_page=end,
|
|
):
|
|
try:
|
|
slice_doc = await converter.convert_file(
|
|
tmp_path, source_uri=source_uri
|
|
)
|
|
except Exception as exc:
|
|
raise ValueError(
|
|
f"Failed to convert slice pages {start}-{end} of {path}: {exc}"
|
|
) from exc
|
|
converted.append(slice_doc)
|
|
finally:
|
|
tmp_path.unlink(missing_ok=True)
|
|
finally:
|
|
# Close the generator under the pdfium lock so src.close() runs
|
|
# even when we abort mid-stream (slice failure, cancellation).
|
|
# Off the event loop because the close path acquires the lock.
|
|
await asyncio.to_thread(it.close)
|
|
|
|
return DoclingDocument.concatenate(converted)
|