danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #386 from ggozad/feat/pdf-containers

Browse source 
Support for embedded documents inside pdfs.
This commit is contained in:
Yiorgis Gozadinos 2026-05-29 10:25:05 +03:00 committed by GitHub
commit c0c83d8037
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 787 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
CHANGELOG.md
View file

@ -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=<percent-encoded-name>` 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)`.

32
docs/configuration/processing.md
View file

@ -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

4
docs/python.md
View file

@ -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:

13
haiku_rag_slim/haiku/rag/client/__init__.py
View file

@ -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(

130
haiku_rag_slim/haiku/rag/client/documents.py
View file

@ -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,

4
haiku_rag_slim/haiku/rag/config/models.py
View file

@ -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(

139
tests/test_cascade_delete.py Normal file
View file

@ -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

464
tests/test_pdf_attachments.py Normal file
View file

@ -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() == []