"""Tests for generic heading/layout multimodal splitting.""" from __future__ import annotations import tempfile import unittest from pathlib import Path from unittest.mock import patch from rag_cut.models import Block, BlockType, SplitConfig, SplitMode from rag_cut.pipeline import chunk_document from rag_cut.splitters.pdf_semantic import CHUNK_STRATEGY, split_pdf_semantic def h(text: str, page: int = 1, level: int = 0, y: int = 100) -> Block: return Block( type=BlockType.HEADING if level else BlockType.PARAGRAPH, text=text, level=level, meta={"page": page, "bbox": [72, y, 520, y + 20], "page_height": 800}, ) def p(text: str, page: int = 1, y: int = 130) -> Block: return Block( type=BlockType.PARAGRAPH, text=text, meta={"page": page, "bbox": [72, y, 520, y + 20], "page_height": 800}, ) def img(name: str, page: int = 1, y: int = 180) -> Block: return Block( type=BlockType.IMAGE, image_id=name, image_path=f"assets/test/{name}.png", ocr_text="Login Submit", meta={"page": page, "bbox": [100, y, 500, y + 180], "page_height": 800}, ) def table(page: int = 1, y: int = 220) -> Block: return Block( type=BlockType.TABLE, markdown="| Field | Description |\n| --- | --- |\n| status | Account status |", image_id="table1", image_path="assets/test/table1.png", meta={ "page": page, "pages": [page], "bbox": [72, y, 520, y + 120], "bboxes": [{"page": page, "bbox": [72, y, 520, y + 120]}], "row_count": 2, "col_count": 2, "cross_page": False, "crop_path": "assets/test/table1.png", }, ) class HeadingLayoutMultimodalTest(unittest.TestCase): def test_numbered_headings_create_same_level_boundaries(self) -> None: chunks = split_pdf_semantic( [ h("1 Overview"), p("Intro text."), h("1.1 Account Status", y=180), p("Account status details.", y=210), h("1.2 Trade Detail", y=260), p("Trade detail body.", y=290), ], SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000), ) headings = [c.meta["heading"] for c in chunks] self.assertEqual(headings, ["1 Overview", "1.1 Account Status", "1.2 Trade Detail"]) self.assertTrue(all(c.meta["chunk_strategy"] == CHUNK_STRATEGY for c in chunks)) def test_letter_subheadings_bind_nearest_images(self) -> None: chunks = split_pdf_semantic( [ h("1 Getting Started"), h("A. Log in", y=140), p("Open the login page.", y=170), img("login", y=210), h("B) Authentication", y=430), p("Approve the authentication request.", y=460), img("auth", y=500), ], SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000), ) self.assertEqual([c.meta["heading"] for c in chunks], ["A. Log in", "B) Authentication"]) self.assertEqual(chunks[0].meta["heading_path"], ["1 Getting Started", "A. Log in"]) self.assertEqual(chunks[0].meta["images"][0]["image_id"], "login") self.assertEqual(chunks[1].meta["images"][0]["image_id"], "auth") def test_table_stays_atomic_with_markdown_and_position(self) -> None: chunks = split_pdf_semantic( [h("2 Settings"), p("The fields are listed below."), table(y=180)], SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000), ) self.assertEqual(len(chunks), 1) self.assertIn("| Field | Description |", chunks[0].content) self.assertEqual(chunks[0].meta["tables"][0]["row_count"], 2) self.assertEqual(chunks[0].meta["tables"][0]["bbox"], [72, 180, 520, 300]) def test_unlabelled_image_only_chunk_is_preview_only(self) -> None: cover_logo = img("cover-logo") cover_logo.ocr_text = "" chunks = split_pdf_semantic( [cover_logo, h("1 Overview", page=2), p("Useful body.", page=2)], SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000), ) self.assertEqual(len(chunks), 2) self.assertFalse(chunks[0].meta["retrieval"]) self.assertTrue(chunks[1].meta["retrieval"]) self.assertIn("cover-logo", chunks[0].content) def test_image_only_chunk_with_ocr_remains_retrievable(self) -> None: chunks = split_pdf_semantic( [img("workflow-screenshot")], SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000), ) self.assertEqual(len(chunks), 1) self.assertTrue(chunks[0].meta["retrieval"]) def test_noise_blocks_are_skipped(self) -> None: blocks = [] for page in range(1, 5): blocks.append(p("Product Manual", page=page, y=20)) blocks.append(p(str(page), page=page, y=760)) blocks.extend( [ p("Contents"), p("1 Intro ........ 3"), p("2 Setup ........ 8"), h("1 Intro", page=2), p("Useful body.", page=2), ] ) chunks = split_pdf_semantic(blocks, SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000)) content = "\n".join(c.content for c in chunks) self.assertIn("Useful body.", content) self.assertNotIn("Product Manual", content) self.assertNotIn("Contents", content) self.assertNotIn("........", content) self.assertNotIn("2 Setup", content) def test_cross_page_same_heading_is_one_chunk(self) -> None: chunks = split_pdf_semantic( [h("3 Client"), p("Page one text.", page=1), p("Page two continuation.", page=2)], SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000), ) self.assertEqual(len(chunks), 1) self.assertEqual(chunks[0].meta["pages"], [1, 2]) self.assertIn("Page two continuation.", chunks[0].content) def test_instructional_steps_stay_in_parent_section(self) -> None: """Step 1/2/3 lines must not become headings that get flushed away.""" chunks = split_pdf_semantic( [ h("5.7 CRS SUBMISSION TO IRD", level=2), h("5.7.1 Submission at 1st time", level=3, y=140), p( "For CRS submission at first time by G3SB, prior consent has to be obtained " "from IRD by submitting test data file to the AEOI Portal for validation.", y=170, ), img("workflow", y=210), p("Step 1 Users prepare the XML as per sections 5.1– 5.5.", y=420), p('Step 2 Set parameter "Export in test data format?" value to "Y".', y=450), p("Step 3 Check XML file tag DocTypeIndic was using OECD11.", y=480), p( "Step 4 Go to Registration and login page of AEOI Portal: " "https://aeoi1.ird.gov.hk/portal/landing/", y=510, ), p("Download the encryption tools to encrypt the XML file.", y=540), ], SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=4000), ) content = "\n".join(c.content for c in chunks) self.assertIn("Step 1 Users prepare the XML", content) self.assertIn("Export in test data format", content) self.assertIn("OECD11", content) self.assertIn("Step 4 Go to Registration", content) self.assertTrue(any(c.meta["heading"] == "5.7.1 Submission at 1st time" for c in chunks)) def test_oversized_section_splits_by_length_without_parent(self) -> None: blocks = [h("4 Reports")] + [p(f"Long paragraph {i} " + "x" * 80, y=130 + i) for i in range(8)] chunks = split_pdf_semantic(blocks, SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=260)) self.assertGreater(len(chunks), 1) self.assertTrue(all(c.meta.get("retrieval", True) for c in chunks)) self.assertFalse(any(c.meta.get("is_section_parent") for c in chunks)) self.assertFalse(any(c.meta.get("is_sub_chunk") for c in chunks)) self.assertFalse(any(c.meta.get("parent_chunk_id") is not None for c in chunks)) def test_catalog_titles_fold_into_section_not_dropped(self) -> None: """TOC-style '1.xxx' lines between catalog headings must not vanish on flush.""" chunks = split_pdf_semantic( [ h("14.W. Whitney George", level=1, y=100), p("Famous fund manager bio.", y=130), Block( type=BlockType.HEADING, text="形態指標", level=1, meta={ "page": 1, "bbox": [72, 200, 200, 220], "page_height": 800, "font_size": 16, "body_font_size": 12, }, ), p("共包含以下11個形態指標說明：", y=240), h("1.頭肩頂形態", level=1, y=260), h("2.頭肩底形態", level=1, y=280), h("1.頭肩頂形態", level=1, page=2, y=120), p("整體介紹：頂部反轉形態說明。", page=2, y=150), ], SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=4000), ) morph = next(c for c in chunks if c.meta.get("heading") == "形態指標") self.assertEqual(morph.meta.get("heading_path"), ["形態指標"]) self.assertIn("共包含以下11個形態指標說明：", morph.content) self.assertIn("1.頭肩頂形態", morph.content) self.assertIn("2.頭肩底形態", morph.content) detail = next( c for c in chunks if c.meta.get("heading") == "1.頭肩頂形態" and "整體介紹" in c.content ) self.assertIn("整體介紹", detail.content) def test_empty_explicit_section_does_not_pollute_previous_chunk(self) -> None: empty_heading = h("2 REFERENCES", level=1, page=3, y=100) empty_heading.meta["source_heading"] = True chunks = split_pdf_semantic( [ h("1 GLOSSARY", level=1, page=2, y=100), table(page=2, y=140), empty_heading, h("3 DOCUMENT HISTORY", level=1, page=3, y=180), p("History body.", page=3, y=220), ], SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=4000), ) glossary = next(c for c in chunks if c.meta.get("heading") == "1 GLOSSARY") self.assertNotIn("2 REFERENCES", glossary.content) self.assertFalse(any("2 REFERENCES" in c.content for c in chunks)) class DummyParser: def parse(self, path: Path, assets_dir: Path) -> list[Block]: return [h("1 Overview"), p("Pipeline body.")] class PipelineSemanticIntegrationTest(unittest.TestCase): def test_pdf_default_mode_uses_heading_layout_multimodal(self) -> None: with tempfile.TemporaryDirectory() as tmp: tmp_root = Path(tmp) source = tmp_root / "sample.pdf" source.write_bytes(b"%PDF-1.4\n") with patch("rag_cut.pipeline.get_parser", return_value=DummyParser()): result = chunk_document( source, config=SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000), storage_root=tmp_root / "storage", ) self.assertEqual(result.chunk_count, 1) self.assertEqual(result.split_config["chunk_strategy"], CHUNK_STRATEGY) self.assertEqual(result.chunks[0].meta["chunk_strategy"], CHUNK_STRATEGY) if __name__ == "__main__": unittest.main()