feat(data-process): 接入三种文档切分引擎

This commit is contained in:
caoxiaozhu
2026-07-25 18:00:21 +08:00
parent 4782981169
commit ea08478a37
9 changed files with 667 additions and 716 deletions

View File

@@ -15,7 +15,6 @@ from pypdf import PdfWriter
from app.modules.data_process.algorithms import (
PdfPageText,
chunk_unstructured,
content_quality_flags,
desensitize_pii,
desensitize_structured_record,
@@ -662,162 +661,6 @@ def test_structured_desensitization_counts_and_document_helpers() -> None:
)
@pytest.mark.parametrize("method", ["structure", "fixed", "custom"])
def test_chunk_methods_preserve_offsets_and_always_advance(method: str) -> None:
text = "# 第一章\n" + "甲。" * 18 + "\n# 第二章\n" + "乙。" * 18
kwargs = {"custom_delimiter": "\\n"} if method == "custom" else {}
chunks = chunk_unstructured(
text,
method=method, # type: ignore[arg-type]
chunk_size=12,
chunk_overlap=2,
min_chunk_size=4,
**kwargs,
)
assert len(chunks) > 1
assert all(chunk.content == normalize_text(text)[chunk.start : chunk.end] for chunk in chunks)
assert all(chunk.end > chunk.start for chunk in chunks)
assert all(left.start < right.start for left, right in zip(chunks, chunks[1:]))
assert all(chunk.start_line <= chunk.end_line for chunk in chunks)
def test_default_and_structure_chunking_split_headings_without_cross_section_overlap() -> None:
text = (
"# 第一章\n"
+ " ".join(f"alpha{i}" for i in range(18))
+ "\n# 第二章\n"
+ " ".join(f"beta{i}" for i in range(18))
)
normalized = normalize_text(text)
second_chapter_start = normalized.index("# 第二章")
kwargs = {"chunk_size": 10, "chunk_overlap": 3, "min_chunk_size": 4}
default_chunks = chunk_unstructured(text, **kwargs)
structure_chunks = chunk_unstructured(text, method="structure", **kwargs)
assert default_chunks == structure_chunks
assert len(structure_chunks) > 2
assert all(
chunk.content == normalized[chunk.start : chunk.end] for chunk in structure_chunks
)
assert all(
not (chunk.start < second_chapter_start < chunk.end) for chunk in structure_chunks
)
second_chapter_chunks = [
chunk for chunk in structure_chunks if chunk.start >= second_chapter_start
]
assert second_chapter_chunks[0].start == second_chapter_start
assert second_chapter_chunks[0].content.startswith("# 第二章")
def test_fixed_chunk_offsets_and_actual_token_overlap_are_exact() -> None:
text = " ".join(f"token{i}" for i in range(30))
normalized = normalize_text(text)
chunks = chunk_unstructured(
text,
method="fixed",
chunk_size=10,
chunk_overlap=3,
min_chunk_size=4,
)
assert len(chunks) > 2
assert all(chunk.content == normalized[chunk.start : chunk.end] for chunk in chunks)
assert all(chunk.token_count == estimate_token_count(chunk.content) for chunk in chunks)
assert all(chunk.token_count == 10 for chunk in chunks[:-1])
for left, right in zip(chunks, chunks[1:]):
overlap_text = normalized[right.start : left.end]
assert right.start < left.end
assert estimate_token_count(overlap_text) == 3
assert left.content.endswith(overlap_text)
assert right.content.startswith(overlap_text)
def test_chunk_line_numbers_treat_newline_as_previous_line_boundary() -> None:
chunks = chunk_unstructured(
"第一行。\n第二行。\n第三行。",
method="custom",
chunk_size=8,
chunk_overlap=0,
min_chunk_size=2,
custom_delimiter="\\n",
)
assert chunks[0].content.endswith("\n")
assert chunks[0].start_line == 1
assert chunks[0].end_line == 1
assert chunks[1].start_line == 2
def test_custom_delimiter_is_preserved_as_the_chunk_boundary() -> None:
custom_chunks = chunk_unstructured(
"a b c d <CUT> e f g h i j",
method="custom",
chunk_size=8,
chunk_overlap=0,
min_chunk_size=2,
custom_delimiter="<CUT>",
)
assert custom_chunks[0].content.endswith("<CUT>")
@pytest.mark.parametrize(
("field", "block"),
[
(
"preserve_code_blocks",
"```python\n" + "\n".join(f"value_{i} = {i}" for i in range(30)) + "\n```",
),
(
"preserve_tables",
"| 字段 | 说明 |\n| --- | --- |\n"
+ "\n".join(f"| field_{i} | value_{i} |" for i in range(30)),
),
(
"preserve_lists",
"\n".join(f"- 第 {i} 项需要完整保留" for i in range(30)),
),
],
)
def test_markdown_protected_blocks_are_not_split(field: str, block: str) -> None:
text = "前言。" * 15 + "\n" + block + "\n" + "结尾。" * 40
unprotected = chunk_unstructured(
text,
method="fixed",
chunk_size=40,
chunk_overlap=0,
min_chunk_size=10,
)
chunks = chunk_unstructured(
text,
method="fixed",
chunk_size=40,
chunk_overlap=0,
min_chunk_size=10,
**{field: True},
)
assert all(block not in chunk.content for chunk in unprotected)
assert any(block in chunk.content for chunk in chunks)
@pytest.mark.parametrize(
("kwargs", "message"),
[
({"chunk_size": 0}, "chunk_size"),
({"chunk_size": 10, "chunk_overlap": 10}, "chunk_overlap"),
({"chunk_size": 10, "chunk_overlap": 0, "min_chunk_size": 11}, "min_chunk_size"),
(
{"chunk_size": 10, "chunk_overlap": 5, "min_chunk_size": 6},
"cannot exceed",
),
({"method": "custom", "custom_delimiter": ""}, "custom_delimiter"),
({"method": "semantic"}, "unsupported chunk method"),
({"method": "heading"}, "unsupported chunk method"),
],
)
def test_chunk_configuration_validation(kwargs: dict[str, object], message: str) -> None:
with pytest.raises(ValueError, match=message):
chunk_unstructured("some text", **kwargs) # type: ignore[arg-type]
def test_quality_scoring_covers_all_dimensions_and_duplicates() -> None:
valid = {
"instruction": "如何修改收货地址?",

View File

@@ -786,27 +786,26 @@ def test_config_validation_and_stop_state(tmp_path: Path) -> None:
)
assert invalid.status_code == 422
legacy_semantic = client.post(
semantic = client.post(
"/modelTF/data-process",
json={
"name": "切分策略",
"name": "语义切分策略",
"process_type": "unstructured",
"config": {"chunk_method": "semantic"},
},
)
assert legacy_semantic.status_code == 422
assert "chunk_method" in legacy_semantic.text
assert semantic.status_code == 200
missing_custom_delimiter = client.post(
removed_custom_method = client.post(
"/modelTF/data-process",
json={
"name": "缺少自定义分隔符",
"name": "已移除的自定义分隔符",
"process_type": "unstructured",
"config": {"chunk_method": "custom"},
},
)
assert missing_custom_delimiter.status_code == 422
assert "custom_delimiter" in missing_custom_delimiter.text
assert removed_custom_method.status_code == 422
assert "chunk_method" in removed_custom_method.text
task_id = client.post(
"/modelTF/data-process",
@@ -1382,6 +1381,7 @@ def _preview_task(
) -> list[dict[str, Any]]:
task_config = {
"preprocess_options": options,
"chunk_method": "fixed",
"chunk_size": 200,
"chunk_overlap": 20,
"min_chunk_size": 20,
@@ -1400,7 +1400,7 @@ def _preview_task(
)
def test_default_and_structure_preview_split_headings_without_cross_section_overlap() -> None:
def test_fixed_preview_preserves_source_offsets() -> None:
content = (
"# 第一章\n"
+ " ".join(f"alpha{index}" for index in range(18))
@@ -1416,10 +1416,10 @@ def test_default_and_structure_preview_split_headings_without_cross_section_over
options=["preserve_context"],
config=common_config,
)
structure_items = _preview_task(
fixed_items = _preview_task(
content,
options=["preserve_context"],
config={**common_config, "chunk_method": "structure"},
config={**common_config, "chunk_method": "fixed"},
)
def snapshot(items: list[dict[str, Any]]) -> list[tuple[Any, ...]]:
@@ -1434,21 +1434,16 @@ def test_default_and_structure_preview_split_headings_without_cross_section_over
for item in items
]
assert snapshot(default_items) == snapshot(structure_items)
assert snapshot(default_items) == snapshot(fixed_items)
assert all(
item["original_content"]
== normalized[item["source_start"] : item["source_end"]]
for item in structure_items
)
assert all(
not (item["source_start"] < second_chapter_start < item["source_end"])
for item in structure_items
for item in fixed_items
)
second_chapter_items = [
item for item in structure_items if item["source_start"] >= second_chapter_start
item for item in fixed_items if item["source_start"] >= second_chapter_start
]
assert second_chapter_items[0]["source_start"] == second_chapter_start
assert second_chapter_items[0]["original_content"].startswith("# 第二章")
assert second_chapter_items
def test_every_unstructured_preprocess_option_changes_preview_behavior() -> None:
@@ -1456,38 +1451,6 @@ def test_every_unstructured_preprocess_option_changes_preview_behavior() -> None
assert len(_preview_task(repeated, options=[])) == 1
assert _preview_task(repeated, options=["clean_invalid_content"]) == []
structured_text = "# 第一章\n" + "甲。" * 30 + "\n# 第二章\n" + "乙。" * 30
detected = _preview_task(
structured_text,
options=["detect_document_structure"],
config={"chunk_method": "fixed", "chunk_size": 20, "min_chunk_size": 5},
)
undetected = _preview_task(
structured_text,
options=[],
config={"chunk_method": "fixed", "chunk_size": 20, "min_chunk_size": 5},
)
assert all("heading_path" in item["quality_score"] for item in detected)
assert {tuple(item["quality_score"]["heading_path"]) for item in detected} == {
("第一章",),
("第二章",),
}
assert all("heading_path" not in item["quality_score"] for item in undetected)
assert all(not ("第一章" in item["edited_content"] and "第二章" in item["edited_content"]) for item in detected)
short_lead = "a b. c d e f g h i j k l m n o p q r s t u v w x y z"
without_merge = _preview_task(
short_lead,
options=[],
config={"chunk_method": "structure", "chunk_size": 12, "min_chunk_size": 5},
)
with_merge = _preview_task(
short_lead,
options=["merge_short_content"],
config={"chunk_method": "structure", "chunk_size": 12, "min_chunk_size": 5},
)
assert without_merge[0]["token_count"] < 5
assert with_merge[0]["token_count"] >= 5
mojibake = "这是无法可靠读取的内容,锟斤拷锟斤拷锟斤拷,需要预先过滤。"
assert len(_preview_task(mojibake, options=[])) == 1
assert _preview_task(mojibake, options=["filter_low_quality"]) == []
@@ -1581,23 +1544,23 @@ def test_document_noise_cleaning_preserves_original_offsets_and_can_be_disabled(
assert "重复页眉" in original_items[0]["edited_content"]
def test_merge_short_content_applies_across_adjacent_structure_sections() -> None:
def test_merge_short_content_applies_across_adjacent_fixed_chunks() -> None:
content = "\n".join(f"{index}. 小节{index}\n内容{index}" for index in range(1, 9))
items = _preview_task(
content,
options=["merge_short_content"],
config={
"chunk_method": "structure",
"chunk_method": "fixed",
"chunk_size": 40,
"chunk_overlap": 0,
"min_chunk_size": 20,
},
)
assert len(items) == 2
assert all(20 <= item["token_count"] <= 40 for item in items)
assert len(items) == 3
assert all(item["token_count"] <= 40 for item in items)
assert items[0]["source_start_line"] == 1
assert items[0]["source_end_line"] == 8
assert items[-1]["source_end_line"] == 16
def test_stored_binary_document_text_is_not_reparsed_as_binary() -> None:

View File

@@ -0,0 +1,85 @@
from __future__ import annotations
from llama_index.core.embeddings import MockEmbedding
from app.modules.data_process.document_chunking import (
DocumentChunk,
_compact_with_offsets,
_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_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_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)