89 lines
3.3 KiB
Python
89 lines
3.3 KiB
Python
|
import unittest
|
||
|
|
|
||
|
|
from rag.chunk_config import ChunkConfig
|
||
|
|
from rag.chunking import chunk_document
|
||
|
|
|
||
|
|
|
||
|
|
class DefaultChunkingTests(unittest.TestCase):
|
||
|
|
def test_keeps_short_paragraphs_as_separate_chunks(self):
|
||
|
|
result = chunk_document(
|
||
|
|
"第一段内容。\n\n第二段内容。",
|
||
|
|
"default",
|
||
|
|
config=ChunkConfig(size=64, overlap=4),
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertEqual([chunk.text for chunk in result.chunks], ["第一段内容。", "第二段内容。"])
|
||
|
|
self.assertEqual(result.actual_strategy, "default")
|
||
|
|
self.assertFalse(result.degraded)
|
||
|
|
|
||
|
|
def test_splits_long_paragraphs_with_overlap(self):
|
||
|
|
result = chunk_document(
|
||
|
|
"abcdefghij" * 3,
|
||
|
|
"default",
|
||
|
|
config=ChunkConfig(size=10, overlap=2),
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertGreater(len(result.chunks), 1)
|
||
|
|
self.assertTrue(all(len(chunk.text) <= 10 for chunk in result.chunks))
|
||
|
|
self.assertEqual(result.chunks[0].text[-2:], result.chunks[1].text[:2])
|
||
|
|
|
||
|
|
|
||
|
|
class QaPairChunkingTests(unittest.TestCase):
|
||
|
|
def test_keeps_each_question_and_answer_together(self):
|
||
|
|
result = chunk_document(
|
||
|
|
"Q: 什么是净值？\nA: 净值是基金单位价值。\nQ: 如何申购？\nA: 通过交易页面申购。",
|
||
|
|
"qa_pair",
|
||
|
|
config=ChunkConfig(size=64, overlap=4),
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertEqual(len(result.chunks), 2)
|
||
|
|
self.assertIn("question: 什么是净值？", result.chunks[0].text)
|
||
|
|
self.assertIn("answer: 净值是基金单位价值。", result.chunks[0].text)
|
||
|
|
|
||
|
|
def test_rejects_missing_qa_markers(self):
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
chunk_document("普通文本，没有问答标记", "qa_pair")
|
||
|
|
|
||
|
|
def test_rejects_the_whole_document_when_one_pair_is_too_long(self):
|
||
|
|
text = "Q: 短问题\nA: 短答案\nQ: 长问题\nA: " + ("很长" * 20)
|
||
|
|
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
chunk_document(text, "qa_pair", config=ChunkConfig(size=20, overlap=2))
|
||
|
|
|
||
|
|
|
||
|
|
class ChapterChunkingTests(unittest.TestCase):
|
||
|
|
def test_preserves_nested_heading_path(self):
|
||
|
|
result = chunk_document(
|
||
|
|
"# 产品说明\n## 风险揭示\n风险内容。",
|
||
|
|
"chapter_semantic",
|
||
|
|
config=ChunkConfig(size=64, overlap=4),
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertEqual(result.actual_strategy, "chapter_semantic")
|
||
|
|
self.assertEqual(result.chunks[0].section_title, "产品说明 > 风险揭示")
|
||
|
|
self.assertTrue(result.chunks[0].text.startswith("【章节：产品说明 > 风险揭示】"))
|
||
|
|
|
||
|
|
def test_splits_only_inside_a_chapter(self):
|
||
|
|
result = chunk_document(
|
||
|
|
"# 第一章\n" + ("甲" * 20) + "\n# 第二章\n" + ("乙" * 20),
|
||
|
|
"chapter_semantic",
|
||
|
|
config=ChunkConfig(size=16, overlap=2),
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertTrue(all("甲" not in chunk.text or "乙" not in chunk.text for chunk in result.chunks))
|
||
|
|
|
||
|
|
def test_degrades_to_default_without_headings(self):
|
||
|
|
result = chunk_document(
|
||
|
|
"没有Markdown标题的普通内容。",
|
||
|
|
"chapter_semantic",
|
||
|
|
config=ChunkConfig(size=64, overlap=4),
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertEqual(result.actual_strategy, "default")
|
||
|
|
self.assertTrue(result.degraded)
|
||
|
|
self.assertTrue(result.warning)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|