diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index 0a326aa2..a667213d 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -361,30 +361,33 @@ async def _reconcile_pdf_attachments( "Cannot scan %s for embedded attachments: %s", parent_doc.uri, exc ) return - attachment_count = pdf.count_attachments() + 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 + 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) + 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} diff --git a/tests/test_pdf_attachments.py b/tests/test_pdf_attachments.py index 1446cf8c..a838e473 100644 --- a/tests/test_pdf_attachments.py +++ b/tests/test_pdf_attachments.py @@ -1,5 +1,4 @@ import io -import logging 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 -async def test_nested_pdf_attachments_recurse_up_to_cap( - temp_db_path, monkeypatch, caplog -): +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 with a - # warning logged at the depth boundary. + # 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)]) @@ -214,8 +210,7 @@ async def test_nested_pdf_attachments_recurse_up_to_cap( 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) + 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" @@ -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(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): @@ -233,7 +227,7 @@ async def test_config_off_skips_extraction(temp_db_path, monkeypatch): fake_ingest_fetch_result, ) 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" pdf_bytes = build_pdf([("a.txt", b"A")]) parent = await _make_parent(client, parent_uri, pdf_bytes)