diff --git a/apps/worker/app/services/document_agent/calibration/phase1.py b/apps/worker/app/services/document_agent/calibration/phase1.py index a651c90d6..176e74434 100644 --- a/apps/worker/app/services/document_agent/calibration/phase1.py +++ b/apps/worker/app/services/document_agent/calibration/phase1.py @@ -5,10 +5,15 @@ and scans forward until one is confirmed; that single confirmation fixes the regime's candidate offset. Phase-2 owns tail verification and bulk anchoring. -A probe scans from ``max(toc_range end + 1, printed)``: a printed label never -resolves to a physical page before itself, so the offset a scan can yield is -structurally non-negative. Scanning below the printed page would let a section -divider that repeats the heading confirm ahead of the numbered body page. +A probe scans from ``toc_range end + 1`` (body after the TOC). That can yield a +negative ``physical - printed`` offset (e.g. journal reprints whose printed +labels exceed the PDF length). Phase-2 already accepts negative offsets and +prunes ``printed + offset`` outside ``1..page_count``. + +To restore the Sydney anti-false-positive floor (non-negative offset), switch +``start_page`` back to ``max(region_scan_start, probe.printed)`` below and +invert ``test_scan_starts_after_toc_so_negative_offset_is_allowed`` to expect +scan start ``= printed`` and offset ``0``. """ from __future__ import annotations @@ -150,7 +155,10 @@ def run_calibration_phase1( scan = scan_title_forward( ctx=ctx, title=probe.title, - start_page=max(region_scan_start, probe.printed), + # Scan from after TOC. Optional floor (disabled): force + # non-negative offset / skip early divider hits — restore with: + # start_page=max(region_scan_start, probe.printed), + start_page=region_scan_start, page_count=resolved_page_count, ) scans.append(scan) diff --git a/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py b/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py index db48fae36..0d991453b 100644 --- a/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py +++ b/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py @@ -19,7 +19,6 @@ ) from app.services.document_agent.registry import register_tool from app.services.document_agent.tools.vlm_toc_extractor import ( - TOC_VLM_MAX_TOKENS, BatchPageResult, vlm_entries_to_toc_hierarchies, ) @@ -35,6 +34,8 @@ TOC_VLM_CONCURRENCY = 10 MAX_BOUNDARY_ROUNDS = 6 MAX_TOC_PAGES = BOUNDARY_STEP_PAGES * MAX_BOUNDARY_ROUNDS # 30 +# Confirm only returns is_toc_start + brief reason per page (≤ BOUNDARY_STEP_PAGES). +TOC_ANCHOR_CONFIRM_MAX_TOKENS = 512 _CONFIRM_PROMPT = ( "You are a document structure analysis expert. " @@ -183,7 +184,7 @@ def _confirm_anchor_chunk( messages=messages, model=resolved, temperature=0.1, - max_tokens=TOC_VLM_MAX_TOKENS, + max_tokens=TOC_ANCHOR_CONFIRM_MAX_TOKENS, response_format={"type": "json_object"}, usage_task="document_agent.toc_anchor_confirm", ) diff --git a/apps/worker/app/services/document_agent/tools/probe_page_features.py b/apps/worker/app/services/document_agent/tools/probe_page_features.py index 7369da44e..f9af0ed14 100644 --- a/apps/worker/app/services/document_agent/tools/probe_page_features.py +++ b/apps/worker/app/services/document_agent/tools/probe_page_features.py @@ -29,6 +29,11 @@ # Near-full-page sparse stroke frames are treated as borders, not figures. _FIGURE_FULLPAGE_AREA_RATIO = 0.92 _FIGURE_FULLPAGE_MAX_PATHS = 25 +# Skip find_tables + drawing clustering on vector/image-dense pages. +# Clustering ~1e5 strokes is quadratic in a crowded grid cell and exceeds +# the 300s probe child timeout (EAS lipoprotein consensus PDF page 2). +_DENSE_DRAWING_LIMIT = 5000 +_DENSE_IMAGE_LIMIT = 200 # MuPDF ``page.get_texttrace()`` type matching PDF text rendering mode ``3 Tr``. _INVISIBLE_TEXT_TRACE_TYPE = 3 @@ -257,6 +262,10 @@ def _figure_bboxes_from_drawings( return figures +def _is_dense_visual_page(*, image_count: int, drawings_count: int) -> bool: + return image_count > _DENSE_IMAGE_LIMIT or drawings_count > _DENSE_DRAWING_LIMIT + + def _probe_visual_assets( page: Any, page_area: float, @@ -295,42 +304,46 @@ def _probe_visual_assets( image_area += max(box[2] - box[0], 0.0) * max(box[3] - box[1], 0.0) bboxes.append({"kind": "image", "bbox": box}) - table_count = 0 - try: - finder = page.find_tables() - tables = getattr(finder, "tables", []) or [] - table_count = len(tables) - for table in tables: - raw = getattr(table, "bbox", None) - if raw is None: - continue - box = _clip_bbox( - float(raw[0]), - float(raw[1]), - float(raw[2]), - float(raw[3]), - clip=page_rect, - ) - if _valid_bbox(box): - bboxes.append({"kind": "table", "bbox": box}) - except Exception: - table_count = 0 - try: drawings = page.get_drawings() or [] except Exception: drawings = [] drawings_count = len(drawings) - bboxes.extend( - _figure_bboxes_from_drawings( - drawings, - page_rect=page_rect, - page_area=page_area, - header_y=header_y, - footer_y=footer_y, - ) + dense_visual = _is_dense_visual_page( + image_count=image_count, drawings_count=drawings_count ) + table_count = 0 + if not dense_visual: + try: + finder = page.find_tables() + tables = getattr(finder, "tables", []) or [] + table_count = len(tables) + for table in tables: + raw = getattr(table, "bbox", None) + if raw is None: + continue + box = _clip_bbox( + float(raw[0]), + float(raw[1]), + float(raw[2]), + float(raw[3]), + clip=page_rect, + ) + if _valid_bbox(box): + bboxes.append({"kind": "table", "bbox": box}) + except Exception: + table_count = 0 + bboxes.extend( + _figure_bboxes_from_drawings( + drawings, + page_rect=page_rect, + page_area=page_area, + header_y=header_y, + footer_y=footer_y, + ) + ) + coverage = min(image_area / page_area, 1.0) if page_area > 0 else 0.0 return { "image_coverage": round(coverage, 4), @@ -338,7 +351,7 @@ def _probe_visual_assets( "table_count": table_count, "drawings_count": drawings_count, # Geometry is only used to derive the gate; do not persist bboxes. - "has_asset": bool(bboxes), + "has_asset": bool(bboxes) or dense_visual, } diff --git a/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py b/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py index cc331b290..eeed005bc 100644 --- a/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py +++ b/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py @@ -8,8 +8,13 @@ from dataclasses import dataclass from typing import Any, cast -# Shared completion budget for TOC VLM calls (confirm batches + extract batches). -TOC_VLM_MAX_TOKENS = 8192 +# Completion budget for TOC VLM calls: known pages in the call × per-page cap. +TOC_VLM_MAX_TOKENS_PER_PAGE = 3000 + + +def toc_vlm_max_tokens(page_count: int) -> int: + """Return ``max_tokens`` for a TOC VLM call covering ``page_count`` pages.""" + return max(1, int(page_count)) * TOC_VLM_MAX_TOKENS_PER_PAGE # --------------------------------------------------------------------------- @@ -199,7 +204,7 @@ def vlm_extract_toc_batch( messages=cast(Any, [{"role": "user", "content": content_parts}]), model=model, temperature=0.1, - max_tokens=TOC_VLM_MAX_TOKENS, + max_tokens=toc_vlm_max_tokens(len(page_pngs)), response_format={"type": "json_object"}, usage_task="document_agent.vlm_toc_batch", ) diff --git a/apps/worker/app/services/document_ingestion/processing_billing.py b/apps/worker/app/services/document_ingestion/processing_billing.py index 6fd6cc321..51c658ea6 100644 --- a/apps/worker/app/services/document_ingestion/processing_billing.py +++ b/apps/worker/app/services/document_ingestion/processing_billing.py @@ -4,7 +4,10 @@ from datetime import datetime from app.services.document_ingestion.page_estimator import WorkloadEstimate -from app.services.document_ingestion.processing_context import ParseJobContext +from app.services.document_ingestion.processing_context import ( + ParseJobContext, + persist_job_metadata_updates, +) from loguru import logger from sqlalchemy import select @@ -145,13 +148,8 @@ def record_processing_start( ) if extra_metadata is not None: metadata_updates.update(extra_metadata) - with get_sync_db_context() as db: - job_result = db.execute(select(Job).where(Job.job_id == job_id).with_for_update()) - job = job_result.scalar_one_or_none() - if job: - job.job_metadata = { - **dict(job.job_metadata or {}), - **metadata_updates, - } - job_context.metadata_service.update_metadata(job_id, metadata_updates) - job_context.job_metadata.update(metadata_updates) + persist_job_metadata_updates( + job_id=job_id, + job_context=job_context, + metadata_updates=metadata_updates, + ) diff --git a/apps/worker/app/services/document_ingestion/processing_context.py b/apps/worker/app/services/document_ingestion/processing_context.py index 9dfe87945..0d8e51728 100644 --- a/apps/worker/app/services/document_ingestion/processing_context.py +++ b/apps/worker/app/services/document_ingestion/processing_context.py @@ -27,6 +27,31 @@ class ParseJobContext: s3_key: str +def persist_job_metadata_updates( + *, + job_id: str, + job_context: ParseJobContext, + metadata_updates: dict[str, object], +) -> None: + """Merge metadata into Redis and the durable Job row. + + Redis remains the live working copy during processing. The Job row is + the durable record after Redis TTL expires, so token usage and other + stage fields must be written here as well. + """ + if not metadata_updates: + return + with get_sync_db_context() as db: + job = _select_job_row_for_update(db, job_id) + if job is not None: + job.job_metadata = { + **dict(job.job_metadata or {}), + **metadata_updates, + } + job_context.metadata_service.update_metadata(job_id, metadata_updates) + job_context.job_metadata.update(metadata_updates) + + def load_parse_job_context( job_id: str, requested_user_id: str | None, @@ -111,3 +136,9 @@ def _load_job_row(job_id: str) -> Job | None: def _select_job_row(db: Session, job_id: str) -> Job | None: return db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() + + +def _select_job_row_for_update(db: Session, job_id: str) -> Job | None: + return db.execute( + select(Job).where(Job.job_id == job_id).with_for_update() + ).scalar_one_or_none() diff --git a/apps/worker/app/services/document_ingestion/processing_run.py b/apps/worker/app/services/document_ingestion/processing_run.py index cd5bcfce0..4a35f645b 100644 --- a/apps/worker/app/services/document_ingestion/processing_run.py +++ b/apps/worker/app/services/document_ingestion/processing_run.py @@ -17,6 +17,7 @@ from app.services.document_ingestion.processing_context import ( ParseJobContext, load_parse_job_context, + persist_job_metadata_updates, ) from app.services.document_ingestion.source_preparation import prepare_source_file from app.services.document_ingestion.success_finalization import finalize_parse_success @@ -199,10 +200,21 @@ def _run_parse_job( result_storage_factory=get_result_storage, ) finally: - job_context.job_metadata["stages"] = { + stages = { "timing_ms": dict(stage_timing_dict), "token_usage": dict(token_usage_dict), } + job_context.job_metadata["stages"] = stages + try: + persist_job_metadata_updates( + job_id=job_id, + job_context=job_context, + metadata_updates={"stages": stages}, + ) + except Exception as exc: + logger.warning( + f"Failed to persist processing stages: job_id={job_id}, error={exc}" + ) cleanup_llm_overrides() cleanup_token_tracker() cleanup_stage_tracker() diff --git a/apps/worker/app/services/document_ingestion/success_finalization.py b/apps/worker/app/services/document_ingestion/success_finalization.py index 9acc0e3ab..42330dbcc 100644 --- a/apps/worker/app/services/document_ingestion/success_finalization.py +++ b/apps/worker/app/services/document_ingestion/success_finalization.py @@ -18,7 +18,10 @@ build_generated_result_package, ) from app.services.document_ingestion.artifact_refs import collect_referenced_artifact_refs -from app.services.document_ingestion.processing_context import ParseJobContext +from app.services.document_ingestion.processing_context import ( + ParseJobContext, + persist_job_metadata_updates, +) from loguru import logger from shared.models.schemas.job_metadata import JobMetadataHelper @@ -210,8 +213,11 @@ def _record_processing_completion( _refresh_processing_stages(job_context) if "stages" in job_context.job_metadata: processing_timing_updates["stages"] = job_context.job_metadata["stages"] - job_context.metadata_service.update_metadata(job_id, processing_timing_updates) - job_context.job_metadata.update(processing_timing_updates) + persist_job_metadata_updates( + job_id=job_id, + job_context=job_context, + metadata_updates=processing_timing_updates, + ) def _generate_result_package( diff --git a/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py b/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py index fdcf64d35..0964825b0 100644 --- a/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py +++ b/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py @@ -43,6 +43,8 @@ CHILD_EXIT_GRACE_SECONDS = 5 POST_RESULT_EXIT_GRACE_SECONDS = 5 POST_KILL_JOIN_GRACE_SECONDS = 1 +# Child can exit 0 before multiprocessing.Queue's feeder flushes; retry once. +EMPTY_QUEUE_EXIT_RETRIES = 1 PROCESS_POOL_SIZE = read_pymupdf_max_concurrent() PROCESS_POOL_CONTEXT = multiprocessing.get_context("spawn") @@ -130,10 +132,66 @@ def _close_result_queue(result_queue: MultiprocessingQueue) -> None: logger.debug(f"Failed to join PyMuPDF result queue thread: {exc}") +def _is_empty_queue_exit(exc: Exception) -> bool: + if not isinstance(exc, PDFParsingException): + return False + if exc.details.get("reason") != "SUBPROCESS_CRASH": + return False + return "exited with code=0 and no result" in (exc.internal_message or "") + + +def _flush_child_result_queue(queue: object) -> None: + """Block until the child's queue feeder has written the payload. + + ``Queue.put()`` returns after buffering on a feeder thread. If the child + process exits before that thread flushes, the parent sees exitcode=0 and + an empty queue (production CRASH on otherwise-healthy PDFs). + """ + close = getattr(queue, "close", None) + if not callable(close): + return + try: + close() + except Exception: + return + join_thread = getattr(queue, "join_thread", None) + if not callable(join_thread): + return + try: + join_thread() + except Exception as exc: + logger.debug(f"Failed to join PyMuPDF child result queue feeder thread: {exc}") + + def _run_worker_in_spawned_process( worker_fn, args: tuple, timeout: int, +) -> dict: + """Spawn an isolated child, retrying once if the result queue is dropped.""" + last_empty_exit: PDFParsingException | None = None + attempts = 1 + EMPTY_QUEUE_EXIT_RETRIES + for attempt in range(1, attempts + 1): + try: + return _run_worker_in_spawned_process_once(worker_fn, args, timeout) + except PDFParsingException as exc: + if attempt < attempts and _is_empty_queue_exit(exc): + logger.warning( + f"[pymupdf-subprocess] retrying fn={worker_fn.__name__} " + f"after empty-queue exit 0 (attempt {attempt}/{attempts})" + ) + last_empty_exit = exc + continue + raise + if last_empty_exit is not None: + raise last_empty_exit + raise RuntimeError("pymupdf child retry loop exited without a result") + + +def _run_worker_in_spawned_process_once( + worker_fn, + args: tuple, + timeout: int, ) -> dict: """Spawn one isolated child process, but only after a pooled slot is available.""" ctx = PROCESS_POOL_CONTEXT @@ -267,6 +325,8 @@ def wrapped(queue, *args): "error_msg": str(exc), } ) + finally: + _flush_child_result_queue(queue) return wrapped diff --git a/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py b/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py index a4f323d42..d13565813 100644 --- a/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py +++ b/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py @@ -12,6 +12,9 @@ if TYPE_CHECKING: from app.services.document_agent.manifest import Shard +_SOURCE_OBJECT_OUT_OF_RANGE = "source object number out of range" +_REWRITTEN_SOURCE_NAME = "_rewritten_source.pdf" + @dataclass class MergedShard: @@ -46,6 +49,11 @@ def map_agent_shards( ] +def _is_source_object_out_of_range(exc: BaseException) -> bool: + """True only for MuPDF xref copy failure during insert_pdf.""" + return isinstance(exc, RuntimeError) and _SOURCE_OBJECT_OUT_OF_RANGE in str(exc) + + def split_pdf( pdf_path: str, shards: list[MergedShard], @@ -69,6 +77,36 @@ def split_pdf( ``None`` when no pages are excluded. """ doc = pymupdf.open(pdf_path) + rewritten_path: str | None = None + try: + try: + return _extract_shards(doc, shards, work_dir, exclude_pages) + except RuntimeError as exc: + if not _is_source_object_out_of_range(exc): + raise + rewritten_path = os.path.join(work_dir, _REWRITTEN_SOURCE_NAME) + logger.warning( + "PDF split hit source object number out of range; " + "rewriting a clean copy and retrying once" + ) + doc.save(rewritten_path, garbage=4, deflate=True) + finally: + doc.close() + + assert rewritten_path is not None + rewritten = pymupdf.open(rewritten_path) + try: + return _extract_shards(rewritten, shards, work_dir, exclude_pages) + finally: + rewritten.close() + + +def _extract_shards( + doc: pymupdf.Document, + shards: list[MergedShard], + work_dir: str, + exclude_pages: set[int] | None, +) -> tuple[list[str], dict[int, int] | None]: paths: list[str] = [] page_remap: dict[int, int] | None = None @@ -79,10 +117,10 @@ def split_pdf( f"{sorted(exclude_pages)}" ) - try: - global_new_idx = 0 # running counter across all shards - for shard in shards: - sub_doc = pymupdf.open() + global_new_idx = 0 # running counter across all shards + for shard in shards: + sub_doc = pymupdf.open() + try: shard_included = 0 for page_num in range(shard.page_start, shard.page_end + 1): if exclude_pages and page_num in exclude_pages: @@ -105,15 +143,14 @@ def split_pdf( logger.warning( f" ⚠️ shard_{shard.shard_index}: all pages excluded, skipping" ) + finally: sub_doc.close() - excluded_in_shard = shard.page_count - shard_included - logger.info( - f" ✂️ shard_{shard.shard_index}: " - f"pages {shard.page_start}-{shard.page_end} " - f"({shard_included} included" - f"{f', {excluded_in_shard} excluded' if excluded_in_shard else ''})" - ) - finally: - doc.close() + excluded_in_shard = shard.page_count - shard_included + logger.info( + f" ✂️ shard_{shard.shard_index}: " + f"pages {shard.page_start}-{shard.page_end} " + f"({shard_included} included" + f"{f', {excluded_in_shard} excluded' if excluded_in_shard else ''})" + ) return paths, page_remap diff --git a/apps/worker/tests/contract/test_calibration_phase1_contract.py b/apps/worker/tests/contract/test_calibration_phase1_contract.py index fe649f35e..2cb747a69 100644 --- a/apps/worker/tests/contract/test_calibration_phase1_contract.py +++ b/apps/worker/tests/contract/test_calibration_phase1_contract.py @@ -96,8 +96,8 @@ def test_offset_is_found_page_minus_printed(patch_scan) -> None: assert result.status == "ok" assert [(r.kind, r.offset) for r in result.regimes] == [("decimal", 5)] - # toc_range=[1, 3], printed=10 → scan starts at the printed page. - assert fake.calls == [("Chapter 1", 10)] + # toc_range=[1, 3] → scan starts at toc_end + 1 (not printed). + assert fake.calls == [("Chapter 1", 4)] def test_first_hit_stops_the_regime(patch_scan) -> None: @@ -111,7 +111,7 @@ def test_first_hit_stops_the_regime(patch_scan) -> None: run_calibration_phase1(ctx=_ctx(), toc_hierarchies=hierarchies, page_count=60) - assert fake.calls == [("Chapter 1", 10)] + assert fake.calls == [("Chapter 1", 4)] def test_second_probe_runs_when_the_first_misses(patch_scan) -> None: @@ -128,7 +128,7 @@ def test_second_probe_runs_when_the_first_misses(patch_scan) -> None: ctx=_ctx(), toc_hierarchies=hierarchies, page_count=60 ) - assert fake.calls == [("Chapter 1", 10), ("Chapter 2", 20)] + assert fake.calls == [("Chapter 1", 4), ("Chapter 2", 4)] assert [r.offset for r in result.regimes] == [5] @@ -146,7 +146,7 @@ def test_probes_use_distinct_printed_pages(patch_scan) -> None: ctx=_ctx(), toc_hierarchies=hierarchies, page_count=60 ) - assert fake.calls == [("Chapter 1", 10), ("Chapter 2", 20)] + assert fake.calls == [("Chapter 1", 4), ("Chapter 2", 4)] assert [r.offset for r in result.regimes] == [5] @@ -164,7 +164,7 @@ def test_probe_prefers_leaf_within_a_printed_page(patch_scan) -> None: ctx=_ctx(), toc_hierarchies=hierarchies, page_count=60 ) - assert fake.calls == [("A1 Purpose", 10)] + assert fake.calls == [("A1 Purpose", 4)] assert [r.offset for r in result.regimes] == [5] @@ -205,14 +205,18 @@ def test_roman_and_decimal_regimes_calibrate_independently(patch_scan) -> None: ("roman", 2), ("decimal", 5), } - # printed=2 sits inside the TOC range → floor at toc end + 1; printed=10 wins. - assert fake.calls == [("Preface", 4), ("Chapter 1", 10)] + # Both regimes scan from toc_end + 1 (printed floor disabled). + assert fake.calls == [("Preface", 4), ("Chapter 1", 4)] -def test_scan_floor_is_the_printed_page_so_offset_is_never_negative( +def test_scan_starts_after_toc_so_negative_offset_is_allowed( monkeypatch: pytest.MonkeyPatch, ) -> None: - """A divider repeating the heading cannot confirm ahead of the printed page.""" + """Confirming on the first post-TOC page can yield physical < printed. + + Restore ``max(region_scan_start, probe.printed)`` in phase1 to bring back + ``test_scan_floor_is_the_printed_page_so_offset_is_never_negative`` behavior. + """ class _ConfirmFirstPage: def __init__(self) -> None: @@ -246,8 +250,8 @@ def __call__( ctx=_ctx(), toc_hierarchies=hierarchies, page_count=60 ) - assert fake.calls == [("Chapter 1", 10)] - assert [r.offset for r in result.regimes] == [0] + assert fake.calls == [("Chapter 1", 4)] + assert [r.offset for r in result.regimes] == [-6] def test_confirmed_anchor_is_reported_as_a_sample(patch_scan) -> None: @@ -275,7 +279,7 @@ def test_entries_without_a_parseable_printed_page_are_skipped(patch_scan) -> Non run_calibration_phase1(ctx=_ctx(), toc_hierarchies=hierarchies, page_count=60) - assert fake.calls == [("Chapter 1", 10)] + assert fake.calls == [("Chapter 1", 4)] def test_empty_toc_fails_without_scanning(patch_scan) -> None: diff --git a/apps/worker/tests/unit/test_pdf_split_object_out_of_range.py b/apps/worker/tests/unit/test_pdf_split_object_out_of_range.py new file mode 100644 index 000000000..0b0fd35a3 --- /dev/null +++ b/apps/worker/tests/unit/test_pdf_split_object_out_of_range.py @@ -0,0 +1,131 @@ +"""PDF split retries once after MuPDF xref copy failure.""" + +from __future__ import annotations + +from pathlib import Path + +import pymupdf +import pytest + +from app.services.document_parser.formats.pdf.shard_splitter import ( + MergedShard, + _is_source_object_out_of_range, + split_pdf, +) + +_OBJECT_OUT_OF_RANGE = RuntimeError("code=4: source object number out of range") +_CUSTOMER_PDF = Path( + "/home/suguan/.cursor/Exercise_Prescription_in_Cardiac_Rehabilitation.pdf" +) + + +def _write_pages(path: Path, page_count: int) -> None: + doc = pymupdf.open() + for index in range(page_count): + page = doc.new_page(width=72, height=72) + page.insert_text((12, 24), f"page-{index + 1}") + doc.save(path) + doc.close() + + +def _page_count(path: str) -> int: + doc = pymupdf.open(path) + try: + return doc.page_count + finally: + doc.close() + + +def test_source_object_matcher_is_specific() -> None: + assert _is_source_object_out_of_range(_OBJECT_OUT_OF_RANGE) is True + assert ( + _is_source_object_out_of_range(RuntimeError("code=7: syntax error")) is False + ) + assert _is_source_object_out_of_range(ValueError(_OBJECT_OUT_OF_RANGE.args[0])) is False + + +def test_split_pdf_copies_pages(tmp_path: Path) -> None: + source = tmp_path / "source.pdf" + _write_pages(source, 3) + work_dir = tmp_path / "shards" + work_dir.mkdir() + + paths, remap = split_pdf( + str(source), + [MergedShard(0, 1, 2), MergedShard(1, 3, 3)], + str(work_dir), + ) + + assert remap is None + assert [Path(path).name for path in paths] == ["shard_0.pdf", "shard_1.pdf"] + assert _page_count(paths[0]) == 2 + assert _page_count(paths[1]) == 1 + assert not (work_dir / "_rewritten_source.pdf").exists() + + +def test_split_pdf_rewrites_once_on_object_out_of_range( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + source = tmp_path / "source.pdf" + _write_pages(source, 2) + work_dir = tmp_path / "shards" + work_dir.mkdir() + + real_insert = pymupdf.Document.insert_pdf + calls = {"n": 0} + + def _insert_pdf(self: pymupdf.Document, *args: object, **kwargs: object) -> None: + calls["n"] += 1 + if calls["n"] == 1: + raise _OBJECT_OUT_OF_RANGE + real_insert(self, *args, **kwargs) + + monkeypatch.setattr(pymupdf.Document, "insert_pdf", _insert_pdf) + + paths, remap = split_pdf( + str(source), + [MergedShard(0, 1, 2)], + str(work_dir), + ) + + assert remap is None + assert calls["n"] == 3 # fail once, then copy both pages from rewritten source + assert (work_dir / "_rewritten_source.pdf").exists() + assert len(paths) == 1 + assert _page_count(paths[0]) == 2 + + +def test_split_pdf_does_not_catch_other_runtime_errors( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + source = tmp_path / "source.pdf" + _write_pages(source, 1) + work_dir = tmp_path / "shards" + work_dir.mkdir() + + def _insert_pdf(self: pymupdf.Document, *args: object, **kwargs: object) -> None: + raise RuntimeError("code=7: syntax error") + + monkeypatch.setattr(pymupdf.Document, "insert_pdf", _insert_pdf) + + with pytest.raises(RuntimeError, match="syntax error"): + split_pdf(str(source), [MergedShard(0, 1, 1)], str(work_dir)) + + assert not (work_dir / "_rewritten_source.pdf").exists() + + +@pytest.mark.skipif(not _CUSTOMER_PDF.exists(), reason="local Quartz incremental PDF not present") +def test_split_incremental_quartz_pdf_around_failing_page(tmp_path: Path) -> None: + work_dir = tmp_path / "shards" + work_dir.mkdir() + + paths, remap = split_pdf( + str(_CUSTOMER_PDF), + [MergedShard(0, 40, 50)], + str(work_dir), + ) + + assert remap is None + assert len(paths) == 1 + assert _page_count(paths[0]) == 11 + assert (work_dir / "_rewritten_source.pdf").exists() diff --git a/apps/worker/tests/unit/test_probe_dense_visual_assets.py b/apps/worker/tests/unit/test_probe_dense_visual_assets.py new file mode 100644 index 000000000..d02deef86 --- /dev/null +++ b/apps/worker/tests/unit/test_probe_dense_visual_assets.py @@ -0,0 +1,88 @@ +"""Dense PDF pages must skip table/figure extraction instead of hanging.""" + +from __future__ import annotations + +from types import SimpleNamespace + +from app.services.document_agent.tools.probe_page_features import ( + _DENSE_DRAWING_LIMIT, + _DENSE_IMAGE_LIMIT, + _probe_visual_assets, + _rect_area, +) + + +class _FakeRect: + def __init__(self, x0: float, y0: float, x1: float, y1: float) -> None: + self.x0 = x0 + self.y0 = y0 + self.x1 = x1 + self.y1 = y1 + self.width = x1 - x0 + self.height = y1 - y0 + + +class _FakePage: + def __init__( + self, + *, + image_count: int = 0, + drawings_count: int = 0, + ) -> None: + self.rect = _FakeRect(0, 0, 600, 800) + self._images = [(index, 0, 0, 0, 0, 0, 0, 0) for index in range(image_count)] + self._drawings = [ + {"rect": _FakeRect(10, 10, 20, 20)} for _ in range(drawings_count) + ] + self.find_tables_calls = 0 + self.get_drawings_calls = 0 + + def get_images(self, full: bool = True) -> list[tuple[int, ...]]: + return self._images + + def get_image_rects(self, _xref: int) -> list[_FakeRect]: + return [] + + def find_tables(self) -> SimpleNamespace: + self.find_tables_calls += 1 + return SimpleNamespace(tables=[]) + + def get_drawings(self) -> list[dict[str, _FakeRect]]: + self.get_drawings_calls += 1 + return self._drawings + + +def test_dense_drawings_skip_tables_and_still_flag_asset() -> None: + page = _FakePage(drawings_count=_DENSE_DRAWING_LIMIT + 1) + result = _probe_visual_assets( + page, _rect_area(page.rect), header_y=None, footer_y=None + ) + + assert page.find_tables_calls == 0 + assert result["drawings_count"] == _DENSE_DRAWING_LIMIT + 1 + assert result["table_count"] == 0 + assert result["has_asset"] is True + + +def test_dense_images_skip_tables() -> None: + page = _FakePage(image_count=_DENSE_IMAGE_LIMIT + 1) + result = _probe_visual_assets( + page, _rect_area(page.rect), header_y=None, footer_y=None + ) + + assert page.find_tables_calls == 0 + assert result["image_count"] == _DENSE_IMAGE_LIMIT + 1 + assert result["table_count"] == 0 + assert result["has_asset"] is True + + +def test_normal_page_still_runs_table_finder() -> None: + page = _FakePage(image_count=2, drawings_count=8) + result = _probe_visual_assets( + page, _rect_area(page.rect), header_y=None, footer_y=None + ) + + assert page.find_tables_calls == 1 + assert result["image_count"] == 2 + assert result["drawings_count"] == 8 + assert result["has_asset"] is False diff --git a/apps/worker/tests/unit/test_processing_metadata_persist.py b/apps/worker/tests/unit/test_processing_metadata_persist.py new file mode 100644 index 000000000..4cd3d8200 --- /dev/null +++ b/apps/worker/tests/unit/test_processing_metadata_persist.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +import os +from types import SimpleNamespace +from unittest.mock import Mock + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +import app.services.document_ingestion.processing_context as processing_context # noqa: E402 +import app.services.document_ingestion.success_finalization as success_finalization # noqa: E402 + + +class _FakeDbContext: + def __init__(self, session: object) -> None: + self.session = session + + def __enter__(self) -> object: + return self.session + + def __exit__(self, *_args: object) -> bool: + return False + + +def _job_context(*, metadata: dict[str, object] | None = None) -> processing_context.ParseJobContext: + return processing_context.ParseJobContext( + job_metadata=metadata or {"namespace": "default"}, + job_user_id="user-1", + metadata_service=Mock(), + redis_service=object(), + s3_key="uploads/job.pdf", + ) + + +def test_persist_job_metadata_updates_merges_stages_into_job_row( + monkeypatch: object, +) -> None: + job = SimpleNamespace(job_metadata={"namespace": "default", "page_count": 12}) + session = Mock() + monkeypatch.setattr( + processing_context, + "get_sync_db_context", + lambda: _FakeDbContext(session), + ) + monkeypatch.setattr( + processing_context, + "_select_job_row_for_update", + lambda _db, _job_id: job, + ) + job_context = _job_context() + stages = { + "token_usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14, "calls": 1}, + "timing_ms": {"worker.parse.document": 1200}, + } + + processing_context.persist_job_metadata_updates( + job_id="job_abc", + job_context=job_context, + metadata_updates={"stages": stages}, + ) + + assert job.job_metadata["namespace"] == "default" + assert job.job_metadata["page_count"] == 12 + assert job.job_metadata["stages"]["token_usage"]["total_tokens"] == 14 + job_context.metadata_service.update_metadata.assert_called_once_with( + "job_abc", + {"stages": stages}, + ) + assert job_context.job_metadata["stages"] == stages + + +def test_record_processing_completion_persists_token_usage_to_job_row( + monkeypatch: object, +) -> None: + persist = Mock() + monkeypatch.setattr(success_finalization, "persist_job_metadata_updates", persist) + monkeypatch.setattr( + success_finalization, + "get_current_token_tracker", + lambda: {"prompt_tokens": 8, "completion_tokens": 2, "total_tokens": 10, "calls": 1}, + ) + monkeypatch.setattr( + success_finalization, + "get_current_stage_tracker", + lambda: {"worker.parse.document": 900}, + ) + from datetime import datetime, timezone + + job_context = _job_context() + started = datetime(2026, 8, 24, 10, 0, tzinfo=timezone.utc) + + success_finalization._record_processing_completion( + job_id="job_abc", + job_context=job_context, + processing_started_at=started, + ) + + persist.assert_called_once() + updates = persist.call_args.kwargs["metadata_updates"] + assert updates["stages"]["token_usage"]["total_tokens"] == 10 + assert "processing_completed_at" in updates + assert "processing_duration_ms" in updates diff --git a/apps/worker/tests/unit/test_pymupdf_subprocess_queue.py b/apps/worker/tests/unit/test_pymupdf_subprocess_queue.py new file mode 100644 index 000000000..1626daf66 --- /dev/null +++ b/apps/worker/tests/unit/test_pymupdf_subprocess_queue.py @@ -0,0 +1,93 @@ +"""Child workers must flush the multiprocessing queue before exiting.""" + +from __future__ import annotations + +import app.services.document_parser.formats.pdf.pymupdf_subprocess as subprocess_mod +from shared.core.exceptions.domain_exceptions import PDFParsingException + + +class _FakeQueue: + def __init__(self) -> None: + self.items: list[object] = [] + self.closed = False + self.joined = False + + def put(self, item: object) -> None: + self.items.append(item) + + def close(self) -> None: + self.closed = True + + def join_thread(self) -> None: + self.joined = True + + +def test_worker_decorator_flushes_queue_after_success() -> None: + @subprocess_mod.worker + def _ok(queue: _FakeQueue, value: str) -> None: + queue.put({"ok": True, "value": value}) + + queue = _FakeQueue() + _ok(queue, "payload") + + assert queue.items == [{"ok": True, "value": "payload"}] + assert queue.closed is True + assert queue.joined is True + + +def test_worker_decorator_flushes_queue_after_failure() -> None: + @subprocess_mod.worker + def _boom(queue: _FakeQueue) -> None: + raise RuntimeError("child failed") + + queue = _FakeQueue() + _boom(queue) + + assert queue.items[0]["ok"] is False + assert queue.items[0]["error_type"] == "RuntimeError" + assert queue.closed is True + assert queue.joined is True + + +def test_empty_queue_exit_is_retryable() -> None: + exc = PDFParsingException( + user_message="Failed to process your document. Please try again.", + reason="SUBPROCESS_CRASH", + internal_message=( + "pymupdf child exited with code=0 and no result: " + "fn=_probe_assets_worker, pid=2049" + ), + ) + assert subprocess_mod._is_empty_queue_exit(exc) is True + + +def test_nonzero_crash_is_not_retryable() -> None: + exc = PDFParsingException( + user_message="Failed to process your document. Please try again.", + reason="SUBPROCESS_CRASH", + internal_message="pymupdf child exited with code=-9 and no result: fn=x, pid=1", + ) + assert subprocess_mod._is_empty_queue_exit(exc) is False + + +def test_empty_queue_exit_retries_once(monkeypatch: object) -> None: + calls = {"n": 0} + crash = PDFParsingException( + user_message="Failed to process your document. Please try again.", + reason="SUBPROCESS_CRASH", + internal_message=( + "pymupdf child exited with code=0 and no result: " + "fn=_probe_assets_worker, pid=2049" + ), + ) + + def _once(_worker_fn: object, _args: tuple, _timeout: int) -> dict: + calls["n"] += 1 + if calls["n"] == 1: + raise crash + return {"ok": True, "assets": []} + + monkeypatch.setattr(subprocess_mod, "_run_worker_in_spawned_process_once", _once) + result = subprocess_mod._run_worker_in_spawned_process(lambda: None, (), 10) + assert result == {"ok": True, "assets": []} + assert calls["n"] == 2 diff --git a/packages/shared-python/shared/models/schemas/job_metadata.py b/packages/shared-python/shared/models/schemas/job_metadata.py index 4efb808b3..0162a257a 100644 --- a/packages/shared-python/shared/models/schemas/job_metadata.py +++ b/packages/shared-python/shared/models/schemas/job_metadata.py @@ -36,6 +36,10 @@ class JobMetadataBase(BaseModel): page_memory_config: Optional[Dict[str, Any]] = Field( None, description="Resolved page-memory worker configuration" ) + stages: Optional[Dict[str, Any]] = Field( + None, + description="Worker processing stages, including token_usage and timing_ms", + ) # result_mode was removed and is no longer supported. # Source-file fields.