wehub-resource-sync
365 lines
12 KiB
Python
365 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
|
|
sys.path.append(os.getcwd())
|
|
|
|
from yuxi.knowledge.chunking.ragflow_like.dispatcher import chunk_markdown
|
|
from yuxi.knowledge.chunking.ragflow_like.nlp import bullets_category, count_tokens
|
|
from yuxi.knowledge.chunking.ragflow_like.utils.semantic_utils import split_sentences_chinese
|
|
from yuxi.knowledge.chunking.ragflow_like.presets import (
|
|
CHUNK_ENGINE_VERSION,
|
|
CHUNK_PRESET_IDS,
|
|
CHUNK_PRESETS,
|
|
get_chunk_preset_options,
|
|
get_default_chunk_parser_config,
|
|
map_to_internal_parser_id,
|
|
resolve_chunk_processing_params,
|
|
)
|
|
from yuxi.knowledge.utils.kb_utils import resolve_processing_params, sanitize_processing_params
|
|
|
|
|
|
def test_general_maps_to_naive() -> None:
|
|
assert map_to_internal_parser_id("general") == "naive"
|
|
|
|
|
|
def test_resolve_chunk_processing_params_priority() -> None:
|
|
resolved = resolve_chunk_processing_params(
|
|
kb_additional_params={
|
|
"chunk_preset_id": "book",
|
|
"chunk_parser_config": {"chunk_token_num": 300, "delimiter": "\\n"},
|
|
},
|
|
file_processing_params={
|
|
"chunk_preset_id": "qa",
|
|
"chunk_parser_config": {"delimiter": "###", "overlapped_percent": 5},
|
|
},
|
|
request_params={
|
|
"chunk_preset_id": "laws",
|
|
"chunk_parser_config": {"chunk_token_num": 666},
|
|
},
|
|
)
|
|
|
|
assert resolved["chunk_preset_id"] == "laws"
|
|
assert resolved["chunk_engine_version"] == CHUNK_ENGINE_VERSION
|
|
assert resolved["chunk_parser_config"] == {
|
|
"chunk_token_num": 666,
|
|
"delimiter": "###",
|
|
"overlapped_percent": 5,
|
|
}
|
|
|
|
|
|
def test_resolve_chunk_processing_params_returns_only_nested_keys() -> None:
|
|
resolved = resolve_chunk_processing_params(
|
|
kb_additional_params={"chunk_parser_config": {"chunk_token_num": 300}},
|
|
file_processing_params={},
|
|
request_params={},
|
|
)
|
|
|
|
assert resolved["chunk_parser_config"] == {"chunk_token_num": 300}
|
|
assert resolved["chunk_preset_id"] == "general"
|
|
assert resolved["chunk_engine_version"] == CHUNK_ENGINE_VERSION
|
|
assert len(resolved) == 3
|
|
|
|
|
|
def test_qa_chunking_from_markdown_headings() -> None:
|
|
content = """
|
|
# 问题一
|
|
这是答案一。
|
|
|
|
## 子问题
|
|
这是答案二。
|
|
""".strip()
|
|
|
|
chunks = chunk_markdown(
|
|
markdown_content=content,
|
|
file_id="file_1",
|
|
filename="faq.md",
|
|
processing_params={"chunk_preset_id": "qa", "chunk_parser_config": {}},
|
|
)
|
|
|
|
assert len(chunks) >= 1
|
|
assert "问题：" in chunks[0]["content"]
|
|
assert "回答：" in chunks[0]["content"]
|
|
|
|
|
|
def test_chunk_records_include_reserved_position_fields() -> None:
|
|
content = "第一段内容。\n\n第二段内容。"
|
|
|
|
chunks = chunk_markdown(
|
|
markdown_content=content,
|
|
file_id="file_pos",
|
|
filename="pos.md",
|
|
processing_params={
|
|
"chunk_preset_id": "separator",
|
|
"chunk_parser_config": {"delimiter": "\\n\\n"},
|
|
},
|
|
)
|
|
|
|
assert chunks[0]["start_char_pos"] == 0
|
|
assert chunks[0]["end_char_pos"] == len("第一段内容。")
|
|
assert chunks[0]["start_token_pos"] is None
|
|
assert chunks[0]["end_token_pos"] is None
|
|
assert "start_char_pos" in chunks[1]
|
|
|
|
|
|
def test_book_chunking_hierarchical_merge() -> None:
|
|
content = """
|
|
第一章 总则
|
|
第一节 适用范围
|
|
本规范适用于测试场景。
|
|
第二节 基本原则
|
|
应当遵循最小改动原则。
|
|
""".strip()
|
|
|
|
chunks = chunk_markdown(
|
|
markdown_content=content,
|
|
file_id="file_2",
|
|
filename="book.txt",
|
|
processing_params={"chunk_preset_id": "book", "chunk_parser_config": {"chunk_token_num": 256}},
|
|
)
|
|
|
|
assert len(chunks) >= 1
|
|
assert any("第一章" in ck["content"] for ck in chunks)
|
|
|
|
|
|
def test_book_chunking_should_apply_overlength_protection() -> None:
|
|
content = "\n".join(
|
|
[
|
|
"第一章 总则",
|
|
"第一节 适用范围",
|
|
"超长正文" * 1200,
|
|
"第二节 基本原则",
|
|
"应当遵循最小改动原则。",
|
|
]
|
|
)
|
|
max_chunk_tokens = 180
|
|
|
|
chunks = chunk_markdown(
|
|
markdown_content=content,
|
|
file_id="file_book_long",
|
|
filename="book.txt",
|
|
processing_params={
|
|
"chunk_preset_id": "book",
|
|
"chunk_parser_config": {"chunk_token_num": max_chunk_tokens, "delimiter": "\\n"},
|
|
},
|
|
)
|
|
|
|
assert len(chunks) > 1
|
|
assert max(count_tokens(ck["content"]) for ck in chunks) <= max_chunk_tokens
|
|
|
|
|
|
def test_split_sentences_chinese_should_keep_quote_boundary() -> None:
|
|
text = '他说：“你好。”然后问：“你在吗？”最后结束！'
|
|
sentences = split_sentences_chinese(text)
|
|
|
|
assert sentences == ["他说：“你好。”", "然后问：“你在吗？”", "最后结束！"]
|
|
|
|
|
|
def test_markdown_heading_has_higher_weight_in_bullet_category() -> None:
|
|
sections = [
|
|
"# 3.2 个人所得项目及计税、申报方式概括",
|
|
"一、关于季节工、临时工等费用税前扣除问题，以下规定继续执行。",
|
|
"二、根据现行规定，补贴收入应并入工资薪金所得。",
|
|
"（一）从超出国家规定比例支付的补贴，不属于免税福利费。",
|
|
]
|
|
|
|
# 命中 markdown 标题模式（BULLET_PATTERN 下标 4）时，应该优先选中该组。
|
|
assert bullets_category(sections) == 4
|
|
|
|
|
|
def test_mid_sentence_bullet_marker_should_not_be_treated_as_heading() -> None:
|
|
sections = [
|
|
"根据前述规则：一、这里是句中枚举，不是章节标题，不能被当成层级。",
|
|
"延续上文：（二）这里同样是正文中的枚举表达，不是独立标题。",
|
|
"## 3.4 交通补贴的个税处理",
|
|
]
|
|
assert bullets_category(sections) == 4
|
|
|
|
|
|
def test_chunk_preset_options_include_description() -> None:
|
|
options = get_chunk_preset_options()
|
|
assert [option["value"] for option in options] == list(CHUNK_PRESETS)
|
|
assert {option["value"] for option in options} == CHUNK_PRESET_IDS
|
|
assert all(isinstance(option.get("description"), str) and option["description"] for option in options)
|
|
|
|
|
|
def test_chunk_preset_defaults_only_include_strategy_specific_fields() -> None:
|
|
for preset_id in CHUNK_PRESET_IDS:
|
|
assert get_default_chunk_parser_config(preset_id) == {}
|
|
|
|
|
|
def test_laws_chunking_should_apply_overlength_protection() -> None:
|
|
lines = ["#### 中华人民共和国企业所得税法实施条例", "##### 微信扫一扫：分享"]
|
|
lines.extend(
|
|
[f"第{i}条 企业所得税法实施细则说明，适用于测试场景，确保条文长度足够用于验证分块策略。" for i in range(1, 260)]
|
|
)
|
|
content = "\n".join(lines)
|
|
|
|
max_chunk_tokens = 180
|
|
chunks = chunk_markdown(
|
|
markdown_content=content,
|
|
file_id="file_laws_long",
|
|
filename="laws.docx",
|
|
processing_params={
|
|
"chunk_preset_id": "laws",
|
|
"chunk_parser_config": {
|
|
"chunk_token_num": max_chunk_tokens,
|
|
"overlapped_percent": 20,
|
|
"delimiter": "\\n",
|
|
},
|
|
},
|
|
)
|
|
|
|
assert len(chunks) > 1
|
|
assert max(count_tokens(ck["content"]) for ck in chunks) <= max_chunk_tokens
|
|
|
|
|
|
def test_laws_chunking_should_prefer_sentence_boundary_split() -> None:
|
|
line = "第一条 企业所得税法实施细则用于测试分块语义边界。"
|
|
content = line * 120
|
|
|
|
chunks = chunk_markdown(
|
|
markdown_content=content,
|
|
file_id="file_laws_sentence",
|
|
filename="laws.docx",
|
|
processing_params={
|
|
"chunk_preset_id": "laws",
|
|
"chunk_parser_config": {
|
|
"chunk_token_num": 120,
|
|
"overlapped_percent": 0,
|
|
"delimiter": "\\n",
|
|
},
|
|
},
|
|
)
|
|
|
|
assert len(chunks) > 1
|
|
for ck in chunks:
|
|
text = ck["content"].strip()
|
|
assert text
|
|
assert count_tokens(text) <= 120
|
|
|
|
|
|
def test_laws_chunking_should_prefer_article_level_before_item_level() -> None:
|
|
content = """
|
|
第六章 特别纳税调整
|
|
第一百零六条 企业所得税法第三十八条规定的可以指定扣缴义务人的情形，包括：
|
|
(一)在资金、经营、购销等方面存在直接或者间接的控制关系;
|
|
(二)可以代表企业实施其他具有约束力的行为。
|
|
第一百零七条 税务机关可以依法核定应纳税所得额。
|
|
""".strip()
|
|
|
|
chunks = chunk_markdown(
|
|
markdown_content=content,
|
|
file_id="file_laws_article",
|
|
filename="laws.docx",
|
|
processing_params={
|
|
"chunk_preset_id": "laws",
|
|
"chunk_parser_config": {
|
|
"chunk_token_num": 1000,
|
|
"overlapped_percent": 0,
|
|
"delimiter": "\\n",
|
|
},
|
|
},
|
|
)
|
|
|
|
# 只要条下款项没有被拆成独立碎片，即可满足“条级优先”的目标。
|
|
target_chunks = [ck["content"] for ck in chunks if "第一百零六条" in ck["content"]]
|
|
assert target_chunks
|
|
assert any("(一)" in chunk and "(二)" in chunk for chunk in target_chunks)
|
|
|
|
|
|
def test_laws_markdown_articles_should_not_collapse_into_chapter_chunk() -> None:
|
|
content = """
|
|
## 第一章 总则
|
|
- **第一条** 为了规范担保活动，保障债权实现，制定本法。
|
|
- **第二条** 在借贷活动中，当事人可以依法设定担保。
|
|
- **第三条** 担保活动应当遵循平等、自愿、公平和诚实信用原则。
|
|
""".strip()
|
|
|
|
chunks = chunk_markdown(
|
|
markdown_content=content,
|
|
file_id="file_laws_markdown_article",
|
|
filename="laws.md",
|
|
processing_params={
|
|
"chunk_preset_id": "laws",
|
|
"chunk_parser_config": {
|
|
"chunk_token_num": 120,
|
|
"overlapped_percent": 0,
|
|
"delimiter": "\\n",
|
|
},
|
|
},
|
|
)
|
|
|
|
first_article_chunks = [ck["content"] for ck in chunks if "第一条" in ck["content"]]
|
|
assert first_article_chunks
|
|
# 条级切分时，第一条与第二条不应被合并到同一块。
|
|
assert all("第二条" not in chunk for chunk in first_article_chunks)
|
|
assert max(count_tokens(ck["content"]) for ck in chunks) <= 120
|
|
|
|
|
|
def test_sanitize_processing_params_should_drop_non_persistent_fields() -> None:
|
|
sanitized = sanitize_processing_params(
|
|
{
|
|
"chunk_preset_id": "general",
|
|
"chunk_parser_config": {"chunk_token_num": 300},
|
|
"ocr_engine": "mineru_ocr",
|
|
"ocr_engine_config": {},
|
|
"auto_index": True,
|
|
"content_hashes": {"a.md": "hash-a"},
|
|
"enable_ocr": "mineru_ocr",
|
|
"_preprocessed_map": {"a.md": {"path": "/tmp/a.md"}},
|
|
}
|
|
)
|
|
|
|
assert sanitized == {
|
|
"chunk_preset_id": "general",
|
|
"chunk_parser_config": {"chunk_token_num": 300},
|
|
"ocr_engine": "mineru_ocr",
|
|
"ocr_engine_config": {},
|
|
}
|
|
|
|
|
|
def test_resolve_processing_params_keeps_ocr_fields_and_chunk_params() -> None:
|
|
resolved = resolve_processing_params(
|
|
kb_additional_params={
|
|
"chunk_preset_id": "book",
|
|
"chunk_parser_config": {"delimiter": "\n", "chunk_token_num": 300},
|
|
},
|
|
file_processing_params={
|
|
"ocr_engine": "mineru_ocr",
|
|
"ocr_engine_config": {"backend": "pipeline"},
|
|
"chunk_preset_id": "qa",
|
|
"chunk_parser_config": {"overlapped_percent": 10},
|
|
"content_hashes": {"a.md": "hash-a"},
|
|
},
|
|
request_params={
|
|
"auto_index": True,
|
|
"chunk_preset_id": "laws",
|
|
"chunk_parser_config": {"chunk_token_num": 666},
|
|
},
|
|
)
|
|
|
|
assert resolved["ocr_engine"] == "mineru_ocr"
|
|
assert resolved["ocr_engine_config"] == {"backend": "pipeline"}
|
|
assert resolved["chunk_preset_id"] == "laws"
|
|
assert resolved["chunk_parser_config"] == {
|
|
"delimiter": "\n",
|
|
"chunk_token_num": 666,
|
|
"overlapped_percent": 10,
|
|
}
|
|
assert "content_hashes" not in resolved
|
|
assert "enable_ocr" not in resolved
|
|
assert "auto_index" not in resolved
|
|
|
|
|
|
def test_resolve_processing_params_defaults_ocr_fields() -> None:
|
|
resolved = resolve_processing_params(
|
|
kb_additional_params={},
|
|
file_processing_params={"ocr_engine_config": "invalid", "enable_ocr": "mineru_ocr"},
|
|
)
|
|
|
|
assert resolved["ocr_engine"] == "rapid_ocr"
|
|
assert resolved["ocr_engine_config"] == {}
|
|
assert "enable_ocr" not in resolved
|