f47611020a
- 正文抽取下钻 SDT 内容控件，目录等内容不再整段丢失 - 切片投影匹配剥离序列化插入的列表自动编号，新增行锚点兜底与游标防回退 - fixed/semantic 切分路径定位失败时保留切片，不再静默丢弃 - Word 预览按段落大纲级别识别标题，未套标题样式的小节正常渲染 - 本地嵌入模型抽为共享单例，供语义分块与质量评分共用
213 lines
7.9 KiB
Python
213 lines
7.9 KiB
Python
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)
|