86 lines
2.9 KiB
Python
86 lines
2.9 KiB
Python
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)
|