mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 08:17:34 +00:00
feat: attachment provenance
This commit is contained in:
1 parent
daf8a5aa26
commit
0a15ce8fbb
15 files changed
+659
-36
No files matched your search
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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()})
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
+129
-18
@@ -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",
|
||||
|
||||
@@ -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
|
||||
@@ -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):
|
||||
|
||||
@@ -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])
|
||||
@@ -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"
|
||||
@@ -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:
|
||||
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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"]
|
||||
Reference in new issue
Block a user