diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index e1defd7e..0a326aa2 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/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, diff --git a/tests/test_pdf_attachments.py b/tests/test_pdf_attachments.py new file mode 100644 index 00000000..fab807ba --- /dev/null +++ b/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() == []