RAG-CUT/backend/tests/test_heading_layout_multimodal.py

"""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_instructions_remain_body_text(self) -> None:
        chunks = split_pdf_semantic(
            [
                h("1 Data Platform", level=1),
                p("1. 用户点击发布则会显示对应的数据发布弹窗", y=140),
                p("(1) 发布路径：下拉选择菜单，只可单选", y=170),
                h("1.1 发布管理", level=2, y=220),
                p("发布管理正文。", y=250),
            ],
            SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
        )

        self.assertEqual([chunk.meta["heading"] for chunk in chunks], ["1 Data Platform", "1.1 发布管理"])
        self.assertIn("1. 用户点击发布", chunks[0].content)
        self.assertNotIn("# 1. 用户点击发布", chunks[0].content)

    def test_numbered_option_description_is_not_a_heading(self) -> None:
        option = h(
            "10.4 Default: This option enables you to set default Qty, Account and/or Best Price.",
            level=2,
            y=160,
        )
        option.meta["source_heading"] = True
        chunks = split_pdf_semantic(
            [h("10 Option settings", level=1), option, p("Following option details.", y=200)],
            SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
        )

        self.assertEqual(len(chunks), 1)
        self.assertEqual(chunks[0].meta["heading"], "10 Option settings")
        self.assertIn("10.4 Default: This option enables", chunks[0].content)
        self.assertNotIn("## 10.4 Default", chunks[0].content)

    def test_pure_numeric_heading_keeps_its_real_depth(self) -> None:
        numeric = h("1.1.1", level=2, page=2)
        numeric.meta["source_heading"] = True
        next_section = h("1.2 ACCOUNT STATUS CODE MASTER", level=2, page=5)
        next_section.meta["source_heading"] = True
        chunks = split_pdf_semantic(
            [numeric, table(page=2), next_section, p("Status body.", page=5)],
            SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=4000),
        )

        self.assertEqual([chunk.meta["heading"] for chunk in chunks], ["1.1.1", "1.2 ACCOUNT STATUS CODE MASTER"])
        self.assertNotIn("1.1.1", chunks[1].content)
        self.assertEqual(chunks[1].meta["pages"], [5])

    def test_clock_text_is_not_a_heading_or_ancestor(self) -> None:
        clock = h("15:18:53", level=2, page=8)
        clock.meta["source_heading"] = True
        chunks = split_pdf_semantic(
            [clock, h("第三节 买入/卖出序", level=2, page=9), p("Useful body.", page=9)],
            SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
        )

        self.assertEqual(len(chunks), 1)
        self.assertEqual(chunks[0].meta["heading"], "第三节 买入/卖出序")
        self.assertNotIn("15:18:53", chunks[0].content)
        self.assertEqual(chunks[0].meta["pages"], [9])

    def test_parent_heading_is_metadata_not_repeated_page_block(self) -> None:
        parent = h("1 Parent", level=1, page=1)
        parent.meta["source_heading"] = True
        child = h("1.1 Child", level=2, page=2)
        child.meta["source_heading"] = True
        chunks = split_pdf_semantic(
            [parent, child, p("Child body.", page=2)],
            SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
        )

        self.assertEqual(len(chunks), 1)
        self.assertEqual(chunks[0].meta["heading_path"], ["1 Parent", "1.1 Child"])
        self.assertNotIn("# 1 Parent", chunks[0].content)
        self.assertEqual(chunks[0].meta["pages"], [2])

    def test_oversized_table_keeps_nearest_heading_and_intro(self) -> None:
        big_table = table(page=1, y=180)
        big_table.markdown += "\n" + "| value | description |\n" * 80
        chunks = split_pdf_semantic(
            [
                h("2 Settings", level=1),
                p("The fields are listed below.", y=140),
                big_table,
                p("Following explanation " + "x" * 180, y=340),
                p("Final paragraph " + "y" * 180, y=380),
            ],
            SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=260),
        )

        table_chunks = [chunk for chunk in chunks if "table" in chunk.block_types]
        self.assertEqual(len(table_chunks), 1)
        self.assertIn("2 Settings", table_chunks[0].content)
        self.assertIn("The fields are listed below.", table_chunks[0].content)
        self.assertFalse(any(chunk.block_types and all(t == "heading" for t in chunk.block_types) for chunk in chunks))
    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()