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

Fix CI flakes for PDF attachment extraction

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-05-28 16:58:16 +03:00
parent 6eee09743b
commit 402957d3e5
No known key found for this signature in database
2 changed files with 29 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -361,30 +361,33 @@ async def _reconcile_pdf_attachments(
"Cannot scan %s for embedded attachments: %s", parent_doc.uri, exc "Cannot scan %s for embedded attachments: %s", parent_doc.uri, exc
) )
return return
attachment_count = pdf.count_attachments() try:
attachment_count = pdf.count_attachments()
if depth + 1 >= MAX_ATTACHMENT_DEPTH: if depth + 1 >= MAX_ATTACHMENT_DEPTH:
if attachment_count > 0: if attachment_count > 0:
logger.warning( logger.warning(
"Attachment depth cap (%d) reached at %s; skipping %d nested " "Attachment depth cap (%d) reached at %s; skipping %d nested "
"attachment(s).", "attachment(s).",
MAX_ATTACHMENT_DEPTH, MAX_ATTACHMENT_DEPTH,
parent_doc.uri, parent_doc.uri,
attachment_count, attachment_count,
) )
return return
new_attachments: dict[str, tuple[str, bytes, str, str]] = {} new_attachments: dict[str, tuple[str, bytes, str, str]] = {}
for i in range(attachment_count): for i in range(attachment_count):
att = pdf.get_attachment(i) att = pdf.get_attachment(i)
name = att.get_name() name = att.get_name()
if not name: if not name:
continue continue
data = bytes(att.get_data()) data = bytes(att.get_data())
child_uri = f"{parent_doc.uri}#attachment={quote(name, safe='')}" child_uri = f"{parent_doc.uri}#attachment={quote(name, safe='')}"
content_type = mimetypes.guess_type(name)[0] or "application/octet-stream" content_type = mimetypes.guess_type(name)[0] or "application/octet-stream"
content_hash = hashlib.md5(data, usedforsecurity=False).hexdigest() content_hash = hashlib.md5(data, usedforsecurity=False).hexdigest()
new_attachments[child_uri] = (name, data, content_type, content_hash) 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 = 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} existing_by_uri: dict[str, Document] = {d.uri: d for d in existing if d.uri}

14
tests/test_pdf_attachments.py
View file

@ -1,5 +1,4 @@
import io import io
import logging
import pypdfium2 as pdfium import pypdfium2 as pdfium
@ -195,16 +194,13 @@ async def test_reingest_adds_new_attachment(temp_db_path, monkeypatch):
assert f"{parent_uri}#attachment=c.txt" in names assert f"{parent_uri}#attachment=c.txt" in names
async def test_nested_pdf_attachments_recurse_up_to_cap( async def test_nested_pdf_attachments_recurse_up_to_cap(temp_db_path, monkeypatch):
temp_db_path, monkeypatch, caplog
):
monkeypatch.setattr( monkeypatch.setattr(
"haiku.rag.client.documents._ingest_fetch_result", "haiku.rag.client.documents._ingest_fetch_result",
fake_ingest_fetch_result, fake_ingest_fetch_result,
) )
# Build a 4-deep chain: root -> L1 -> L2 -> L3. MAX_ATTACHMENT_DEPTH=3 # 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 # means root + L1 + L2 ingested (3 PDFs total); L3 is skipped at the cap.
# warning logged at the depth boundary.
assert MAX_ATTACHMENT_DEPTH == 3 assert MAX_ATTACHMENT_DEPTH == 3
l3 = build_pdf([("leaf.txt", b"deepest")]) l3 = build_pdf([("leaf.txt", b"deepest")])
l2 = build_pdf([("l3.pdf", l3)]) l2 = build_pdf([("l3.pdf", l3)])
@ -214,8 +210,7 @@ async def test_nested_pdf_attachments_recurse_up_to_cap(
async with HaikuRAG(temp_db_path, create=True) as client: async with HaikuRAG(temp_db_path, create=True) as client:
root_uri = "file:///fixtures/root.pdf" root_uri = "file:///fixtures/root.pdf"
parent = await _make_parent(client, root_uri, root) 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)
await _reconcile_pdf_attachments(client, parent, root, depth=0)
l1_uri = f"{root_uri}#attachment=l1.pdf" l1_uri = f"{root_uri}#attachment=l1.pdf"
l2_uri = f"{l1_uri}#attachment=l2.pdf" l2_uri = f"{l1_uri}#attachment=l2.pdf"
@ -224,7 +219,6 @@ async def test_nested_pdf_attachments_recurse_up_to_cap(
assert await client.get_document_by_uri(l1_uri) is not None 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(l2_uri) is not None
assert await client.get_document_by_uri(l3_uri) is 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): async def test_config_off_skips_extraction(temp_db_path, monkeypatch):
@ -233,7 +227,7 @@ async def test_config_off_skips_extraction(temp_db_path, monkeypatch):
fake_ingest_fetch_result, fake_ingest_fetch_result,
) )
async with HaikuRAG(temp_db_path, create=True) as client: async with HaikuRAG(temp_db_path, create=True) as client:
client._config.processing.extract_pdf_attachments = False monkeypatch.setattr(client._config.processing, "extract_pdf_attachments", False)
parent_uri = "file:///fixtures/parent.pdf" parent_uri = "file:///fixtures/parent.pdf"
pdf_bytes = build_pdf([("a.txt", b"A")]) pdf_bytes = build_pdf([("a.txt", b"A")])
parent = await _make_parent(client, parent_uri, pdf_bytes) parent = await _make_parent(client, parent_uri, pdf_bytes)