Extract PDF /EmbeddedFiles attachments as child documents

haiku_rag_slim/haiku/rag/client/documents.py

@@ -1,8 +1,11 @@
+import hashlib
 import json
+import logging
+import mimetypes
 import tempfile
 from pathlib import Path
 from typing import TYPE_CHECKING
-from urllib.parse import unquote, urlparse
+from urllib.parse import quote, unquote, urlparse
 
 from haiku.rag.client.exceptions import UnsupportedSourceError
 from haiku.rag.client.processing import (
@@ -27,6 +30,14 @@ if TYPE_CHECKING:
     from haiku.rag.client import HaikuRAG
     from haiku.rag.ingester.sources.base import Source
 
+logger = logging.getLogger(__name__)
+
+# Maximum length of an attachment chain rooted at a top-level ingest. With
+# value 3, a PDF whose attachments contain PDFs which themselves contain
+# PDFs is fully ingested (3 levels); a fourth nested level logs a warning
+# and is skipped.
+MAX_ATTACHMENT_DEPTH = 3
+
 
 def parent_uri_filter(parent_uri: str) -> str:
     """SQL `WHERE` clause matching documents whose ``metadata.parent_uri``
@@ -232,9 +243,11 @@ async def _ingest_fetch_result(
     user_metadata: dict,
     stored_uri: str,
     existing_doc: Document | None,
+    depth: int = 0,
 ) -> Document:
     """Convert / chunk / embed / store a fetched document. Replaces an
-    existing document if one is supplied."""
+    existing document if one is supplied. ``depth`` tracks position in an
+    attachment chain so the reconciliation step can bound recursion."""
     from haiku.rag.embeddings import embed_chunks
 
     converter = get_converter(client._config)
@@ -296,6 +309,7 @@ async def _ingest_fetch_result(
                 client, existing_doc, embedded_chunks, docling_document
             )
             store_span.set_attribute("document_id", updated.id)
+        await _reconcile_pdf_attachments(client, updated, result.body, depth=depth)
         return updated
 
     if title is None:
@@ -312,9 +326,107 @@ async def _ingest_fetch_result(
             client, document, embedded_chunks, docling_document
         )
         store_span.set_attribute("document_id", created.id)
+    await _reconcile_pdf_attachments(client, created, result.body, depth=depth)
     return created
 
 
+async def _reconcile_pdf_attachments(
+    client: "HaikuRAG",
+    parent_doc: Document,
+    parent_body: bytes,
+    *,
+    depth: int,
+) -> None:
+    """Diff the parent PDF's ``/EmbeddedFiles`` table against any children
+    already linked via ``metadata.parent_uri`` and bring the child set in line:
+    ingest additions, update changed bytes, cascade-delete removed names.
+
+    Re-uses ``_ingest_fetch_result`` for each child so the standard conversion
+    path runs uniformly — child PDFs recurse into this helper one level deeper,
+    bounded by ``MAX_ATTACHMENT_DEPTH``.
+    """
+    if not client._config.processing.extract_pdf_attachments:
+        return
+    if not parent_doc.uri:
+        return
+    if (parent_doc.metadata or {}).get("content_type") != "application/pdf":
+        return
+
+    import pypdfium2 as pdfium
+
+    try:
+        pdf = pdfium.PdfDocument(parent_body)
+    except pdfium.PdfiumError as exc:
+        logger.warning(
+            "Cannot scan %s for embedded attachments: %s", parent_doc.uri, exc
+        )
+        return
+    attachment_count = pdf.count_attachments()
+
+    if depth + 1 >= MAX_ATTACHMENT_DEPTH:
+        if attachment_count > 0:
+            logger.warning(
+                "Attachment depth cap (%d) reached at %s; skipping %d nested "
+                "attachment(s).",
+                MAX_ATTACHMENT_DEPTH,
+                parent_doc.uri,
+                attachment_count,
+            )
+        return
+
+    new_attachments: dict[str, tuple[str, bytes, str, str]] = {}
+    for i in range(attachment_count):
+        att = pdf.get_attachment(i)
+        name = att.get_name()
+        if not name:
+            continue
+        data = bytes(att.get_data())
+        child_uri = f"{parent_doc.uri}#attachment={quote(name, safe='')}"
+        content_type = mimetypes.guess_type(name)[0] or "application/octet-stream"
+        content_hash = hashlib.md5(data, usedforsecurity=False).hexdigest()
+        new_attachments[child_uri] = (name, data, content_type, content_hash)
+
+    existing = await client.list_documents(filter=parent_uri_filter(parent_doc.uri))
+    existing_by_uri: dict[str, Document] = {d.uri: d for d in existing if d.uri}
+
+    for child_uri, (name, data, content_type, content_hash) in new_attachments.items():
+        existing_child = existing_by_uri.get(child_uri)
+        if (
+            existing_child
+            and (existing_child.metadata or {}).get("md5") == content_hash
+        ):
+            continue
+
+        child_fr = FetchResult(
+            uri=child_uri,
+            body=data,
+            content_type=content_type,
+            content_hash=content_hash,
+            extra_metadata={"parent_uri": parent_doc.uri},
+        )
+        try:
+            await _ingest_fetch_result(
+                client,
+                child_fr,
+                title=None,
+                user_metadata={},
+                stored_uri=child_uri,
+                existing_doc=existing_child,
+                depth=depth + 1,
+            )
+        except UnsupportedSourceError:
+            logger.warning(
+                "Skipping attachment %r in %s: unsupported content type %r",
+                name,
+                parent_doc.uri,
+                content_type,
+            )
+
+    for child_uri, child in existing_by_uri.items():
+        if child_uri not in new_attachments and child.id:
+            await client.delete_document(child.id)
+
+
 async def create_document_from_source(
     client: "HaikuRAG",
     source: str | Path,

tests/test_pdf_attachments.py

@@ -0,0 +1,295 @@
+import io
+import logging
+
+import pypdfium2 as pdfium
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.client.documents import (
+    MAX_ATTACHMENT_DEPTH,
+    _reconcile_pdf_attachments,
+    parent_uri_filter,
+)
+from haiku.rag.ingester.sources import FetchResult
+from haiku.rag.store.models.document import Document
+
+
+def build_pdf(attachments: list[tuple[str, bytes]]) -> bytes:
+    """Build a minimal one-page PDF with the given (name, bytes) attachments."""
+    pdf = pdfium.PdfDocument.new()
+    pdf.new_page(200, 200)
+    for name, data in attachments:
+        att = pdf.new_attachment(name)
+        att.set_data(data)
+    buf = io.BytesIO()
+    pdf.save(buf)
+    return buf.getvalue()
+
+
+async def fake_ingest_fetch_result(
+    client,
+    result: FetchResult,
+    *,
+    title,
+    user_metadata,
+    stored_uri,
+    existing_doc,
+    depth=0,
+):
+    """A stand-in for ``_ingest_fetch_result`` that skips docling/embedder
+    entirely: it writes the document with content_type/md5/parent_uri set
+    correctly, then defers to the real ``_reconcile_pdf_attachments`` so
+    recursive logic stays under test."""
+    final_metadata = {
+        **(user_metadata or {}),
+        "content_type": result.content_type,
+        "md5": result.content_hash,
+        **result.extra_metadata,
+    }
+    if result.revision is not None:
+        final_metadata["source_revision"] = result.revision
+
+    if existing_doc:
+        existing_doc.content = ""
+        existing_doc.metadata = final_metadata
+        if title is not None:
+            existing_doc.title = title
+        doc = await client.document_repository.update(existing_doc)
+    else:
+        doc = await client.document_repository.create(
+            Document(
+                content="",
+                uri=stored_uri,
+                title=title,
+                metadata=final_metadata,
+            )
+        )
+    await _reconcile_pdf_attachments(client, doc, result.body, depth=depth)
+    return doc
+
+
+async def _make_parent(
+    client: HaikuRAG,
+    uri: str,
+    body: bytes,
+    *,
+    content_type: str = "application/pdf",
+) -> Document:
+    """Insert a parent Document directly + invoke reconciliation for its body."""
+    import hashlib
+
+    md5 = hashlib.md5(body, usedforsecurity=False).hexdigest()
+    parent = await client.document_repository.create(
+        Document(
+            content="",
+            uri=uri,
+            metadata={"content_type": content_type, "md5": md5},
+        )
+    )
+    return parent
+
+
+async def test_first_ingest_creates_one_doc_per_attachment(temp_db_path, monkeypatch):
+    monkeypatch.setattr(
+        "haiku.rag.client.documents._ingest_fetch_result",
+        fake_ingest_fetch_result,
+    )
+    leaf_pdf = build_pdf([])
+    pdf_bytes = build_pdf([("a.pdf", leaf_pdf), ("notes.txt", b"plain text payload")])
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        parent_uri = "file:///fixtures/parent.pdf"
+        parent = await _make_parent(client, parent_uri, pdf_bytes)
+
+        await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
+
+        children = await client.list_documents(filter=parent_uri_filter(parent_uri))
+        assert len(children) == 2
+        by_uri = {c.uri: c for c in children}
+        assert f"{parent_uri}#attachment=a.pdf" in by_uri
+        assert f"{parent_uri}#attachment=notes.txt" in by_uri
+
+        a = by_uri[f"{parent_uri}#attachment=a.pdf"]
+        assert a.metadata["parent_uri"] == parent_uri
+        assert a.metadata["content_type"] == "application/pdf"
+        assert "md5" in a.metadata
+
+        txt = by_uri[f"{parent_uri}#attachment=notes.txt"]
+        assert txt.metadata["content_type"] == "text/plain"
+
+
+async def test_attachment_with_spaces_in_name_is_percent_encoded(
+    temp_db_path, monkeypatch
+):
+    monkeypatch.setattr(
+        "haiku.rag.client.documents._ingest_fetch_result",
+        fake_ingest_fetch_result,
+    )
+    pdf_bytes = build_pdf([("memo with spaces.pdf", b"payload")])
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        parent_uri = "file:///fixtures/parent.pdf"
+        parent = await _make_parent(client, parent_uri, pdf_bytes)
+        await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
+
+        children = await client.list_documents(filter=parent_uri_filter(parent_uri))
+        assert len(children) == 1
+        assert children[0].uri == f"{parent_uri}#attachment=memo%20with%20spaces.pdf"
+
+
+async def test_reingest_removes_dropped_attachment(temp_db_path, monkeypatch):
+    monkeypatch.setattr(
+        "haiku.rag.client.documents._ingest_fetch_result",
+        fake_ingest_fetch_result,
+    )
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        parent_uri = "file:///fixtures/parent.pdf"
+        first = build_pdf([("a.txt", b"A"), ("b.txt", b"B")])
+        parent = await _make_parent(client, parent_uri, first)
+        await _reconcile_pdf_attachments(client, parent, first, depth=0)
+        assert (
+            len(await client.list_documents(filter=parent_uri_filter(parent_uri))) == 2
+        )
+
+        second = build_pdf([("a.txt", b"A")])
+        await _reconcile_pdf_attachments(client, parent, second, depth=0)
+        remaining = await client.list_documents(filter=parent_uri_filter(parent_uri))
+        assert len(remaining) == 1
+        assert remaining[0].uri == f"{parent_uri}#attachment=a.txt"
+
+
+async def test_reingest_updates_changed_attachment_in_place(temp_db_path, monkeypatch):
+    monkeypatch.setattr(
+        "haiku.rag.client.documents._ingest_fetch_result",
+        fake_ingest_fetch_result,
+    )
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        parent_uri = "file:///fixtures/parent.pdf"
+        first = build_pdf([("a.txt", b"original")])
+        parent = await _make_parent(client, parent_uri, first)
+        await _reconcile_pdf_attachments(client, parent, first, depth=0)
+        before = (await client.list_documents(filter=parent_uri_filter(parent_uri)))[0]
+
+        second = build_pdf([("a.txt", b"different")])
+        await _reconcile_pdf_attachments(client, parent, second, depth=0)
+        after = (await client.list_documents(filter=parent_uri_filter(parent_uri)))[0]
+
+        assert after.id == before.id
+        assert after.metadata["md5"] != before.metadata["md5"]
+
+
+async def test_reingest_adds_new_attachment(temp_db_path, monkeypatch):
+    monkeypatch.setattr(
+        "haiku.rag.client.documents._ingest_fetch_result",
+        fake_ingest_fetch_result,
+    )
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        parent_uri = "file:///fixtures/parent.pdf"
+        first = build_pdf([("a.txt", b"A")])
+        parent = await _make_parent(client, parent_uri, first)
+        await _reconcile_pdf_attachments(client, parent, first, depth=0)
+
+        second = build_pdf([("a.txt", b"A"), ("c.txt", b"C")])
+        await _reconcile_pdf_attachments(client, parent, second, depth=0)
+        children = await client.list_documents(filter=parent_uri_filter(parent_uri))
+        assert len(children) == 2
+        names = {c.uri for c in children}
+        assert f"{parent_uri}#attachment=a.txt" in names
+        assert f"{parent_uri}#attachment=c.txt" in names
+
+
+async def test_nested_pdf_attachments_recurse_up_to_cap(
+    temp_db_path, monkeypatch, caplog
+):
+    monkeypatch.setattr(
+        "haiku.rag.client.documents._ingest_fetch_result",
+        fake_ingest_fetch_result,
+    )
+    # Build a 4-deep chain: root -> L1 -> L2 -> L3. MAX_ATTACHMENT_DEPTH=3
+    # means root + L1 + L2 ingested (3 PDFs total); L3 is skipped with a
+    # warning logged at the depth boundary.
+    assert MAX_ATTACHMENT_DEPTH == 3
+    l3 = build_pdf([("leaf.txt", b"deepest")])
+    l2 = build_pdf([("l3.pdf", l3)])
+    l1 = build_pdf([("l2.pdf", l2)])
+    root = build_pdf([("l1.pdf", l1)])
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        root_uri = "file:///fixtures/root.pdf"
+        parent = await _make_parent(client, root_uri, root)
+        with caplog.at_level(logging.WARNING, logger="haiku.rag.client.documents"):
+            await _reconcile_pdf_attachments(client, parent, root, depth=0)
+
+        l1_uri = f"{root_uri}#attachment=l1.pdf"
+        l2_uri = f"{l1_uri}#attachment=l2.pdf"
+        l3_uri = f"{l2_uri}#attachment=l3.pdf"
+
+        assert await client.get_document_by_uri(l1_uri) is not None
+        assert await client.get_document_by_uri(l2_uri) is not None
+        assert await client.get_document_by_uri(l3_uri) is None
+        assert any("depth cap" in r.message for r in caplog.records)
+
+
+async def test_config_off_skips_extraction(temp_db_path, monkeypatch):
+    monkeypatch.setattr(
+        "haiku.rag.client.documents._ingest_fetch_result",
+        fake_ingest_fetch_result,
+    )
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        client._config.processing.extract_pdf_attachments = False
+        parent_uri = "file:///fixtures/parent.pdf"
+        pdf_bytes = build_pdf([("a.txt", b"A")])
+        parent = await _make_parent(client, parent_uri, pdf_bytes)
+        await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
+
+        assert await client.list_documents(filter=parent_uri_filter(parent_uri)) == []
+
+
+async def test_non_pdf_parent_is_ignored(temp_db_path, monkeypatch):
+    """A non-PDF document with a PDF blob would be a logic bug, but the helper
+    must short-circuit on content_type alone — never call pypdfium2."""
+    monkeypatch.setattr(
+        "haiku.rag.client.documents._ingest_fetch_result",
+        fake_ingest_fetch_result,
+    )
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        parent_uri = "file:///fixtures/parent.txt"
+        parent = await _make_parent(
+            client,
+            parent_uri,
+            b"not a pdf",
+            content_type="text/plain",
+        )
+        # Even with PDF bytes, content_type=text/plain blocks extraction.
+        pdf_bytes = build_pdf([("a.txt", b"A")])
+        await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
+        assert await client.list_documents(filter=parent_uri_filter(parent_uri)) == []
+
+
+async def test_parent_without_uri_is_skipped(temp_db_path, monkeypatch):
+    monkeypatch.setattr(
+        "haiku.rag.client.documents._ingest_fetch_result",
+        fake_ingest_fetch_result,
+    )
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        parent = Document(
+            content="",
+            uri=None,
+            metadata={"content_type": "application/pdf", "md5": "abc"},
+        )
+        pdf_bytes = build_pdf([("a.txt", b"A")])
+        await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
+        assert await client.list_documents() == []
+
+
+async def test_cascade_delete_removes_reconciled_children(temp_db_path, monkeypatch):
+    monkeypatch.setattr(
+        "haiku.rag.client.documents._ingest_fetch_result",
+        fake_ingest_fetch_result,
+    )
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        parent_uri = "file:///fixtures/parent.pdf"
+        pdf_bytes = build_pdf([("a.txt", b"A"), ("b.txt", b"B")])
+        parent = await _make_parent(client, parent_uri, pdf_bytes)
+        await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
+        assert len(await client.list_documents()) == 3
+
+        await client.delete_document(parent.id)
+        assert await client.list_documents() == []