diff --git a/CHANGELOG.md b/CHANGELOG.md index 7040da22..53c47954 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # Changelog ## [Unreleased] +### Added + +- PDF `/EmbeddedFiles` attachments are ingested as separate Documents linked to the wrapper through `metadata.parent_uri`. Child URIs use a `#attachment=` fragment on the parent URI. Re-ingest reconciles the child set (add / update / delete) against the wrapper's current attachments; `delete_document` cascades through `parent_uri`. Nested chains are bounded at 3 levels. Toggle with `processing.extract_pdf_attachments` (default `true`). + ### Changed - A successful DELETE job auto-prunes dead jobs with the same `(source_id, uri)`. New `JobRepo.prune_dead(source_id, uri)`. diff --git a/docs/configuration/processing.md b/docs/configuration/processing.md index 8508016e..f8ddf80e 100644 --- a/docs/configuration/processing.md +++ b/docs/configuration/processing.md @@ -23,6 +23,9 @@ processing: chunking_merge_peers: true # Merge undersized successive chunks chunking_use_markdown_tables: false # Use markdown tables vs narrative format + # PDF /EmbeddedFiles attachments + extract_pdf_attachments: true # Ingest embedded files as separate Documents + # Automatic title generation auto_title: false # Auto-generate titles on ingestion title_model: # LLM for title generation (fallback) @@ -368,6 +371,35 @@ Explicit titles passed via `title=` parameter always take precedence and are nev To generate titles for existing untitled documents, use [`rebuild --title-only`](../cli.md#rebuild-database). +### PDF Embedded Attachments + +A PDF can carry other files inside it via the `/EmbeddedFiles` table (signed memos, appendices, supporting documents). With `extract_pdf_attachments: true` (the default), each embedded file is ingested as a separate Document linked to the wrapper through `metadata.parent_uri`: + +```yaml +processing: + extract_pdf_attachments: true +``` + +```python +# After ingesting a PDF with two attachments: +parent = await client.create_document_from_source("/path/to/parent.pdf") +children = await client.list_documents( + filter=f"metadata LIKE '%\"parent_uri\": \"{parent.uri}\"%'" +) +# children: 2 Documents, each with parent.uri in metadata.parent_uri, +# URIs like file:///path/to/parent.pdf#attachment=memo.pdf +``` + +Behavior: + +- Children inherit the standard ingest metadata (`content_type`, `md5`, `source_revision`) plus `parent_uri`. +- Re-ingesting the wrapper reconciles its current attachment set against existing children: new files are added, changed bytes update in place, and dropped names are deleted. +- `delete_document(parent_id)` cascades through `parent_uri` and removes all children. +- Nested attachments (a PDF whose attachment is itself a PDF with attachments) recurse up to 3 levels. Deeper chains log a warning and skip. +- Attachments whose extension or content type the converter does not support log a warning and are skipped without aborting the rest of the set. + +Set `extract_pdf_attachments: false` to ingest only the wrapper. + ## Continuous ingestion For automatic ingestion of local directories, S3 buckets, or HTTP diff --git a/docs/python.md b/docs/python.md index 1380ed77..229599d7 100644 --- a/docs/python.md +++ b/docs/python.md @@ -80,6 +80,8 @@ doc = await client.create_document_from_source( ) ``` +PDFs that carry attachments via the `/EmbeddedFiles` table are split into one Document per attachment, linked to the wrapper through `metadata.parent_uri`. See [PDF Embedded Attachments](configuration/processing.md#pdf-embedded-attachments). + ### Retrieving Documents By ID: @@ -167,6 +169,8 @@ await client.update_document(document_id=doc.id, chunks=custom_chunks) await client.delete_document(doc.id) ``` +Deleting a document also removes any child Documents linked to it via `metadata.parent_uri` (PDF attachment children, primarily). The cascade is transitive. + ## Searching Documents The search method performs native hybrid search (vector + full-text) using LanceDB with optional reranking for improved relevance: diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 3946d4b3..13b7643e 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -310,7 +310,18 @@ class HaikuRAG: return None async def delete_document(self, document_id: str) -> bool: - """Delete a document by its ID.""" + """Delete a document by its ID. Cascades to children linked via + ``metadata.parent_uri``.""" + from haiku.rag.client.documents import parent_uri_filter + + doc = await self.get_document_by_id(document_id) + if doc is None: + return False + if doc.uri: + children = await self.list_documents(filter=parent_uri_filter(doc.uri)) + for child in children: + if child.id and child.id != document_id: + await self.delete_document(child.id) return await self.document_repository.delete(document_id) async def list_documents( diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index 574a309e..a667213d 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -1,7 +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 ( @@ -26,6 +30,24 @@ 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`` + equals ``parent_uri``. ``metadata`` is stored as a JSON string produced by + the standard library's ``json.dumps`` (which inserts ``": "`` between key + and value), so the match is a substring search over that serialized form — + escape JSON-meaningful chars in the URI, then SQL-escape single quotes.""" + json_fragment = json.dumps(parent_uri)[1:-1].replace("'", "''") + return f'metadata LIKE \'%"parent_uri": "{json_fragment}"%\'' + async def _store_document_with_chunks( client: "HaikuRAG", @@ -221,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) @@ -285,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: @@ -301,9 +326,110 @@ 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 + try: + 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) + finally: + pdf.close() + + 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/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 2719623d..b6102143 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -182,6 +182,10 @@ class ProcessingConfig(BaseModel): - ``"image"``: docling generates picture images and stores them in ``document_items.picture_data``; no VLM runs at ingest. """ + extract_pdf_attachments: bool = True + """When a PDF carries `/EmbeddedFiles`, ingest each attachment as a separate + Document linked back to the wrapper via ``metadata.parent_uri``. Cap depth + at 3 to bound nested-attachment recursion.""" auto_title: bool = False title_model: ModelConfig = Field( default_factory=lambda: ModelConfig( diff --git a/tests/test_cascade_delete.py b/tests/test_cascade_delete.py new file mode 100644 index 00000000..3963ffbf --- /dev/null +++ b/tests/test_cascade_delete.py @@ -0,0 +1,139 @@ +from haiku.rag.client import HaikuRAG +from haiku.rag.client.documents import parent_uri_filter +from haiku.rag.store.models.document import Document + + +def test_parent_uri_filter_simple(): + f = parent_uri_filter("file:///path/to/parent.pdf") + assert f == 'metadata LIKE \'%"parent_uri": "file:///path/to/parent.pdf"%\'' + + +def test_parent_uri_filter_escapes_single_quote(): + f = parent_uri_filter("file:///x's.pdf") + assert "''" in f + + +def test_parent_uri_filter_escapes_backslash(): + f = parent_uri_filter("file:///x\\y.pdf") + assert "\\\\" in f + + +async def test_delete_cascades_to_children(temp_db_path): + async with HaikuRAG(temp_db_path, create=True) as client: + parent_uri = "file:///path/to/parent.pdf" + parent = await client.document_repository.create( + Document(content="parent body", uri=parent_uri, metadata={}) + ) + child_a = await client.document_repository.create( + Document( + content="child A body", + uri=f"{parent_uri}#attachment=a.pdf", + metadata={"parent_uri": parent_uri}, + ) + ) + child_b = await client.document_repository.create( + Document( + content="child B body", + uri=f"{parent_uri}#attachment=b.pdf", + metadata={"parent_uri": parent_uri}, + ) + ) + + deleted = await client.delete_document(parent.id) + assert deleted is True + + assert await client.get_document_by_id(parent.id) is None + assert await client.get_document_by_id(child_a.id) is None + assert await client.get_document_by_id(child_b.id) is None + + +async def test_delete_leaves_unrelated_documents(temp_db_path): + async with HaikuRAG(temp_db_path, create=True) as client: + parent_uri = "file:///path/to/parent.pdf" + parent = await client.document_repository.create( + Document(content="parent", uri=parent_uri, metadata={}) + ) + child = await client.document_repository.create( + Document( + content="child", + uri=f"{parent_uri}#attachment=a.pdf", + metadata={"parent_uri": parent_uri}, + ) + ) + unrelated = await client.document_repository.create( + Document( + content="unrelated", + uri="file:///path/to/other.pdf", + metadata={}, + ) + ) + + await client.delete_document(parent.id) + + assert await client.get_document_by_id(child.id) is None + survivor = await client.get_document_by_id(unrelated.id) + assert survivor is not None + assert survivor.id == unrelated.id + + +async def test_delete_cascades_recursively(temp_db_path): + async with HaikuRAG(temp_db_path, create=True) as client: + gp_uri = "file:///path/to/grandparent.pdf" + parent_uri = f"{gp_uri}#attachment=parent.pdf" + + grandparent = await client.document_repository.create( + Document(content="gp", uri=gp_uri, metadata={}) + ) + parent = await client.document_repository.create( + Document( + content="p", + uri=parent_uri, + metadata={"parent_uri": gp_uri}, + ) + ) + child = await client.document_repository.create( + Document( + content="c", + uri=f"{parent_uri}#attachment=leaf.pdf", + metadata={"parent_uri": parent_uri}, + ) + ) + + await client.delete_document(grandparent.id) + + assert await client.get_document_by_id(grandparent.id) is None + assert await client.get_document_by_id(parent.id) is None + assert await client.get_document_by_id(child.id) is None + + +async def test_delete_nonexistent_returns_false(temp_db_path): + async with HaikuRAG(temp_db_path, create=True) as client: + result = await client.delete_document("does-not-exist") + assert result is False + + +async def test_delete_handles_self_referential_parent(temp_db_path): + """A document whose metadata.parent_uri points at its own uri must not + cascade into infinite recursion.""" + async with HaikuRAG(temp_db_path, create=True) as client: + uri = "file:///path/to/self.pdf" + doc = await client.document_repository.create( + Document(content="self-loop", uri=uri, metadata={"parent_uri": uri}) + ) + + deleted = await client.delete_document(doc.id) + assert deleted is True + assert await client.get_document_by_id(doc.id) is None + + +def test_processing_config_extract_pdf_attachments_default_true(): + from haiku.rag.config.models import ProcessingConfig + + assert ProcessingConfig().extract_pdf_attachments is True + + +def test_processing_config_extract_pdf_attachments_overridable(): + from haiku.rag.config.models import ProcessingConfig + + cfg = ProcessingConfig(extract_pdf_attachments=False) + assert cfg.extract_pdf_attachments is False diff --git a/tests/test_pdf_attachments.py b/tests/test_pdf_attachments.py new file mode 100644 index 00000000..1feb1b49 --- /dev/null +++ b/tests/test_pdf_attachments.py @@ -0,0 +1,464 @@ +import io + +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): + 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 at the cap. + 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) + 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 + + +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: + monkeypatch.setattr(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_malformed_pdf_logs_warning_and_skips(temp_db_path, monkeypatch, caplog): + """A non-PDF body labelled as application/pdf must not crash the helper — + pypdfium2's open raises PdfiumError, which we log and return from.""" + import logging + + 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/junk.pdf" + garbage = b"this is not a pdf at all" + parent = await _make_parent(client, parent_uri, garbage) + with caplog.at_level(logging.WARNING, logger="haiku.rag.client.documents"): + await _reconcile_pdf_attachments(client, parent, garbage, depth=0) + assert await client.list_documents(filter=parent_uri_filter(parent_uri)) == [] + + +async def test_unsupported_attachment_continues_loop(temp_db_path, monkeypatch): + """One attachment whose ingest raises UnsupportedSourceError must not + prevent siblings from being ingested. The unsupported attachment is + skipped with a warning; the others land.""" + from haiku.rag.client.exceptions import UnsupportedSourceError + + async def picky_fake( + client, + result, + *, + title, + user_metadata, + stored_uri, + existing_doc, + depth=0, + ): + if stored_uri.endswith("unsupported.xyz"): + raise UnsupportedSourceError("nope") + return await fake_ingest_fetch_result( + client, + result, + title=title, + user_metadata=user_metadata, + stored_uri=stored_uri, + existing_doc=existing_doc, + depth=depth, + ) + + monkeypatch.setattr("haiku.rag.client.documents._ingest_fetch_result", picky_fake) + async with HaikuRAG(temp_db_path, create=True) as client: + parent_uri = "file:///fixtures/parent.pdf" + pdf_bytes = build_pdf([("ok.txt", b"keep me"), ("unsupported.xyz", b"data")]) + 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 {c.uri for c in children} == {f"{parent_uri}#attachment=ok.txt"} + + +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() == [] + + +async def test_create_document_from_source_extracts_attachments( + tmp_path, temp_db_path, monkeypatch +): + """The full create_document_from_source path — the same entry point the + ingester worker uses for an UPSERT job — produces parent + child docs + when the source is a PDF with embedded files on disk.""" + monkeypatch.setattr( + "haiku.rag.client.documents._ingest_fetch_result", + fake_ingest_fetch_result, + ) + pdf_path = tmp_path / "parent.pdf" + pdf_path.write_bytes( + build_pdf([("notes.txt", b"plain text"), ("data.txt", b"more data")]) + ) + + async with HaikuRAG(temp_db_path, create=True) as client: + parent = await client.create_document_from_source(pdf_path) + + async with HaikuRAG(temp_db_path) as client: + assert isinstance(parent, Document) + children = await client.list_documents(filter=parent_uri_filter(parent.uri)) + assert len(children) == 2 + assert {c.metadata["parent_uri"] for c in children} == {parent.uri} + + +async def test_create_document_from_source_reingest_after_attachment_edit( + tmp_path, temp_db_path, monkeypatch +): + """Mutate the parent PDF's attachments and re-ingest. The md5 short-circuit + must NOT fire (parent bytes changed); reconciliation diffs children to add, + update, and delete in one pass while leaving unrelated children untouched.""" + monkeypatch.setattr( + "haiku.rag.client.documents._ingest_fetch_result", + fake_ingest_fetch_result, + ) + pdf_path = tmp_path / "parent.pdf" + pdf_path.write_bytes( + build_pdf( + [ + ("stable.txt", b"unchanged across runs"), + ("changed.txt", b"old contents"), + ("removed.txt", b"goes away"), + ] + ) + ) + + async with HaikuRAG(temp_db_path, create=True) as client: + parent = await client.create_document_from_source(pdf_path) + assert isinstance(parent, Document) + before_children = { + c.uri: c + for c in await client.list_documents(filter=parent_uri_filter(parent.uri)) + } + assert set(before_children) == { + f"{parent.uri}#attachment=stable.txt", + f"{parent.uri}#attachment=changed.txt", + f"{parent.uri}#attachment=removed.txt", + } + stable_id_before = before_children[f"{parent.uri}#attachment=stable.txt"].id + changed_id_before = before_children[f"{parent.uri}#attachment=changed.txt"].id + + pdf_path.write_bytes( + build_pdf( + [ + ("stable.txt", b"unchanged across runs"), + ("changed.txt", b"new contents"), + ("added.txt", b"brand new"), + ] + ) + ) + + async with HaikuRAG(temp_db_path) as client: + await client.create_document_from_source(pdf_path) + after_children = { + c.uri: c + for c in await client.list_documents(filter=parent_uri_filter(parent.uri)) + } + + assert set(after_children) == { + f"{parent.uri}#attachment=stable.txt", + f"{parent.uri}#attachment=changed.txt", + f"{parent.uri}#attachment=added.txt", + } + stable = after_children[f"{parent.uri}#attachment=stable.txt"] + changed = after_children[f"{parent.uri}#attachment=changed.txt"] + assert stable.id == stable_id_before + assert changed.id == changed_id_before + assert ( + stable.metadata["md5"] + == before_children[f"{parent.uri}#attachment=stable.txt"].metadata["md5"] + ) + assert ( + changed.metadata["md5"] + != before_children[f"{parent.uri}#attachment=changed.txt"].metadata["md5"] + ) + + +async def test_create_document_from_source_delete_cascades( + tmp_path, temp_db_path, monkeypatch +): + """The ingester worker's DELETE path is just client.delete_document(doc.id). + A parent ingested via the full pipeline must cascade to its children when + that path runs — mirrors what happens when a watched file is removed.""" + monkeypatch.setattr( + "haiku.rag.client.documents._ingest_fetch_result", + fake_ingest_fetch_result, + ) + pdf_path = tmp_path / "parent.pdf" + pdf_path.write_bytes(build_pdf([("a.txt", b"A"), ("b.txt", b"B")])) + + async with HaikuRAG(temp_db_path, create=True) as client: + parent = await client.create_document_from_source(pdf_path) + assert isinstance(parent, Document) + assert len(await client.list_documents()) == 3 + + await client.delete_document(parent.id) + assert await client.list_documents() == []