From 0a15ce8fbb70d24d6f31f51225b5757e649c4c52 Mon Sep 17 00:00:00 2001 From: Alex Date: Thu, 13 Aug 2026 14:30:28 +0100 Subject: [PATCH 1/2] feat: attachment provenance --- application/api/user/tasks.py | 8 +- application/llm/handlers/base.py | 36 ++- application/parser/chunking.py | 6 +- application/parser/chunking_strategies.py | 6 +- application/parser/document_reader.py | 9 +- application/utils.py | 5 +- application/vectorstore/base.py | 2 +- application/worker.py | 147 ++++++++-- tests/llm/handlers/test_llm_handlers.py | 87 ++++++ tests/parser/file/test_markdown_parser.py | 3 + tests/parser/test_chunking.py | 21 ++ tests/parser/test_document_reader.py | 58 ++++ tests/test_utils.py | 25 ++ tests/worker/test_attachment_worker.py | 16 +- .../test_attachment_worker_provenance.py | 266 ++++++++++++++++++ 15 files changed, 659 insertions(+), 36 deletions(-) create mode 100644 tests/worker/test_attachment_worker_provenance.py diff --git a/application/api/user/tasks.py b/application/api/user/tasks.py index d85e38f8..2229f804 100644 --- a/application/api/user/tasks.py +++ b/application/api/user/tasks.py @@ -219,7 +219,9 @@ def _emit_attachment_poison_event(task_name, bound): Mirrors ``_emit_ingest_poison_event``: the guard returns before the worker runs, so ``attachment_worker``'s own events never fire and the - upload toast would otherwise spin on "processing" forever. + upload toast would otherwise spin on "processing" forever. Also writes + the failure row the worker never got to write, so the poisoned upload + stays visible to a DB scan and not just to whoever saw the toast. """ user = bound.get("user") file_info = bound.get("file_info") or {} @@ -227,7 +229,11 @@ def _emit_attachment_poison_event(task_name, bound): if not user or not attachment_id: return from application.events.publisher import publish_user_event + from application.worker import record_attachment_failure + record_attachment_failure( + user, file_info, "Attachment processing stopped after repeated failures." + ) publish_user_event( user, "attachment.failed", diff --git a/application/llm/handlers/base.py b/application/llm/handlers/base.py index 5983f796..ef379823 100644 --- a/application/llm/handlers/base.py +++ b/application/llm/handlers/base.py @@ -316,11 +316,39 @@ class LLMHandler(ABC): attachment_texts = [] for attachment in attachments: - logger.info(f"Adding attachment {attachment.get('id')} to context") - if "content" in attachment: - attachment_texts.append( - f"Attached file content:\n\n{attachment['content']}" + # ``metadata.extraction`` records what parsing actually did. + # Rows predating it have no ``extraction`` key and pass; rows + # whose extraction failed must not reach the prompt (a PG row + # always has a ``content`` key, so key membership is no gate). + extraction = (attachment.get("metadata") or {}).get("extraction") or {} + status = extraction.get("status") + if status is not None and status != "ok": + logger.info( + f"Skipping attachment {attachment.get('id')}: extraction status {status}" ) + continue + content = attachment.get("content") + if content is None: + logger.info( + f"Skipping attachment {attachment.get('id')}: no extracted content" + ) + continue + logger.info(f"Adding attachment {attachment.get('id')} to context") + note = "" + if extraction.get("truncated"): + stored = extraction.get("stored_tokens") + original = extraction.get("original_tokens") + counts = ( + f"only the first {stored:,} of ~{original:,} tokens are included" + if isinstance(stored, int) and isinstance(original, int) + else "part of the document is missing" + ) + name = attachment.get("filename") or "This file" + note = ( + f'[NOTE: "{name}" was truncated during extraction — {counts}. ' + "Scope any whole-document claims to this portion.]\n\n" + ) + attachment_texts.append(f"Attached file content:\n\n{note}{content}") if attachment_texts: combined_text = "\n\n".join(attachment_texts) diff --git a/application/parser/chunking.py b/application/parser/chunking.py index 322c356b..32366b4f 100644 --- a/application/parser/chunking.py +++ b/application/parser/chunking.py @@ -44,8 +44,8 @@ class Chunker: def split_document(self, doc: Document) -> List[Document]: split_docs = [] header, body = self.separate_header_and_body(doc.text) - header_tokens = self.encoding.encode(header) if header else [] - body_tokens = self.encoding.encode(body) + header_tokens = self.encoding.encode_ordinary(header) if header else [] + body_tokens = self.encoding.encode_ordinary(body) current_position = 0 part_index = 0 @@ -71,7 +71,7 @@ class Chunker: i = 0 while i < len(documents): doc = documents[i] - tokens = self.encoding.encode(doc.text) + tokens = self.encoding.encode_ordinary(doc.text) token_count = len(tokens) if self.min_tokens <= token_count <= self.max_tokens: diff --git a/application/parser/chunking_strategies.py b/application/parser/chunking_strategies.py index 6e4a0dcb..933ec7f0 100644 --- a/application/parser/chunking_strategies.py +++ b/application/parser/chunking_strategies.py @@ -44,11 +44,11 @@ class _BaseStrategyChunker: self.encoding = get_encoding() def _token_count(self, text: str) -> int: - return len(self.encoding.encode(text)) + return len(self.encoding.encode_ordinary(text)) def _split_by_tokens(self, text: str) -> List[str]: """Split ``text`` into pieces no larger than ``max_tokens`` tokens.""" - tokens = self.encoding.encode(text) + tokens = self.encoding.encode_ordinary(text) pieces = [] for start in range(0, len(tokens), self.max_tokens): chunk_tokens = tokens[start:start + self.max_tokens] @@ -188,7 +188,7 @@ class ParentChildChunker(_BaseStrategyChunker): processed: List[Document] = [] child_size = self._child_size() for doc in documents: - tokens = self.encoding.encode(doc.text) + tokens = self.encoding.encode_ordinary(doc.text) part_index = 0 for p_start in range(0, len(tokens), self.max_tokens): parent_tokens = tokens[p_start:p_start + self.max_tokens] diff --git a/application/parser/document_reader.py b/application/parser/document_reader.py index 978e29bd..57bb2a98 100644 --- a/application/parser/document_reader.py +++ b/application/parser/document_reader.py @@ -263,8 +263,15 @@ def _docling_structured(path: Path, *, ocr_enabled: bool, include_tables: bool, if include_tables: for tbl in getattr(doc, "tables", []) or []: try: + from application.parser.file.tabular_parser import cell_to_text + + # ``astype(str)`` leaks pandas-3 NaNs and mangles gapped int + # columns ("1001.0") — same contract as the tabular parser. df = tbl.export_to_dataframe() - tables.append({"columns": [str(c) for c in df.columns], "rows": df.astype(str).values.tolist()}) + tables.append({ + "columns": [cell_to_text(c) for c in df.columns], + "rows": [[cell_to_text(v) for v in row] for row in df.values.tolist()], + }) except Exception: try: tables.append({"markdown": tbl.export_to_markdown()}) diff --git a/application/utils.py b/application/utils.py index 16dad7b8..80f21b07 100644 --- a/application/utils.py +++ b/application/utils.py @@ -105,7 +105,10 @@ def truncate_to_line_boundary(data: bytes) -> bytes: def num_tokens_from_string(string: str) -> int: encoding = get_encoding() if isinstance(string, str): - num_tokens = len(encoding.encode(string)) + # encode_ordinary: plain ``encode()`` raises ValueError on literal + # special-token text like <|endoftext|>, which user documents + # legitimately contain. + num_tokens = len(encoding.encode_ordinary(string)) return num_tokens else: return 0 diff --git a/application/vectorstore/base.py b/application/vectorstore/base.py index e58fb4d9..1a999f26 100644 --- a/application/vectorstore/base.py +++ b/application/vectorstore/base.py @@ -51,7 +51,7 @@ class RemoteEmbeddings: def clip(text): if not isinstance(text, str): return text - tokens = encoding.encode(text) + tokens = encoding.encode_ordinary(text) if len(tokens) <= limit: return text logging.warning( diff --git a/application/worker.py b/application/worker.py index c152c25b..4f00da42 100755 --- a/application/worker.py +++ b/application/worker.py @@ -57,6 +57,7 @@ from application.storage.db.source_config import SourceConfig from application.storage.storage_creator import StorageCreator from application.utils import ( count_tokens_docs, + get_encoding, num_tokens_from_string, safe_filename, truncate_to_line_boundary, @@ -67,6 +68,9 @@ from application.utils import ( MIN_TOKENS = 150 MAX_TOKENS = 1250 +# Attachment content stored for prompting is capped in tokens — the same +# unit the gate is expressed in — so a stored row can never exceed the cap. +ATTACHMENT_MAX_TOKENS = 100_000 RECURSION_DEPTH = 2 INGEST_HEARTBEAT_INTERVAL_SECONDS = 30 @@ -1539,6 +1543,93 @@ def _bounded_attachment_copy(local_path: str) -> tuple[str, bool]: return tmp.name, True +def _upsert_attachment_row( + user, filename, relative_path, *, mime_type, content, token_count, metadata, attachment_id +): + """Create or update the attachment row for one upload handle. + + The upload route mints a UUID-shaped ``attachment_id`` (stored in the + storage path); the PG ``attachments.id`` is DB-generated, so the handle + lives in ``legacy_mongo_id``. Task retries share the handle — an earlier + attempt's failure row is updated in place, never duplicated. + """ + with db_session() as conn: + repo = AttachmentsRepository(conn) + existing = repo.get_by_legacy_id(str(attachment_id), user) + if existing: + repo.update( + existing["id"], + user, + { + "filename": filename, + "upload_path": relative_path, + "mime_type": mime_type, + "content": content, + "token_count": token_count, + "metadata": metadata, + }, + ) + else: + repo.create( + user, + filename, + relative_path, + mime_type=mime_type, + content=content, + token_count=token_count, + metadata=metadata, + legacy_mongo_id=str(attachment_id), + ) + + +def record_attachment_failure(user, file_info, error, parser=None): + """Persist a failure row so a broken parse is visible to a DB scan. + + Called on the worker's error path and by the poison guard; never raises — + the original exception must keep propagating, and the parser's error text + goes into ``metadata.extraction``, never into ``content``. + """ + attachment_id = file_info.get("attachment_id") + filename = file_info.get("filename") or "" + if not attachment_id: + return + try: + # An exception after the success write (event publishing, state + # updates) must not replace durable extracted content with a + # NULL-content failed row. + with db_readonly() as conn: + existing = AttachmentsRepository(conn).get_by_legacy_id( + str(attachment_id), user + ) + if existing is not None and existing.get("content") is not None: + return + metadata = { + **(file_info.get("metadata") or {}), + "extraction": { + "status": "failed", + "parser": parser, + "truncated": False, + "error": str(error)[:1024], + }, + } + _upsert_attachment_row( + user, + filename, + file_info.get("path") or "", + mime_type=mimetypes.guess_type(filename)[0] or "application/octet-stream", + content=None, + token_count=None, + metadata=metadata, + attachment_id=attachment_id, + ) + except Exception: + logging.error( + f"Failed to record failure row for attachment {attachment_id}", + extra={"user": user}, + exc_info=True, + ) + + def attachment_worker(self, file_info, user): """ Process and store a single attachment without vectorization. @@ -1548,6 +1639,7 @@ def attachment_worker(self, file_info, user): attachment_id = file_info["attachment_id"] relative_path = file_info["path"] metadata = file_info.get("metadata", {}) + parser_name = None publish_user_event( user, @@ -1578,6 +1670,9 @@ def attachment_worker(self, file_info, user): file_extractor = get_default_file_extractor( ocr_enabled=settings.DOCLING_OCR_ATTACHMENTS_ENABLED ) + _parser = file_extractor.get(os.path.splitext(filename)[1].lower()) + parser_name = type(_parser).__name__ if _parser is not None else "SimpleDirectoryReader" + def _parse_local_file(local_path: str, **kwargs) -> Document: _reject_attachment_zip_bomb(local_path) parse_path, is_temp_copy = _bounded_attachment_copy(local_path) @@ -1606,10 +1701,29 @@ def attachment_worker(self, file_info, user): if parser_metadata: metadata = {**metadata, **parser_metadata} - token_count = num_tokens_from_string(content) - if token_count > 100000: - content = content[:250000] - token_count = num_tokens_from_string(content) + # Gate and cut in the same unit. The old form gated on tokens but cut + # at 250k *chars*, which for dense scripts (CJK ~1.4 tokens/char) + # stored 300k+ tokens while looking like a clean extraction. + encoding = get_encoding() + tokens = encoding.encode_ordinary(content) + original_tokens = len(tokens) + truncated = original_tokens > ATTACHMENT_MAX_TOKENS + if truncated: + content = encoding.decode(tokens[:ATTACHMENT_MAX_TOKENS]) + token_count = ATTACHMENT_MAX_TOKENS + else: + token_count = original_tokens + + metadata = { + **metadata, + "extraction": { + "status": "ok", + "parser": parser_name, + "truncated": truncated, + "original_tokens": original_tokens, + "stored_tokens": token_count, + }, + } self.update_state( state="PROGRESS", meta={"current": 80, "status": "Storing in database"} @@ -1628,20 +1742,16 @@ def attachment_worker(self, file_info, user): mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream" - # The upload route produces a UUID-shaped ``attachment_id`` (stored - # in the storage path) but the PG ``attachments.id`` is generated - # by the DB. Keep ``attachment_id`` as the caller-visible handle - # used for the storage path, and stash it in ``legacy_mongo_id`` - # so the attachment row is resolvable via that handle too. - with db_session() as conn: - AttachmentsRepository(conn).create( - user, filename, relative_path, - mime_type=mime_type, - content=content, - token_count=token_count, - metadata=metadata, - legacy_mongo_id=str(attachment_id), - ) + _upsert_attachment_row( + user, + filename, + relative_path, + mime_type=mime_type, + content=content, + token_count=token_count, + metadata=metadata, + attachment_id=attachment_id, + ) logging.info( f"Stored attachment with ID: {attachment_id}", extra={"user": user} @@ -1675,6 +1785,7 @@ def attachment_worker(self, file_info, user): extra={"user": user}, exc_info=True, ) + record_attachment_failure(user, file_info, e, parser=parser_name) publish_user_event( user, "attachment.failed", diff --git a/tests/llm/handlers/test_llm_handlers.py b/tests/llm/handlers/test_llm_handlers.py index a480b27e..1b068eba 100644 --- a/tests/llm/handlers/test_llm_handlers.py +++ b/tests/llm/handlers/test_llm_handlers.py @@ -2105,3 +2105,90 @@ class TestHandleToolCallsCompressionSuccess: if isinstance(e, dict) and e.get("data", {}).get("status") == "skipped" ] assert len(skip_events) == 1 # Only second call skipped + + +# --------------------------------------------------------------------------- +# _append_unsupported_attachments — extraction-provenance gate +# --------------------------------------------------------------------------- + + +class TestAttachmentExtractionGate: + """Rows now carry ``metadata.extraction`` provenance; the append gate + must honor it. Legacy rows (no ``extraction`` key) keep working.""" + + def test_failed_extraction_row_is_skipped(self): + handler = ConcreteHandler() + messages = [{"role": "system", "content": "sys"}] + attachments = [{ + "id": "a1", + "mime_type": "application/vnd.ms-excel", + "content": None, + "metadata": {"extraction": {"status": "failed", "error": "docling boom"}}, + }] + + result = handler._append_unsupported_attachments(messages, attachments) + + assert result[0]["content"] == "sys" + assert "None" not in result[0]["content"] + assert "docling boom" not in result[0]["content"] + + def test_null_content_row_is_skipped_even_without_extraction(self): + # A PG row always has the ``content`` key; key membership is not a + # signal. ``content=None`` must never append the string "None". + handler = ConcreteHandler() + messages = [{"role": "system", "content": "sys"}] + attachments = [{"id": "a1", "mime_type": "text/plain", "content": None}] + + result = handler._append_unsupported_attachments(messages, attachments) + + assert result[0]["content"] == "sys" + + def test_legacy_row_without_extraction_still_appended(self): + handler = ConcreteHandler() + messages = [{"role": "system", "content": "sys"}] + attachments = [{ + "id": "a1", + "content": "legacy file text", + "metadata": {"storage_type": "local"}, + }] + + result = handler._append_unsupported_attachments(messages, attachments) + + assert "legacy file text" in result[0]["content"] + + def test_ok_row_appended_without_truncation_notice(self): + handler = ConcreteHandler() + messages = [{"role": "system", "content": "sys"}] + attachments = [{ + "id": "a1", + "content": "full file text", + "metadata": {"extraction": {"status": "ok", "truncated": False}}, + }] + + result = handler._append_unsupported_attachments(messages, attachments) + + assert "full file text" in result[0]["content"] + assert "truncated" not in result[0]["content"] + + def test_truncated_row_appended_with_partial_document_notice(self): + handler = ConcreteHandler() + messages = [{"role": "system", "content": "sys"}] + attachments = [{ + "id": "a1", + "filename": "big.pdf", + "content": "partial file text", + "metadata": {"extraction": { + "status": "ok", + "truncated": True, + "original_tokens": 250000, + "stored_tokens": 100000, + }}, + }] + + result = handler._append_unsupported_attachments(messages, attachments) + + prompt = result[0]["content"] + assert "partial file text" in prompt + assert "truncated" in prompt + assert "100,000" in prompt + assert "250,000" in prompt diff --git a/tests/parser/file/test_markdown_parser.py b/tests/parser/file/test_markdown_parser.py index c3cb61ea..bed76ce1 100644 --- a/tests/parser/file/test_markdown_parser.py +++ b/tests/parser/file/test_markdown_parser.py @@ -11,6 +11,9 @@ class _Enc: def encode(self, s: str): return list(s) + def encode_ordinary(self, s: str): + return list(s) + @pytest.fixture(autouse=True) def _patch_tokenizer(monkeypatch): diff --git a/tests/parser/test_chunking.py b/tests/parser/test_chunking.py index 97126f58..55e2e587 100644 --- a/tests/parser/test_chunking.py +++ b/tests/parser/test_chunking.py @@ -281,3 +281,24 @@ class TestChunkerIntegration: assert doc.extra_info is not None assert "token_count" in doc.extra_info assert doc.extra_info["token_count"] > 0 + + +# ===================================================================== +# Special-token text must chunk as ordinary text, not raise +# ===================================================================== + + +@pytest.mark.unit +class TestSpecialTokenText: + + def test_chunking_document_containing_special_token_markers(self): + # A document ABOUT LLMs contains literal <|endoftext|>; plain + # ``encode()`` raises ValueError on it and destroyed the ingest. + chunker = Chunker(chunking_strategy="classic_chunk", max_tokens=50, min_tokens=0) + doc = Document( + text="The <|endoftext|> marker separates documents. " * 30, + doc_id="d1", + ) + chunks = chunker.chunk([doc]) + assert chunks + assert all("<|endoftext|>" in c.text for c in chunks[:1]) diff --git a/tests/parser/test_document_reader.py b/tests/parser/test_document_reader.py index 83d2e435..e73ef3a3 100644 --- a/tests/parser/test_document_reader.py +++ b/tests/parser/test_document_reader.py @@ -564,3 +564,61 @@ def test_vanilla_converter_applies_inference_settings(monkeypatch, tmp_path): assert out["markdown"] == "# md" assert called == [True] + + +# --------------------------------------------------------------------------- +# _docling_structured tables: blank cells must not leak NaN or mangle ints +# --------------------------------------------------------------------------- +@pytest.mark.unit +def test_structured_tables_render_blank_cells_as_strings(monkeypatch, tmp_path): + """pandas 3 preserves missing values through ``astype(str)``, so a blank + spreadsheet cell leaked a float NaN into the tables payload and a blank + in an integer column mangled every id to ``1001.0``. Cells must come out + as join-safe strings (same contract as the tabular parser's + ``cell_to_text``).""" + import sys + import types + + import pandas as pd + + df = pd.DataFrame({"id": [1001, None, 1003], "name": ["a", "b", None]}) + + class _Table: + def export_to_dataframe(self): + return df + + class _Doc: + tables = [_Table()] + pages = {"1": {}} + + def export_to_markdown(self): + return "# md" + + def export_to_dict(self): + return {"texts": [], "tables": [{}], "pages": {"1": {}}} + + class _Converter: + def convert(self, *a, **k): + return types.SimpleNamespace(document=_Doc()) + + monkeypatch.setattr( + "application.parser.file.docling_parser._apply_inference_settings", + lambda: None, + ) + fake_docling = types.ModuleType("docling") + fake_dc_module = types.ModuleType("docling.document_converter") + fake_dc_module.DocumentConverter = _Converter + fake_docling.document_converter = fake_dc_module + monkeypatch.setitem(sys.modules, "docling", fake_docling) + monkeypatch.setitem(sys.modules, "docling.document_converter", fake_dc_module) + + path = tmp_path / "doc.xlsx" + path.write_bytes(b"stub") + + out = dr._docling_structured(path, ocr_enabled=False, include_tables=True, parser=None) + + [table] = out["tables"] + cells = [c for row in table["rows"] for c in row] + assert all(isinstance(c, str) for c in cells), f"non-string cells: {cells!r}" + assert "1001" in cells and "1001.0" not in cells + assert "" in cells # blank cell renders as empty string, not "nan"/"None" diff --git a/tests/test_utils.py b/tests/test_utils.py index 16ec3852..c789ab00 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -102,6 +102,31 @@ class TestNumTokens: def test_empty_string(self): assert num_tokens_from_string("") == 0 + @pytest.mark.unit + def test_special_token_text_counts_instead_of_raising(self): + # User documents legitimately contain literal marker text like + # <|endoftext|> (any paper about LLMs). Plain ``encode()`` raises + # ValueError on it; counting must treat it as ordinary text. + count = num_tokens_from_string( + "This paper discusses the <|endoftext|> token in GPT models." + ) + assert count > 0 + + @pytest.mark.unit + def test_unbroken_cjk_run_counts_in_bounded_time(self): + # An unbroken letter run is a single BPE piece; tiktoken < 0.13.0 + # merged it quadratically (~26s for this input, minutes for real + # uploads). tiktoken 0.13.0 does it in milliseconds. Guards + # against a dependency downgrade re-introducing the hang. + import time + + dense = "統計資料表格內容分析報告書類文書處理系統設計開發" * 4000 # 96k chars, one piece + start = time.monotonic() + count = num_tokens_from_string(dense) + elapsed = time.monotonic() - start + assert count > 100000 + assert elapsed < 10, f"pathological encode took {elapsed:.1f}s — tiktoken downgraded?" + class TestNumTokensFromObjectOrList: diff --git a/tests/worker/test_attachment_worker.py b/tests/worker/test_attachment_worker.py index 1f1e8624..4172a944 100644 --- a/tests/worker/test_attachment_worker.py +++ b/tests/worker/test_attachment_worker.py @@ -67,7 +67,7 @@ class TestAttachmentWorker: assert row["content"] == "hello world" assert row["user_id"] == "user1" - def test_parse_failure_stores_nothing_and_tells_the_user( + def test_parse_failure_stores_failure_row_and_tells_the_user( self, pg_conn, patch_worker_db, task_self, monkeypatch ): """A parse failure must fail loudly, not store the error as content. @@ -76,6 +76,9 @@ class TestAttachmentWorker: parser returned its own traceback as the "document", and the worker stored it and published ``attachment.completed`` — so the upload looked fine and the model was handed an error message as the PDF. + The failure now leaves a queryable row whose ``metadata.extraction`` + records the error; ``content`` stays NULL so the model can never + read a traceback as the document. """ from application import worker from application.parser.file.base_parser import DocumentParseError @@ -109,12 +112,17 @@ class TestAttachmentWorker: with pytest.raises(DocumentParseError): worker.attachment_worker(task_self, file_info, "user1") - # Nothing may be persisted — an attachment whose text is a traceback - # is worse than no attachment, because the model will read it. + # The error lives in metadata, never in content — an attachment whose + # text is a traceback is worse than no attachment, because the model + # will read it. row = AttachmentsRepository(pg_conn).get_by_legacy_id( file_info["attachment_id"], "user1" ) - assert row is None, "a failed parse must not insert an attachment row" + assert row is not None, "a failed parse must leave a queryable row" + assert row["content"] is None + extraction = row["metadata"]["extraction"] + assert extraction["status"] == "failed" + assert "InvalidCxxCompiler" in extraction["error"] # The user is told it failed, and never told it completed. events = [event for event, _ in published] diff --git a/tests/worker/test_attachment_worker_provenance.py b/tests/worker/test_attachment_worker_provenance.py new file mode 100644 index 00000000..95b49ec8 --- /dev/null +++ b/tests/worker/test_attachment_worker_provenance.py @@ -0,0 +1,266 @@ +"""Extraction-provenance tests for ``attachment_worker``. + +Every parse attempt must leave a queryable trace in ``attachments``: +``metadata.extraction`` records what happened (status, parser, truncation, +token counts, error) on success, truncation, and terminal failure alike. +Runs against a real ephemeral Postgres so the rows the worker writes are +asserted as stored, not as mocked calls. +""" + +from __future__ import annotations + +import uuid +from unittest.mock import patch + +import pytest + +import application.storage.db.engine as engine_module +from application.parser.file.base_parser import DocumentParseError +from application.storage.db.repositories.attachments import AttachmentsRepository +from application.storage.db.session import db_readonly +from application.utils import get_encoding + + +class _StubTask: + def update_state(self, *args, **kwargs): + pass + + +class _Doc: + def __init__(self, text): + self.text = text + self.extra_info = {} + + +@pytest.fixture() +def wired_engine(pg_engine, monkeypatch): + """Point the app's module-level engine cache at the ephemeral DB.""" + monkeypatch.setattr(engine_module, "_engine", None) + eng = engine_module.get_engine() + yield eng + eng.dispose() + monkeypatch.setattr(engine_module, "_engine", None) + + +@pytest.fixture() +def storage_dir(tmp_path, monkeypatch): + """LocalStorage rooted at tmp_path, patched into the worker.""" + from application.storage.local import LocalStorage + + storage = LocalStorage(base_dir=str(tmp_path)) + monkeypatch.setattr( + "application.storage.storage_creator.StorageCreator.get_storage", + classmethod(lambda cls: storage), + ) + return tmp_path + + +def _run_worker(file_info, user="prov-user"): + from application.worker import attachment_worker + + return attachment_worker(_StubTask(), file_info, user) + + +def _file_info(storage_dir, filename="doc.txt", content=b"hello attachment"): + attachment_id = str(uuid.uuid4()) + rel_path = f"inputs/prov-user/attachments/{attachment_id}/{filename}" + full = storage_dir / rel_path + full.parent.mkdir(parents=True, exist_ok=True) + full.write_bytes(content) + return { + "filename": filename, + "attachment_id": attachment_id, + "path": rel_path, + "metadata": {"storage_type": "local"}, + } + + +def _fetch(attachment_id, user="prov-user"): + with db_readonly() as conn: + return AttachmentsRepository(conn).get_by_legacy_id(str(attachment_id), user) + + +def _count(user="prov-user"): + with db_readonly() as conn: + return len(AttachmentsRepository(conn).list_for_user(user)) + + +@pytest.mark.usefixtures("wired_engine") +class TestSuccessProvenance: + def test_ok_row_records_extraction_metadata(self, storage_dir): + info = _file_info(storage_dir) + + result = _run_worker(info) + + row = _fetch(info["attachment_id"]) + assert row is not None + extraction = (row["metadata"] or {}).get("extraction") + assert extraction is not None + assert extraction["status"] == "ok" + assert extraction["truncated"] is False + assert extraction["original_tokens"] == extraction["stored_tokens"] + assert extraction["stored_tokens"] == row["token_count"] + assert extraction.get("parser") + assert "hello attachment" in row["content"] + assert result["token_count"] == row["token_count"] + + def test_upload_metadata_keys_preserved(self, storage_dir): + info = _file_info(storage_dir) + + _run_worker(info) + + row = _fetch(info["attachment_id"]) + # The storage-level metadata written at upload time must survive + # the extraction merge. + assert row["metadata"]["storage_type"] == "local" + + +@pytest.mark.usefixtures("wired_engine") +class TestTruncationProvenance: + def test_over_gate_content_truncated_in_token_units(self, storage_dir, monkeypatch): + # Dense CJK: ~1.4 tokens/char, so a char-unit cut (the old + # behavior) would store far more than 100k tokens. + dense = "統計資料表格內容分析、報告書類文書處理系統。設計開發運用管理。\n" * 10000 + info = _file_info(storage_dir) + monkeypatch.setattr( + "application.worker.SimpleDirectoryReader", + lambda **kwargs: type("R", (), {"load_data": lambda self: [_Doc(dense)]})(), + ) + + _run_worker(info) + + row = _fetch(info["attachment_id"]) + extraction = row["metadata"]["extraction"] + assert extraction["status"] == "ok" + assert extraction["truncated"] is True + assert extraction["stored_tokens"] == 100000 + assert extraction["original_tokens"] > 100000 + assert row["token_count"] == 100000 + # The stored content really is ~100k tokens, not a 250k-char cut + # still worth 300k+ tokens. + enc = get_encoding() + assert len(enc.encode_ordinary(row["content"])) <= 100100 + + def test_under_gate_content_not_truncated(self, storage_dir, monkeypatch): + text = "plain short attachment content" + info = _file_info(storage_dir) + monkeypatch.setattr( + "application.worker.SimpleDirectoryReader", + lambda **kwargs: type("R", (), {"load_data": lambda self: [_Doc(text)]})(), + ) + + _run_worker(info) + + row = _fetch(info["attachment_id"]) + extraction = row["metadata"]["extraction"] + assert extraction["truncated"] is False + assert row["content"] == text + + +@pytest.mark.usefixtures("wired_engine") +class TestFailureProvenance: + def test_terminal_parse_failure_writes_failed_row(self, storage_dir, monkeypatch): + info = _file_info(storage_dir, filename="broken.xlsx") + + def _raise(**kwargs): + raise DocumentParseError("Failed to parse broken.xlsx with docling: boom") + + monkeypatch.setattr("application.worker.SimpleDirectoryReader", _raise) + + with pytest.raises(DocumentParseError): + _run_worker(info) + + row = _fetch(info["attachment_id"]) + assert row is not None, "terminal failure must leave a queryable row" + assert row["content"] is None + extraction = row["metadata"]["extraction"] + assert extraction["status"] == "failed" + assert "boom" in extraction["error"] + # The parser's exception text must never masquerade as document + # content. + assert row["token_count"] is None + + def test_retry_then_success_upserts_single_row(self, storage_dir, monkeypatch): + info = _file_info(storage_dir) + calls = {"n": 0} + + def _flaky(**kwargs): + calls["n"] += 1 + if calls["n"] == 1: + raise RuntimeError("transient blip") + return type("R", (), {"load_data": lambda self: [_Doc("recovered fine")]})() + + monkeypatch.setattr("application.worker.SimpleDirectoryReader", _flaky) + + with pytest.raises(RuntimeError): + _run_worker(info) + failed_row = _fetch(info["attachment_id"]) + assert failed_row["metadata"]["extraction"]["status"] == "failed" + + _run_worker(info) + + assert _count() == 1, "retry success must update the failure row, not add one" + row = _fetch(info["attachment_id"]) + assert row["metadata"]["extraction"]["status"] == "ok" + assert row["content"] == "recovered fine" + + def test_post_store_failure_does_not_clobber_ok_row(self, storage_dir): + # An exception after the success write (e.g. event publishing) must + # not replace stored content with a NULL-content failed row; the + # extraction result is already durable. + from application.worker import record_attachment_failure + + info = _file_info(storage_dir) + _run_worker(info) + + record_attachment_failure("prov-user", info, RuntimeError("post-store blip")) + + row = _fetch(info["attachment_id"]) + assert row["metadata"]["extraction"]["status"] == "ok" + assert "hello attachment" in row["content"] + + def test_failure_row_write_failure_does_not_mask_original_error( + self, storage_dir, monkeypatch + ): + info = _file_info(storage_dir) + + def _raise(**kwargs): + raise DocumentParseError("original parse error") + + monkeypatch.setattr("application.worker.SimpleDirectoryReader", _raise) + monkeypatch.setattr( + "application.worker.db_session", + _raising_db_session, + ) + + with pytest.raises(DocumentParseError, match="original parse error"): + _run_worker(info) + + +def _raising_db_session(): + raise ConnectionError("db down") + + +@pytest.mark.usefixtures("wired_engine") +class TestPoisonProvenance: + def test_poison_guard_writes_failed_row(self): + from application.api.user.tasks import _emit_attachment_poison_event + + attachment_id = str(uuid.uuid4()) + bound = { + "user": "prov-user", + "file_info": { + "filename": "poison.pdf", + "attachment_id": attachment_id, + "path": "inputs/prov-user/attachments/x/poison.pdf", + }, + } + + with patch("application.events.publisher.publish_user_event"): + _emit_attachment_poison_event("store_attachment", bound) + + row = _fetch(attachment_id) + assert row is not None, "poison-guard trip must leave a queryable row" + extraction = row["metadata"]["extraction"] + assert extraction["status"] == "failed" + assert "repeated failures" in extraction["error"] From 7ad2f385120d4f5ef75b4c307fc0fa56c82b720e Mon Sep 17 00:00:00 2001 From: Alex Date: Thu, 13 Aug 2026 16:04:45 +0100 Subject: [PATCH 2/2] feat: faster attachment processing --- application/core/settings.py | 11 + application/parser/file/bulk.py | 47 +++- application/parser/file/pdfium_parser.py | 261 ++++++++++++++++++ application/requirements.txt | 1 + .../storage/db/repositories/attachments.py | 35 +++ application/worker.py | 50 ++-- tests/parser/file/test_pdfium_parser.py | 184 ++++++++++++ .../db/repositories/test_attachments.py | 40 +++ tests/worker/test_attachment_worker.py | 4 +- tests/worker/test_worker_sse_publishes.py | 4 +- 10 files changed, 611 insertions(+), 26 deletions(-) create mode 100644 application/parser/file/pdfium_parser.py create mode 100644 tests/parser/file/test_pdfium_parser.py diff --git a/application/core/settings.py b/application/core/settings.py index 2892da80..b8135fe6 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -108,6 +108,17 @@ class Settings(BaseSettings): DOCLING_COMPILE_TORCH_MODELS: bool = False DOCLING_TABULAR_MAX_BYTES: int = 2_000_000 DOCLING_MARKUP_MAX_BYTES: int = 8_000_000 + # Read PDF *attachments* via their embedded text layer (pypdfium2) instead + # of docling, falling back to docling when there is no text layer to read. + # Attachments go into a prompt, so docling's structural markdown earns far + # less than the tens of seconds per file it costs; source ingestion is + # unaffected and always uses docling, because chunking and retrieval do + # depend on that structure. + ATTACHMENT_PDF_TEXT_FAST_PATH: bool = True + # Median characters per sampled page below which a PDF is treated as a scan + # and handed to docling. Measured separation on real uploads: scans at + # 0-17 chars/page, text-layer documents at 433-6834. + ATTACHMENT_PDF_TEXT_MIN_MEDIAN_CHARS: int = 32 ATTACHMENT_TEXT_MAX_BYTES: int = 5_000_000 AGENT_IMAGE_MAX_BYTES: int = 5_000_000 AGENT_IMAGE_MAX_PIXELS: int = 16_777_216 diff --git a/application/parser/file/bulk.py b/application/parser/file/bulk.py index 9008b2a1..d8e5ae8d 100644 --- a/application/parser/file/bulk.py +++ b/application/parser/file/bulk.py @@ -25,13 +25,50 @@ def _build_audio_parser_mapping() -> Dict[str, BaseParser]: return {extension: AudioParser() for extension in SUPPORTED_AUDIO_EXTENSIONS} +def _wrap_pdf_fast_path(pdf_parser: BaseParser) -> BaseParser: + """Put the pypdfium2 text-layer parser in front of ``pdf_parser``. + + Args: + pdf_parser: Parser to fall back to for PDFs with no usable text layer. + + Returns: + BaseParser: The fast-path parser, or ``pdf_parser`` unchanged when + pypdfium2 is unavailable. + """ + try: + from application.parser.file.pdfium_parser import PdfiumTextParser + except ImportError: + logging.warning( + "pypdfium2 is not installed; PDF attachments will use %s", + type(pdf_parser).__name__, + ) + return pdf_parser + return PdfiumTextParser( + fallback_parser=pdf_parser, + min_median_chars=settings.ATTACHMENT_PDF_TEXT_MIN_MEDIAN_CHARS, + ) + + def get_default_file_extractor( ocr_enabled: Optional[bool] = None, + pdf_text_fast_path: bool = False, ) -> Dict[str, BaseParser]: """Get the default file extractor. Uses docling parsers by default for advanced document processing. Falls back to standard parsers if docling is not installed. + + Args: + ocr_enabled: Enable OCR in the docling parsers. Defaults to + ``settings.DOCLING_OCR_ENABLED``. + pdf_text_fast_path: Read PDFs via their embedded text layer + (pypdfium2) instead of docling, falling back to docling for scans. + Off by default: it trades docling's structural markdown for speed, + which suits attachments (read into a prompt) but not source + ingestion (chunked and embedded for retrieval). + + Returns: + Dict[str, BaseParser]: Parser keyed by lower-case file suffix. """ try: from application.parser.file.docling_parser import ( @@ -48,9 +85,12 @@ def get_default_file_extractor( ) if ocr_enabled is None: ocr_enabled = settings.DOCLING_OCR_ENABLED + pdf_parser: BaseParser = DoclingPDFParser(ocr_enabled=ocr_enabled) + if pdf_text_fast_path: + pdf_parser = _wrap_pdf_fast_path(pdf_parser) return { # Documents - ".pdf": DoclingPDFParser(ocr_enabled=ocr_enabled), + ".pdf": pdf_parser, ".docx": DoclingDocxParser(), ".pptx": DoclingPPTXParser(), ".xlsx": DoclingXLSXParser(), @@ -88,8 +128,11 @@ def get_default_file_extractor( "For advanced document parsing, install with: pip install docling" ) # Fallback to standard parsers + legacy_pdf: BaseParser = PDFParser() + if pdf_text_fast_path: + legacy_pdf = _wrap_pdf_fast_path(legacy_pdf) return { - ".pdf": PDFParser(), + ".pdf": legacy_pdf, ".docx": DocxParser(), ".csv": PandasCSVParser(), ".xlsx": ExcelParser(), diff --git a/application/parser/file/pdfium_parser.py b/application/parser/file/pdfium_parser.py new file mode 100644 index 00000000..a0ddb0cd --- /dev/null +++ b/application/parser/file/pdfium_parser.py @@ -0,0 +1,261 @@ +"""Direct PDF text-layer extraction with pypdfium2, for the attachment path. + +Attachments are read straight into a chat prompt; sources are chunked and +embedded for retrieval. That difference is what justifies two parsers. Docling +produces structured markdown — headings, table markup, image placeholders — +which is worth its cost when the output feeds chunking, and mostly wasted when +the output is concatenated into a prompt. + +The cost is not marginal. Measured on eight real uploads drawn from the p90 of +production attachment-ingest latency (docling 2.119.0, OCR off, warm process): + + docling total 155.6 s pypdfium2 total 0.33 s (~470x) + token yield within +/-15% per file, three files yielding *more* text + +What is lost is structure: docling emitted 20-187 headings and 44-88 table rows +per file where this parser emits none, and text arrives in content-stream order +rather than reading order. Acceptable for prompt context; not acceptable for +retrieval, which is why `get_default_file_extractor` keeps docling by default +and only swaps in this parser when a caller opts in. + +A scanned PDF has no text layer, so this parser would return an empty document +for exactly the files that need docling most. It therefore probes first and +delegates to `fallback_parser` when the text layer is too thin to be real, +leaving scanned-PDF behaviour unchanged. +""" + +import logging +import statistics +from pathlib import Path +from typing import Any, Dict, List, Optional, Union + +from application.parser.file.base_parser import BaseParser, DocumentParseError + +logger = logging.getLogger(__name__) + +# Pages sampled by the text-layer probe. Sampling a spread rather than the +# first N matters: scanned documents routinely carry a text cover page, and +# born-digital ones sometimes open with a full-page image. +DEFAULT_SAMPLE_PAGES = 16 + +# Median characters per sampled page below which the file is treated as having +# no usable text layer. Measured separation on the production corpus was stark: +# scanned files sat at 0 and 17 characters per page, text-layer files at +# 433-6834, so anything in the low hundreds separates them. 32 is deliberately +# conservative — it keeps sparse-but-real documents on docling. +DEFAULT_MIN_MEDIAN_CHARS = 32 + + +class PdfiumTextParser(BaseParser): + """Extract a PDF's embedded text layer with pypdfium2. + + Falls back to ``fallback_parser`` when the file has no usable text layer + (a scan), cannot be opened, or when pypdfium2 raises. The fallback is the + parser this one displaces, so delegating restores the previous behaviour + exactly rather than degrading it. + + Attributes: + fallback_parser: Parser used when the text layer is unusable. + min_median_chars: Median chars/page required to use the fast path. + sample_pages: Maximum number of pages inspected by the probe. + last_engine: Name of the engine that produced the most recent parse, + ``"pypdfium2"`` or the fallback parser's class name. + """ + + def __init__( + self, + fallback_parser: Optional[BaseParser] = None, + min_median_chars: int = DEFAULT_MIN_MEDIAN_CHARS, + sample_pages: int = DEFAULT_SAMPLE_PAGES, + parser_config: Optional[Dict] = None, + ) -> None: + """Initialize the parser. + + Args: + fallback_parser: Parser to delegate to when the text layer is + unusable. When ``None`` such files raise ``DocumentParseError``. + min_median_chars: Median characters per sampled page required to + take the fast path. + sample_pages: Maximum pages the probe inspects. + parser_config: Optional pre-built parser config. + """ + super().__init__(parser_config) + self.fallback_parser = fallback_parser + self.min_median_chars = min_median_chars + self.sample_pages = max(1, sample_pages) + self.last_engine: Optional[str] = None + self._last_metadata: Dict[str, Any] = {} + + def _init_parser(self) -> Dict: + """Initialize the parser config. + + Returns: + Dict: Static config describing the probe thresholds. + """ + return { + "engine": "pypdfium2", + "min_median_chars": self.min_median_chars, + "sample_pages": self.sample_pages, + } + + def _sample_indices(self, page_count: int) -> List[int]: + """Return a spread of page indices to probe.""" + if page_count <= self.sample_pages: + return list(range(page_count)) + step = (page_count - 1) / (self.sample_pages - 1) + return sorted({round(i * step) for i in range(self.sample_pages)}) + + def _probe_text_layer(self, file: Path) -> Dict[str, Any]: + """Measure how much extractable text the PDF carries. + + Args: + file: Path to the PDF. + + Returns: + Dict: ``page_count`` and ``median_chars`` across sampled pages. + + Raises: + Exception: Propagated from pypdfium2 when the file cannot be read. + """ + import pypdfium2 as pdfium + + pdf = pdfium.PdfDocument(str(file)) + try: + page_count = len(pdf) + counts = [] + for index in self._sample_indices(page_count): + page = pdf[index] + textpage = page.get_textpage() + try: + counts.append(textpage.count_chars()) + finally: + textpage.close() + page.close() + finally: + pdf.close() + return { + "page_count": page_count, + "median_chars": int(statistics.median(counts)) if counts else 0, + } + + def _extract_text(self, file: Path) -> str: + """Extract every page's text layer. + + Args: + file: Path to the PDF. + + Returns: + str: Page texts joined by blank lines. + """ + import pypdfium2 as pdfium + + pdf = pdfium.PdfDocument(str(file)) + pages: List[str] = [] + try: + for index in range(len(pdf)): + page = pdf[index] + textpage = page.get_textpage() + try: + pages.append(textpage.get_text_range()) + finally: + textpage.close() + page.close() + finally: + pdf.close() + return "\n\n".join(pages) + + def _delegate(self, file: Path, errors: str, reason: str) -> Union[str, List[str]]: + """Hand the file to the fallback parser. + + Args: + file: Path to the PDF. + errors: Decoding leniency passed through to the fallback. + reason: Why the fast path was skipped, for the log line. + + Returns: + The fallback parser's output. + + Raises: + DocumentParseError: When no fallback parser is configured. + """ + if self.fallback_parser is None: + raise DocumentParseError( + f"{file.name} has no usable PDF text layer ({reason}) and no " + "fallback parser is configured." + ) + logger.info( + "PDF text-layer fast path skipped for %s (%s); using %s", + file.name, + reason, + type(self.fallback_parser).__name__, + ) + if not self.fallback_parser.parser_config_set: + self.fallback_parser.init_parser() + self.last_engine = type(self.fallback_parser).__name__ + result = self.fallback_parser.parse_file(file, errors=errors) + self._last_metadata = { + **self.fallback_parser.get_file_metadata(file), + "parse_engine": self.last_engine, + } + return result + + def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]: + """Parse a PDF, preferring its embedded text layer. + + Args: + file: Path to the PDF. + errors: Decoding leniency, passed through to the fallback parser. + + Returns: + The extracted text, or the fallback parser's output. + + Raises: + DocumentParseError: When the text layer is unusable and no fallback + parser is configured. + """ + file = Path(file) + try: + probe = self._probe_text_layer(file) + except Exception as exc: # noqa: BLE001 - any read failure defers to docling + return self._delegate(file, errors, f"unreadable by pypdfium2: {exc}") + + if probe["page_count"] == 0: + return self._delegate(file, errors, "no pages") + if probe["median_chars"] < self.min_median_chars: + return self._delegate( + file, + errors, + f"median {probe['median_chars']} chars/page below " + f"{self.min_median_chars}", + ) + + try: + text = self._extract_text(file) + except Exception as exc: # noqa: BLE001 - as above + return self._delegate(file, errors, f"extraction failed: {exc}") + + self.last_engine = "pypdfium2" + self._last_metadata = { + "parse_engine": "pypdfium2", + "pdf_pages": probe["page_count"], + } + logger.info( + "Parsed %s via pypdfium2 text layer: %d pages, %d chars", + file.name, + probe["page_count"], + len(text), + ) + return text + + def get_file_metadata(self, file: Path) -> Dict: + """Return metadata for the most recently parsed file. + + Args: + file: Path to the PDF (unused; state is per-parse). + + Returns: + Dict: Includes ``parse_engine`` so a stored attachment records + which parser actually produced its content. + """ + _ = file + return dict(self._last_metadata) diff --git a/application/requirements.txt b/application/requirements.txt index 061648d0..bd368613 100644 --- a/application/requirements.txt +++ b/application/requirements.txt @@ -78,6 +78,7 @@ praw==8.0.2 pymilvus==3.0.1 milvus-lite==3.2.0; sys_platform != "win32" pypdf==6.15.0 +pypdfium2==5.12.1 qdrant-client==1.19.0 python-dateutil==2.9.0.post0 python-dotenv diff --git a/application/storage/db/repositories/attachments.py b/application/storage/db/repositories/attachments.py index eeddd5e9..3f3752f3 100644 --- a/application/storage/db/repositories/attachments.py +++ b/application/storage/db/repositories/attachments.py @@ -270,3 +270,38 @@ class AttachmentsRepository: params, ) return result.rowcount > 0 + + def update_metadata_if_content_null( + self, legacy_mongo_id: str, user_id: str, metadata: Any + ) -> bool: + """Failure-provenance write that refuses to touch stored content. + + The ``content IS NULL`` predicate makes the no-clobber guard atomic: + a concurrent success write (task redelivery, the poison guard racing + a still-live attempt) can commit between any prior check and this + statement, and the row it filled is left alone. + + Returns: + True when a NULL-content row was updated. + """ + if user_id is None: + return False + result = self._conn.execute( + text( + "UPDATE attachments SET metadata = CAST(:metadata AS jsonb) " + "WHERE legacy_mongo_id = :legacy_id AND user_id = :user_id " + "AND content IS NULL" + ), + { + "metadata": ( + json.dumps(strip_null_bytes(metadata)) + if metadata is not None + else None + ), + "legacy_id": ( + str(legacy_mongo_id) if legacy_mongo_id is not None else None + ), + "user_id": user_id, + }, + ) + return result.rowcount > 0 diff --git a/application/worker.py b/application/worker.py index 4f00da42..f4149c99 100755 --- a/application/worker.py +++ b/application/worker.py @@ -1594,15 +1594,6 @@ def record_attachment_failure(user, file_info, error, parser=None): if not attachment_id: return try: - # An exception after the success write (event publishing, state - # updates) must not replace durable extracted content with a - # NULL-content failed row. - with db_readonly() as conn: - existing = AttachmentsRepository(conn).get_by_legacy_id( - str(attachment_id), user - ) - if existing is not None and existing.get("content") is not None: - return metadata = { **(file_info.get("metadata") or {}), "extraction": { @@ -1612,16 +1603,27 @@ def record_attachment_failure(user, file_info, error, parser=None): "error": str(error)[:1024], }, } - _upsert_attachment_row( - user, - filename, - file_info.get("path") or "", - mime_type=mimetypes.guess_type(filename)[0] or "application/octet-stream", - content=None, - token_count=None, - metadata=metadata, - attachment_id=attachment_id, - ) + with db_session() as conn: + repo = AttachmentsRepository(conn) + # The conditional update carries the no-clobber guard in its own + # WHERE clause: a success row committed by a concurrent duplicate + # execution (broker redelivery, the poison guard racing a live + # attempt) or earlier in this attempt is never overwritten with a + # NULL-content failure. + if repo.update_metadata_if_content_null(str(attachment_id), user, metadata): + return + if repo.get_by_legacy_id(str(attachment_id), user) is not None: + return + repo.create( + user, + filename, + file_info.get("path") or "", + mime_type=mimetypes.guess_type(filename)[0] or "application/octet-stream", + content=None, + token_count=None, + metadata=metadata, + legacy_mongo_id=str(attachment_id), + ) except Exception: logging.error( f"Failed to record failure row for attachment {attachment_id}", @@ -1667,8 +1669,13 @@ def attachment_worker(self, file_info, user): scope={"kind": "attachment", "id": str(attachment_id)}, ) + # Attachments only: PDFs are read via their text layer where one + # exists (see ``ATTACHMENT_PDF_TEXT_FAST_PATH``). Source ingestion + # calls SimpleDirectoryReader without a ``file_extractor`` and so keeps + # the docling default, which is what retrieval quality depends on. file_extractor = get_default_file_extractor( - ocr_enabled=settings.DOCLING_OCR_ATTACHMENTS_ENABLED + ocr_enabled=settings.DOCLING_OCR_ATTACHMENTS_ENABLED, + pdf_text_fast_path=settings.ATTACHMENT_PDF_TEXT_FAST_PATH, ) _parser = file_extractor.get(os.path.splitext(filename)[1].lower()) parser_name = type(_parser).__name__ if _parser is not None else "SimpleDirectoryReader" @@ -1693,6 +1700,9 @@ def attachment_worker(self, file_info, user): attachment_document = storage.process_file(relative_path, _parse_local_file) content = attachment_document.text + # A fast-path parser may have delegated to its fallback for this file, + # so record the engine that actually ran rather than the one selected. + parser_name = getattr(_parser, "last_engine", None) or parser_name parser_metadata = { key: value for key, value in (attachment_document.extra_info or {}).items() diff --git a/tests/parser/file/test_pdfium_parser.py b/tests/parser/file/test_pdfium_parser.py new file mode 100644 index 00000000..62e92491 --- /dev/null +++ b/tests/parser/file/test_pdfium_parser.py @@ -0,0 +1,184 @@ +"""Tests for the pypdfium2 text-layer PDF parser used on the attachment path. + +Two behaviours matter here and are asserted separately: + +* a PDF with a real text layer is read directly by pypdfium2, skipping docling + entirely (this is the whole point — docling costs tens of seconds per file); +* a PDF with no text layer (a scan) is handed to the fallback parser rather + than returning an empty document, so the scanned-PDF path is unchanged. + +The extractor-wiring test is the guard that source ingestion keeps docling: +the fast path must only appear when a caller explicitly asks for it. +""" + +import pytest + +from application.parser.file.base_parser import BaseParser, DocumentParseError + +pypdfium2 = pytest.importorskip("pypdfium2") + + +def _text_pdf(path, pages=3, text="The quick brown fox jumps over the lazy dog."): + """Write a PDF that has a genuine text layer.""" + from reportlab.lib.pagesizes import letter + from reportlab.pdfgen import canvas + + c = canvas.Canvas(str(path), pagesize=letter) + for page in range(pages): + # Several lines per page so the per-page character count clears the + # probe threshold the way a real document does. + for line in range(12): + c.drawString(72, 720 - line * 18, f"page {page} line {line} {text}") + c.showPage() + c.save() + return path + + +def _scanned_pdf(path, pages=3): + """Write a PDF with pages but no text layer, i.e. what a scan looks like.""" + from pypdf import PdfWriter + + writer = PdfWriter() + for _ in range(pages): + writer.add_blank_page(width=612, height=792) + with open(path, "wb") as fh: + writer.write(fh) + return path + + +class _RecordingFallback(BaseParser): + """Stands in for the docling parser so delegation is observable.""" + + def __init__(self): + super().__init__(parser_config={}) + self.calls = [] + + def _init_parser(self): + return {} + + def parse_file(self, file, errors="ignore"): + self.calls.append(file) + return "FALLBACK OUTPUT" + + def get_file_metadata(self, file): + return {"fallback_meta": True} + + +@pytest.fixture +def parser_cls(): + from application.parser.file.pdfium_parser import PdfiumTextParser + + return PdfiumTextParser + + +def test_text_layer_pdf_is_read_by_pypdfium2(tmp_path, parser_cls): + pdf = _text_pdf(tmp_path / "text.pdf") + fallback = _RecordingFallback() + parser = parser_cls(fallback_parser=fallback) + parser.init_parser() + + out = parser.parse_file(pdf) + + assert "quick brown fox" in out + assert "page 2" in out, "every page should be extracted, not just the first" + assert fallback.calls == [], "a text-layer PDF must not reach docling" + assert parser.last_engine == "pypdfium2" + + +def test_scanned_pdf_delegates_to_fallback(tmp_path, parser_cls): + pdf = _scanned_pdf(tmp_path / "scan.pdf") + fallback = _RecordingFallback() + parser = parser_cls(fallback_parser=fallback) + parser.init_parser() + + out = parser.parse_file(pdf) + + assert out == "FALLBACK OUTPUT" + assert fallback.calls == [pdf] + assert parser.last_engine == "_RecordingFallback" + + +def test_scanned_pdf_without_fallback_raises(tmp_path, parser_cls): + pdf = _scanned_pdf(tmp_path / "scan.pdf") + parser = parser_cls(fallback_parser=None) + parser.init_parser() + + with pytest.raises(DocumentParseError): + parser.parse_file(pdf) + + +def test_unopenable_pdf_delegates_to_fallback(tmp_path, parser_cls): + pdf = tmp_path / "broken.pdf" + pdf.write_bytes(b"not a pdf at all") + fallback = _RecordingFallback() + parser = parser_cls(fallback_parser=fallback) + parser.init_parser() + + assert parser.parse_file(pdf) == "FALLBACK OUTPUT" + assert fallback.calls == [pdf] + + +def test_metadata_reports_engine_and_page_count(tmp_path, parser_cls): + pdf = _text_pdf(tmp_path / "text.pdf", pages=4) + parser = parser_cls(fallback_parser=_RecordingFallback()) + parser.init_parser() + parser.parse_file(pdf) + + meta = parser.get_file_metadata(pdf) + + assert meta["parse_engine"] == "pypdfium2" + assert meta["pdf_pages"] == 4 + + +def test_fallback_metadata_is_passed_through(tmp_path, parser_cls): + pdf = _scanned_pdf(tmp_path / "scan.pdf") + parser = parser_cls(fallback_parser=_RecordingFallback()) + parser.init_parser() + parser.parse_file(pdf) + + meta = parser.get_file_metadata(pdf) + + assert meta["fallback_meta"] is True + assert meta["parse_engine"] == "_RecordingFallback" + + +def test_threshold_sends_sparse_text_to_fallback(tmp_path, parser_cls): + """A PDF whose text layer is a stray character is a scan, not a document.""" + pdf = _text_pdf(tmp_path / "sparse.pdf", pages=2, text="") + fallback = _RecordingFallback() + parser = parser_cls(fallback_parser=fallback, min_median_chars=10_000) + parser.init_parser() + + assert parser.parse_file(pdf) == "FALLBACK OUTPUT" + + +# --- extractor wiring: the guard that source ingestion keeps docling --------- + + +def test_extractor_defaults_to_docling_for_pdf(): + """Sources must be unaffected: no fast path unless explicitly requested.""" + from application.parser.file.bulk import get_default_file_extractor + + pdf_parser = get_default_file_extractor()[".pdf"] + + assert type(pdf_parser).__name__ == "DoclingPDFParser" + + +def test_extractor_uses_fast_path_when_requested(): + from application.parser.file.bulk import get_default_file_extractor + from application.parser.file.pdfium_parser import PdfiumTextParser + + pdf_parser = get_default_file_extractor(pdf_text_fast_path=True)[".pdf"] + + assert isinstance(pdf_parser, PdfiumTextParser) + assert type(pdf_parser.fallback_parser).__name__ == "DoclingPDFParser" + + +def test_fast_path_does_not_change_non_pdf_parsers(): + from application.parser.file.bulk import get_default_file_extractor + + plain = get_default_file_extractor() + fast = get_default_file_extractor(pdf_text_fast_path=True) + + for suffix in (".docx", ".xlsx", ".csv", ".html"): + assert type(fast[suffix]).__name__ == type(plain[suffix]).__name__ diff --git a/tests/storage/db/repositories/test_attachments.py b/tests/storage/db/repositories/test_attachments.py index ad2a2626..c9373b02 100644 --- a/tests/storage/db/repositories/test_attachments.py +++ b/tests/storage/db/repositories/test_attachments.py @@ -115,3 +115,43 @@ class TestListForUser: repo = _repo(pg_conn) results = repo.list_for_user("nonexistent") assert results == [] + + +class TestUpdateMetadataIfContentNull: + """Atomic no-clobber write for failure provenance: the ``content IS + NULL`` predicate lives in the UPDATE itself, so a success row committed + by a concurrent duplicate execution can never be overwritten between a + check and the write.""" + + def test_updates_row_without_stored_content(self, pg_conn): + repo = _repo(pg_conn) + repo.create("u", "f.xlsx", "/p", content=None, + metadata={"extraction": {"status": "failed", "error": "first"}}, + legacy_mongo_id="handle-1") + + updated = repo.update_metadata_if_content_null( + "handle-1", "u", {"extraction": {"status": "failed", "error": "second"}} + ) + + assert updated is True + row = repo.get_by_legacy_id("handle-1", "u") + assert row["metadata"]["extraction"]["error"] == "second" + + def test_refuses_row_with_stored_content(self, pg_conn): + repo = _repo(pg_conn) + repo.create("u", "f.txt", "/p", content="durable text", token_count=2, + metadata={"extraction": {"status": "ok"}}, + legacy_mongo_id="handle-2") + + updated = repo.update_metadata_if_content_null( + "handle-2", "u", {"extraction": {"status": "failed", "error": "late loser"}} + ) + + assert updated is False + row = repo.get_by_legacy_id("handle-2", "u") + assert row["metadata"]["extraction"]["status"] == "ok" + assert row["content"] == "durable text" + + def test_missing_row_returns_false(self, pg_conn): + repo = _repo(pg_conn) + assert repo.update_metadata_if_content_null("nope", "u", {"x": 1}) is False diff --git a/tests/worker/test_attachment_worker.py b/tests/worker/test_attachment_worker.py index 4172a944..93e34d97 100644 --- a/tests/worker/test_attachment_worker.py +++ b/tests/worker/test_attachment_worker.py @@ -38,7 +38,7 @@ class TestAttachmentWorker: # Stub the parser selection so the docling import path isn't taken. monkeypatch.setattr( - worker, "get_default_file_extractor", lambda ocr_enabled=False: {} + worker, "get_default_file_extractor", lambda ocr_enabled=False, pdf_text_fast_path=False: {} ) file_info = { @@ -99,7 +99,7 @@ class TestAttachmentWorker: worker.StorageCreator, "get_storage", lambda: fake_storage ) monkeypatch.setattr( - worker, "get_default_file_extractor", lambda ocr_enabled=False: {} + worker, "get_default_file_extractor", lambda ocr_enabled=False, pdf_text_fast_path=False: {} ) file_info = { diff --git a/tests/worker/test_worker_sse_publishes.py b/tests/worker/test_worker_sse_publishes.py index 5dd9f107..5b2fe37e 100644 --- a/tests/worker/test_worker_sse_publishes.py +++ b/tests/worker/test_worker_sse_publishes.py @@ -598,7 +598,7 @@ class TestAttachmentWorkerPublishes: worker.StorageCreator, "get_storage", lambda: fake_storage ) monkeypatch.setattr( - worker, "get_default_file_extractor", lambda ocr_enabled=False: {} + worker, "get_default_file_extractor", lambda ocr_enabled=False, pdf_text_fast_path=False: {} ) file_info = { @@ -654,7 +654,7 @@ class TestAttachmentWorkerPublishes: worker.StorageCreator, "get_storage", lambda: fake_storage ) monkeypatch.setattr( - worker, "get_default_file_extractor", lambda ocr_enabled=False: {} + worker, "get_default_file_extractor", lambda ocr_enabled=False, pdf_text_fast_path=False: {} ) file_info = {