YG_FT/backend/tests/test_document_chunking.py

from __future__ import annotations

import sys
from types import ModuleType, SimpleNamespace

from llama_index.core.embeddings import MockEmbedding

from app.modules.data_process.document_chunking import (
    DocumentChunk,
    _compact_with_offsets,
    _document_converter,
    _nodes_to_chunks,
    _project_layout_span,
    chunk_fixed_text,
    chunk_semantic_text,
    merge_short_chunks,
)


def test_fixed_splitter_preserves_offsets_and_token_limit() -> None:
    text = "第一段说明苹果。第二段说明香蕉。\n第三段说明数据库。第四段说明索引。"
    chunks = chunk_fixed_text(text, chunk_size=20, chunk_overlap=0)

    assert len(chunks) > 1
    assert all(chunk.source_start is not None for chunk in chunks)
    assert all(chunk.source_end is not None for chunk in chunks)
    assert all(
        chunk.original_content == text[chunk.source_start : chunk.source_end]
        for chunk in chunks
        if chunk.source_start is not None and chunk.source_end is not None
    )
    assert all(chunk.token_count <= 20 for chunk in chunks)


def test_semantic_splitter_uses_llamaindex_and_reapplies_maximum_size() -> None:
    text = "第一段讨论水果。第二段继续讨论香蕉。第三段讨论数据库。第四段讨论索引。"
    chunks = chunk_semantic_text(
        text,
        chunk_size=30,
        chunk_overlap=0,
        breakpoint_percentile_threshold=95,
        embed_model=MockEmbedding(embed_dim=8),
    )

    assert len(chunks) >= 2
    assert all(chunk.token_count <= 30 for chunk in chunks)
    assert "".join(chunk.original_content for chunk in chunks) == text


def test_layout_converter_disables_ocr(monkeypatch) -> None:
    class FakePipelineOptions:
        def __init__(self) -> None:
            self.do_ocr = True

    class FakePdfFormatOption:
        def __init__(self, *, pipeline_options) -> None:
            self.pipeline_options = pipeline_options

    class FakeDocumentConverter:
        def __init__(self, *, format_options) -> None:
            self.format_options = format_options

    docling_module = ModuleType("docling")
    docling_module.__path__ = []
    document_converter_module = ModuleType("docling.document_converter")
    document_converter_module.DocumentConverter = FakeDocumentConverter
    document_converter_module.PdfFormatOption = FakePdfFormatOption
    datamodel_module = ModuleType("docling.datamodel")
    datamodel_module.__path__ = []
    base_models_module = ModuleType("docling.datamodel.base_models")
    base_models_module.InputFormat = SimpleNamespace(PDF="pdf")
    pipeline_options_module = ModuleType("docling.datamodel.pipeline_options")
    pipeline_options_module.PdfPipelineOptions = FakePipelineOptions
    for name, module in {
        "docling": docling_module,
        "docling.document_converter": document_converter_module,
        "docling.datamodel": datamodel_module,
        "docling.datamodel.base_models": base_models_module,
        "docling.datamodel.pipeline_options": pipeline_options_module,
    }.items():
        monkeypatch.setitem(sys.modules, name, module)

    _document_converter.cache_clear()
    try:
        converter = _document_converter()
        options = converter.format_options["pdf"].pipeline_options
        assert options.do_ocr is False
    finally:
        _document_converter.cache_clear()


def test_layout_projection_ignores_layout_whitespace_but_keeps_source_lines() -> None:
    source = "标题\n第一条 这是正文。\n第二条 后续正文。"
    compact_source, offsets = _compact_with_offsets(source)
    start, end, cursor = _project_layout_span(
        source,
        "第一条\n这是正文。",
        compact_source=compact_source,
        source_offsets=offsets,
        compact_start=0,
    )

    assert source[start:end] == "第一条 这是正文。"
    assert cursor > 0


def test_layout_projection_tolerates_list_numbers_inserted_by_serializer() -> None:
    # Word 自动编号存放在 numbering.xml，python-docx 抽取的正文没有编号，
    # 而 Docling 序列化切片时会补上 "1. " 前缀，投影不能因此失败。
    source = "接入方式说明\n结构化数据接入需要先配置连接地址。\n非结构化接入需要上传文档。"
    compact_source, offsets = _compact_with_offsets(source)
    start, end, cursor = _project_layout_span(
        source,
        "1. 结构化数据接入需要先配置连接地址。\n2. 非结构化接入需要上传文档。",
        compact_source=compact_source,
        source_offsets=offsets,
        compact_start=0,
    )

    assert start is not None and end is not None
    assert source[start:end] == "结构化数据接入需要先配置连接地址。\n非结构化接入需要上传文档。"
    assert cursor > 0


def test_layout_projection_never_moves_cursor_backwards() -> None:
    source = "重复段落内容。\n中间正文。\n重复段落内容。"
    compact_source, offsets = _compact_with_offsets(source)
    # 重复内容回退匹配命中已消费的更早位置时，游标必须保持不退。
    _, _, cursor = _project_layout_span(
        source,
        "重复段落内容。",
        compact_source=compact_source,
        source_offsets=offsets,
        compact_start=compact_source.index("中间正文"),
    )
    assert cursor >= compact_source.index("中间正文")


def test_layout_projection_falls_back_to_line_anchors_for_inserted_content() -> None:
    # 表格跨切片时 Docling 会在续片中重复表头，正文不再是连续子串；
    # 按行锚点匹配仍应定位到表头所在行到末行数据之间的连续区间。
    source = "表头甲\t表头乙\n第一行数据\t说明一\n第二行数据\t说明二"
    compact_source, offsets = _compact_with_offsets(source)
    start, end, _ = _project_layout_span(
        source,
        "表头甲 表头乙\n第二行数据 说明二",
        compact_source=compact_source,
        source_offsets=offsets,
        compact_start=0,
    )

    assert start is not None and end is not None
    assert source[start:end] == (
        "表头甲\t表头乙\n第一行数据\t说明一\n第二行数据\t说明二"
    )


def test_layout_projection_refuses_low_coverage_anchor_match() -> None:
    source = "完全无关的正文内容甲。\n完全无关的正文内容乙。"
    compact_source, offsets = _compact_with_offsets(source)
    start, end, cursor = _project_layout_span(
        source,
        "找不到的数据行内容\n另一条找不到的数据行内容",
        compact_source=compact_source,
        source_offsets=offsets,
        compact_start=0,
    )

    assert start is None
    assert end is None
    assert cursor == 0


def test_text_splitter_keeps_chunks_that_cannot_be_located() -> None:
    class FakeNode:
        def get_content(self) -> str:
            return "这段文本在源文本中不存在。"

    chunks = _nodes_to_chunks([FakeNode()], "完全不同的源文本。")

    assert len(chunks) == 1
    assert chunks[0].original_content == "这段文本在源文本中不存在。"
    assert chunks[0].source_start is None
    assert chunks[0].source_start_line is None


def test_short_layout_chunk_merges_with_neighbor_and_keeps_page_provenance() -> None:
    source = "短标题\n这是一段足够长的正文内容，用于测试相邻切片合并。"
    chunks = [
        DocumentChunk("短标题", "短标题", 0, 3, 1, 1, 2, source_pages=(1,)),
        DocumentChunk(
            "这是一段足够长的正文内容，用于测试相邻切片合并。",
            "这是一段足够长的正文内容，用于测试相邻切片合并。",
            4,
            len(source),
            2,
            2,
            20,
            source_pages=(1, 2),
        ),
    ]

    merged = merge_short_chunks(
        chunks,
        source_text=source,
        min_token_count=10,
        max_token_count=100,
    )

    assert len(merged) == 1
    assert merged[0].original_content == source
    assert merged[0].source_pages == (1, 2)