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)