mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-12 14:14:50 +00:00
feat: faster attachment processing
This commit is contained in:
1 parent
0a15ce8fbb
commit
7ad2f38512
10 files changed
+611
-26
No files matched your search
@@ -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
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
+30
-20
@@ -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()
|
||||
|
||||
@@ -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__
|
||||
@@ -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
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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 = {
|
||||
|
||||
Reference in new issue
Block a user