55 lines
2.1 KiB
Python
55 lines
2.1 KiB
Python
|
import tempfile
|
||
|
|
import unittest
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from rag.preview import preview_document
|
||
|
|
|
||
|
|
|
||
|
|
class PreviewTests(unittest.TestCase):
|
||
|
|
def test_previews_manually_selected_strategy_with_custom_config(self):
|
||
|
|
with tempfile.TemporaryDirectory() as tmp:
|
||
|
|
path = Path(tmp) / "faq.md"
|
||
|
|
path.write_text("Q: 什么是净值？\nA: 基金单位价值。", encoding="utf-8")
|
||
|
|
|
||
|
|
preview = preview_document(path, strategy="qa_pair", chunk_size=64, chunk_overlap=4)
|
||
|
|
|
||
|
|
self.assertEqual(preview["strategy"], "qa_pair")
|
||
|
|
self.assertEqual(preview["actual_strategy"], "qa_pair")
|
||
|
|
self.assertEqual(len(preview["chunks"]), 1)
|
||
|
|
self.assertEqual(preview["chunks"][0]["section_title"], None)
|
||
|
|
|
||
|
|
def test_rejects_auto_strategy(self):
|
||
|
|
with tempfile.TemporaryDirectory() as tmp:
|
||
|
|
path = Path(tmp) / "faq.md"
|
||
|
|
path.write_text("普通文本", encoding="utf-8")
|
||
|
|
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
preview_document(path, strategy="auto")
|
||
|
|
|
||
|
|
def test_preview_reports_chapter_fallback(self):
|
||
|
|
with tempfile.TemporaryDirectory() as tmp:
|
||
|
|
path = Path(tmp) / "notice.md"
|
||
|
|
path.write_text("没有标题的普通文本", encoding="utf-8")
|
||
|
|
|
||
|
|
preview = preview_document(path, strategy="chapter_semantic")
|
||
|
|
|
||
|
|
self.assertEqual(preview["actual_strategy"], "default")
|
||
|
|
self.assertTrue(preview["degraded"])
|
||
|
|
self.assertTrue(preview["warning"])
|
||
|
|
|
||
|
|
def test_preview_exposes_cleaning_warnings_and_uses_cleaned_text(self):
|
||
|
|
with tempfile.TemporaryDirectory() as tmp:
|
||
|
|
path = Path(tmp) / "notice.md"
|
||
|
|
path.write_text("\ufeff第一段\r\n\r\n\r\n第二段\t内容", encoding="utf-8")
|
||
|
|
|
||
|
|
preview = preview_document(path, strategy="default")
|
||
|
|
|
||
|
|
self.assertTrue(preview["cleaning_changed"])
|
||
|
|
self.assertTrue(preview["cleaning_warnings"])
|
||
|
|
self.assertEqual(preview["chunks"][0]["text"], "第一段")
|
||
|
|
self.assertEqual(preview["chunks"][1]["text"], "第二段 内容")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|