From 35aed6e24bed1e227b62d321ae28c9b2b0a5df52 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 16 Oct 2026 15:29:00 +0800 Subject: [PATCH 1/5] feat: enhance markdown and pdf parsing with image region metadata - Added support for extra metadata in the MarkdownParseState to store image region information. - Updated the append_row method to accept additional metadata and store it in a new dictionary. - Modified the parse_md function to include image region metadata when processing markdown content. - Integrated image region loading in the PDF parsing workflow to ensure accurate page number mapping after TOC removal. - Adjusted tests to validate the new metadata structure and ensure consistency across chunk and zip formats. --- .../formats/markdown/layout_bbox_index.py | 127 ++++++++ .../formats/markdown/parse_state.py | 19 +- .../formats/markdown/parser.py | 12 +- .../document_parser/formats/pdf/parser.py | 18 ++ .../formats/pdf/shard_splitter.py | 22 +- .../app/services/page_memory/page_assets.py | 2 + apps/worker/scripts/debug_text_track.py | 26 +- ...est_page_memory_node_assembler_contract.py | 11 +- .../test_page_memory_output_contract.py | 19 +- .../contract/test_parse_task_contract.py | 189 ++++++++++++ .../test_text_track_image_bbox_contract.py | 288 ++++++++++++++++++ 11 files changed, 706 insertions(+), 27 deletions(-) create mode 100644 apps/worker/app/services/document_parser/formats/markdown/layout_bbox_index.py create mode 100644 apps/worker/tests/contract/test_text_track_image_bbox_contract.py diff --git a/apps/worker/app/services/document_parser/formats/markdown/layout_bbox_index.py b/apps/worker/app/services/document_parser/formats/markdown/layout_bbox_index.py new file mode 100644 index 000000000..88d1c5458 --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/markdown/layout_bbox_index.py @@ -0,0 +1,127 @@ +"""Index MinerU ``layout.json`` figure regions by image file name. + +MinerU writes one ``layout.json`` per parsed PDF. Every figure span +(``image`` / ``chart``) carries its bounding box and the bare ``{hash}.jpg`` +name of the extracted image, which is the same name the markdown references +before ``build_markdown_image_asset`` renames the file. This module turns that +into the ``asset_region`` metadata contract shared with the page track. +""" + +from __future__ import annotations + +import json +import os +from collections.abc import Iterator, Mapping, Sequence +from typing import Any + +from loguru import logger + +AssetRegion = dict[str, Any] + +COORD_SYSTEM_MINERU_LAYOUT = "mineru_layout" + +_FIGURE_SPAN_TYPES = frozenset({"image", "chart"}) + + +def load_layout_image_regions( + layout_json_path: str, + page_numbers: Sequence[int] | None = None, +) -> dict[str, AssetRegion]: + """Return ``image file name -> asset_region`` for one ``layout.json``. + + Args: + layout_json_path: MinerU ``layout.json`` of a single MinerU request. + page_numbers: Original 1-based page number of each page of that request, + indexed by the layout's 0-based ``page_idx``. ``None`` means the + request parsed the whole original PDF, so ``page_idx + 1`` is the + original page number. + + A missing ``layout.json`` yields an empty index. A figure whose page cannot + be resolved is skipped (and logged) rather than given a guessed page. + """ + if not os.path.exists(layout_json_path): + logger.debug("layout.json not found, image bbox will not be attached") + return {} + + with open(layout_json_path, "r", encoding="utf-8") as layout_file: + layout = json.load(layout_file) + + regions: dict[str, AssetRegion] = {} + for page in layout["pdf_info"]: + page_idx = page["page_idx"] + page_width, page_height = page["page_size"] + for span in _iter_figure_spans(page["para_blocks"]): + image_name = os.path.basename(span.get("image_path") or "") + if not image_name: + continue + page_num = _resolve_page_num(page_idx, page_numbers) + if page_num is None: + logger.warning( + "Skip image bbox for {}: page_idx={} is outside the request page map", + image_name, + page_idx, + ) + continue + regions.setdefault( + image_name, + { + "page_num": page_num, + "bbox": list(span["bbox"]), + "coord_origin": "top_left", + "coord_system": COORD_SYSTEM_MINERU_LAYOUT, + "page_width": page_width, + "page_height": page_height, + # The ZIP/API asset_region contract requires this field. + # MinerU crops the figure at its layout bbox, no margin. + "crop_margin_px": 0, + }, + ) + return regions + + +def load_shard_image_regions( + shard_output_dirs: Sequence[str | None], + shard_page_numbers: Sequence[Sequence[int] | None], +) -> dict[str, AssetRegion]: + """Merge the per-shard indexes of a sharded PDF into one index. + + ``shard_page_numbers[i]`` is the page map of ``shard_output_dirs[i]`` (see + ``load_layout_image_regions``). When the same image file name appears in + several shards, the earliest shard wins, matching document order. + """ + merged: dict[str, AssetRegion] = {} + for shard_dir, page_numbers in zip(shard_output_dirs, shard_page_numbers): + if shard_dir is None: + continue + shard_regions = load_layout_image_regions( + os.path.join(shard_dir, "layout.json"), page_numbers + ) + for image_name, region in shard_regions.items(): + merged.setdefault(image_name, region) + return merged + + +def image_region_metadata( + image_regions: Mapping[str, AssetRegion], + image_path: str, +) -> dict[str, AssetRegion] | None: + """``extra_metadata`` for the image row of a markdown image reference.""" + region = image_regions.get(os.path.basename(image_path)) + return {"asset_region": region} if region is not None else None + + +def _resolve_page_num(page_idx: int, page_numbers: Sequence[int] | None) -> int | None: + if page_numbers is None: + return page_idx + 1 + if 0 <= page_idx < len(page_numbers): + return page_numbers[page_idx] + return None + + +def _iter_figure_spans(blocks: Sequence[dict[str, Any]]) -> Iterator[dict[str, Any]]: + for block in blocks: + for line in block.get("lines", []): + for span in line["spans"]: + if span["type"] in _FIGURE_SPAN_TYPES: + yield span + yield from _iter_figure_spans(block.get("blocks", [])) diff --git a/apps/worker/app/services/document_parser/formats/markdown/parse_state.py b/apps/worker/app/services/document_parser/formats/markdown/parse_state.py index a8d709d38..9fa3c76a9 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/parse_state.py +++ b/apps/worker/app/services/document_parser/formats/markdown/parse_state.py @@ -1,6 +1,6 @@ from __future__ import annotations - +import json from collections.abc import Callable from dataclasses import dataclass, field from typing import Any @@ -31,6 +31,7 @@ class MarkdownParseState: timestamp: str row_updater: RowUpdater rows: list[ParserRowValues] = field(default_factory=list) + row_extra_metadata: dict[int, str] = field(default_factory=dict) content_items: list[str] = field(default_factory=list) path_stack: list[tuple[str, int]] = field(default_factory=list) inner_paths: list[str] = field(default_factory=list) @@ -125,7 +126,15 @@ def append_content_item(self, item: str) -> None: def append_plain_text(self, text: str) -> None: self.content_items.append(text.strip()) - def append_row(self, row: ParserRowValues) -> None: + def append_row( + self, + row: ParserRowValues, + extra_metadata: dict[str, Any] | None = None, + ) -> None: + if extra_metadata: + self.row_extra_metadata[len(self.rows)] = json.dumps( + extra_metadata, ensure_ascii=False + ) self.rows.append(row) def schedule_deferred_task(self, task: MarkdownDeferredSummaryTask) -> None: @@ -168,6 +177,10 @@ def to_dataframe(self) -> pd.DataFrame: asset_title=str(row_values[12]) if len(row_values) > 12 else "", ) ) - return process_dup_paths_df(rows_builder.to_dataframe()) + rows_df = rows_builder.to_dataframe() + rows_df["extra_metadata"] = [ + self.row_extra_metadata.get(index, "") for index in range(len(rows_df)) + ] + return process_dup_paths_df(rows_df) diff --git a/apps/worker/app/services/document_parser/formats/markdown/parser.py b/apps/worker/app/services/document_parser/formats/markdown/parser.py index f863ddcb1..1abd5a173 100755 --- a/apps/worker/app/services/document_parser/formats/markdown/parser.py +++ b/apps/worker/app/services/document_parser/formats/markdown/parser.py @@ -14,6 +14,10 @@ build_markdown_image_name, build_markdown_image_asset, ) +from app.services.document_parser.formats.markdown.layout_bbox_index import ( + AssetRegion, + image_region_metadata, +) from app.services.document_parser.formats.markdown.parse_state import MarkdownParseState from app.services.document_parser.formats.markdown.table_asset import ( MarkdownTableAssetRequest, @@ -238,6 +242,7 @@ def parse_md( lines_with_heading=None, is_first_shard=True, skip_toc_detection=False, + image_regions: dict[str, AssetRegion] | None = None, ): if lines_with_heading is not None: # ── Phase A bypass ── @@ -407,7 +412,12 @@ def parse_md( and image_asset.row_values is not None ): parser_state.append_content_item(image_asset.content_item) - parser_state.append_row(image_asset.row_values) + parser_state.append_row( + image_asset.row_values, + extra_metadata=image_region_metadata( + image_regions or {}, img_path + ), + ) if ( image_asset.cache_key is not None and image_asset.cache_entry is not None diff --git a/apps/worker/app/services/document_parser/formats/pdf/parser.py b/apps/worker/app/services/document_parser/formats/pdf/parser.py index a937b1c60..accc74ecf 100755 --- a/apps/worker/app/services/document_parser/formats/pdf/parser.py +++ b/apps/worker/app/services/document_parser/formats/pdf/parser.py @@ -3,6 +3,10 @@ import re import shutil +from app.services.document_parser.formats.markdown.layout_bbox_index import ( + load_layout_image_regions, + load_shard_image_regions, +) from app.services.document_parser.formats.markdown.parser import parse_md from app.services.document_parser.orchestration.oversized_pdf_policy import ( build_oversized_pdf_processing_failed_exception, @@ -86,6 +90,9 @@ def parse_pdfs( base_llm_paras=base_llm_paras, relative_root=relative_root, skip_toc_detection=True, + image_regions=load_layout_image_regions( + os.path.join(output_dir, "layout.json") + ), ) @@ -118,6 +125,7 @@ def _parse_pdf_via_shards( ) from app.services.document_parser.formats.pdf.shard_splitter import ( map_agent_shards, + shard_page_numbers, split_pdf, ) from shared.services.ai.token_tracking import ( @@ -346,6 +354,15 @@ def _predict_shard_headings_with_tracking( if not fast_path_original_pdf: merge_images(shard_output_dirs, output_dir) + # The fast path parsed the whole original PDF; otherwise each shard PDF + # holds only its page range minus the excluded TOC pages. + shard_page_maps = ( + [None] + if fast_path_original_pdf + else [shard_page_numbers(shard, toc_pages) for shard in merged_shards] + ) + image_regions = load_shard_image_regions(shard_output_dirs, shard_page_maps) + logger.info("✅ Shard-first hierarchy complete, entering parse_md Phase B") # 8. parse_md Phase B only (skip TOC detection + heading prediction) @@ -356,6 +373,7 @@ def _predict_shard_headings_with_tracking( base_llm_paras=base_llm_paras, relative_root=relative_root, lines_with_heading=all_lines_with_heading, + image_regions=image_regions, ) finally: _cleanup_temp_shard_s3_assets(temp_shard_s3_keys) 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 d13565813..7bfddb1c5 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 @@ -28,11 +28,6 @@ class MergedShard: def page_count(self) -> int: return self.page_end - self.page_start + 1 - @property - def page_offset(self) -> int: - """Offset to add to MinerU's 0-based page_idx to get original page_idx.""" - return self.page_start - 1 - def map_agent_shards( agent_shards: list["Shard"], @@ -49,6 +44,19 @@ def map_agent_shards( ] +def shard_page_numbers(shard: MergedShard, exclude_pages: set[int]) -> list[int]: + """Original 1-based page numbers physically present in the shard PDF, in order. + + Index ``i`` is the original page of MinerU's 0-based ``page_idx == i`` for + this shard, because excluded (e.g. TOC) pages are removed before MinerU runs. + """ + return [ + page_num + for page_num in range(shard.page_start, shard.page_end + 1) + if page_num not in exclude_pages + ] + + 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) @@ -122,9 +130,7 @@ def _extract_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: - continue + for page_num in shard_page_numbers(shard, exclude_pages or set()): sub_doc.insert_pdf( doc, from_page=page_num - 1, diff --git a/apps/worker/app/services/page_memory/page_assets.py b/apps/worker/app/services/page_memory/page_assets.py index 19c828510..1821c5619 100644 --- a/apps/worker/app/services/page_memory/page_assets.py +++ b/apps/worker/app/services/page_memory/page_assets.py @@ -544,6 +544,8 @@ def build_asset_rows( def _asset_extra_metadata(asset: PageAsset) -> dict[str, Any]: + if asset.kind == "table": + return {} # bbox is the detected (VLM) box on the rendered page image, not the cropped box. return { "asset_region": { diff --git a/apps/worker/scripts/debug_text_track.py b/apps/worker/scripts/debug_text_track.py index d85fc9a99..4b36c1217 100644 --- a/apps/worker/scripts/debug_text_track.py +++ b/apps/worker/scripts/debug_text_track.py @@ -664,10 +664,33 @@ def _stage_full_pdf( filename: str, merged_lines: list[str], model: str | None, + anatomy, ): + from app.services.document_parser.formats.markdown.layout_bbox_index import ( + load_shard_image_regions, + ) from app.services.document_parser.formats.markdown.parser import parse_md + from app.services.document_parser.formats.pdf.shard_splitter import ( + map_agent_shards, + shard_page_numbers, + ) from app.services.document_parser.orchestration.postprocess import apply_parse_postprocess + merged_shards = map_agent_shards(anatomy.shard_plan.shards) + toc_pages: set[int] = set() + if anatomy.toc_result and anatomy.toc_result.toc_pages: + toc_pages = set(anatomy.toc_result.toc_pages) + fast_path = len(merged_shards) == 1 and not toc_pages + shard_page_maps = ( + [None] + if fast_path + else [shard_page_numbers(shard, toc_pages) for shard in merged_shards] + ) + image_regions = load_shard_image_regions( + [str(out_dir / "_shards" / f"shard_{i}") for i in range(len(merged_shards))], + shard_page_maps, + ) + logger.info("=" * 70) logger.info("📦 Stage 4: parse_md Phase B → DataFrame → chunks") logger.info("=" * 70) @@ -687,6 +710,7 @@ def _stage_full_pdf( base_llm_paras=base_llm_paras, relative_root=filename, lines_with_heading=merged_lines, + image_regions=image_regions, ) parsed_df = apply_parse_postprocess(str(out_dir), parsed_df) elapsed_parse = time.time() - t0 @@ -878,7 +902,7 @@ def main() -> int: # Stage 4: Full extraction chunks, full_elapsed = _stage_full_pdf( - out_dir, filename, merged_lines, args.model + out_dir, filename, merged_lines, args.model, anatomy ) trace["stages"]["full"] = { "elapsed_s": round(full_elapsed, 1), diff --git a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py index d0300fb8e..bedac5e1d 100644 --- a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py +++ b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py @@ -422,15 +422,8 @@ def test_build_node_rows_prepends_asset_rows_and_links_page_nodes() -> None: assert "title" not in rows[0]["extra_metadata"] assert "html_uri" not in rows[0]["extra_metadata"] assert "image_uri" not in rows[0]["extra_metadata"] - assert rows[0]["extra_metadata"]["asset_region"] == { - "page_num": 231, - "bbox": [10, 20, 200, 120], - "coord_origin": "top_left", - "coord_system": "page_image_px", - "page_width": 1000, - "page_height": 1400, - "crop_margin_px": 4, - } + assert "asset_region" not in rows[0]["extra_metadata"] + assert rows[0]["page_nums"] == "231" by_path = {row["path"]: row for row in rows} owner = by_path["demo.pdf/3 基本规定/3.1 职责"] diff --git a/apps/worker/tests/contract/test_page_memory_output_contract.py b/apps/worker/tests/contract/test_page_memory_output_contract.py index 1f9a8801f..3bd7684f8 100644 --- a/apps/worker/tests/contract/test_page_memory_output_contract.py +++ b/apps/worker/tests/contract/test_page_memory_output_contract.py @@ -73,14 +73,23 @@ def test_asset_region_and_metadata_match_between_chunks_and_zip() -> None: "page_height": 1400, "crop_margin_px": 4, } - for asset_id in ("asset_figure_1", "asset_table_2"): - chunk_metadata = chunk_by_id[asset_id] - zip_metadata = zip_by_id[asset_id] - assert chunk_metadata["asset_region"] == expected_region + figure_chunk = chunk_by_id["asset_figure_1"] + figure_zip = zip_by_id["asset_figure_1"] + table_chunk = chunk_by_id["asset_table_2"] + table_zip = zip_by_id["asset_table_2"] + assert figure_chunk["asset_region"] == expected_region + assert figure_zip["asset_region"] == expected_region + assert "asset_region" not in table_chunk + assert "asset_region" not in table_zip + assert table_chunk["page_nums"] == [5] + assert table_zip["page_nums"] == [5] + for chunk_metadata, zip_metadata in ( + (figure_chunk, figure_zip), + (table_chunk, table_zip), + ): assert "image_uri" not in chunk_metadata assert set(zip_metadata) <= set(chunk_metadata) for key in ( - "asset_region", "entities", "asset_title", "keywords", diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index 0fdf813b8..0e2249e69 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -1357,3 +1357,192 @@ def _identity_eval_md_headings( "oversized.pdf/Chapter 1", "oversized.pdf/Chapter 2", ] + + +def test_sharded_pdf_image_bbox_page_numbers_are_original_pages_after_toc_removal( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setenv("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") + monkeypatch.setenv("TMP_PATH", str(tmp_path)) + monkeypatch.setenv("S3_BUCKET_NAME", "test-uploads") + monkeypatch.setenv("S3_ACCESS_KEY_ID", "test") + monkeypatch.setenv("S3_SECRET_ACCESS_KEY", "test") + monkeypatch.setenv("S3_TEMP_PATH", str(tmp_path)) + + from PIL import Image + + from app.services.document_agent.manifest import ( + PageAnatomyMap, + Shard, + ShardPlan, + TocResult, + ) + from app.services.document_parser.formats.pdf import parser as pdf_parser + from app.services.document_parser.profiling.profile_model import ParserDocumentProfile + from app.services.document_parser.profiling.taxonomy import PdfRoutingCategory + from shared.core.constants.processing import ProcessingConstants + + pdf_path = tmp_path / "sharded.pdf" + output_dir = tmp_path / "output" + output_dir.mkdir() + _write_blank_pdf(pdf_path, page_count=3) + + class _FakeJobFileStorage: + def upload_source_file(self, local_file_path: str, storage_key: str) -> dict[str, str]: + return {"etag": "test"} + + def verify_upload_exists(self, storage_key: str) -> dict[str, bool]: + return {"exists": True} + + def delete_upload_file(self, storage_key: str) -> bool: + return True + + # Shard 0 covers pages 1-2 with TOC page 1 removed, so MinerU page_idx 0 is + # original page 2. Shard 1 covers page 3 only. + anatomy = PageAnatomyMap( + job_id="job-bbox-pages", + file_path=str(pdf_path), + page_count=3, + page_features=[], + page_labels=[], + toc_result=TocResult(toc_pages=[1], method="vlm_batch"), + shard_plan=ShardPlan( + enabled=True, + reason="too_large", + shards=[ + Shard( + shard_index=0, + page_start=1, + page_end=2, + page_offset=0, + anchor_type="toc_leaf_boundary", + anchor_evidence="Chapter 1", + ), + Shard( + shard_index=1, + page_start=3, + page_end=3, + page_offset=2, + anchor_type="toc_leaf_boundary", + anchor_evidence="Chapter 2", + ), + ], + ), + toc_hierarchies=[{"toc_tree": {"Chapter 1": {}, "Chapter 2": {}}}], + ) + + image_names = {0: "1" * 64 + ".jpg", 1: "2" * 64 + ".jpg"} + image_bboxes = {0: [10, 20, 300, 400], 1: [30, 40, 500, 600]} + + def _fake_split_pdf(pdf_path_arg, shards, work_dir, exclude_pages=None): + paths = [] + for shard in shards: + shard_path = Path(work_dir) / f"shard_{shard.shard_index}.pdf" + shard_path.write_bytes(b"%PDF-1.4\n%%EOF\n") + paths.append(str(shard_path)) + return paths, None + + def _fake_parse_via_full(shard_pdf, shard_filename, shard_out, s3_key=None): + shard_index = 0 if "shard0" in shard_filename else 1 + shard_dir = Path(shard_out) + (shard_dir / "images").mkdir(parents=True, exist_ok=True) + image_name = image_names[shard_index] + width = height = 640 + pixels = bytearray(width * height * 3) + for index in range(len(pixels)): + pixels[index] = (index * (37 + shard_index * 16)) % 256 + image_file = shard_dir / "images" / image_name + Image.frombytes("RGB", (width, height), bytes(pixels)).save( + image_file, format="JPEG", quality=95 + ) + with image_file.open("ab") as handle: + handle.write(b"\x00" * ProcessingConstants.IMG_MIN_SIZE) + (shard_dir / "full.md").write_text( + f"# Chapter {shard_index + 1}\n![](images/{image_name})", + encoding="utf-8", + ) + bbox = image_bboxes[shard_index] + span = {"bbox": bbox, "type": "image", "image_path": image_name} + (shard_dir / "layout.json").write_text( + json.dumps( + { + "pdf_info": [ + { + "page_idx": 0, + "page_size": [960, 540], + "para_blocks": [ + { + "type": "image", + "bbox": bbox, + "blocks": [ + { + "type": "image_body", + "bbox": bbox, + "lines": [{"bbox": bbox, "spans": [span]}], + } + ], + } + ], + } + ] + } + ), + encoding="utf-8", + ) + + def _identity_eval_md_headings( + md_lines, + source_type, + toc_hierarchies=None, + smart_parse=True, + model_name=None, + output_dir=None, + layout_json_path=None, + is_first_shard=True, + ): + return list(md_lines) + + monkeypatch.setattr(pdf_parser.settings, "MAX_PDF_PAGE_LIMIT", 2) + monkeypatch.setattr(pdf_parser.settings, "MINERU_SHARD_CONCURRENCY", 1) + monkeypatch.setattr( + "app.services.document_parser.formats.pdf.shard_splitter.split_pdf", + _fake_split_pdf, + ) + monkeypatch.setattr(pdf_parser, "parse_via_full", _fake_parse_via_full) + monkeypatch.setattr( + "app.services.document_parser.formats.markdown.parser.eval_md_headings", + _identity_eval_md_headings, + ) + monkeypatch.setattr(pdf_parser, "JobFileStorage", _FakeJobFileStorage) + + df = pdf_parser.parse_pdfs( + str(pdf_path), + "sharded.pdf", + str(output_dir), + { + "smart_title_parse": False, + "summary_image": False, + "summary_table": False, + "summary_txt": False, + "stopwords": [], + "model_name": "mock-model", + "hierarchy_model_name": "mock-model", + }, + profile=ParserDocumentProfile( + file_type="pdf", + category="generic document", + routing_category=PdfRoutingCategory.GENERIC, + page_count=3, + anatomy=anatomy, + ), + relative_root="sharded.pdf", + s3_key="uploads/job-bbox-pages.pdf", + job_id="job-bbox-pages", + ) + + image_rows = df[df["type"] == "image"] + regions = [json.loads(value)["asset_region"] for value in image_rows["extra_metadata"]] + assert [region["page_num"] for region in regions] == [2, 3] + assert [region["bbox"] for region in regions] == [image_bboxes[0], image_bboxes[1]] + assert {region["coord_system"] for region in regions} == {"mineru_layout"} diff --git a/apps/worker/tests/contract/test_text_track_image_bbox_contract.py b/apps/worker/tests/contract/test_text_track_image_bbox_contract.py new file mode 100644 index 000000000..2e6113944 --- /dev/null +++ b/apps/worker/tests/contract/test_text_track_image_bbox_contract.py @@ -0,0 +1,288 @@ +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +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") + +from PIL import Image # noqa: E402 + +from app.services.document_parser.formats.markdown.layout_bbox_index import ( # noqa: E402 + load_layout_image_regions, + load_shard_image_regions, +) +from app.services.document_parser.formats.markdown.parser import parse_md # noqa: E402 +from app.services.document_parser.formats.pdf.shard_splitter import ( # noqa: E402 + MergedShard, + shard_page_numbers, +) +from shared.core.constants.processing import ProcessingConstants # noqa: E402 +from shared.services.chunks.dataframe_chunk_converter import ( # noqa: E402 + dataframe_to_chunks, +) +from shared.services.storage.zip_chunk_schema import ZipChunkSchemaBuilder # noqa: E402 + +_PAGE_SIZE = [960, 540] + + +def _figure_block(block_type: str, image_path: str, bbox: list[int]) -> dict[str, Any]: + body_type = f"{block_type}_body" + span = {"bbox": bbox, "type": block_type, "image_path": image_path} + return { + "type": block_type, + "bbox": bbox, + "blocks": [ + { + "type": body_type, + "bbox": bbox, + "lines": [{"bbox": bbox, "spans": [span]}], + } + ], + } + + +def _text_block(bbox: list[int]) -> dict[str, Any]: + span = {"bbox": bbox, "type": "text", "content": "body"} + return {"type": "text", "bbox": bbox, "lines": [{"bbox": bbox, "spans": [span]}]} + + +def _write_layout(path: Path, pages: list[list[dict[str, Any]]]) -> None: + layout = { + "pdf_info": [ + {"page_idx": index, "page_size": _PAGE_SIZE, "para_blocks": blocks} + for index, blocks in enumerate(pages) + ] + } + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(layout), encoding="utf-8") + + +def _region(page_num: int, bbox: list[int]) -> dict[str, Any]: + return { + "page_num": page_num, + "bbox": bbox, + "coord_origin": "top_left", + "coord_system": "mineru_layout", + "page_width": 960, + "page_height": 540, + "crop_margin_px": 0, + } + + +def test_layout_index_collects_image_and_chart_but_not_table_or_text( + tmp_path: Path, +) -> None: + layout_path = tmp_path / "layout.json" + _write_layout( + layout_path, + [ + [ + _figure_block("image", "aaa.jpg", [1, 2, 30, 40]), + _figure_block("chart", "bbb.jpg", [5, 6, 70, 80]), + _figure_block("table", "ccc.jpg", [9, 9, 99, 99]), + _text_block([0, 0, 10, 10]), + ] + ], + ) + + regions = load_layout_image_regions(str(layout_path)) + + assert regions == { + "aaa.jpg": _region(1, [1, 2, 30, 40]), + "bbb.jpg": _region(1, [5, 6, 70, 80]), + } + + +def test_layout_index_maps_shard_pages_through_the_page_map(tmp_path: Path) -> None: + layout_path = tmp_path / "layout.json" + _write_layout( + layout_path, + [ + [_figure_block("image", "p0.jpg", [1, 1, 2, 2])], + [_figure_block("image", "p1.jpg", [3, 3, 4, 4])], + [_figure_block("image", "p2.jpg", [5, 5, 6, 6])], + ], + ) + + # Original page 11 (a TOC page) was removed from this shard PDF. + regions = load_layout_image_regions(str(layout_path), [10, 12]) + + assert regions["p0.jpg"]["page_num"] == 10 + assert regions["p1.jpg"]["page_num"] == 12 + assert "p2.jpg" not in regions + + +def test_layout_index_is_empty_without_layout_json(tmp_path: Path) -> None: + assert load_layout_image_regions(str(tmp_path / "layout.json")) == {} + + +def test_layout_index_keeps_first_region_of_a_repeated_image(tmp_path: Path) -> None: + layout_path = tmp_path / "layout.json" + _write_layout( + layout_path, + [ + [_figure_block("image", "same.jpg", [1, 1, 2, 2])], + [_figure_block("image", "same.jpg", [7, 7, 8, 8])], + ], + ) + + regions = load_layout_image_regions(str(layout_path)) + + assert regions["same.jpg"] == _region(1, [1, 1, 2, 2]) + + +def test_shard_regions_merge_in_shard_order_and_skip_missing_dirs( + tmp_path: Path, +) -> None: + first = tmp_path / "shard_0" + second = tmp_path / "shard_1" + _write_layout( + first / "layout.json", + [[_figure_block("image", "dup.jpg", [1, 1, 2, 2])]], + ) + _write_layout( + second / "layout.json", + [ + [ + _figure_block("image", "dup.jpg", [9, 9, 9, 9]), + _figure_block("image", "only_second.jpg", [3, 3, 4, 4]), + ] + ], + ) + + regions = load_shard_image_regions( + [str(first), None, str(second)], + [[1], [2], [20]], + ) + + assert regions["dup.jpg"] == _region(1, [1, 1, 2, 2]) + assert regions["only_second.jpg"] == _region(20, [3, 3, 4, 4]) + + +def test_shard_page_numbers_skip_excluded_pages() -> None: + shard = MergedShard(shard_index=0, page_start=5, page_end=9) + + assert shard_page_numbers(shard, set()) == [5, 6, 7, 8, 9] + assert shard_page_numbers(shard, {7, 8, 100}) == [5, 6, 9] + + +def _write_jpeg(path: Path, *, seed: int) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + width = height = 640 + pixels = bytearray(width * height * 3) + for index in range(len(pixels)): + pixels[index] = (index * seed) % 256 + Image.frombytes("RGB", (width, height), bytes(pixels)).save( + path, format="JPEG", quality=95 + ) + missing = ProcessingConstants.IMG_MIN_SIZE - path.stat().st_size + if missing > 0: + with path.open("ab") as handle: + handle.write(b"\x00" * missing) + + +def _parse_phase_b( + output_dir: Path, + lines: list[str], + image_regions: dict[str, dict[str, Any]] | None, +): + return parse_md( + str(output_dir), + source_type="md", + base_llm_paras={ + "summary_image": False, + "summary_table": False, + "summary_txt": False, + "stopwords": [], + }, + relative_root="deck.pdf", + lines_with_heading=lines, + image_regions=image_regions, + ) + + +def test_parse_md_attaches_asset_region_to_image_chunks_only(tmp_path: Path) -> None: + hash_name = "a" * 64 + ".jpg" + _write_jpeg(tmp_path / "images" / hash_name, seed=37) + lines = [ + "# Results", + "Revenue grew this quarter.", + f"![](images/{hash_name})", + "
MetricValue
", + ] + region = _region(7, [10, 20, 300, 400]) + + df = _parse_phase_b(tmp_path, lines, {hash_name: region}) + chunks = dataframe_to_chunks(df) + + by_type = {str(chunk["type"]): chunk for chunk in chunks} + assert by_type["image"]["metadata"]["asset_region"] == region + assert "asset_region" not in by_type["table"]["metadata"] + assert "asset_region" not in by_type["text"]["metadata"] + + image_chunk = by_type["image"] + zip_chunks = ZipChunkSchemaBuilder().format_chunks( + chunks, + image_files_map={ + str(image_chunk["chunk_id"]): { + "file_path": image_chunk["metadata"]["file_path"] + } + }, + table_files_map={ + str(by_type["table"]["chunk_id"]): { + "file_path": by_type["table"]["metadata"]["file_path"] + } + }, + ) + zip_by_type = {str(chunk["type"]): chunk for chunk in zip_chunks} + assert zip_by_type["image"]["metadata"]["asset_region"] == region + assert "asset_region" not in zip_by_type["table"]["metadata"] + + +def test_parse_md_without_image_regions_writes_no_asset_region(tmp_path: Path) -> None: + hash_name = "b" * 64 + ".jpg" + _write_jpeg(tmp_path / "images" / hash_name, seed=41) + lines = ["# Results", f"![](images/{hash_name})"] + + df = _parse_phase_b(tmp_path, lines, None) + + assert list(df["extra_metadata"]) == [""] * len(df) + assert all("asset_region" not in chunk["metadata"] for chunk in dataframe_to_chunks(df)) + + +def test_parse_md_looks_up_each_duplicate_image_by_its_own_file_name( + tmp_path: Path, +) -> None: + first_name = "c" * 64 + ".jpg" + second_name = "d" * 64 + ".jpg" + _write_jpeg(tmp_path / "images" / first_name, seed=53) + # Identical pixels, different file name: the second one takes the duplicate path. + (tmp_path / "images" / second_name).write_bytes( + (tmp_path / "images" / first_name).read_bytes() + ) + lines = [ + "# Results", + f"![](images/{first_name})", + f"![](images/{second_name})", + ] + first_region = _region(3, [1, 1, 100, 100]) + second_region = _region(4, [2, 2, 200, 200]) + + df = _parse_phase_b( + tmp_path, + lines, + {first_name: first_region, second_name: second_region}, + ) + + image_rows = df[df["type"] == "image"] + assert [json.loads(value)["asset_region"] for value in image_rows["extra_metadata"]] == [ + first_region, + second_region, + ] From 6ec95961e412f6e682fd96c64dfb5698b0973486 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 16 Oct 2026 15:37:50 +0800 Subject: [PATCH 2/5] feat: add extra metadata support for image assets in markdown and pdf parsing - Introduced `extra_metadata` parameter in `build_image_asset_row` to accommodate additional image region information. - Updated `ParsedRow` and related methods to handle the new `extra_metadata` field. - Modified markdown and pdf parsing functions to integrate image region metadata, ensuring accurate representation in output. - Adjusted tests to validate the inclusion of `extra_metadata` in various parsing scenarios. --- .../document_parser/assets/inline_asset.py | 2 + .../document_parser/formats/image/parser.py | 4 +- .../formats/markdown/image_asset.py | 19 +++++++++ .../formats/markdown/layout_bbox_index.py | 29 +++++++++----- .../formats/markdown/parse_state.py | 20 ++-------- .../formats/markdown/parser.py | 14 ++----- .../formats/markdown/table_embedded_images.py | 4 ++ .../document_parser/formats/pdf/parser.py | 6 +-- .../formats/pdf/shard_splitter.py | 31 +++++++------- .../document_parser/support/parser_rows.py | 13 ++++-- .../services/page_memory/memory_service.py | 8 ++-- .../app/services/page_memory/page_assets.py | 9 ----- apps/worker/scripts/debug_text_track.py | 8 ++-- .../contract/test_parse_task_contract.py | 4 +- .../test_text_track_image_bbox_contract.py | 40 ++++++++++++++++++- .../tests/unit/test_pdf_shard_s3_upload.py | 6 +-- .../test_pdf_split_object_out_of_range.py | 9 ++--- .../services/storage/zip_chunk_schema.py | 13 +++--- 18 files changed, 144 insertions(+), 95 deletions(-) diff --git a/apps/worker/app/services/document_parser/assets/inline_asset.py b/apps/worker/app/services/document_parser/assets/inline_asset.py index 27348a19b..0a28d1ecf 100644 --- a/apps/worker/app/services/document_parser/assets/inline_asset.py +++ b/apps/worker/app/services/document_parser/assets/inline_asset.py @@ -13,6 +13,7 @@ def build_image_asset_row( keywords: str = "", entities: str = "", asset_title: str = "", + extra_metadata: str = "", ) -> ParsedRow: return ParsedRow( content=content, @@ -26,6 +27,7 @@ def build_image_asset_row( addtime=addtime, entities=entities, asset_title=asset_title, + extra_metadata=extra_metadata, ) diff --git a/apps/worker/app/services/document_parser/formats/image/parser.py b/apps/worker/app/services/document_parser/formats/image/parser.py index 5357b10d8..aa9acfa81 100755 --- a/apps/worker/app/services/document_parser/formats/image/parser.py +++ b/apps/worker/app/services/document_parser/formats/image/parser.py @@ -11,7 +11,7 @@ from app.services.document_parser.tables.dataframe_helpers import process_dup_paths_df from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time from app.services.document_parser.support.parser_rows import ( - PARSER_ROW_COLUMNS, + PARSER_DF_COLUMNS, ParsedRow, ParsedRowsBuilder, ) @@ -236,7 +236,7 @@ def parse_image( from shared.core.constants import ProcessingConstants if discard_undersized_image_file(img_path, label=f"image {filename}"): - return pd.DataFrame(columns=list(PARSER_ROW_COLUMNS)) + return pd.DataFrame(columns=list(PARSER_DF_COLUMNS)) # Extract image content client = _get_vision_client() diff --git a/apps/worker/app/services/document_parser/formats/markdown/image_asset.py b/apps/worker/app/services/document_parser/formats/markdown/image_asset.py index 1958bdbf0..fc13cb27d 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/image_asset.py +++ b/apps/worker/app/services/document_parser/formats/markdown/image_asset.py @@ -2,6 +2,7 @@ import os from dataclasses import dataclass +from collections.abc import Mapping from pathlib import Path from typing import cast @@ -15,6 +16,10 @@ ImageDeferredSummaryTask, MarkdownDeferredSummaryTask, ) +from app.services.document_parser.formats.markdown.layout_bbox_index import ( + AssetRegion, + image_region_extra_metadata, +) from app.services.document_parser.formats.markdown.parse_state import ParserRowValues from loguru import logger @@ -48,6 +53,7 @@ class MarkdownImageAssetRequest: summary_image: bool row_index: int rename_on_summary: bool = True + image_regions: Mapping[str, AssetRegion] | None = None def build_markdown_image_asset( @@ -72,11 +78,19 @@ def build_markdown_image_asset( with open(source_path, "rb") as image_file: image_binary_hash = perceptual_hash(image_file.read()) + # Look up the region by this reference's original hash file name, before the + # file is renamed. A duplicate reuses the first image's content but keeps its + # own region. + extra_metadata = image_region_extra_metadata( + request.image_regions, request.image_path + ) + if image_binary_hash in request.seen_images: return _build_duplicate_image_asset( source_path=source_path, cache_entry=request.seen_images[image_binary_hash], timestamp=request.timestamp, + extra_metadata=extra_metadata, ) relative_image_path = f"images/{request.image_name}{image_suffix}" @@ -102,6 +116,7 @@ def build_markdown_image_asset( summary=image_summary_field, know_id=image_know_id, timestamp=request.timestamp, + extra_metadata=extra_metadata, ) cache_entry = { "relative_img_path": relative_image_path, @@ -181,6 +196,7 @@ def _build_duplicate_image_asset( source_path: Path, cache_entry: dict[str, str], timestamp: str, + extra_metadata: str, ) -> MarkdownImageAsset: relative_path = cache_entry["relative_img_path"] row_values = _build_image_row_values( @@ -189,6 +205,7 @@ def _build_duplicate_image_asset( summary=cache_entry["img_summary_field"], know_id=cache_entry["temp_uid"], timestamp=timestamp, + extra_metadata=extra_metadata, ) try: source_path.unlink() @@ -220,6 +237,7 @@ def _build_image_row_values( summary: str, know_id: str, timestamp: str, + extra_metadata: str, ) -> ParserRowValues: image_row = build_image_asset_row( content=content, @@ -227,6 +245,7 @@ def _build_image_row_values( summary=summary, know_id=know_id, addtime=timestamp, + extra_metadata=extra_metadata, ) return cast(ParserRowValues, image_row.to_list()) diff --git a/apps/worker/app/services/document_parser/formats/markdown/layout_bbox_index.py b/apps/worker/app/services/document_parser/formats/markdown/layout_bbox_index.py index 88d1c5458..2b8394647 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/layout_bbox_index.py +++ b/apps/worker/app/services/document_parser/formats/markdown/layout_bbox_index.py @@ -71,9 +71,6 @@ def load_layout_image_regions( "coord_system": COORD_SYSTEM_MINERU_LAYOUT, "page_width": page_width, "page_height": page_height, - # The ZIP/API asset_region contract requires this field. - # MinerU crops the figure at its layout bbox, no margin. - "crop_margin_px": 0, }, ) return regions @@ -87,10 +84,13 @@ def load_shard_image_regions( ``shard_page_numbers[i]`` is the page map of ``shard_output_dirs[i]`` (see ``load_layout_image_regions``). When the same image file name appears in - several shards, the earliest shard wins, matching document order. + several shards, the earliest shard wins, matching document order. The two + sequences must be aligned one-to-one; a length mismatch raises. """ merged: dict[str, AssetRegion] = {} - for shard_dir, page_numbers in zip(shard_output_dirs, shard_page_numbers): + for shard_dir, page_numbers in zip( + shard_output_dirs, shard_page_numbers, strict=True + ): if shard_dir is None: continue shard_regions = load_layout_image_regions( @@ -101,13 +101,22 @@ def load_shard_image_regions( return merged -def image_region_metadata( - image_regions: Mapping[str, AssetRegion], +def image_region_extra_metadata( + image_regions: Mapping[str, AssetRegion] | None, image_path: str, -) -> dict[str, AssetRegion] | None: - """``extra_metadata`` for the image row of a markdown image reference.""" +) -> str: + """Serialized ``extra_metadata`` row value for a markdown image reference. + + ``image_path`` is the reference as written in the markdown (hash file name, + before ``build_markdown_image_asset`` renames the file). Returns ``""`` when + the image has no indexed region. + """ + if not image_regions: + return "" region = image_regions.get(os.path.basename(image_path)) - return {"asset_region": region} if region is not None else None + if region is None: + return "" + return json.dumps({"asset_region": region}, ensure_ascii=False) def _resolve_page_num(page_idx: int, page_numbers: Sequence[int] | None) -> int | None: diff --git a/apps/worker/app/services/document_parser/formats/markdown/parse_state.py b/apps/worker/app/services/document_parser/formats/markdown/parse_state.py index 9fa3c76a9..36d88c942 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/parse_state.py +++ b/apps/worker/app/services/document_parser/formats/markdown/parse_state.py @@ -1,6 +1,6 @@ from __future__ import annotations -import json + from collections.abc import Callable from dataclasses import dataclass, field from typing import Any @@ -31,7 +31,6 @@ class MarkdownParseState: timestamp: str row_updater: RowUpdater rows: list[ParserRowValues] = field(default_factory=list) - row_extra_metadata: dict[int, str] = field(default_factory=dict) content_items: list[str] = field(default_factory=list) path_stack: list[tuple[str, int]] = field(default_factory=list) inner_paths: list[str] = field(default_factory=list) @@ -126,15 +125,7 @@ def append_content_item(self, item: str) -> None: def append_plain_text(self, text: str) -> None: self.content_items.append(text.strip()) - def append_row( - self, - row: ParserRowValues, - extra_metadata: dict[str, Any] | None = None, - ) -> None: - if extra_metadata: - self.row_extra_metadata[len(self.rows)] = json.dumps( - extra_metadata, ensure_ascii=False - ) + def append_row(self, row: ParserRowValues) -> None: self.rows.append(row) def schedule_deferred_task(self, task: MarkdownDeferredSummaryTask) -> None: @@ -175,12 +166,9 @@ def to_dataframe(self) -> pd.DataFrame: page_nums=str(row_values[10]), entities=str(row_values[11]) if len(row_values) > 11 else "", asset_title=str(row_values[12]) if len(row_values) > 12 else "", + extra_metadata=str(row_values[13]) if len(row_values) > 13 else "", ) ) - rows_df = rows_builder.to_dataframe() - rows_df["extra_metadata"] = [ - self.row_extra_metadata.get(index, "") for index in range(len(rows_df)) - ] - return process_dup_paths_df(rows_df) + return process_dup_paths_df(rows_builder.to_dataframe()) diff --git a/apps/worker/app/services/document_parser/formats/markdown/parser.py b/apps/worker/app/services/document_parser/formats/markdown/parser.py index 1abd5a173..d84f7f53a 100755 --- a/apps/worker/app/services/document_parser/formats/markdown/parser.py +++ b/apps/worker/app/services/document_parser/formats/markdown/parser.py @@ -14,10 +14,7 @@ build_markdown_image_name, build_markdown_image_asset, ) -from app.services.document_parser.formats.markdown.layout_bbox_index import ( - AssetRegion, - image_region_metadata, -) +from app.services.document_parser.formats.markdown.layout_bbox_index import AssetRegion from app.services.document_parser.formats.markdown.parse_state import MarkdownParseState from app.services.document_parser.formats.markdown.table_asset import ( MarkdownTableAssetRequest, @@ -405,6 +402,7 @@ def parse_md( seen_images=parser_state.seen_images, summary_image=bool(base_llm_paras["summary_image"]), row_index=len(parser_state.rows), + image_regions=image_regions, ) ) if ( @@ -412,12 +410,7 @@ def parse_md( and image_asset.row_values is not None ): parser_state.append_content_item(image_asset.content_item) - parser_state.append_row( - image_asset.row_values, - extra_metadata=image_region_metadata( - image_regions or {}, img_path - ), - ) + parser_state.append_row(image_asset.row_values) if ( image_asset.cache_key is not None and image_asset.cache_entry is not None @@ -464,6 +457,7 @@ def parse_md( output_dir=output_dir, image_dir=img_dir, summary_image=bool(base_llm_paras["summary_image"]), + image_regions=image_regions, ) for image_asset in embedded.image_assets: if image_asset.row_values is not None: diff --git a/apps/worker/app/services/document_parser/formats/markdown/table_embedded_images.py b/apps/worker/app/services/document_parser/formats/markdown/table_embedded_images.py index ab82543d5..76c44314c 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/table_embedded_images.py +++ b/apps/worker/app/services/document_parser/formats/markdown/table_embedded_images.py @@ -3,6 +3,7 @@ from __future__ import annotations import re +from collections.abc import Mapping from dataclasses import dataclass from bs4 import BeautifulSoup, Tag @@ -14,6 +15,7 @@ build_markdown_image_asset, build_markdown_image_name, ) +from app.services.document_parser.formats.markdown.layout_bbox_index import AssetRegion from app.services.document_parser.formats.markdown.parse_state import MarkdownParseState from shared.utils.chunk_refs import build_chunk_ref @@ -46,6 +48,7 @@ def extract_table_embedded_images( output_dir: str, image_dir: str, summary_image: bool, + image_regions: Mapping[str, AssetRegion] | None = None, ) -> TableEmbeddedImagesResult: """Pull assets out of a table, rename them, and rewrite HTML srcs. @@ -88,6 +91,7 @@ def extract_table_embedded_images( summary_image=summary_image, row_index=len(parser_state.rows) + len(image_assets), rename_on_summary=False, + image_regions=image_regions, ) ) if image_asset.discarded_undersized: diff --git a/apps/worker/app/services/document_parser/formats/pdf/parser.py b/apps/worker/app/services/document_parser/formats/pdf/parser.py index accc74ecf..18456461a 100755 --- a/apps/worker/app/services/document_parser/formats/pdf/parser.py +++ b/apps/worker/app/services/document_parser/formats/pdf/parser.py @@ -125,8 +125,8 @@ def _parse_pdf_via_shards( ) from app.services.document_parser.formats.pdf.shard_splitter import ( map_agent_shards, - shard_page_numbers, split_pdf, + split_pdf_page_numbers, ) from shared.services.ai.token_tracking import ( bind_token_tracker, @@ -190,7 +190,7 @@ def _parse_pdf_via_shards( work_dir = os.path.join(output_dir, "_shards") os.makedirs(work_dir, exist_ok=True) with stage_timer("pdf.split", filename=filename): - shard_pdf_paths, _page_remap = split_pdf( + shard_pdf_paths = split_pdf( pdf_path, merged_shards, work_dir, exclude_pages=toc_pages if toc_pages else None, ) @@ -359,7 +359,7 @@ def _predict_shard_headings_with_tracking( shard_page_maps = ( [None] if fast_path_original_pdf - else [shard_page_numbers(shard, toc_pages) for shard in merged_shards] + else split_pdf_page_numbers(merged_shards, toc_pages) ) image_regions = load_shard_image_regions(shard_output_dirs, shard_page_maps) 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 7bfddb1c5..af47d9e48 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 @@ -57,6 +57,18 @@ def shard_page_numbers(shard: MergedShard, exclude_pages: set[int]) -> list[int] ] +def split_pdf_page_numbers( + shards: list[MergedShard], exclude_pages: set[int] +) -> list[list[int]]: + """Page map of each shard PDF ``split_pdf`` writes, in the same order. + + A shard whose pages are all excluded produces no PDF, so it has no entry + here either; index ``i`` therefore lines up with ``split_pdf``'s ``i``-th path. + """ + page_maps = [shard_page_numbers(shard, exclude_pages) for shard in shards] + return [page_map for page_map in page_maps if page_map] + + 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) @@ -67,7 +79,7 @@ def split_pdf( shards: list[MergedShard], work_dir: str, exclude_pages: set[int] | None = None, -) -> tuple[list[str], dict[int, int] | None]: +) -> list[str]: """Physically split PDF into sub-PDFs using PyMuPDF. Args: @@ -78,11 +90,8 @@ def split_pdf( (e.g. TOC pages detected during PROFILE). Returns: - (shard_paths, page_remap) - - shard_paths: one temp PDF path per shard. - - page_remap: when pages are excluded, maps each shard's local - 0-based page index to the original 1-based page number. - ``None`` when no pages are excluded. + One temp PDF path per shard that kept at least one page. Page ``i`` of + each PDF maps to the original page ``split_pdf_page_numbers`` reports. """ doc = pymupdf.open(pdf_path) rewritten_path: str | None = None @@ -114,18 +123,15 @@ def _extract_shards( shards: list[MergedShard], work_dir: str, exclude_pages: set[int] | None, -) -> tuple[list[str], dict[int, int] | None]: +) -> list[str]: paths: list[str] = [] - page_remap: dict[int, int] | None = None if exclude_pages: - page_remap = {} logger.info( f"📌 Excluding {len(exclude_pages)} pages from PDF: " f"{sorted(exclude_pages)}" ) - global_new_idx = 0 # running counter across all shards for shard in shards: sub_doc = pymupdf.open() try: @@ -136,9 +142,6 @@ def _extract_shards( from_page=page_num - 1, to_page=page_num - 1, ) - if page_remap is not None: - page_remap[global_new_idx] = page_num - global_new_idx += 1 shard_included += 1 shard_path = os.path.join(work_dir, f"shard_{shard.shard_index}.pdf") @@ -159,4 +162,4 @@ def _extract_shards( f"({shard_included} included" f"{f', {excluded_in_shard} excluded' if excluded_in_shard else ''})" ) - return paths, page_remap + return paths diff --git a/apps/worker/app/services/document_parser/support/parser_rows.py b/apps/worker/app/services/document_parser/support/parser_rows.py index 739946a50..414c19b67 100644 --- a/apps/worker/app/services/document_parser/support/parser_rows.py +++ b/apps/worker/app/services/document_parser/support/parser_rows.py @@ -26,6 +26,11 @@ def _resolve_parser_columns() -> tuple[str, ...]: PARSER_ROW_COLUMNS: tuple[str, ...] = _resolve_parser_columns() +# ``extra_metadata`` is the trailing JSON-string column merged into chunk +# metadata by ``dataframe_to_chunks``. It is not part of the env-configurable +# ``ALL_DF_COLS`` list, the same way the page track appends it. +PARSER_DF_COLUMNS: tuple[str, ...] = (*PARSER_ROW_COLUMNS, "extra_metadata") + def serialize_entities(entities: Any) -> str: """Serialize typed entities (§4.4) into the JSON string stored in the row. @@ -65,6 +70,7 @@ class ParsedRow: length: int | None = None entities: str = "" asset_title: str = "" + extra_metadata: str = "" def to_list(self) -> list[object]: content_length = self.length if self.length is not None else len(self.content) @@ -82,10 +88,11 @@ def to_list(self) -> list[object]: self.page_nums, self.entities, self.asset_title, + self.extra_metadata, ] def to_dict(self) -> dict[str, object]: - return dict(zip(PARSER_ROW_COLUMNS, self.to_list())) + return dict(zip(PARSER_DF_COLUMNS, self.to_list())) class ParsedRowsBuilder: @@ -101,7 +108,7 @@ def extend(self, rows: list[ParsedRow]) -> None: def to_dataframe(self) -> pd.DataFrame: return pd.DataFrame( [row.to_list() for row in self._rows], - columns=Index(PARSER_ROW_COLUMNS), + columns=Index(PARSER_DF_COLUMNS), ) @@ -141,5 +148,5 @@ def apply_asset_summary(row: list[Any], result: Any) -> None: def _ensure_row_width(row: list[Any]) -> None: """Pad a legacy 11-column row out to the current schema width in place.""" - while len(row) < len(PARSER_ROW_COLUMNS): + while len(row) < len(PARSER_DF_COLUMNS): row.append("") diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index b18d60101..bd97ddaa0 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -15,7 +15,7 @@ ) from app.services.document_parser.profiling.doc_profiler import profile_document from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time -from app.services.document_parser.support.parser_rows import PARSER_ROW_COLUMNS +from app.services.document_parser.support.parser_rows import PARSER_DF_COLUMNS from app.services.document_parser.support.stage_profiler import stage_timer from app.services.page_memory.normalizer import normalize_to_pdf from app.services.page_memory._utils import ( @@ -231,7 +231,7 @@ def _build_page_dataframe( anatomy = getattr(profile, "anatomy", None) page_count = max(int(profile.page_count or 0), 0) if page_count <= 0: - return pd.DataFrame(columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) + return pd.DataFrame(columns=pd.Index(PARSER_DF_COLUMNS)) asset_extraction_enabled = page_memory_config.asset_extraction_enabled @@ -465,7 +465,7 @@ def _build_page_dataframe( "page_to_node": _page_to_node_map(rows), }, ) - return pd.DataFrame(rows, columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) + return pd.DataFrame(rows, columns=pd.Index(PARSER_DF_COLUMNS)) def _build_hierarchy_scopes( @@ -854,7 +854,7 @@ def _build_whole_doc_dataframe( page_info=page_scope_info(pages), variables={"summary": summary, "verdict": verdict}, ) - return pd.DataFrame([row], columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) + return pd.DataFrame([row], columns=pd.Index(PARSER_DF_COLUMNS)) def _build_summary(*, filename: str, page_count: int, raw_text: str) -> str: diff --git a/apps/worker/app/services/page_memory/page_assets.py b/apps/worker/app/services/page_memory/page_assets.py index 1821c5619..43f05d962 100644 --- a/apps/worker/app/services/page_memory/page_assets.py +++ b/apps/worker/app/services/page_memory/page_assets.py @@ -681,15 +681,6 @@ def _is_meaningful_frame(frame: pd.DataFrame) -> bool: return bool(non_empty.values.sum() >= 2) -def _normalize_keywords(value: object) -> list[str]: - if isinstance(value, list): - return [str(item).strip() for item in value if str(item).strip()] - if isinstance(value, str): - delimiter = ";" if ";" in value else "," - return [item.strip() for item in value.split(delimiter) if item.strip()] - return [] - - def _safe_float(value: object, *, default: float) -> float: try: return float(str(value)) diff --git a/apps/worker/scripts/debug_text_track.py b/apps/worker/scripts/debug_text_track.py index 4b36c1217..8702c04ec 100644 --- a/apps/worker/scripts/debug_text_track.py +++ b/apps/worker/scripts/debug_text_track.py @@ -396,7 +396,7 @@ def _stage_mineru_pdf( parse_via_full(pdf_path, filename, shard_out) shard_output_dirs = [shard_out] else: - shard_pdf_paths, _page_remap = split_pdf( + shard_pdf_paths = split_pdf( pdf_path, merged_shards, work_dir, exclude_pages=toc_pages if toc_pages else None, ) @@ -672,7 +672,7 @@ def _stage_full_pdf( from app.services.document_parser.formats.markdown.parser import parse_md from app.services.document_parser.formats.pdf.shard_splitter import ( map_agent_shards, - shard_page_numbers, + split_pdf_page_numbers, ) from app.services.document_parser.orchestration.postprocess import apply_parse_postprocess @@ -684,10 +684,10 @@ def _stage_full_pdf( shard_page_maps = ( [None] if fast_path - else [shard_page_numbers(shard, toc_pages) for shard in merged_shards] + else split_pdf_page_numbers(merged_shards, toc_pages) ) image_regions = load_shard_image_regions( - [str(out_dir / "_shards" / f"shard_{i}") for i in range(len(merged_shards))], + [str(out_dir / "_shards" / f"shard_{i}") for i in range(len(shard_page_maps))], shard_page_maps, ) diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index 0e2249e69..3d60d8c52 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -1275,7 +1275,7 @@ def _fake_split_pdf(pdf_path_arg, shards, work_dir, exclude_pages=None): shard_path = Path(work_dir) / f"shard_{shard.shard_index}.pdf" shard_path.write_bytes(b"%PDF-1.4\n%%EOF\n") paths.append(str(shard_path)) - return paths, None + return paths def _fake_parse_via_full(shard_pdf, shard_filename, shard_out, s3_key=None): parse_s3_keys.append(s3_key) @@ -1441,7 +1441,7 @@ def _fake_split_pdf(pdf_path_arg, shards, work_dir, exclude_pages=None): shard_path = Path(work_dir) / f"shard_{shard.shard_index}.pdf" shard_path.write_bytes(b"%PDF-1.4\n%%EOF\n") paths.append(str(shard_path)) - return paths, None + return paths def _fake_parse_via_full(shard_pdf, shard_filename, shard_out, s3_key=None): shard_index = 0 if "shard0" in shard_filename else 1 diff --git a/apps/worker/tests/contract/test_text_track_image_bbox_contract.py b/apps/worker/tests/contract/test_text_track_image_bbox_contract.py index 2e6113944..08001e697 100644 --- a/apps/worker/tests/contract/test_text_track_image_bbox_contract.py +++ b/apps/worker/tests/contract/test_text_track_image_bbox_contract.py @@ -19,9 +19,12 @@ load_shard_image_regions, ) from app.services.document_parser.formats.markdown.parser import parse_md # noqa: E402 +import pytest # noqa: E402 + from app.services.document_parser.formats.pdf.shard_splitter import ( # noqa: E402 MergedShard, shard_page_numbers, + split_pdf_page_numbers, ) from shared.core.constants.processing import ProcessingConstants # noqa: E402 from shared.services.chunks.dataframe_chunk_converter import ( # noqa: E402 @@ -72,7 +75,6 @@ def _region(page_num: int, bbox: list[int]) -> dict[str, Any]: "coord_system": "mineru_layout", "page_width": 960, "page_height": 540, - "crop_margin_px": 0, } @@ -173,6 +175,22 @@ def test_shard_page_numbers_skip_excluded_pages() -> None: assert shard_page_numbers(shard, {7, 8, 100}) == [5, 6, 9] +def test_split_pdf_page_numbers_drop_shards_without_a_pdf() -> None: + shards = [ + MergedShard(shard_index=0, page_start=1, page_end=2), + MergedShard(shard_index=1, page_start=3, page_end=4), + MergedShard(shard_index=2, page_start=5, page_end=6), + ] + + # Shard 1 is entirely excluded, so split_pdf writes no PDF for it. + assert split_pdf_page_numbers(shards, {3, 4, 5}) == [[1, 2], [6]] + + +def test_shard_regions_reject_misaligned_page_maps(tmp_path: Path) -> None: + with pytest.raises(ValueError): + load_shard_image_regions([str(tmp_path / "shard_0")], [[1], [2]]) + + def _write_jpeg(path: Path, *, seed: int) -> None: path.parent.mkdir(parents=True, exist_ok=True) width = height = 640 @@ -246,6 +264,26 @@ def test_parse_md_attaches_asset_region_to_image_chunks_only(tmp_path: Path) -> assert "asset_region" not in zip_by_type["table"]["metadata"] +def test_parse_md_attaches_asset_region_to_table_embedded_images( + tmp_path: Path, +) -> None: + hash_name = "e" * 64 + ".jpg" + _write_jpeg(tmp_path / "images" / hash_name, seed=59) + lines = [ + "# Results", + "" + f'
Flow
', + ] + region = _region(9, [5, 6, 200, 300]) + + df = _parse_phase_b(tmp_path, lines, {hash_name: region}) + chunks = dataframe_to_chunks(df) + + by_type = {str(chunk["type"]): chunk for chunk in chunks} + assert by_type["image"]["metadata"]["asset_region"] == region + assert "asset_region" not in by_type["table"]["metadata"] + + def test_parse_md_without_image_regions_writes_no_asset_region(tmp_path: Path) -> None: hash_name = "b" * 64 + ".jpg" _write_jpeg(tmp_path / "images" / hash_name, seed=41) diff --git a/apps/worker/tests/unit/test_pdf_shard_s3_upload.py b/apps/worker/tests/unit/test_pdf_shard_s3_upload.py index 9c638a54a..bc5d219ea 100644 --- a/apps/worker/tests/unit/test_pdf_shard_s3_upload.py +++ b/apps/worker/tests/unit/test_pdf_shard_s3_upload.py @@ -88,7 +88,7 @@ def fake_split_pdf(pdf_path, shards, work_dir, exclude_pages=None): events.append("split") shard_path = Path(work_dir) / "shard_0.pdf" shard_path.write_bytes(b"%PDF-1.4") - return [str(shard_path)], None + return [str(shard_path)] class _FakeStorage: def upload_source_file( @@ -158,7 +158,7 @@ def test_failed_shard_upload_does_not_start_mineru( def fake_split_pdf(pdf_path, shards, work_dir, exclude_pages=None): shard_path = Path(work_dir) / "shard_0.pdf" shard_path.write_bytes(b"%PDF-1.4") - return [str(shard_path)], None + return [str(shard_path)] class _FailingStorage: def upload_source_file( @@ -202,7 +202,7 @@ def test_unverified_shard_upload_does_not_start_mineru( def fake_split_pdf(pdf_path, shards, work_dir, exclude_pages=None): shard_path = Path(work_dir) / "shard_0.pdf" shard_path.write_bytes(b"%PDF-1.4") - return [str(shard_path)], None + return [str(shard_path)] class _InvisibleStorage: def upload_source_file( 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 index 0b0fd35a3..ac0ee4b64 100644 --- 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 @@ -50,13 +50,12 @@ def test_split_pdf_copies_pages(tmp_path: Path) -> None: work_dir = tmp_path / "shards" work_dir.mkdir() - paths, remap = split_pdf( + paths = 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 @@ -82,13 +81,12 @@ def _insert_pdf(self: pymupdf.Document, *args: object, **kwargs: object) -> None monkeypatch.setattr(pymupdf.Document, "insert_pdf", _insert_pdf) - paths, remap = split_pdf( + paths = 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 @@ -119,13 +117,12 @@ def test_split_incremental_quartz_pdf_around_failing_page(tmp_path: Path) -> Non work_dir = tmp_path / "shards" work_dir.mkdir() - paths, remap = split_pdf( + paths = 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/packages/shared-python/shared/services/storage/zip_chunk_schema.py b/packages/shared-python/shared/services/storage/zip_chunk_schema.py index 0449671bd..d9043cb9c 100644 --- a/packages/shared-python/shared/services/storage/zip_chunk_schema.py +++ b/packages/shared-python/shared/services/storage/zip_chunk_schema.py @@ -114,10 +114,6 @@ def format_chunks( metadata["keywords"] = existing_metadata.get("keywords") or chunk.get( "keywords", [] ) - if asset_region := _normalize_asset_region( - existing_metadata.get("asset_region") - ): - metadata["asset_region"] = asset_region metadata["connect_to"] = _build_embed_connect_to( chunk=chunk, chunk_type_str=chunk_type_str, @@ -250,7 +246,6 @@ def _normalize_asset_region(value: Any) -> dict[str, Any] | None: coord_system = str(value.get("coord_system") or "").strip() page_width = _positive_int_or_none(value.get("page_width")) page_height = _positive_int_or_none(value.get("page_height")) - crop_margin_px = _non_negative_int_or_none(value.get("crop_margin_px")) if ( page_num is None or bbox is None @@ -258,18 +253,20 @@ def _normalize_asset_region(value: Any) -> dict[str, Any] | None: or not coord_system or page_width is None or page_height is None - or crop_margin_px is None ): return None - return { + region: dict[str, Any] = { "page_num": page_num, "bbox": bbox, "coord_origin": coord_origin, "coord_system": coord_system, "page_width": page_width, "page_height": page_height, - "crop_margin_px": crop_margin_px, } + # Only the page track crops with a margin; the text track does not crop. + if (crop_margin_px := _non_negative_int_or_none(value.get("crop_margin_px"))) is not None: + region["crop_margin_px"] = crop_margin_px + return region def _bbox_or_none(value: Any) -> list[int] | None: From 271a3d5ca0dcfbd99de457c5799900e45cb110e4 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 16 Oct 2026 17:48:43 +0800 Subject: [PATCH 3/5] feat: enhance document profiling with multi-category support and improved metadata - Updated `DocumentProfile` to support a list of categories and an additional `category_label` for better classification. - Modified the coarse profiling process to extract domain IDs and category labels from input data, ensuring accurate document classification. - Enhanced the parsing instructions to reflect the new structure, requiring domain IDs and category labels for document classification. - Adjusted related functions and tests to accommodate the new multi-category and label features, ensuring consistency across document types. --- AGENTS.md | 7 +- .../coarse_profile/classifier.py | 23 +- .../document_agent/coarse_profile/prompts.py | 15 +- .../app/services/document_agent/manifest.py | 4 +- .../tools/find_toc_anchor_pages.py | 16 + .../formats/markdown/layout_bbox_index.py | 4 +- .../document_parser/formats/pdf/parser.py | 105 +++--- .../formats/pdf/shard_merger.py | 24 +- .../formats/pdf/shard_splitter.py | 103 +++-- .../document_parser/profiling/doc_profiler.py | 5 +- .../profiling/domain_taxonomy.py | 254 +++++++++++++ .../profiling/profile_model.py | 7 +- apps/worker/scripts/debug_text_track.py | 75 ++-- .../scripts/page_memory/_debug_pm_shared.py | 9 +- .../page_memory/debug_pm_stage0_bootstrap.py | 1 + .../test_coarse_profile_protocol_contract.py | 70 +++- .../test_doc_profile_anatomy_contract.py | 43 ++- .../contract/test_domain_taxonomy_contract.py | 67 ++++ .../contract/test_parse_task_contract.py | 355 +++++++++++++++++- .../test_text_track_image_bbox_contract.py | 18 +- .../tests/unit/test_pdf_shard_s3_upload.py | 37 +- .../test_pdf_split_object_out_of_range.py | 49 ++- .../tests/unit/test_toc_anchor_selection.py | 9 + 23 files changed, 1106 insertions(+), 194 deletions(-) create mode 100644 apps/worker/app/services/document_parser/profiling/domain_taxonomy.py create mode 100644 apps/worker/tests/contract/test_domain_taxonomy_contract.py diff --git a/AGENTS.md b/AGENTS.md index 9e82c3f5d..f458a3852 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -129,9 +129,10 @@ This is the typed `ParseOutput` entry for all file types. The parser flow: 1. **Profiles** the document via `profiling.doc_profiler.profile_document()`. PDF profiling uses `document_agent` as the single PyMuPDF feature source, then runs - VLM coarse classification with two fields: open semantic `category` (for example, - `Financial Prospectus`) and routing-only `routing_category` - (`atlas/scanned/slides/generic`). Oversized non-atlas PDFs additionally run the + VLM coarse classification with three fields: `category` (domain codes picked from + the numbered ISIC Rev.5 leaf list in `profiling/domain_taxonomy.py`, multiple + allowed), `category_label` (open short phrase, for example `Financial Prospectus`), + and routing-only `routing_category` (`atlas/scanned/slides/generic`). Oversized non-atlas PDFs additionally run the structural anatomy stage once at the entry point and pass the resulting shard plan to PDF parsing. 2. **Routes** to the appropriate parser based on file extension. diff --git a/apps/worker/app/services/document_agent/coarse_profile/classifier.py b/apps/worker/app/services/document_agent/coarse_profile/classifier.py index 482586ad9..e2e7fe90b 100644 --- a/apps/worker/app/services/document_agent/coarse_profile/classifier.py +++ b/apps/worker/app/services/document_agent/coarse_profile/classifier.py @@ -19,6 +19,10 @@ from app.services.document_agent.coarse_profile.prompts import COARSE_PROFILE_INSTRUCTIONS from app.services.document_agent.manifest import DocumentProfile, ToolContext, ToolResult from app.services.document_agent.visual import render_pages +from app.services.document_parser.profiling.domain_taxonomy import ( + OTHER_CODE, + domain_codes_from_ids, +) from app.services.document_parser.profiling.taxonomy import PdfRoutingCategory PAGE_KIND_DEFINITIONS = { @@ -137,10 +141,15 @@ def _parse_profile(raw: str) -> DocumentProfile: data = json.loads(raw) if not isinstance(data, dict): raise ValueError("coarse profile output must be a JSON object") - category = " ".join(str(data.get("category") or "unknown document").split()[:5]) - routing_category = PdfRoutingCategory.normalize( - data.get("routing_category") or data.get("category") - ).value + try: + category = domain_codes_from_ids(data.get("domain_ids")) + except ValueError as exc: + logger.warning("[document_agent] invalid domain_ids, using OTHER: {}", exc) + category = [OTHER_CODE] + category_label = " ".join( + str(data.get("category_label") or "unknown document").split()[:5] + ) + routing_category = PdfRoutingCategory.normalize(data.get("routing_category")).value raw_is_scanned = data.get("is_scanned") if isinstance(raw_is_scanned, bool): is_scanned = raw_is_scanned @@ -155,7 +164,8 @@ def _parse_profile(raw: str) -> DocumentProfile: footer_y = None return DocumentProfile( is_scanned=is_scanned, - category=category or "unknown document", + category=category, + category_label=category_label or "unknown document", routing_category=routing_category, language=str(data.get("language") or "unknown"), header_y=header_y, @@ -197,7 +207,8 @@ def classify(self) -> tuple[DocumentProfile, ToolResult]: if not model: profile = DocumentProfile( is_scanned=False, - category="unknown document", + category=[OTHER_CODE], + category_label="unknown document", routing_category=PdfRoutingCategory.GENERIC.value, ) return profile, ToolResult( diff --git a/apps/worker/app/services/document_agent/coarse_profile/prompts.py b/apps/worker/app/services/document_agent/coarse_profile/prompts.py index aab5d0994..efa3e45f3 100644 --- a/apps/worker/app/services/document_agent/coarse_profile/prompts.py +++ b/apps/worker/app/services/document_agent/coarse_profile/prompts.py @@ -1,11 +1,17 @@ """Prompts for one-shot coarse document profiling.""" +from app.services.document_parser.profiling.domain_taxonomy import DOMAIN_MENU + COARSE_PROFILE_INSTRUCTIONS = ( "You are a document profile classifier. Use page-feature statistics " "and the provided page screenshots to classify the PDF. " - "Return strict JSON only with keys: is_scanned, category, routing_category, " - "language, header_y, footer_y. " - "category is a concise semantic document type in at most 5 English words. " + "Return strict JSON only with keys: is_scanned, domain_ids, category_label, " + "routing_category, language, header_y, footer_y. " + "domain_ids is a non-empty JSON array of integers taken from the numbered " + "domain list below. Pick every numbered entry whose domain the document's " + "subject matter belongs to; section titles without a number are not " + "selectable. Pick the Other entry alone, and only when no listed domain fits. " + "category_label is a concise semantic document type in at most 5 English words. " "routing_category must be one of atlas, scanned, slides, generic. " "Set routing_category=atlas only when pages are primarily drawing/detail " "sheets rather than prose. " @@ -15,7 +21,8 @@ "you observe (largest y) when any header is present, otherwise null; " "footer_y is the highest footer line you observe (smallest y) when any " "footer is present, otherwise null. When both are set, require " - "header_y < footer_y." + "header_y < footer_y.\n" + "Domain list:\n" + DOMAIN_MENU ) __all__ = ["COARSE_PROFILE_INSTRUCTIONS"] diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index 3e2cbfdc6..93f1e4b42 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -49,7 +49,9 @@ def to_dict(self) -> dict[str, Any]: @dataclass class DocumentProfile: is_scanned: bool - category: str + category: list[str] + """Domain codes picked from the domain taxonomy (e.g. ``["6612"]``).""" + category_label: str routing_category: str = "generic" language: str = "unknown" # Content-band margins as fractions of page height (top origin, y down). diff --git a/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py index c73f0f069..5d3e71d4f 100644 --- a/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py +++ b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py @@ -27,6 +27,7 @@ TOC_KEYWORDS = frozenset({"目录", "目次", "contents", "table of contents"}) MAX_ANCHOR_CANDIDATES = 20 +HEAD_WINDOW_MAX_PAGES = 10 # English keywords must sit at a line edge once candidates overflow; CJK 目录/目次 are exempt. EDGE_ANCHORED_KEYWORDS = frozenset({"contents", "table of contents"}) MIN_CONSECUTIVE_DUPLICATE_RUN = 2 @@ -228,6 +229,13 @@ def _select_anchor_pages( return _match_pages(matches), narrowing +def _head_window_pages(page_count: int) -> set[int]: + """First ``min(page_count, 10)`` pages when keyword scan found nothing.""" + if page_count <= 0: + return set() + return set(range(1, min(page_count, HEAD_WINDOW_MAX_PAGES) + 1)) + + @worker def _render_pages_worker( queue, pdf_path: str, pages: list[int], output_dir: str, dpi: int @@ -274,6 +282,14 @@ def find_toc_anchor_pages(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult page_count=total_pages, ) anchor_pages, narrowing = _select_anchor_pages(keyword_matches) + if not anchor_pages and narrowing.get("raw") == 0: + anchor_pages = _head_window_pages(total_pages) + narrowing["head_window"] = len(anchor_pages) + logger.info( + "[find.toc_anchor_pages] no keyword candidates; " + "using first {} pages as head-window fallback", + len(anchor_pages), + ) logger.info("[find.toc_anchor_pages] candidate narrowing: {}", narrowing) if narrowing.get("gave_up"): logger.warning( diff --git a/apps/worker/app/services/document_parser/formats/markdown/layout_bbox_index.py b/apps/worker/app/services/document_parser/formats/markdown/layout_bbox_index.py index 2b8394647..814e91610 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/layout_bbox_index.py +++ b/apps/worker/app/services/document_parser/formats/markdown/layout_bbox_index.py @@ -77,7 +77,7 @@ def load_layout_image_regions( def load_shard_image_regions( - shard_output_dirs: Sequence[str | None], + shard_output_dirs: Sequence[str], shard_page_numbers: Sequence[Sequence[int] | None], ) -> dict[str, AssetRegion]: """Merge the per-shard indexes of a sharded PDF into one index. @@ -91,8 +91,6 @@ def load_shard_image_regions( for shard_dir, page_numbers in zip( shard_output_dirs, shard_page_numbers, strict=True ): - if shard_dir is None: - continue shard_regions = load_layout_image_regions( os.path.join(shard_dir, "layout.json"), page_numbers ) diff --git a/apps/worker/app/services/document_parser/formats/pdf/parser.py b/apps/worker/app/services/document_parser/formats/pdf/parser.py index 18456461a..ab4ef4055 100755 --- a/apps/worker/app/services/document_parser/formats/pdf/parser.py +++ b/apps/worker/app/services/document_parser/formats/pdf/parser.py @@ -7,6 +7,7 @@ load_layout_image_regions, load_shard_image_regions, ) +from app.services.document_parser.formats.pdf.shard_splitter import LiveShard from app.services.document_parser.formats.markdown.parser import parse_md from app.services.document_parser.orchestration.oversized_pdf_policy import ( build_oversized_pdf_processing_failed_exception, @@ -125,8 +126,8 @@ def _parse_pdf_via_shards( ) from app.services.document_parser.formats.pdf.shard_splitter import ( map_agent_shards, + select_live_shards, split_pdf, - split_pdf_page_numbers, ) from shared.services.ai.token_tracking import ( bind_token_tracker, @@ -176,13 +177,15 @@ def _parse_pdf_via_shards( # 4. Parse via MinerU. The 1-shard/no-TOC case keeps the original # PDF/S3 object to avoid temporary split/upload churn. - shard_output_dirs: list[str | None] + live_shards: list[LiveShard] + shard_output_dirs: list[str] concurrency = settings.MINERU_SHARD_CONCURRENCY if fast_path_original_pdf: logger.info("📄 Single shard without TOC pages; using original PDF fast path") with stage_timer("pdf.extract.single_shard_fast", filename=filename): parse_via_full(pdf_path, filename, output_dir, s3_key=s3_key) + live_shards = select_live_shards(merged_shards) shard_output_dirs = [output_dir] else: # Physically split PDF when TOC pages must be excluded or multiple @@ -190,12 +193,12 @@ def _parse_pdf_via_shards( work_dir = os.path.join(output_dir, "_shards") os.makedirs(work_dir, exist_ok=True) with stage_timer("pdf.split", filename=filename): - shard_pdf_paths = split_pdf( + live_shards = split_pdf( pdf_path, merged_shards, work_dir, exclude_pages=toc_pages if toc_pages else None, ) - if not shard_pdf_paths: + if not live_shards: raise RuntimeError( f"PDF shard split for {filename} produced no shard PDFs" ) @@ -205,48 +208,61 @@ def _parse_pdf_via_shards( source_s3_key=s3_key, job_id=job_id, filename=filename, - shard_index=shard_index, + shard_index=live.shard_index, ) - for shard_index, _shard_pdf_path in enumerate(shard_pdf_paths) + for live in live_shards ] with stage_timer("pdf.upload_shards", filename=filename): - _upload_temp_shard_pdfs(shard_pdf_paths, temp_shard_s3_keys) + _upload_temp_shard_pdfs(live_shards, temp_shard_s3_keys) # 5. Parse each shard via MinerU (parallel) only after S3 has the # split objects. URL mode HeadObject then reuses those keys. - shard_output_dirs = [None] * len(shard_pdf_paths) + parsed_output_dirs: list[str | None] = [None] * len(live_shards) - def _parse_single_shard(shard_idx, shard_pdf): + def _parse_single_shard(live: LiveShard, shard_s3_key: str) -> str: assert work_dir is not None - shard_out = os.path.join(work_dir, f"shard_{shard_idx}_output") + shard_out = os.path.join( + work_dir, f"shard_{live.shard_index}_output" + ) os.makedirs(shard_out, exist_ok=True) shard_filename = ( - f"{os.path.splitext(filename)[0]}_shard{shard_idx}.pdf" + f"{os.path.splitext(filename)[0]}_shard{live.shard_index}.pdf" ) - shard_s3_key = temp_shard_s3_keys[shard_idx] logger.info( - f" 🔄 MinerU shard_{shard_idx}: parsing via S3 URL " + f" 🔄 MinerU shard_{live.shard_index}: parsing via S3 URL " f"({shard_s3_key})" ) parse_via_full( - shard_pdf, shard_filename, shard_out, s3_key=shard_s3_key + live.pdf_path, shard_filename, shard_out, s3_key=shard_s3_key ) return shard_out with stage_timer( "pdf.mineru_parallel", filename=filename, - shard_count=len(shard_pdf_paths), + shard_count=len(live_shards), ): with ThreadPoolExecutor(max_workers=concurrency) as executor: futures = { - executor.submit(_parse_single_shard, i, shard_pdf_path): i - for i, shard_pdf_path in enumerate(shard_pdf_paths) + executor.submit( + _parse_single_shard, live, s3_key + ): index + for index, (live, s3_key) in enumerate( + zip(live_shards, temp_shard_s3_keys, strict=True) + ) } for future in as_completed(futures): idx = futures[future] - shard_output_dirs[idx] = future.result() + parsed_output_dirs[idx] = future.result() + + shard_output_dirs = [] + for index, shard_dir in enumerate(parsed_output_dirs): + if shard_dir is None: + raise RuntimeError( + f"Missing MinerU output for shard_{live_shards[index].shard_index}" + ) + shard_output_dirs.append(shard_dir) # 6. Per-shard heading prediction (parallel) @dataclass @@ -262,21 +278,22 @@ class ShardHeadingResult: ) def _predict_shard_headings( - shard_idx: int, + live: LiveShard, shard_out_dir: str, ) -> ShardHeadingResult: """Run full heading prediction pipeline on a single shard's full.md.""" md_path = os.path.join(shard_out_dir, "full.md") if not os.path.exists(md_path): - raise FileNotFoundError(f"shard_{shard_idx}: full.md not found") + raise FileNotFoundError( + f"shard_{live.shard_index}: full.md not found" + ) with open(md_path, "r", encoding="utf-8") as f: md_lines = f.readlines() md_lines = [line.strip() for line in md_lines if line.strip() != ""] md_lines = merge_html_tables(md_lines) - is_first_shard = shard_idx == 0 - agent_shard = agent_shards[shard_idx] + agent_shard = agent_shards[live.shard_index] shard_toc = getattr(agent_shard, "toc_hierarchies", None) lines_with_heading = eval_md_headings( @@ -291,41 +308,42 @@ def _predict_shard_headings( if os.path.exists(os.path.join(shard_out_dir, "layout.json")) else None ), - is_first_shard=is_first_shard, + is_first_shard=live.is_first_shard, ) heading_count = sum(1 for line in lines_with_heading if line.startswith("#")) logger.info( - f" ✅ shard_{shard_idx}: {heading_count} headings identified " + f" ✅ shard_{live.shard_index}: {heading_count} headings identified " f"from {len(lines_with_heading)} lines" ) return ShardHeadingResult( - shard_index=shard_idx, + shard_index=live.shard_index, lines_with_heading=lines_with_heading, heading_count=heading_count, ) def _predict_shard_headings_with_tracking( - shard_idx: int, + live: LiveShard, shard_out_dir: str, ) -> ShardHeadingResult: with bind_token_tracker(token_tracker_root_id): - return _predict_shard_headings(shard_idx, shard_out_dir) + return _predict_shard_headings(live, shard_out_dir) shard_heading_results: list[ShardHeadingResult | None] = [None] * len( - shard_output_dirs + live_shards ) with stage_timer( - "pdf.shard_headings", filename=filename, shard_count=len(shard_output_dirs) + "pdf.shard_headings", filename=filename, shard_count=len(live_shards) ): with ThreadPoolExecutor(max_workers=concurrency) as executor: futures = { executor.submit( - _predict_shard_headings_with_tracking, i, shard_dir - ): i - for i, shard_dir in enumerate(shard_output_dirs) - if shard_dir is not None + _predict_shard_headings_with_tracking, live, shard_dir + ): index + for index, (live, shard_dir) in enumerate( + zip(live_shards, shard_output_dirs, strict=True) + ) } for future in as_completed(futures): idx = futures[future] @@ -335,11 +353,16 @@ def _predict_shard_headings_with_tracking( complete_heading_results: list[ShardHeadingResult] = [] for index, result in enumerate(shard_heading_results): if result is None: - raise RuntimeError(f"Missing heading result for shard_{index}") + raise RuntimeError( + f"Missing heading result for shard_{live_shards[index].shard_index}" + ) complete_heading_results.append(result) all_lines_with_heading: list[str] = merge_shard_lines( [result.lines_with_heading for result in complete_heading_results], + shard_indexes=[ + result.shard_index for result in complete_heading_results + ], ) total_headings = sum( 1 for line in all_lines_with_heading if line.startswith("#") @@ -359,7 +382,7 @@ def _predict_shard_headings_with_tracking( shard_page_maps = ( [None] if fast_path_original_pdf - else split_pdf_page_numbers(merged_shards, toc_pages) + else [live.page_numbers for live in live_shards] ) image_regions = load_shard_image_regions(shard_output_dirs, shard_page_maps) @@ -381,18 +404,16 @@ def _predict_shard_headings_with_tracking( def _upload_temp_shard_pdfs( - shard_pdf_paths: list[str], + live_shards: list[LiveShard], shard_s3_keys: list[str], ) -> None: """Upload split shards to S3 and confirm each object exists before MinerU.""" storage = JobFileStorage() - for shard_index, (shard_path, shard_s3_key) in enumerate( - zip(shard_pdf_paths, shard_s3_keys, strict=True) - ): + for live, shard_s3_key in zip(live_shards, shard_s3_keys, strict=True): logger.info( - f" ☁️ Uploading MinerU shard_{shard_index} to S3 ({shard_s3_key})" + f" ☁️ Uploading MinerU shard_{live.shard_index} to S3 ({shard_s3_key})" ) - storage.upload_source_file(shard_path, shard_s3_key) + storage.upload_source_file(live.pdf_path, shard_s3_key) existing_file = storage.verify_upload_exists(shard_s3_key) if not existing_file.get("exists"): raise StorageServiceException( @@ -402,7 +423,7 @@ def _upload_temp_shard_pdfs( ), operation="verify_source_object", ) - logger.info(f" ✅ Uploaded MinerU shard_{shard_index}: {shard_s3_key}") + logger.info(f" ✅ Uploaded MinerU shard_{live.shard_index}: {shard_s3_key}") def _build_temp_shard_s3_key( diff --git a/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py b/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py index 78ed71a0a..aa6225236 100644 --- a/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py +++ b/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py @@ -5,6 +5,7 @@ import os import re import shutil +from collections.abc import Sequence from loguru import logger @@ -19,6 +20,8 @@ def _extract_heading_key(line: str) -> tuple[int, str] | None: def merge_shard_lines( shard_lines_list: list[list[str]], + *, + shard_indexes: Sequence[int] | None = None, ) -> list[str]: """Concatenate per-shard lines_with_heading in order, removing boundary duplicates. @@ -30,17 +33,28 @@ def merge_shard_lines( shard's output has the same (level, text) as the *first heading* in shard N+1's output, strip the trailing heading (and any non-heading lines that follow it, i.e. the divider page content) from shard N's lines. + + ``shard_indexes`` is the original plan index of each entry (same identity as + ``LiveShard.shard_index``). When omitted, logs fall back to list position. """ if not shard_lines_list: return [] + if shard_indexes is not None and len(shard_indexes) != len(shard_lines_list): + raise ValueError( + "shard_indexes must align one-to-one with shard_lines_list " + f"(got {len(shard_indexes)} indexes for {len(shard_lines_list)} shards)" + ) result: list[str] = [] - for shard_idx, lines in enumerate(shard_lines_list): + for position, lines in enumerate(shard_lines_list): if not lines: continue + log_index = ( + shard_indexes[position] if shard_indexes is not None else position + ) next_first_heading: tuple[int, str] | None = None - for future_idx in range(shard_idx + 1, len(shard_lines_list)): + for future_idx in range(position + 1, len(shard_lines_list)): future_lines = shard_lines_list[future_idx] for next_line in future_lines: key = _extract_heading_key(next_line) @@ -66,7 +80,7 @@ def merge_shard_lines( and last_heading_key == next_first_heading ): logger.info( - f"🔗 shard_{shard_idx}: removing trailing boundary heading " + f"🔗 shard_{log_index}: removing trailing boundary heading " f"'{last_heading_key[1]}' (L{last_heading_key[0]}) duplicated " f"at start of next shard ({len(lines_to_add) - last_heading_pos} lines trimmed)" ) @@ -77,14 +91,12 @@ def merge_shard_lines( return result -def merge_images(shard_dirs: list[str], target_dir: str) -> None: +def merge_images(shard_dirs: Sequence[str], target_dir: str) -> None: """Copy all images from shard images/ dirs into target_dir/images/.""" target_img_dir = os.path.join(target_dir, "images") os.makedirs(target_img_dir, exist_ok=True) total = 0 for shard_dir in shard_dirs: - if shard_dir is None: - continue img_dir = os.path.join(shard_dir, "images") if not os.path.isdir(img_dir): continue 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 af47d9e48..1a368dfc5 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 @@ -3,7 +3,8 @@ from __future__ import annotations import os -from dataclasses import dataclass +from collections.abc import Sequence +from dataclasses import dataclass, replace from typing import TYPE_CHECKING import pymupdf @@ -29,6 +30,26 @@ def page_count(self) -> int: return self.page_end - self.page_start + 1 +@dataclass(frozen=True) +class LiveShard: + """A planned shard that still has pages after TOC exclusion. + + ``shard_index`` is the original plan index (same as ``agent_shards``), not + the compressed position among surviving PDFs. Downstream MinerU, heading + prediction, S3 keys, and bbox page maps must use this identity. + """ + + shard_index: int + page_start: int + page_end: int + page_numbers: list[int] + pdf_path: str = "" + + @property + def is_first_shard(self) -> bool: + return self.shard_index == 0 + + def map_agent_shards( agent_shards: list["Shard"], ) -> list[MergedShard]: @@ -57,16 +78,30 @@ def shard_page_numbers(shard: MergedShard, exclude_pages: set[int]) -> list[int] ] -def split_pdf_page_numbers( - shards: list[MergedShard], exclude_pages: set[int] -) -> list[list[int]]: - """Page map of each shard PDF ``split_pdf`` writes, in the same order. +def select_live_shards( + shards: Sequence[MergedShard], + exclude_pages: set[int] | None = None, +) -> list[LiveShard]: + """Planned shards that still have at least one page after exclusion. - A shard whose pages are all excluded produces no PDF, so it has no entry - here either; index ``i`` therefore lines up with ``split_pdf``'s ``i``-th path. + Order matches document order. Skipped shards are omitted, not padded. + ``pdf_path`` is empty until ``split_pdf`` writes the file. """ - page_maps = [shard_page_numbers(shard, exclude_pages) for shard in shards] - return [page_map for page_map in page_maps if page_map] + excluded = exclude_pages or set() + live: list[LiveShard] = [] + for shard in shards: + page_numbers = shard_page_numbers(shard, excluded) + if not page_numbers: + continue + live.append( + LiveShard( + shard_index=shard.shard_index, + page_start=shard.page_start, + page_end=shard.page_end, + page_numbers=page_numbers, + ) + ) + return live def _is_source_object_out_of_range(exc: BaseException) -> bool: @@ -79,7 +114,7 @@ def split_pdf( shards: list[MergedShard], work_dir: str, exclude_pages: set[int] | None = None, -) -> list[str]: +) -> list[LiveShard]: """Physically split PDF into sub-PDFs using PyMuPDF. Args: @@ -90,8 +125,8 @@ def split_pdf( (e.g. TOC pages detected during PROFILE). Returns: - One temp PDF path per shard that kept at least one page. Page ``i`` of - each PDF maps to the original page ``split_pdf_page_numbers`` reports. + One live shard per planned shard that kept at least one page. Page + ``i`` of each PDF is original page ``live.page_numbers[i]``. """ doc = pymupdf.open(pdf_path) rewritten_path: str | None = None @@ -123,43 +158,47 @@ def _extract_shards( shards: list[MergedShard], work_dir: str, exclude_pages: set[int] | None, -) -> list[str]: - paths: list[str] = [] - +) -> list[LiveShard]: if exclude_pages: logger.info( f"📌 Excluding {len(exclude_pages)} pages from PDF: " f"{sorted(exclude_pages)}" ) + live_plan = select_live_shards(shards, exclude_pages) + live_indexes = {live.shard_index for live in live_plan} for shard in shards: + if shard.shard_index not in live_indexes: + logger.warning( + f" ⚠️ shard_{shard.shard_index}: all pages excluded, skipping" + ) + logger.info( + f" ✂️ shard_{shard.shard_index}: " + f"pages {shard.page_start}-{shard.page_end} " + f"(0 included, {shard.page_count} excluded)" + ) + + live_shards: list[LiveShard] = [] + for live in live_plan: sub_doc = pymupdf.open() try: - shard_included = 0 - for page_num in shard_page_numbers(shard, exclude_pages or set()): + for page_num in live.page_numbers: sub_doc.insert_pdf( doc, from_page=page_num - 1, to_page=page_num - 1, ) - shard_included += 1 - - shard_path = os.path.join(work_dir, f"shard_{shard.shard_index}.pdf") - if shard_included > 0: - sub_doc.save(shard_path) - paths.append(shard_path) - else: - logger.warning( - f" ⚠️ shard_{shard.shard_index}: all pages excluded, skipping" - ) + shard_path = os.path.join(work_dir, f"shard_{live.shard_index}.pdf") + sub_doc.save(shard_path) + live_shards.append(replace(live, pdf_path=shard_path)) finally: sub_doc.close() - excluded_in_shard = shard.page_count - shard_included + excluded_in_shard = (live.page_end - live.page_start + 1) - len(live.page_numbers) logger.info( - f" ✂️ shard_{shard.shard_index}: " - f"pages {shard.page_start}-{shard.page_end} " - f"({shard_included} included" + f" ✂️ shard_{live.shard_index}: " + f"pages {live.page_start}-{live.page_end} " + f"({len(live.page_numbers)} included" f"{f', {excluded_in_shard} excluded' if excluded_in_shard else ''})" ) - return paths + return live_shards diff --git a/apps/worker/app/services/document_parser/profiling/doc_profiler.py b/apps/worker/app/services/document_parser/profiling/doc_profiler.py index 57f733169..8534793ce 100644 --- a/apps/worker/app/services/document_parser/profiling/doc_profiler.py +++ b/apps/worker/app/services/document_parser/profiling/doc_profiler.py @@ -15,6 +15,7 @@ build_oversized_pdf_processing_failed_exception, raise_if_oversized_pdf_not_supported, ) +from app.services.document_parser.profiling.domain_taxonomy import OTHER_CODE from app.services.document_parser.profiling.profile_model import ( ParserDocumentProfile, ParserTocProfile, @@ -81,7 +82,8 @@ def profile_document( return ParserDocumentProfile( file_type=ext.lstrip("."), - category=f"{ext.lstrip('.') or 'unknown'} document", + category=[OTHER_CODE], + category_label=f"{ext.lstrip('.') or 'unknown'} document", routing_category=PdfRoutingCategory.GENERIC, ) @@ -146,6 +148,7 @@ def _profile_pdf_with_db( profile = ParserDocumentProfile( file_type="pdf", category=agent_profile.category, + category_label=agent_profile.category_label, routing_category=routing_category, is_scanned=agent_profile.is_scanned, page_count=coordinator.blackboard.page_count, diff --git a/apps/worker/app/services/document_parser/profiling/domain_taxonomy.py b/apps/worker/app/services/document_parser/profiling/domain_taxonomy.py new file mode 100644 index 000000000..74f64c4d1 --- /dev/null +++ b/apps/worker/app/services/document_parser/profiling/domain_taxonomy.py @@ -0,0 +1,254 @@ +"""Document domain taxonomy based on UN ISIC Rev.5. + +Leaves are ISIC divisions (2-digit); finance, professional/scientific and +health are expanded to groups/classes because documents there are common and +dense. The model picks leaf numbers from the rendered menu; only the stable ISIC +code of each pick is persisted, so renumbering the menu never changes stored data. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +OTHER_CODE = "OTHER" + + +@dataclass(frozen=True) +class DomainLeaf: + code: str + title: str + + +@dataclass(frozen=True) +class DomainSection: + title: str + leaves: tuple[DomainLeaf, ...] + + +def _section(title: str, *leaves: tuple[str, str]) -> DomainSection: + return DomainSection(title, tuple(DomainLeaf(code, name) for code, name in leaves)) + + +DOMAIN_SECTIONS: tuple[DomainSection, ...] = ( + _section( + "Agriculture, forestry and fishing", + ("01", "Crop and animal production"), + ("02", "Forestry and logging"), + ("03", "Fishing and aquaculture"), + ), + _section( + "Mining and quarrying", + ("05", "Coal and lignite"), + ("06", "Crude petroleum and natural gas"), + ("07", "Metal ores"), + ("08", "Other mining and quarrying"), + ("09", "Mining support services"), + ), + _section( + "Manufacturing", + ("10", "Food products"), + ("11", "Beverages"), + ("12", "Tobacco products"), + ("13", "Textiles"), + ("14", "Wearing apparel"), + ("15", "Leather and related products"), + ("16", "Wood and cork products"), + ("17", "Paper and paper products"), + ("18", "Printing and recorded media"), + ("19", "Coke and refined petroleum products"), + ("20", "Chemicals and chemical products"), + ("21", "Pharmaceutical products"), + ("22", "Rubber and plastic products"), + ("23", "Non-metallic mineral products"), + ("24", "Basic metals"), + ("25", "Fabricated metal products"), + ("26", "Computer, electronic and optical products"), + ("27", "Electrical equipment"), + ("28", "Machinery and equipment"), + ("29", "Motor vehicles, trailers and semi-trailers"), + ("30", "Other transport equipment"), + ("31", "Furniture"), + ("32", "Other manufacturing"), + ("33", "Repair and installation of machinery and equipment"), + ), + _section( + "Electricity, gas, steam and air conditioning supply", + ("35", "Electricity, gas, steam and air conditioning supply"), + ), + _section( + "Water supply, sewerage, waste management and remediation", + ("36", "Water collection, treatment and supply"), + ("37", "Sewerage"), + ("38", "Waste collection, treatment, disposal and recovery"), + ("39", "Remediation and other waste management services"), + ), + _section( + "Construction", + ("41", "Construction of buildings"), + ("42", "Civil engineering"), + ("43", "Specialized construction activities"), + ), + _section( + "Wholesale and retail trade", + ("46", "Wholesale trade"), + ("47", "Retail trade"), + ), + _section( + "Transportation and storage", + ("49", "Land transport and pipelines"), + ("50", "Water transport"), + ("51", "Air transport"), + ("52", "Warehousing and transport support activities"), + ("53", "Postal and courier activities"), + ), + _section( + "Accommodation and food service", + ("55", "Accommodation"), + ("56", "Food and beverage service"), + ), + _section( + "Publishing, broadcasting and content production", + ("58", "Publishing"), + ("59", "Film, video, television and music production"), + ("60", "Broadcasting, news agencies and content distribution"), + ), + _section( + "Telecommunications, computer programming and information services", + ("61", "Telecommunications"), + ("62", "Computer programming and consultancy"), + ("63", "Computing infrastructure, data processing and information services"), + ), + _section( + "Financial and insurance activities", + ("641", "Monetary intermediation and banking"), + ("642", "Holding companies and financing conduits"), + ("643", "Trusts, funds and similar financial entities"), + ("649", "Other financial services (credit, leasing, factoring)"), + ("651", "Insurance"), + ("652", "Reinsurance"), + ("653", "Pension funding"), + ("6611", "Administration of financial markets"), + ("6612", "Securities and commodity contracts brokerage"), + ("6619", "Other activities auxiliary to financial services"), + ("662", "Activities auxiliary to insurance and pension funding"), + ("663", "Fund management"), + ), + _section( + "Real estate activities", + ("68", "Real estate activities"), + ), + _section( + "Professional, scientific and technical activities", + ("691", "Legal activities"), + ("692", "Accounting, auditing and tax consultancy"), + ("701", "Head offices"), + ("702", "Management consultancy"), + ("711", "Architecture, engineering and technical consultancy"), + ("712", "Technical testing and analysis"), + ("721", "Research and development in natural sciences and engineering"), + ("722", "Research and development in social sciences and humanities"), + ("731", "Advertising"), + ("732", "Market research and public opinion polling"), + ("741", "Specialized design"), + ("742", "Photography"), + ("743", "Translation and interpretation"), + ("749", "Other professional, scientific and technical activities"), + ("75", "Veterinary activities"), + ), + _section( + "Administrative and support services", + ("77", "Rental and leasing"), + ("78", "Employment activities"), + ("79", "Travel agencies and tour operators"), + ("80", "Investigation and security"), + ("81", "Services to buildings and landscape"), + ("82", "Office administration and business support"), + ), + _section( + "Public administration and defence", + ("84", "Public administration, defence and compulsory social security"), + ), + _section( + "Education", + ("85", "Education"), + ), + _section( + "Human health and social work", + ("861", "Hospitals"), + ("862", "Medical and dental practice"), + ("869", "Other human health activities"), + ("87", "Residential care"), + ("88", "Social work without accommodation"), + ), + _section( + "Arts, sports and recreation", + ("90", "Arts creation and performing arts"), + ("91", "Libraries, archives, museums and cultural activities"), + ("92", "Gambling and betting"), + ("93", "Sports, amusement and recreation"), + ), + _section( + "Other service activities", + ("94", "Membership organizations"), + ("95", "Repair of computers, personal and household goods, and vehicles"), + ("96", "Personal services"), + ), + _section( + "Other", + (OTHER_CODE, "Other (no listed domain fits)"), + ), +) + +_CODE_BY_ID: dict[int, str] = { + number: leaf.code + for number, leaf in enumerate( + (leaf for section in DOMAIN_SECTIONS for leaf in section.leaves), + start=1, + ) +} + + +def _render_menu() -> str: + lines: list[str] = [] + number = 1 + for section in DOMAIN_SECTIONS: + lines.append(section.title) + for leaf in section.leaves: + lines.append(f"{number}. {leaf.title}") + number += 1 + return "\n".join(lines) + + +DOMAIN_MENU = _render_menu() + + +def domain_codes_from_ids(ids: object) -> list[str]: + """Map model-picked menu numbers to stable domain codes. + + Raises ``ValueError`` for anything that is not a non-empty list of known, + distinct numbers, or when ``OTHER`` is combined with other picks. + """ + if not isinstance(ids, list) or not ids: + raise ValueError("domain_ids must be a non-empty list") + codes: list[str] = [] + for number in ids: + if isinstance(number, bool) or not isinstance(number, int): + raise ValueError(f"domain id must be an integer: {number!r}") + if number not in _CODE_BY_ID: + raise ValueError(f"unknown domain id: {number}") + codes.append(_CODE_BY_ID[number]) + if len(set(codes)) != len(codes): + raise ValueError("domain_ids contains duplicates") + if OTHER_CODE in codes and len(codes) > 1: + raise ValueError(f"{OTHER_CODE} cannot be combined with other domains") + return codes + + +__all__ = [ + "DOMAIN_MENU", + "DOMAIN_SECTIONS", + "OTHER_CODE", + "DomainLeaf", + "DomainSection", + "domain_codes_from_ids", +] diff --git a/apps/worker/app/services/document_parser/profiling/profile_model.py b/apps/worker/app/services/document_parser/profiling/profile_model.py index 7c7e79687..b47810bb1 100644 --- a/apps/worker/app/services/document_parser/profiling/profile_model.py +++ b/apps/worker/app/services/document_parser/profiling/profile_model.py @@ -3,6 +3,7 @@ from dataclasses import asdict, dataclass, field from typing import Any +from app.services.document_parser.profiling.domain_taxonomy import OTHER_CODE from app.services.document_parser.profiling.taxonomy import PdfRoutingCategory @@ -33,7 +34,8 @@ class ParserDocumentProfile: """Parser-entry document profile used for routing and PDF anatomy reuse.""" file_type: str - category: str = "unknown document" + category: list[str] = field(default_factory=lambda: [OTHER_CODE]) + category_label: str = "unknown document" routing_category: PdfRoutingCategory = PdfRoutingCategory.GENERIC is_scanned: bool = False page_count: int = 0 @@ -67,7 +69,8 @@ def to_dict(self) -> dict[str, Any]: def summary(self) -> str: parts = ( - f"[{self.file_type.upper()}] category={self.category}, " + f"[{self.file_type.upper()}] category={self.category} " + f"({self.category_label}), " f"routing={self.routing_category.value}, " f"scanned={self.is_scanned}, pages={self.page_count}" ) diff --git a/apps/worker/scripts/debug_text_track.py b/apps/worker/scripts/debug_text_track.py index 8702c04ec..dd79d3010 100644 --- a/apps/worker/scripts/debug_text_track.py +++ b/apps/worker/scripts/debug_text_track.py @@ -207,8 +207,9 @@ def _stage_profile(pdf_path: str, filename: str, out_dir: Path, model: str | Non logger.info(" profile done in {:.1f}s", elapsed) logger.info( - " category={!r} routing={} header_y={} footer_y={}", + " category={!r} label={!r} routing={} header_y={} footer_y={}", getattr(profile, "category", None), + getattr(profile, "category_label", None), getattr(profile, "routing_category", None), header_y, footer_y, @@ -396,21 +397,23 @@ def _stage_mineru_pdf( parse_via_full(pdf_path, filename, shard_out) shard_output_dirs = [shard_out] else: - shard_pdf_paths = split_pdf( + live_shards = split_pdf( pdf_path, merged_shards, work_dir, exclude_pages=toc_pages if toc_pages else None, ) logger.info(" split into {} shard PDFs (excluded {} TOC pages)", - len(shard_pdf_paths), len(toc_pages)) + len(live_shards), len(toc_pages)) - shard_output_dirs: list[str] = [None] * len(shard_pdf_paths) # type: ignore[list-item] - for i, shard_pdf in enumerate(shard_pdf_paths): - shard_out = os.path.join(work_dir, f"shard_{i}") + shard_output_dirs = [] + for live in live_shards: + shard_out = os.path.join(work_dir, f"shard_{live.shard_index}") os.makedirs(shard_out, exist_ok=True) - shard_filename = f"{os.path.splitext(filename)[0]}_shard{i}.pdf" - logger.info(" 🔄 MinerU shard_{}: parsing...", i) - parse_via_full(shard_pdf, shard_filename, shard_out) - shard_output_dirs[i] = shard_out + shard_filename = ( + f"{os.path.splitext(filename)[0]}_shard{live.shard_index}.pdf" + ) + logger.info(" 🔄 MinerU shard_{}: parsing...", live.shard_index) + parse_via_full(live.pdf_path, shard_filename, shard_out) + shard_output_dirs.append(shard_out) elapsed = time.time() - t0 logger.info(" ✅ MinerU extraction done in {:.1f}s → {} shard dirs", elapsed, len(shard_output_dirs)) @@ -433,7 +436,10 @@ def _stage_hierarchy_pdf( merge_images, merge_shard_lines, ) - from app.services.document_parser.formats.pdf.shard_splitter import map_agent_shards + from app.services.document_parser.formats.pdf.shard_splitter import ( + map_agent_shards, + select_live_shards, + ) logger.info("=" * 70) logger.info("🔬 Stage 3: Per-shard heading prediction → merged hierarchy") @@ -443,31 +449,39 @@ def _stage_hierarchy_pdf( agent_shards = anatomy.shard_plan.shards merged_shards = map_agent_shards(agent_shards) + toc_pages: set[int] = set() + if anatomy.toc_result and anatomy.toc_result.toc_pages: + toc_pages = set(anatomy.toc_result.toc_pages) + live_shards = select_live_shards(merged_shards, toc_pages) work_dir = out_dir / "_shards" - # Discover shard output dirs + # Discover live shard output dirs by original shard_index, not compressed position. shard_output_dirs: list[str] = [] - for i in range(len(merged_shards)): - shard_dir = str(work_dir / f"shard_{i}") + for live in live_shards: + shard_dir = str(work_dir / f"shard_{live.shard_index}") if not os.path.isdir(shard_dir): - raise FileNotFoundError(f"shard_{i} dir not found: {shard_dir} (run --stop-at mineru first)") + raise FileNotFoundError( + f"shard_{live.shard_index} dir not found: {shard_dir} " + "(run --stop-at mineru first)" + ) shard_output_dirs.append(shard_dir) hierarchy_model = model or os.environ.get("NORMOL_MODEL") - def _predict_shard(shard_idx: int, shard_out_dir: str) -> list[str]: + def _predict_shard(live, shard_out_dir: str) -> list[str]: md_path = os.path.join(shard_out_dir, "full.md") if not os.path.exists(md_path): - raise FileNotFoundError(f"shard_{shard_idx}: full.md not found at {md_path}") + raise FileNotFoundError( + f"shard_{live.shard_index}: full.md not found at {md_path}" + ) with open(md_path, "r", encoding="utf-8") as f: md_lines = f.readlines() md_lines = [line.strip() for line in md_lines if line.strip()] md_lines = merge_html_tables(md_lines) - is_first = shard_idx == 0 - agent_shard = agent_shards[shard_idx] + agent_shard = agent_shards[live.shard_index] shard_toc = getattr(agent_shard, "toc_hierarchies", None) lines_with_heading = eval_md_headings( @@ -482,12 +496,12 @@ def _predict_shard(shard_idx: int, shard_out_dir: str) -> list[str]: if os.path.exists(os.path.join(shard_out_dir, "layout.json")) else None ), - is_first_shard=is_first, + is_first_shard=live.is_first_shard, ) heading_count = sum(1 for line in lines_with_heading if line.startswith("#")) logger.info(" ✅ shard_{}: {} headings from {} lines", - shard_idx, heading_count, len(lines_with_heading)) + live.shard_index, heading_count, len(lines_with_heading)) _write_json( Path(shard_out_dir) / "lines_with_heading.json", @@ -496,12 +510,15 @@ def _predict_shard(shard_idx: int, shard_out_dir: str) -> list[str]: return lines_with_heading all_shard_lines: list[list[str]] = [] - for i, shard_dir in enumerate(shard_output_dirs): - lines = _predict_shard(i, shard_dir) + for live, shard_dir in zip(live_shards, shard_output_dirs, strict=True): + lines = _predict_shard(live, shard_dir) all_shard_lines.append(lines) # Merge (boundary-heading dedup; signature matches production parser) - merged_lines = merge_shard_lines(all_shard_lines) + merged_lines = merge_shard_lines( + all_shard_lines, + shard_indexes=[live.shard_index for live in live_shards], + ) # Always copy shard images into the package root. Unlike production's MinerU # fast path (which writes directly to output_dir), this debug script always @@ -672,7 +689,7 @@ def _stage_full_pdf( from app.services.document_parser.formats.markdown.parser import parse_md from app.services.document_parser.formats.pdf.shard_splitter import ( map_agent_shards, - split_pdf_page_numbers, + select_live_shards, ) from app.services.document_parser.orchestration.postprocess import apply_parse_postprocess @@ -681,13 +698,17 @@ def _stage_full_pdf( if anatomy.toc_result and anatomy.toc_result.toc_pages: toc_pages = set(anatomy.toc_result.toc_pages) fast_path = len(merged_shards) == 1 and not toc_pages + live_shards = select_live_shards(merged_shards, toc_pages) shard_page_maps = ( [None] if fast_path - else split_pdf_page_numbers(merged_shards, toc_pages) + else [live.page_numbers for live in live_shards] ) image_regions = load_shard_image_regions( - [str(out_dir / "_shards" / f"shard_{i}") for i in range(len(shard_page_maps))], + [ + str(out_dir / "_shards" / f"shard_{live.shard_index}") + for live in live_shards + ], shard_page_maps, ) diff --git a/apps/worker/scripts/page_memory/_debug_pm_shared.py b/apps/worker/scripts/page_memory/_debug_pm_shared.py index aa6b598cd..900d5d122 100644 --- a/apps/worker/scripts/page_memory/_debug_pm_shared.py +++ b/apps/worker/scripts/page_memory/_debug_pm_shared.py @@ -358,11 +358,13 @@ def _build_debug_coordinator( def _document_profile_from_dict(data: dict[str, Any] | None): from app.services.document_agent.manifest import DocumentProfile + from app.services.document_parser.profiling.domain_taxonomy import OTHER_CODE raw = data or {} return DocumentProfile( is_scanned=bool(raw.get("is_scanned", False)), - category=str(raw.get("category") or "unknown"), + category=list(raw.get("category") or [OTHER_CODE]), + category_label=str(raw.get("category_label") or "unknown"), routing_category=str(raw.get("routing_category") or "generic"), language=str(raw.get("language") or "unknown"), header_y=raw.get("header_y"), @@ -537,6 +539,8 @@ def run_stage0_bootstrap( payload={ "page_count": coordinator.blackboard.page_count, "is_scanned": bool(getattr(profile, "is_scanned", False)), + "category": getattr(profile, "category", None), + "category_label": getattr(profile, "category_label", None), "routing_category": getattr(profile, "routing_category", None), "text_pages": len(coordinator.blackboard.page_full_text_cache or {}), "assets_probed": bool( @@ -551,8 +555,9 @@ def run_stage0_bootstrap( logger.info(f" stage0 done in {time.time() - t0:.1f}s → {state_path}") logger.info( - " category={} routing={} page_count={} is_scanned={} has_asset_pages={}", + " category={} label={!r} routing={} page_count={} is_scanned={} has_asset_pages={}", getattr(profile, "category", None), + getattr(profile, "category_label", None), getattr(profile, "routing_category", None), coordinator.blackboard.page_count, getattr(profile, "is_scanned", None), diff --git a/apps/worker/scripts/page_memory/debug_pm_stage0_bootstrap.py b/apps/worker/scripts/page_memory/debug_pm_stage0_bootstrap.py index 45ea96318..ec3d6394e 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage0_bootstrap.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage0_bootstrap.py @@ -74,6 +74,7 @@ def main() -> int: variables={ "source": "stop_after_asset_probe", "category": getattr(profile, "category", None), + "category_label": getattr(profile, "category_label", None), "routing_category": getattr(profile, "routing_category", None), "is_scanned": bool(getattr(profile, "is_scanned", False)), "text_pages": text_pages, diff --git a/apps/worker/tests/contract/test_coarse_profile_protocol_contract.py b/apps/worker/tests/contract/test_coarse_profile_protocol_contract.py index a85561e81..72771260e 100644 --- a/apps/worker/tests/contract/test_coarse_profile_protocol_contract.py +++ b/apps/worker/tests/contract/test_coarse_profile_protocol_contract.py @@ -5,6 +5,8 @@ import json import os +import pytest + 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") @@ -20,13 +22,29 @@ PageLabel, TocResult, ) +from app.services.document_parser.profiling.domain_taxonomy import ( + DOMAIN_SECTIONS, + OTHER_CODE, +) + + +def _menu_id(code: str) -> int: + return next( + number + for number, leaf_code in enumerate( + (leaf.code for section in DOMAIN_SECTIONS for leaf in section.leaves), + start=1, + ) + if leaf_code == code + ) def test_parse_profile_reads_classification_fields() -> None: raw = json.dumps( { "is_scanned": True, - "category": "Feasibility Study Report", + "domain_ids": [_menu_id("6612"), _menu_id("29")], + "category_label": "Feasibility Study Report", "routing_category": "generic", "language": "zh", "header_y": None, @@ -35,16 +53,59 @@ def test_parse_profile_reads_classification_fields() -> None: ) profile = _parse_profile(raw) assert profile.is_scanned is True - assert profile.category == "Feasibility Study Report" + assert profile.category == ["6612", "29"] + assert profile.category_label == "Feasibility Study Report" assert profile.routing_category == "generic" assert profile.language == "zh" +@pytest.mark.parametrize( + "domain_ids", + [ + None, + [], + "6612", + [0], + [10_000], + ["12"], + [True], + [1, 1], + [_menu_id(OTHER_CODE), 1], + ], +) +def test_parse_profile_falls_back_to_other_for_invalid_domain_ids(domain_ids) -> None: + raw = json.dumps( + { + "is_scanned": False, + "domain_ids": domain_ids, + "category_label": "Report", + "routing_category": "generic", + "language": "en", + } + ) + profile = _parse_profile(raw) + assert profile.category == [OTHER_CODE] + assert profile.category_label == "Report" + + +def test_parse_profile_does_not_route_from_category_label() -> None: + raw = json.dumps( + { + "is_scanned": False, + "domain_ids": [_menu_id(OTHER_CODE)], + "category_label": "Engineering Atlas", + "language": "en", + } + ) + assert _parse_profile(raw).routing_category == "generic" + + def test_parse_profile_rejects_invalid_header_footer_order() -> None: raw = json.dumps( { "is_scanned": False, - "category": "Report", + "domain_ids": [_menu_id(OTHER_CODE)], + "category_label": "Report", "routing_category": "generic", "language": "en", "header_y": 0.8, @@ -84,7 +145,8 @@ def _seed_pages(coordinator: ProfileCoordinator, page_count: int) -> None: blackboard.toc_result = TocResult(method="none", notes="no toc") blackboard.document_profile = DocumentProfile( is_scanned=True, - category="Report", + category=[OTHER_CODE], + category_label="Report", routing_category="generic", ) diff --git a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py index 2882f6c1b..a19b51d7f 100644 --- a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py +++ b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py @@ -149,7 +149,7 @@ def test_run_lightweight_anatomy_builds_single_shard_without_coarse_vlm( _seed_preprobed_pages(coordinator, page_count=2) coordinator.blackboard.document_profile = DocumentProfile( is_scanned=False, - category="Research Report", + category=["OTHER"], category_label="Research Report", routing_category=PdfRoutingCategory.GENERIC.value, ) coordinator.blackboard.toc_result = TocResult(method="none") @@ -196,7 +196,7 @@ def fake_classify(_self): return ( DocumentProfile( is_scanned=False, - category="Research Report", + category=["OTHER"], category_label="Research Report", routing_category=PdfRoutingCategory.GENERIC.value, ), ToolResult(status="ok", payload={}), @@ -227,7 +227,8 @@ def fake_toc() -> None: profile = coordinator.run_coarse() - assert profile.category == "Research Report" + assert profile.category == ["OTHER"] + assert profile.category_label == "Research Report" assert calls == [ "coarse_profile", "text_scan", @@ -243,7 +244,8 @@ def test_parse_run_recorder_doc_profile_uses_final_anatomy_toc() -> None: recorder.set_doc_profile( ParserDocumentProfile( file_type="pdf", - category="Financial Prospectus", + category=["6612"], + category_label="Financial Prospectus", routing_category=PdfRoutingCategory.GENERIC, page_count=0, ) @@ -266,7 +268,8 @@ def test_parse_run_recorder_doc_profile_uses_final_anatomy_toc() -> None: doc_profile = recorder._doc_profile_for_plan() # noqa: SLF001 assert doc_profile is not None - assert doc_profile["category"] == "Financial Prospectus" + assert doc_profile["category"] == ["6612"] + assert doc_profile["category_label"] == "Financial Prospectus" assert doc_profile["routing_category"] == PdfRoutingCategory.GENERIC.value assert doc_profile["page_count"] == 2 assert doc_profile["toc"]["method"] == "vlm_batch" @@ -293,14 +296,16 @@ def rollback(self) -> None: recorder.persist_doc_profile( ParserDocumentProfile( file_type="pdf", - category="Atlas", + category=["OTHER"], + category_label="Atlas", routing_category=PdfRoutingCategory.ATLAS, page_count=12, ) ) assert len(db.rows) == 1 - assert db.rows[0].doc_profile["category"] == "Atlas" + assert db.rows[0].doc_profile["category"] == ["OTHER"] + assert db.rows[0].doc_profile["category_label"] == "Atlas" assert db.rows[0].shard_plan is None recorder._anatomy = PageAnatomyMap( # noqa: SLF001 @@ -334,7 +339,7 @@ def test_run_structural_retries_transient_confirm_failed_toc_result( _seed_preprobed_pages(coordinator, page_count=3, pages=[1, 2]) coordinator.blackboard.document_profile = DocumentProfile( is_scanned=False, - category="Prospectus", + category=["OTHER"], category_label="Prospectus", routing_category=PdfRoutingCategory.GENERIC.value, ) coordinator.blackboard.toc_result = TocResult( @@ -419,7 +424,7 @@ def test_run_structural_skip_shard_plan_uses_placeholder_without_finalize( _seed_preprobed_pages(coordinator, page_count=4) coordinator.blackboard.document_profile = DocumentProfile( is_scanned=False, - category="Prospectus", + category=["OTHER"], category_label="Prospectus", routing_category=PdfRoutingCategory.GENERIC.value, ) coordinator.blackboard.toc_result = TocResult( @@ -482,7 +487,7 @@ def test_run_structural_trusts_rejected_all_toc_and_fails_open( _seed_preprobed_pages(coordinator, page_count=3, pages=[1, 2]) coordinator.blackboard.document_profile = DocumentProfile( is_scanned=False, - category="Prospectus", + category=["OTHER"], category_label="Prospectus", routing_category=PdfRoutingCategory.GENERIC.value, ) coordinator.blackboard.toc_result = TocResult( @@ -586,7 +591,7 @@ def fake_classify(_self): return ( DocumentProfile( is_scanned=False, - category="Prospectus", + category=["OTHER"], category_label="Prospectus", routing_category=PdfRoutingCategory.GENERIC.value, ), ToolResult(status="ok", payload={}), @@ -642,7 +647,7 @@ def test_run_text_scan_native_uses_read_page_text_bands( coordinator.blackboard.page_count = 2 coordinator.blackboard.document_profile = DocumentProfile( is_scanned=False, - category="Report", + category=["OTHER"], category_label="Report", routing_category=PdfRoutingCategory.GENERIC.value, header_y=0.06, footer_y=0.95, @@ -697,7 +702,7 @@ def test_run_text_scan_scanned_dispatches_ocr_pages( coordinator.blackboard.page_count = 2 coordinator.blackboard.document_profile = DocumentProfile( is_scanned=True, - category="Report", + category=["OTHER"], category_label="Report", routing_category=PdfRoutingCategory.GENERIC.value, ) @@ -791,7 +796,7 @@ def run_coarse(self) -> DocumentProfile: self.blackboard.toc_result = TocResult(method="none") return DocumentProfile( is_scanned=False, - category="Research Report", + category=["OTHER"], category_label="Research Report", routing_category=PdfRoutingCategory.GENERIC.value, ) @@ -842,7 +847,7 @@ def run_coarse(self) -> DocumentProfile: self.calls.append("run_coarse") return DocumentProfile( is_scanned=False, - category="Research Report", + category=["OTHER"], category_label="Research Report", routing_category=PdfRoutingCategory.GENERIC.value, ) @@ -899,7 +904,7 @@ def run_coarse(self) -> DocumentProfile: self.blackboard.toc_result = TocResult(method="none") return DocumentProfile( is_scanned=False, - category="Research Report", + category=["OTHER"], category_label="Research Report", routing_category=PdfRoutingCategory.GENERIC.value, ) @@ -950,7 +955,7 @@ def run_coarse(self) -> DocumentProfile: self.calls.append("run_coarse") return DocumentProfile( is_scanned=False, - category="Research Report", + category=["OTHER"], category_label="Research Report", routing_category=PdfRoutingCategory.GENERIC.value, ) @@ -1013,7 +1018,7 @@ def run_coarse(self) -> DocumentProfile: ] return DocumentProfile( is_scanned=False, - category="Research Report", + category=["OTHER"], category_label="Research Report", routing_category=PdfRoutingCategory.GENERIC.value, ) @@ -1084,7 +1089,7 @@ def run_coarse(self) -> DocumentProfile: ] return DocumentProfile( is_scanned=False, - category="Engineering Atlas", + category=["OTHER"], category_label="Engineering Atlas", routing_category=PdfRoutingCategory.ATLAS.value, ) diff --git a/apps/worker/tests/contract/test_domain_taxonomy_contract.py b/apps/worker/tests/contract/test_domain_taxonomy_contract.py new file mode 100644 index 000000000..f0a357357 --- /dev/null +++ b/apps/worker/tests/contract/test_domain_taxonomy_contract.py @@ -0,0 +1,67 @@ +"""Contract tests for the document domain taxonomy and its prompt rendering.""" + +from __future__ import annotations + +import os + +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 pytest + +from app.services.document_agent.coarse_profile.prompts import ( + COARSE_PROFILE_INSTRUCTIONS, +) +from app.services.document_parser.profiling.domain_taxonomy import ( + DOMAIN_MENU, + DOMAIN_SECTIONS, + OTHER_CODE, + domain_codes_from_ids, +) + +_LEAVES = [leaf for section in DOMAIN_SECTIONS for leaf in section.leaves] + + +def test_leaf_codes_are_unique_and_digit_or_other() -> None: + codes = [leaf.code for leaf in _LEAVES] + assert len(codes) == len(set(codes)) + assert all(code == OTHER_CODE or code.isdigit() for code in codes) + + +def test_every_section_has_leaves_and_other_is_last() -> None: + assert all(section.leaves for section in DOMAIN_SECTIONS) + assert _LEAVES[-1].code == OTHER_CODE + assert [leaf.code for leaf in _LEAVES].count(OTHER_CODE) == 1 + + +def test_menu_numbers_leaves_contiguously_and_leaves_sections_unnumbered() -> None: + numbered = [line for line in DOMAIN_MENU.splitlines() if line[:1].isdigit()] + assert len(numbered) == len(_LEAVES) + for expected, line in enumerate(numbered, start=1): + assert line.startswith(f"{expected}. ") + headers = [line for line in DOMAIN_MENU.splitlines() if not line[:1].isdigit()] + assert headers == [section.title for section in DOMAIN_SECTIONS] + + +def test_ids_map_to_menu_order_codes() -> None: + assert domain_codes_from_ids([1]) == [_LEAVES[0].code] + assert domain_codes_from_ids([len(_LEAVES)]) == [OTHER_CODE] + assert domain_codes_from_ids([3, 1]) == [_LEAVES[2].code, _LEAVES[0].code] + + +@pytest.mark.parametrize( + "ids", + [None, [], "1", [0], [10_000], ["1"], [True], [1, 1], [len(_LEAVES), 1]], +) +def test_invalid_ids_raise(ids) -> None: + with pytest.raises(ValueError): + domain_codes_from_ids(ids) + + +def test_coarse_prompt_carries_menu_and_json_word() -> None: + assert DOMAIN_MENU in COARSE_PROFILE_INSTRUCTIONS + assert "json" in COARSE_PROFILE_INSTRUCTIONS.lower() diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index 3d60d8c52..0661ae452 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -35,6 +35,21 @@ _SAMPLE_XLSX_PATH: Path = _FIXTURES_ROOT / "sample_100rows.xlsx" +def _fake_split_to_live_shards(shards, work_dir, exclude_pages=None): + from dataclasses import replace + + from app.services.document_parser.formats.pdf.shard_splitter import ( + select_live_shards, + ) + + live = [] + for planned in select_live_shards(shards, exclude_pages): + shard_path = Path(work_dir) / f"shard_{planned.shard_index}.pdf" + shard_path.write_bytes(b"%PDF-1.4\n%%EOF\n") + live.append(replace(planned, pdf_path=str(shard_path))) + return live + + def _write_blank_pdf(file_path: Path, page_count: int) -> None: from pypdf import PdfWriter @@ -1158,7 +1173,7 @@ def _fail_oversized_parse(*args, **kwargs): {}, profile=ParserDocumentProfile( file_type="pdf", - category="generic document", + category=["OTHER"], category_label="generic document", routing_category=PdfRoutingCategory.GENERIC, page_count=2, anatomy=PageAnatomyMap( @@ -1270,12 +1285,7 @@ def delete_upload_file(self, storage_key: str) -> bool: def _fake_split_pdf(pdf_path_arg, shards, work_dir, exclude_pages=None): calls["exclude_pages"] = exclude_pages - paths = [] - for shard in shards: - shard_path = Path(work_dir) / f"shard_{shard.shard_index}.pdf" - shard_path.write_bytes(b"%PDF-1.4\n%%EOF\n") - paths.append(str(shard_path)) - return paths + return _fake_split_to_live_shards(shards, work_dir, exclude_pages) def _fake_parse_via_full(shard_pdf, shard_filename, shard_out, s3_key=None): parse_s3_keys.append(s3_key) @@ -1331,7 +1341,7 @@ def _identity_eval_md_headings( }, profile=ParserDocumentProfile( file_type="pdf", - category="generic document", + category=["OTHER"], category_label="generic document", routing_category=PdfRoutingCategory.GENERIC, page_count=3, anatomy=anatomy, @@ -1436,12 +1446,7 @@ def delete_upload_file(self, storage_key: str) -> bool: image_bboxes = {0: [10, 20, 300, 400], 1: [30, 40, 500, 600]} def _fake_split_pdf(pdf_path_arg, shards, work_dir, exclude_pages=None): - paths = [] - for shard in shards: - shard_path = Path(work_dir) / f"shard_{shard.shard_index}.pdf" - shard_path.write_bytes(b"%PDF-1.4\n%%EOF\n") - paths.append(str(shard_path)) - return paths + return _fake_split_to_live_shards(shards, work_dir, exclude_pages) def _fake_parse_via_full(shard_pdf, shard_filename, shard_out, s3_key=None): shard_index = 0 if "shard0" in shard_filename else 1 @@ -1531,7 +1536,7 @@ def _identity_eval_md_headings( }, profile=ParserDocumentProfile( file_type="pdf", - category="generic document", + category=["OTHER"], category_label="generic document", routing_category=PdfRoutingCategory.GENERIC, page_count=3, anatomy=anatomy, @@ -1546,3 +1551,323 @@ def _identity_eval_md_headings( assert [region["page_num"] for region in regions] == [2, 3] assert [region["bbox"] for region in regions] == [image_bboxes[0], image_bboxes[1]] assert {region["coord_system"] for region in regions} == {"mineru_layout"} + + +def test_skipped_middle_shard_uses_original_shard_identity_for_toc_and_paths( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setenv("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") + monkeypatch.setenv("TMP_PATH", str(tmp_path)) + monkeypatch.setenv("S3_BUCKET_NAME", "test-uploads") + monkeypatch.setenv("S3_ACCESS_KEY_ID", "test") + monkeypatch.setenv("S3_SECRET_ACCESS_KEY", "test") + monkeypatch.setenv("S3_TEMP_PATH", str(tmp_path)) + + from app.services.document_agent.manifest import ( + PageAnatomyMap, + Shard, + ShardPlan, + TocResult, + ) + from app.services.document_parser.formats.pdf import parser as pdf_parser + from app.services.document_parser.profiling.profile_model import ParserDocumentProfile + from app.services.document_parser.profiling.taxonomy import PdfRoutingCategory + + pdf_path = tmp_path / "identity.pdf" + output_dir = tmp_path / "output" + output_dir.mkdir() + _write_blank_pdf(pdf_path, page_count=6) + + uploaded: list[str] = [] + mineru_filenames: list[str] = [] + mineru_dirs: list[str] = [] + heading_calls: list[dict[str, object]] = [] + + class _FakeJobFileStorage: + def upload_source_file(self, local_file_path: str, storage_key: str) -> dict[str, str]: + uploaded.append(storage_key) + return {"etag": "test"} + + def verify_upload_exists(self, storage_key: str) -> dict[str, bool]: + return {"exists": True} + + def delete_upload_file(self, storage_key: str) -> bool: + return True + + toc_by_shard = { + 0: [{"toc_tree": {"Chapter A": {}}}], + 1: [{"toc_tree": {"Skipped Middle": {}}}], + 2: [{"toc_tree": {"Chapter C": {}}}], + } + anatomy = PageAnatomyMap( + job_id="job-identity", + file_path=str(pdf_path), + page_count=6, + page_features=[], + page_labels=[], + toc_result=TocResult(toc_pages=[3, 4], method="vlm_batch"), + shard_plan=ShardPlan( + enabled=True, + reason="too_large", + shards=[ + Shard( + shard_index=0, + page_start=1, + page_end=2, + page_offset=0, + anchor_type="toc_leaf_boundary", + anchor_evidence="Chapter A", + toc_hierarchies=toc_by_shard[0], + ), + Shard( + shard_index=1, + page_start=3, + page_end=4, + page_offset=2, + anchor_type="toc_leaf_boundary", + anchor_evidence="Skipped Middle", + toc_hierarchies=toc_by_shard[1], + ), + Shard( + shard_index=2, + page_start=5, + page_end=6, + page_offset=4, + anchor_type="toc_leaf_boundary", + anchor_evidence="Chapter C", + toc_hierarchies=toc_by_shard[2], + ), + ], + ), + ) + + def _fake_parse_via_full(shard_pdf, shard_filename, shard_out, s3_key=None): + mineru_filenames.append(shard_filename) + mineru_dirs.append(shard_out) + Path(shard_out).mkdir(parents=True, exist_ok=True) + chapter = "A" if "shard0" in shard_filename else "C" + (Path(shard_out) / "full.md").write_text( + f"# Chapter {chapter}\nBody {chapter}.", + encoding="utf-8", + ) + + def _identity_eval_md_headings( + md_lines, + source_type, + toc_hierarchies=None, + smart_parse=True, + model_name=None, + output_dir=None, + layout_json_path=None, + is_first_shard=True, + ): + heading_calls.append( + { + "toc_hierarchies": toc_hierarchies, + "is_first_shard": is_first_shard, + "output_dir": output_dir, + } + ) + return list(md_lines) + + monkeypatch.setattr(pdf_parser.settings, "MAX_PDF_PAGE_LIMIT", 2) + monkeypatch.setattr(pdf_parser.settings, "MINERU_SHARD_CONCURRENCY", 1) + monkeypatch.setattr( + "app.services.document_parser.formats.pdf.shard_splitter.split_pdf", + lambda pdf_path_arg, shards, work_dir, exclude_pages=None: _fake_split_to_live_shards( + shards, work_dir, exclude_pages + ), + ) + monkeypatch.setattr(pdf_parser, "parse_via_full", _fake_parse_via_full) + monkeypatch.setattr( + "app.services.document_parser.formats.markdown.parser.eval_md_headings", + _identity_eval_md_headings, + ) + monkeypatch.setattr(pdf_parser, "JobFileStorage", _FakeJobFileStorage) + + df = pdf_parser.parse_pdfs( + str(pdf_path), + "identity.pdf", + str(output_dir), + { + "smart_title_parse": False, + "summary_image": False, + "summary_table": False, + "summary_txt": False, + "stopwords": [], + "model_name": "mock-model", + "hierarchy_model_name": "mock-model", + }, + profile=ParserDocumentProfile( + file_type="pdf", + category=["OTHER"], category_label="generic document", + routing_category=PdfRoutingCategory.GENERIC, + page_count=6, + anatomy=anatomy, + ), + relative_root="identity.pdf", + s3_key="uploads/job-identity.pdf", + job_id="job-identity", + ) + + assert mineru_filenames == [ + "identity_shard0.pdf", + "identity_shard2.pdf", + ] + assert all("shard_1" not in path for path in mineru_dirs) + assert uploaded == [ + "tmp/mineru-shards/job-identity/shard_0.pdf", + "tmp/mineru-shards/job-identity/shard_2.pdf", + ] + assert [call["is_first_shard"] for call in heading_calls] == [True, False] + assert heading_calls[0]["toc_hierarchies"] == toc_by_shard[0] + assert heading_calls[1]["toc_hierarchies"] == toc_by_shard[2] + assert "shard_0_output" in str(heading_calls[0]["output_dir"]) + assert "shard_2_output" in str(heading_calls[1]["output_dir"]) + assert list(df["path"]) == [ + "identity.pdf/Chapter A", + "identity.pdf/Chapter C", + ] + + +def test_skipped_first_shard_is_not_treated_as_document_front_matter( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setenv("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") + monkeypatch.setenv("TMP_PATH", str(tmp_path)) + monkeypatch.setenv("S3_BUCKET_NAME", "test-uploads") + monkeypatch.setenv("S3_ACCESS_KEY_ID", "test") + monkeypatch.setenv("S3_SECRET_ACCESS_KEY", "test") + monkeypatch.setenv("S3_TEMP_PATH", str(tmp_path)) + + from app.services.document_agent.manifest import ( + PageAnatomyMap, + Shard, + ShardPlan, + TocResult, + ) + from app.services.document_parser.formats.pdf import parser as pdf_parser + from app.services.document_parser.profiling.profile_model import ParserDocumentProfile + from app.services.document_parser.profiling.taxonomy import PdfRoutingCategory + + pdf_path = tmp_path / "front.pdf" + output_dir = tmp_path / "output" + output_dir.mkdir() + _write_blank_pdf(pdf_path, page_count=4) + + heading_calls: list[dict[str, object]] = [] + + class _FakeJobFileStorage: + def upload_source_file(self, local_file_path: str, storage_key: str) -> dict[str, str]: + return {"etag": "test"} + + def verify_upload_exists(self, storage_key: str) -> dict[str, bool]: + return {"exists": True} + + def delete_upload_file(self, storage_key: str) -> bool: + return True + + body_toc = [{"toc_tree": {"Body": {}}}] + anatomy = PageAnatomyMap( + job_id="job-front", + file_path=str(pdf_path), + page_count=4, + page_features=[], + page_labels=[], + toc_result=TocResult(toc_pages=[1, 2], method="vlm_batch"), + shard_plan=ShardPlan( + enabled=True, + reason="too_large", + shards=[ + Shard( + shard_index=0, + page_start=1, + page_end=2, + page_offset=0, + anchor_type="toc_leaf_boundary", + anchor_evidence="TOC", + toc_hierarchies=[{"toc_tree": {"Cover": {}}}], + ), + Shard( + shard_index=1, + page_start=3, + page_end=4, + page_offset=2, + anchor_type="toc_leaf_boundary", + anchor_evidence="Body", + toc_hierarchies=body_toc, + ), + ], + ), + ) + + def _fake_parse_via_full(shard_pdf, shard_filename, shard_out, s3_key=None): + Path(shard_out).mkdir(parents=True, exist_ok=True) + (Path(shard_out) / "full.md").write_text( + "# Body\nBody text.", + encoding="utf-8", + ) + + def _identity_eval_md_headings( + md_lines, + source_type, + toc_hierarchies=None, + smart_parse=True, + model_name=None, + output_dir=None, + layout_json_path=None, + is_first_shard=True, + ): + heading_calls.append( + { + "toc_hierarchies": toc_hierarchies, + "is_first_shard": is_first_shard, + } + ) + return list(md_lines) + + monkeypatch.setattr(pdf_parser.settings, "MAX_PDF_PAGE_LIMIT", 2) + monkeypatch.setattr(pdf_parser.settings, "MINERU_SHARD_CONCURRENCY", 1) + monkeypatch.setattr( + "app.services.document_parser.formats.pdf.shard_splitter.split_pdf", + lambda pdf_path_arg, shards, work_dir, exclude_pages=None: _fake_split_to_live_shards( + shards, work_dir, exclude_pages + ), + ) + monkeypatch.setattr(pdf_parser, "parse_via_full", _fake_parse_via_full) + monkeypatch.setattr( + "app.services.document_parser.formats.markdown.parser.eval_md_headings", + _identity_eval_md_headings, + ) + monkeypatch.setattr(pdf_parser, "JobFileStorage", _FakeJobFileStorage) + + pdf_parser.parse_pdfs( + str(pdf_path), + "front.pdf", + str(output_dir), + { + "smart_title_parse": False, + "summary_image": False, + "summary_table": False, + "summary_txt": False, + "stopwords": [], + "model_name": "mock-model", + "hierarchy_model_name": "mock-model", + }, + profile=ParserDocumentProfile( + file_type="pdf", + category=["OTHER"], category_label="generic document", + routing_category=PdfRoutingCategory.GENERIC, + page_count=4, + anatomy=anatomy, + ), + relative_root="front.pdf", + s3_key="uploads/job-front.pdf", + job_id="job-front", + ) + + assert len(heading_calls) == 1 + assert heading_calls[0]["is_first_shard"] is False + assert heading_calls[0]["toc_hierarchies"] == body_toc diff --git a/apps/worker/tests/contract/test_text_track_image_bbox_contract.py b/apps/worker/tests/contract/test_text_track_image_bbox_contract.py index 08001e697..305b8bf42 100644 --- a/apps/worker/tests/contract/test_text_track_image_bbox_contract.py +++ b/apps/worker/tests/contract/test_text_track_image_bbox_contract.py @@ -23,8 +23,8 @@ from app.services.document_parser.formats.pdf.shard_splitter import ( # noqa: E402 MergedShard, + select_live_shards, shard_page_numbers, - split_pdf_page_numbers, ) from shared.core.constants.processing import ProcessingConstants # noqa: E402 from shared.services.chunks.dataframe_chunk_converter import ( # noqa: E402 @@ -140,7 +140,7 @@ def test_layout_index_keeps_first_region_of_a_repeated_image(tmp_path: Path) -> assert regions["same.jpg"] == _region(1, [1, 1, 2, 2]) -def test_shard_regions_merge_in_shard_order_and_skip_missing_dirs( +def test_shard_regions_merge_in_shard_order_keeps_first_duplicate( tmp_path: Path, ) -> None: first = tmp_path / "shard_0" @@ -160,8 +160,8 @@ def test_shard_regions_merge_in_shard_order_and_skip_missing_dirs( ) regions = load_shard_image_regions( - [str(first), None, str(second)], - [[1], [2], [20]], + [str(first), str(second)], + [[1], [20]], ) assert regions["dup.jpg"] == _region(1, [1, 1, 2, 2]) @@ -175,15 +175,19 @@ def test_shard_page_numbers_skip_excluded_pages() -> None: assert shard_page_numbers(shard, {7, 8, 100}) == [5, 6, 9] -def test_split_pdf_page_numbers_drop_shards_without_a_pdf() -> None: +def test_select_live_shards_drops_fully_excluded_middle_shard() -> None: shards = [ MergedShard(shard_index=0, page_start=1, page_end=2), MergedShard(shard_index=1, page_start=3, page_end=4), MergedShard(shard_index=2, page_start=5, page_end=6), ] - # Shard 1 is entirely excluded, so split_pdf writes no PDF for it. - assert split_pdf_page_numbers(shards, {3, 4, 5}) == [[1, 2], [6]] + live = select_live_shards(shards, {3, 4, 5}) + + assert [item.shard_index for item in live] == [0, 2] + assert [item.page_numbers for item in live] == [[1, 2], [6]] + assert live[0].is_first_shard is True + assert live[1].is_first_shard is False def test_shard_regions_reject_misaligned_page_maps(tmp_path: Path) -> None: diff --git a/apps/worker/tests/unit/test_pdf_shard_s3_upload.py b/apps/worker/tests/unit/test_pdf_shard_s3_upload.py index bc5d219ea..df5ce66d5 100644 --- a/apps/worker/tests/unit/test_pdf_shard_s3_upload.py +++ b/apps/worker/tests/unit/test_pdf_shard_s3_upload.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from dataclasses import replace from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock, call @@ -23,6 +24,10 @@ TocResult, ) from app.services.document_parser.formats.pdf import parser as pdf_parser +from app.services.document_parser.formats.pdf.shard_splitter import ( + LiveShard, + select_live_shards, +) from shared.core.exceptions.domain_exceptions import StorageServiceException @@ -34,7 +39,10 @@ def test_upload_temp_shard_pdfs_writes_each_local_file( monkeypatch.setattr(pdf_parser, "JobFileStorage", Mock(return_value=storage)) pdf_parser._upload_temp_shard_pdfs( - ["/tmp/shard_0.pdf", "/tmp/shard_1.pdf"], + [ + LiveShard(0, 1, 1, [1], "/tmp/shard_0.pdf"), + LiveShard(1, 2, 2, [2], "/tmp/shard_1.pdf"), + ], [ "tmp/mineru-shards/job_abc/shard_0.pdf", "tmp/mineru-shards/job_abc/shard_1.pdf", @@ -86,9 +94,12 @@ def test_shard_pipeline_uploads_before_mineru( def fake_split_pdf(pdf_path, shards, work_dir, exclude_pages=None): events.append("split") - shard_path = Path(work_dir) / "shard_0.pdf" - shard_path.write_bytes(b"%PDF-1.4") - return [str(shard_path)] + live = [] + for planned in select_live_shards(shards, exclude_pages): + shard_path = Path(work_dir) / f"shard_{planned.shard_index}.pdf" + shard_path.write_bytes(b"%PDF-1.4") + live.append(replace(planned, pdf_path=str(shard_path))) + return live class _FakeStorage: def upload_source_file( @@ -156,9 +167,12 @@ def test_failed_shard_upload_does_not_start_mineru( parse_via_full = Mock() def fake_split_pdf(pdf_path, shards, work_dir, exclude_pages=None): - shard_path = Path(work_dir) / "shard_0.pdf" - shard_path.write_bytes(b"%PDF-1.4") - return [str(shard_path)] + live = [] + for planned in select_live_shards(shards, exclude_pages): + shard_path = Path(work_dir) / f"shard_{planned.shard_index}.pdf" + shard_path.write_bytes(b"%PDF-1.4") + live.append(replace(planned, pdf_path=str(shard_path))) + return live class _FailingStorage: def upload_source_file( @@ -200,9 +214,12 @@ def test_unverified_shard_upload_does_not_start_mineru( parse_via_full = Mock() def fake_split_pdf(pdf_path, shards, work_dir, exclude_pages=None): - shard_path = Path(work_dir) / "shard_0.pdf" - shard_path.write_bytes(b"%PDF-1.4") - return [str(shard_path)] + live = [] + for planned in select_live_shards(shards, exclude_pages): + shard_path = Path(work_dir) / f"shard_{planned.shard_index}.pdf" + shard_path.write_bytes(b"%PDF-1.4") + live.append(replace(planned, pdf_path=str(shard_path))) + return live class _InvisibleStorage: def upload_source_file( 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 index ac0ee4b64..649e6de3b 100644 --- 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 @@ -50,15 +50,18 @@ def test_split_pdf_copies_pages(tmp_path: Path) -> None: work_dir = tmp_path / "shards" work_dir.mkdir() - paths = split_pdf( + live = split_pdf( str(source), [MergedShard(0, 1, 2), MergedShard(1, 3, 3)], str(work_dir), ) - 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 [Path(item.pdf_path).name for item in live] == ["shard_0.pdf", "shard_1.pdf"] + assert [item.shard_index for item in live] == [0, 1] + assert live[0].page_numbers == [1, 2] + assert live[1].page_numbers == [3] + assert _page_count(live[0].pdf_path) == 2 + assert _page_count(live[1].pdf_path) == 1 assert not (work_dir / "_rewritten_source.pdf").exists() @@ -81,7 +84,7 @@ def _insert_pdf(self: pymupdf.Document, *args: object, **kwargs: object) -> None monkeypatch.setattr(pymupdf.Document, "insert_pdf", _insert_pdf) - paths = split_pdf( + live = split_pdf( str(source), [MergedShard(0, 1, 2)], str(work_dir), @@ -89,8 +92,9 @@ def _insert_pdf(self: pymupdf.Document, *args: object, **kwargs: object) -> 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 + assert len(live) == 1 + assert live[0].shard_index == 0 + assert _page_count(live[0].pdf_path) == 2 def test_split_pdf_does_not_catch_other_runtime_errors( @@ -117,12 +121,37 @@ def test_split_incremental_quartz_pdf_around_failing_page(tmp_path: Path) -> Non work_dir = tmp_path / "shards" work_dir.mkdir() - paths = split_pdf( + live = split_pdf( str(_CUSTOMER_PDF), [MergedShard(0, 40, 50)], str(work_dir), ) - assert len(paths) == 1 - assert _page_count(paths[0]) == 11 + assert len(live) == 1 + assert live[0].shard_index == 0 + assert _page_count(live[0].pdf_path) == 11 assert (work_dir / "_rewritten_source.pdf").exists() + + +def test_split_pdf_skips_fully_excluded_middle_shard(tmp_path: Path) -> None: + source = tmp_path / "source.pdf" + _write_pages(source, 6) + work_dir = tmp_path / "shards" + work_dir.mkdir() + + live = split_pdf( + str(source), + [ + MergedShard(0, 1, 2), + MergedShard(1, 3, 4), + MergedShard(2, 5, 6), + ], + str(work_dir), + exclude_pages={3, 4}, + ) + + assert [item.shard_index for item in live] == [0, 2] + assert live[0].page_numbers == [1, 2] + assert live[1].page_numbers == [5, 6] + assert [Path(item.pdf_path).name for item in live] == ["shard_0.pdf", "shard_2.pdf"] + assert not (work_dir / "shard_1.pdf").exists() diff --git a/apps/worker/tests/unit/test_toc_anchor_selection.py b/apps/worker/tests/unit/test_toc_anchor_selection.py index 548163bf6..2df82fe9e 100644 --- a/apps/worker/tests/unit/test_toc_anchor_selection.py +++ b/apps/worker/tests/unit/test_toc_anchor_selection.py @@ -9,8 +9,10 @@ line_fingerprint, ) from app.services.document_agent.tools.find_toc_anchor_pages import ( + HEAD_WINDOW_MAX_PAGES, MAX_ANCHOR_CANDIDATES, _filter_edge_anchored, + _head_window_pages, _scan_toc_matches, _select_anchor_pages, ) @@ -139,6 +141,13 @@ def test_under_cap_does_not_apply_edge_or_body() -> None: assert "after_body" not in narrowing +def test_head_window_is_at_most_ten_pages() -> None: + assert _head_window_pages(0) == set() + assert _head_window_pages(5) == {1, 2, 3, 4, 5} + assert _head_window_pages(10) == set(range(1, 11)) + assert _head_window_pages(200) == set(range(1, HEAD_WINDOW_MAX_PAGES + 1)) + + def test_scan_uses_stage0_lines() -> None: bands = { 1: build_page_text_bands((LineDraft(text="General table of contents"),)), From 926d2978aac8081be5d766916029410a376c3358 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 16 Oct 2026 18:06:33 +0800 Subject: [PATCH 4/5] refactor: update document classification to use IPTC Media Topics - Changed document profiling to utilize IPTC Media Topics instead of ISIC codes for categorization. - Updated `DocumentProfile` and related parsing instructions to reflect the new category structure. - Adjusted coarse profiling prompts and functions to align with the IPTC classification system. - Modified tests to validate the new category codes and ensure consistency across document types. --- AGENTS.md | 9 +- .../document_agent/coarse_profile/prompts.py | 6 +- .../app/services/document_agent/manifest.py | 2 +- .../tools/find_toc_anchor_pages.py | 114 ++++++++-- .../profiling/domain_taxonomy.py | 200 +++--------------- .../test_coarse_profile_protocol_contract.py | 8 +- .../test_doc_profile_anatomy_contract.py | 4 +- .../contract/test_domain_taxonomy_contract.py | 7 +- .../tests/unit/test_toc_anchor_selection.py | 33 +++ 9 files changed, 177 insertions(+), 206 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index f458a3852..587b3a673 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -129,10 +129,11 @@ This is the typed `ParseOutput` entry for all file types. The parser flow: 1. **Profiles** the document via `profiling.doc_profiler.profile_document()`. PDF profiling uses `document_agent` as the single PyMuPDF feature source, then runs - VLM coarse classification with three fields: `category` (domain codes picked from - the numbered ISIC Rev.5 leaf list in `profiling/domain_taxonomy.py`, multiple - allowed), `category_label` (open short phrase, for example `Financial Prospectus`), - and routing-only `routing_category` (`atlas/scanned/slides/generic`). Oversized non-atlas PDFs additionally run the + VLM coarse classification with three fields: `category` (IPTC Media Topics + top-level codes from the numbered list in `profiling/domain_taxonomy.py`, + multiple allowed), `category_label` (open short phrase, for example + `Financial Prospectus`), and routing-only `routing_category` + (`atlas/scanned/slides/generic`). Oversized non-atlas PDFs additionally run the structural anatomy stage once at the entry point and pass the resulting shard plan to PDF parsing. 2. **Routes** to the appropriate parser based on file extension. diff --git a/apps/worker/app/services/document_agent/coarse_profile/prompts.py b/apps/worker/app/services/document_agent/coarse_profile/prompts.py index efa3e45f3..ff03a77a6 100644 --- a/apps/worker/app/services/document_agent/coarse_profile/prompts.py +++ b/apps/worker/app/services/document_agent/coarse_profile/prompts.py @@ -8,9 +8,9 @@ "Return strict JSON only with keys: is_scanned, domain_ids, category_label, " "routing_category, language, header_y, footer_y. " "domain_ids is a non-empty JSON array of integers taken from the numbered " - "domain list below. Pick every numbered entry whose domain the document's " - "subject matter belongs to; section titles without a number are not " - "selectable. Pick the Other entry alone, and only when no listed domain fits. " + "IPTC Media Topics list below. Pick every numbered topic that describes what " + "the document is about; section titles without a number are not selectable. " + "Pick the Other entry alone, and only when no listed topic fits. " "category_label is a concise semantic document type in at most 5 English words. " "routing_category must be one of atlas, scanned, slides, generic. " "Set routing_category=atlas only when pages are primarily drawing/detail " diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index 93f1e4b42..7118cfacc 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -50,7 +50,7 @@ def to_dict(self) -> dict[str, Any]: class DocumentProfile: is_scanned: bool category: list[str] - """Domain codes picked from the domain taxonomy (e.g. ``["6612"]``).""" + """IPTC Media Topic codes picked from the domain taxonomy (e.g. ``["04000000"]``).""" category_label: str routing_category: str = "generic" language: str = "unknown" diff --git a/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py index 5d3e71d4f..5ce9559fa 100644 --- a/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py +++ b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py @@ -16,6 +16,9 @@ page_bands_map, ) from app.services.document_agent.registry import has_page_full_text, has_page_labels, register_tool +from app.services.document_parser.profiling.domain_taxonomy import ( + ECONOMY_BUSINESS_FINANCE_CODE, +) from app.services.document_parser.structure.body_boundary import normalize_match_text from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( run_in_child_process, @@ -25,6 +28,10 @@ # CJK and English TOC keywords used for first-pass anchor detection. TOC_KEYWORDS = frozenset({"目录", "目次", "contents", "table of contents"}) +# Domain extras: only when that IPTC category was picked; match line-start only. +DOMAIN_PREFIX_TOC_KEYWORDS: dict[str, frozenset[str]] = { + ECONOMY_BUSINESS_FINANCE_CODE: frozenset({"index to"}), +} MAX_ANCHOR_CANDIDATES = 20 HEAD_WINDOW_MAX_PAGES = 10 @@ -32,28 +39,50 @@ EDGE_ANCHORED_KEYWORDS = frozenset({"contents", "table of contents"}) MIN_CONSECUTIVE_DUPLICATE_RUN = 2 -# Upper bound on consecutive lines joined when repairing a keyword split by -# newlines (e.g. 目\\n录). Derived from the longest keyword character length. -_MAX_KEYWORD_SPLIT_LINES = max(len(keyword) for keyword in TOC_KEYWORDS) +def _prefix_toc_keywords_for_categories(categories: object) -> frozenset[str]: + if not isinstance(categories, list): + return frozenset() + codes = {str(code) for code in categories} + keywords: set[str] = set() + for code, extras in DOMAIN_PREFIX_TOC_KEYWORDS.items(): + if code in codes: + keywords.update(extras) + return frozenset(keywords) + + +def _all_scan_keywords( + *, + prefix_keywords: frozenset[str] = frozenset(), +) -> frozenset[str]: + return TOC_KEYWORDS | prefix_keywords + + +def _max_keyword_split_lines(keywords: frozenset[str]) -> int: + return max((len(keyword) for keyword in keywords), default=1) -def _match_toc_keyword_parts(parts: list[str]) -> str | None: + +def _match_toc_keyword_parts( + parts: list[str], + keywords: frozenset[str], +) -> str | None: candidates = {normalize_match_text(parts[0])} if parts else set() for part in parts[1:]: next_candidates: set[str] = set() for prefix in candidates: for separator in ("", " "): candidate = normalize_match_text(f"{prefix}{separator}{part}") - if any(keyword.startswith(candidate) for keyword in TOC_KEYWORDS): + if any(keyword.startswith(candidate) for keyword in keywords): next_candidates.add(candidate) candidates = next_candidates if not candidates: return None - return next((candidate for candidate in candidates if candidate in TOC_KEYWORDS), None) + return next((candidate for candidate in candidates if candidate in keywords), None) def _merge_keyword_split_lines( lines: list[str], + keywords: frozenset[str], ) -> list[tuple[str, int, int]]: """Merge consecutive lines that together exactly equal one TOC keyword. @@ -62,12 +91,13 @@ def _merge_keyword_split_lines( """ merged: list[tuple[str, int, int]] = [] index = 0 + max_parts = _max_keyword_split_lines(keywords) while index < len(lines): joined: tuple[str, int, int] | None = None - upper = min(_MAX_KEYWORD_SPLIT_LINES, len(lines) - index) + upper = min(max_parts, len(lines) - index) for part_count in range(upper, 1, -1): parts = [lines[index + offset].strip() for offset in range(part_count)] - keyword = _match_toc_keyword_parts(parts) + keyword = _match_toc_keyword_parts(parts, keywords) if keyword is None: continue joined = (keyword, index, index + part_count - 1) @@ -81,29 +111,59 @@ def _merge_keyword_split_lines( return merged -def _match_toc_keyword_in_line(normalized_line: str) -> str | None: - """Return the longest TOC keyword contained in a normalized line.""" +def _match_toc_keyword_in_line( + normalized_line: str, + *, + keywords: frozenset[str], + prefix_keywords: frozenset[str], +) -> str | None: + """Return the longest matching TOC keyword for a normalized line. + + Base keywords use containment; domain extras require a line-start prefix. + """ if not normalized_line: return None - return next( + contained = next( ( keyword - for keyword in sorted(TOC_KEYWORDS, key=len, reverse=True) + for keyword in sorted(keywords - prefix_keywords, key=len, reverse=True) if keyword in normalized_line ), None, ) + prefixed = next( + ( + keyword + for keyword in sorted(prefix_keywords, key=len, reverse=True) + if normalized_line.startswith(keyword) + ), + None, + ) + if contained is None: + return prefixed + if prefixed is None: + return contained + return contained if len(contained) >= len(prefixed) else prefixed -def _find_toc_text_matches(lines: list[str]) -> list[dict[str, Any]]: - """Match TOC keywords as line-level containment after keyword-split repair. +def _find_toc_text_matches( + lines: list[str], + *, + keywords: frozenset[str], + prefix_keywords: frozenset[str] = frozenset(), +) -> list[dict[str, Any]]: + """Match TOC keywords after keyword-split repair. - Cross-line handling stays keyword-internal only (e.g. 目+录). Hit rule is - ``keyword in normalized_line`` (longest match wins), not whole-line equality. + Cross-line handling stays keyword-internal only (e.g. 目+录). Base keywords + use ``keyword in normalized_line``; domain extras require line-start only. """ matches: list[dict[str, Any]] = [] - for raw_line, start_idx, end_idx in _merge_keyword_split_lines(lines): - keyword = _match_toc_keyword_in_line(normalize_match_text(raw_line)) + for raw_line, start_idx, end_idx in _merge_keyword_split_lines(lines, keywords): + keyword = _match_toc_keyword_in_line( + normalize_match_text(raw_line), + keywords=keywords, + prefix_keywords=prefix_keywords, + ) if keyword is None: continue matches.append( @@ -121,14 +181,20 @@ def _scan_toc_matches( bands_by_page: dict[int, PageTextBands], *, page_count: int, + prefix_keywords: frozenset[str] = frozenset(), ) -> list[dict[str, Any]]: + keywords = _all_scan_keywords(prefix_keywords=prefix_keywords) matches: list[dict[str, Any]] = [] for page in range(1, page_count + 1): bands = bands_by_page.get(page) if bands is None: continue texts = [" ".join(line.text.split()) for line in bands.lines] - for match in _find_toc_text_matches(texts): + for match in _find_toc_text_matches( + texts, + keywords=keywords, + prefix_keywords=prefix_keywords, + ): covered = bands.lines[match["line_index"] : match["line_end_index"] + 1] matches.append( { @@ -277,9 +343,19 @@ def find_toc_anchor_pages(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult start = time.monotonic() total_pages = ctx.blackboard.page_count + profile = ctx.blackboard.document_profile + prefix_keywords = _prefix_toc_keywords_for_categories( + getattr(profile, "category", None) + ) + if prefix_keywords: + logger.info( + "[find.toc_anchor_pages] domain prefix TOC keywords enabled: {}", + sorted(prefix_keywords), + ) keyword_matches = _scan_toc_matches( page_bands_map(ctx.blackboard.page_full_text_cache), page_count=total_pages, + prefix_keywords=prefix_keywords, ) anchor_pages, narrowing = _select_anchor_pages(keyword_matches) if not anchor_pages and narrowing.get("raw") == 0: diff --git a/apps/worker/app/services/document_parser/profiling/domain_taxonomy.py b/apps/worker/app/services/document_parser/profiling/domain_taxonomy.py index 74f64c4d1..8cdf70cfb 100644 --- a/apps/worker/app/services/document_parser/profiling/domain_taxonomy.py +++ b/apps/worker/app/services/document_parser/profiling/domain_taxonomy.py @@ -1,9 +1,9 @@ -"""Document domain taxonomy based on UN ISIC Rev.5. +"""Document domain taxonomy based on IPTC Media Topics (top level only). -Leaves are ISIC divisions (2-digit); finance, professional/scientific and -health are expanded to groups/classes because documents there are common and -dense. The model picks leaf numbers from the rendered menu; only the stable ISIC -code of each pick is persisted, so renumbering the menu never changes stored data. +Source: IPTC NewsCodes Media Topics (https://cv.iptc.org/newscodes/mediatopic/). +Only the 17 official top-level topics are listed; the model picks menu numbers and +we persist the stable Media Topic codes (e.g. ``04000000``). ``OTHER`` is our +local escape hatch when none of the listed topics fit. """ from __future__ import annotations @@ -11,6 +11,7 @@ from dataclasses import dataclass OTHER_CODE = "OTHER" +ECONOMY_BUSINESS_FINANCE_CODE = "04000000" @dataclass(frozen=True) @@ -29,173 +30,28 @@ def _section(title: str, *leaves: tuple[str, str]) -> DomainSection: return DomainSection(title, tuple(DomainLeaf(code, name) for code, name in leaves)) +# Official IPTC Media Topics top-level codes (medtop:NNNNNNNN → NNNNNNNN). DOMAIN_SECTIONS: tuple[DomainSection, ...] = ( _section( - "Agriculture, forestry and fishing", - ("01", "Crop and animal production"), - ("02", "Forestry and logging"), - ("03", "Fishing and aquaculture"), - ), - _section( - "Mining and quarrying", - ("05", "Coal and lignite"), - ("06", "Crude petroleum and natural gas"), - ("07", "Metal ores"), - ("08", "Other mining and quarrying"), - ("09", "Mining support services"), - ), - _section( - "Manufacturing", - ("10", "Food products"), - ("11", "Beverages"), - ("12", "Tobacco products"), - ("13", "Textiles"), - ("14", "Wearing apparel"), - ("15", "Leather and related products"), - ("16", "Wood and cork products"), - ("17", "Paper and paper products"), - ("18", "Printing and recorded media"), - ("19", "Coke and refined petroleum products"), - ("20", "Chemicals and chemical products"), - ("21", "Pharmaceutical products"), - ("22", "Rubber and plastic products"), - ("23", "Non-metallic mineral products"), - ("24", "Basic metals"), - ("25", "Fabricated metal products"), - ("26", "Computer, electronic and optical products"), - ("27", "Electrical equipment"), - ("28", "Machinery and equipment"), - ("29", "Motor vehicles, trailers and semi-trailers"), - ("30", "Other transport equipment"), - ("31", "Furniture"), - ("32", "Other manufacturing"), - ("33", "Repair and installation of machinery and equipment"), - ), - _section( - "Electricity, gas, steam and air conditioning supply", - ("35", "Electricity, gas, steam and air conditioning supply"), - ), - _section( - "Water supply, sewerage, waste management and remediation", - ("36", "Water collection, treatment and supply"), - ("37", "Sewerage"), - ("38", "Waste collection, treatment, disposal and recovery"), - ("39", "Remediation and other waste management services"), - ), - _section( - "Construction", - ("41", "Construction of buildings"), - ("42", "Civil engineering"), - ("43", "Specialized construction activities"), - ), - _section( - "Wholesale and retail trade", - ("46", "Wholesale trade"), - ("47", "Retail trade"), - ), - _section( - "Transportation and storage", - ("49", "Land transport and pipelines"), - ("50", "Water transport"), - ("51", "Air transport"), - ("52", "Warehousing and transport support activities"), - ("53", "Postal and courier activities"), - ), - _section( - "Accommodation and food service", - ("55", "Accommodation"), - ("56", "Food and beverage service"), - ), - _section( - "Publishing, broadcasting and content production", - ("58", "Publishing"), - ("59", "Film, video, television and music production"), - ("60", "Broadcasting, news agencies and content distribution"), - ), - _section( - "Telecommunications, computer programming and information services", - ("61", "Telecommunications"), - ("62", "Computer programming and consultancy"), - ("63", "Computing infrastructure, data processing and information services"), - ), - _section( - "Financial and insurance activities", - ("641", "Monetary intermediation and banking"), - ("642", "Holding companies and financing conduits"), - ("643", "Trusts, funds and similar financial entities"), - ("649", "Other financial services (credit, leasing, factoring)"), - ("651", "Insurance"), - ("652", "Reinsurance"), - ("653", "Pension funding"), - ("6611", "Administration of financial markets"), - ("6612", "Securities and commodity contracts brokerage"), - ("6619", "Other activities auxiliary to financial services"), - ("662", "Activities auxiliary to insurance and pension funding"), - ("663", "Fund management"), - ), - _section( - "Real estate activities", - ("68", "Real estate activities"), - ), - _section( - "Professional, scientific and technical activities", - ("691", "Legal activities"), - ("692", "Accounting, auditing and tax consultancy"), - ("701", "Head offices"), - ("702", "Management consultancy"), - ("711", "Architecture, engineering and technical consultancy"), - ("712", "Technical testing and analysis"), - ("721", "Research and development in natural sciences and engineering"), - ("722", "Research and development in social sciences and humanities"), - ("731", "Advertising"), - ("732", "Market research and public opinion polling"), - ("741", "Specialized design"), - ("742", "Photography"), - ("743", "Translation and interpretation"), - ("749", "Other professional, scientific and technical activities"), - ("75", "Veterinary activities"), - ), - _section( - "Administrative and support services", - ("77", "Rental and leasing"), - ("78", "Employment activities"), - ("79", "Travel agencies and tour operators"), - ("80", "Investigation and security"), - ("81", "Services to buildings and landscape"), - ("82", "Office administration and business support"), - ), - _section( - "Public administration and defence", - ("84", "Public administration, defence and compulsory social security"), - ), - _section( - "Education", - ("85", "Education"), - ), - _section( - "Human health and social work", - ("861", "Hospitals"), - ("862", "Medical and dental practice"), - ("869", "Other human health activities"), - ("87", "Residential care"), - ("88", "Social work without accommodation"), - ), - _section( - "Arts, sports and recreation", - ("90", "Arts creation and performing arts"), - ("91", "Libraries, archives, museums and cultural activities"), - ("92", "Gambling and betting"), - ("93", "Sports, amusement and recreation"), - ), - _section( - "Other service activities", - ("94", "Membership organizations"), - ("95", "Repair of computers, personal and household goods, and vehicles"), - ("96", "Personal services"), - ), - _section( - "Other", - (OTHER_CODE, "Other (no listed domain fits)"), + "IPTC Media Topics", + ("01000000", "Arts, culture, entertainment and media"), + ("16000000", "Conflict, war and peace"), + ("02000000", "Crime, law and justice"), + ("03000000", "Disaster, accident and emergency incident"), + (ECONOMY_BUSINESS_FINANCE_CODE, "Economy, business and finance"), + ("05000000", "Education"), + ("06000000", "Environment"), + ("07000000", "Health"), + ("08000000", "Human interest"), + ("09000000", "Labor"), + ("10000000", "Lifestyle and leisure"), + ("11000000", "Politics and government"), + ("12000000", "Religion"), + ("13000000", "Science and technology"), + ("14000000", "Society"), + ("15000000", "Sport"), + ("17000000", "Weather"), + (OTHER_CODE, "Other (no listed topic fits)"), ), ) @@ -212,7 +68,8 @@ def _render_menu() -> str: lines: list[str] = [] number = 1 for section in DOMAIN_SECTIONS: - lines.append(section.title) + if section.title: + lines.append(section.title) for leaf in section.leaves: lines.append(f"{number}. {leaf.title}") number += 1 @@ -247,6 +104,7 @@ def domain_codes_from_ids(ids: object) -> list[str]: __all__ = [ "DOMAIN_MENU", "DOMAIN_SECTIONS", + "ECONOMY_BUSINESS_FINANCE_CODE", "OTHER_CODE", "DomainLeaf", "DomainSection", diff --git a/apps/worker/tests/contract/test_coarse_profile_protocol_contract.py b/apps/worker/tests/contract/test_coarse_profile_protocol_contract.py index 72771260e..ce9ebbe50 100644 --- a/apps/worker/tests/contract/test_coarse_profile_protocol_contract.py +++ b/apps/worker/tests/contract/test_coarse_profile_protocol_contract.py @@ -43,7 +43,7 @@ def test_parse_profile_reads_classification_fields() -> None: raw = json.dumps( { "is_scanned": True, - "domain_ids": [_menu_id("6612"), _menu_id("29")], + "domain_ids": [_menu_id("04000000"), _menu_id("13000000")], "category_label": "Feasibility Study Report", "routing_category": "generic", "language": "zh", @@ -53,7 +53,7 @@ def test_parse_profile_reads_classification_fields() -> None: ) profile = _parse_profile(raw) assert profile.is_scanned is True - assert profile.category == ["6612", "29"] + assert profile.category == ["04000000", "13000000"] assert profile.category_label == "Feasibility Study Report" assert profile.routing_category == "generic" assert profile.language == "zh" @@ -64,10 +64,10 @@ def test_parse_profile_reads_classification_fields() -> None: [ None, [], - "6612", + "04000000", [0], [10_000], - ["12"], + ["5"], [True], [1, 1], [_menu_id(OTHER_CODE), 1], diff --git a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py index a19b51d7f..9cee4b281 100644 --- a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py +++ b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py @@ -244,7 +244,7 @@ def test_parse_run_recorder_doc_profile_uses_final_anatomy_toc() -> None: recorder.set_doc_profile( ParserDocumentProfile( file_type="pdf", - category=["6612"], + category=["04000000"], category_label="Financial Prospectus", routing_category=PdfRoutingCategory.GENERIC, page_count=0, @@ -268,7 +268,7 @@ def test_parse_run_recorder_doc_profile_uses_final_anatomy_toc() -> None: doc_profile = recorder._doc_profile_for_plan() # noqa: SLF001 assert doc_profile is not None - assert doc_profile["category"] == ["6612"] + assert doc_profile["category"] == ["04000000"] assert doc_profile["category_label"] == "Financial Prospectus" assert doc_profile["routing_category"] == PdfRoutingCategory.GENERIC.value assert doc_profile["page_count"] == 2 diff --git a/apps/worker/tests/contract/test_domain_taxonomy_contract.py b/apps/worker/tests/contract/test_domain_taxonomy_contract.py index f0a357357..94a62b5c8 100644 --- a/apps/worker/tests/contract/test_domain_taxonomy_contract.py +++ b/apps/worker/tests/contract/test_domain_taxonomy_contract.py @@ -26,10 +26,13 @@ _LEAVES = [leaf for section in DOMAIN_SECTIONS for leaf in section.leaves] -def test_leaf_codes_are_unique_and_digit_or_other() -> None: +def test_leaf_codes_are_unique_iptc_or_other() -> None: codes = [leaf.code for leaf in _LEAVES] assert len(codes) == len(set(codes)) - assert all(code == OTHER_CODE or code.isdigit() for code in codes) + assert all( + code == OTHER_CODE or (code.isdigit() and len(code) == 8) for code in codes + ) + assert len(_LEAVES) == 18 # 17 IPTC top-level + OTHER def test_every_section_has_leaves_and_other_is_last() -> None: diff --git a/apps/worker/tests/unit/test_toc_anchor_selection.py b/apps/worker/tests/unit/test_toc_anchor_selection.py index 2df82fe9e..f3df57529 100644 --- a/apps/worker/tests/unit/test_toc_anchor_selection.py +++ b/apps/worker/tests/unit/test_toc_anchor_selection.py @@ -13,9 +13,13 @@ MAX_ANCHOR_CANDIDATES, _filter_edge_anchored, _head_window_pages, + _prefix_toc_keywords_for_categories, _scan_toc_matches, _select_anchor_pages, ) +from app.services.document_parser.profiling.domain_taxonomy import ( + ECONOMY_BUSINESS_FINANCE_CODE, +) def _line( @@ -156,3 +160,32 @@ def test_scan_uses_stage0_lines() -> None: matches = _scan_toc_matches(bands, page_count=2) assert [item["page"] for item in matches] == [1] assert matches[0]["region"] == "body" + + +def test_economy_category_enables_index_to_prefix_keyword() -> None: + assert _prefix_toc_keywords_for_categories( + [ECONOMY_BUSINESS_FINANCE_CODE] + ) == frozenset({"index to"}) + assert _prefix_toc_keywords_for_categories(["13000000"]) == frozenset() + assert _prefix_toc_keywords_for_categories(None) == frozenset() + + +def test_index_to_matches_only_at_line_start_when_enabled() -> None: + bands = { + 1: build_page_text_bands( + (LineDraft(text="INDEX TO FINANCIAL STATEMENTS"),) + ), + 2: build_page_text_bands( + (LineDraft(text="See the index to financial statements below."),) + ), + 3: build_page_text_bands((LineDraft(text="No keyword here"),)), + } + without = _scan_toc_matches(bands, page_count=3) + assert without == [] + with_prefix = _scan_toc_matches( + bands, + page_count=3, + prefix_keywords=frozenset({"index to"}), + ) + assert [item["page"] for item in with_prefix] == [1] + assert with_prefix[0]["match_kind"] == "keyword:index to" From 29648c8267a686d610e7771c759abd18984ba3a1 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 16 Oct 2026 18:29:12 +0800 Subject: [PATCH 5/5] feat: scan exhibit index for all domains as TOC candidates Move exhibit index to the global line-start keyword set, and keep finance-only index to/of prefixes for SEC-style statement indexes. Co-authored-by: Cursor --- .../tools/find_toc_anchor_pages.py | 19 +++++---- .../test_doc_profile_anatomy_contract.py | 11 ++++-- .../tests/unit/test_toc_anchor_selection.py | 39 +++++++++++++------ 3 files changed, 47 insertions(+), 22 deletions(-) diff --git a/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py index 5ce9559fa..d14b9b59e 100644 --- a/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py +++ b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py @@ -28,9 +28,13 @@ # CJK and English TOC keywords used for first-pass anchor detection. TOC_KEYWORDS = frozenset({"目录", "目次", "contents", "table of contents"}) +# Strict extras for every domain: match line-start only. +# Exhibit lists appear in filings, contracts, and legal documents alike. +GLOBAL_PREFIX_TOC_KEYWORDS = frozenset({"exhibit index"}) # Domain extras: only when that IPTC category was picked; match line-start only. DOMAIN_PREFIX_TOC_KEYWORDS: dict[str, frozenset[str]] = { - ECONOMY_BUSINESS_FINANCE_CODE: frozenset({"index to"}), + # SEC filings: "index to/of …" covers financial statements and exhibits. + ECONOMY_BUSINESS_FINANCE_CODE: frozenset({"index to", "index of"}), } MAX_ANCHOR_CANDIDATES = 20 @@ -41,10 +45,10 @@ def _prefix_toc_keywords_for_categories(categories: object) -> frozenset[str]: + keywords: set[str] = set(GLOBAL_PREFIX_TOC_KEYWORDS) if not isinstance(categories, list): - return frozenset() + return frozenset(keywords) codes = {str(code) for code in categories} - keywords: set[str] = set() for code, extras in DOMAIN_PREFIX_TOC_KEYWORDS.items(): if code in codes: keywords.update(extras) @@ -347,11 +351,10 @@ def find_toc_anchor_pages(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult prefix_keywords = _prefix_toc_keywords_for_categories( getattr(profile, "category", None) ) - if prefix_keywords: - logger.info( - "[find.toc_anchor_pages] domain prefix TOC keywords enabled: {}", - sorted(prefix_keywords), - ) + logger.info( + "[find.toc_anchor_pages] line-start TOC keywords enabled: {}", + sorted(prefix_keywords), + ) keyword_matches = _scan_toc_matches( page_bands_map(ctx.blackboard.page_full_text_cache), page_count=total_pages, diff --git a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py index 9cee4b281..226b3a31c 100644 --- a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py +++ b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py @@ -94,10 +94,15 @@ def test_toc_anchor_text_scan_line_contains_keyword_and_split_repair() -> None: "The basic contents of a typical contract document are shown below:", ] - late_matches = toc_anchor_tool._find_toc_text_matches(late_lines) # noqa: SLF001 - split_matches = toc_anchor_tool._find_toc_text_matches(split_lines) # noqa: SLF001 + keywords = toc_anchor_tool.TOC_KEYWORDS + late_matches = toc_anchor_tool._find_toc_text_matches( # noqa: SLF001 + late_lines, keywords=keywords + ) + split_matches = toc_anchor_tool._find_toc_text_matches( # noqa: SLF001 + split_lines, keywords=keywords + ) contains_matches = toc_anchor_tool._find_toc_text_matches( # noqa: SLF001 - contains_lines + contains_lines, keywords=keywords ) assert late_matches[0]["line_index"] == 60 diff --git a/apps/worker/tests/unit/test_toc_anchor_selection.py b/apps/worker/tests/unit/test_toc_anchor_selection.py index f3df57529..8813cf669 100644 --- a/apps/worker/tests/unit/test_toc_anchor_selection.py +++ b/apps/worker/tests/unit/test_toc_anchor_selection.py @@ -162,30 +162,47 @@ def test_scan_uses_stage0_lines() -> None: assert matches[0]["region"] == "body" -def test_economy_category_enables_index_to_prefix_keyword() -> None: +def test_economy_category_adds_finance_prefix_toc_keywords() -> None: assert _prefix_toc_keywords_for_categories( [ECONOMY_BUSINESS_FINANCE_CODE] - ) == frozenset({"index to"}) - assert _prefix_toc_keywords_for_categories(["13000000"]) == frozenset() - assert _prefix_toc_keywords_for_categories(None) == frozenset() + ) == frozenset({"index to", "index of", "exhibit index"}) -def test_index_to_matches_only_at_line_start_when_enabled() -> None: +def test_exhibit_index_prefix_keyword_applies_to_every_domain() -> None: + assert _prefix_toc_keywords_for_categories(["13000000"]) == frozenset( + {"exhibit index"} + ) + assert _prefix_toc_keywords_for_categories(None) == frozenset({"exhibit index"}) + + +def test_finance_prefix_keywords_match_only_at_line_start_when_enabled() -> None: bands = { 1: build_page_text_bands( (LineDraft(text="INDEX TO FINANCIAL STATEMENTS"),) ), 2: build_page_text_bands( + (LineDraft(text="INDEX OF FINANCIAL STATEMENTS"),) + ), + 3: build_page_text_bands((LineDraft(text="EXHIBIT INDEX"),)), + 4: build_page_text_bands((LineDraft(text="INDEX TO EXHIBITS"),)), + 5: build_page_text_bands( (LineDraft(text="See the index to financial statements below."),) ), - 3: build_page_text_bands((LineDraft(text="No keyword here"),)), + 6: build_page_text_bands( + (LineDraft(text="See the exhibit index below."),) + ), + 7: build_page_text_bands((LineDraft(text="No keyword here"),)), } - without = _scan_toc_matches(bands, page_count=3) + without = _scan_toc_matches(bands, page_count=7) assert without == [] with_prefix = _scan_toc_matches( bands, - page_count=3, - prefix_keywords=frozenset({"index to"}), + page_count=7, + prefix_keywords=frozenset({"index to", "index of", "exhibit index"}), ) - assert [item["page"] for item in with_prefix] == [1] - assert with_prefix[0]["match_kind"] == "keyword:index to" + assert [(item["page"], item["match_kind"]) for item in with_prefix] == [ + (1, "keyword:index to"), + (2, "keyword:index of"), + (3, "keyword:exhibit index"), + (4, "keyword:index to"), + ]