85 lines
2.9 KiB
Python
85 lines
2.9 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""树脂 RAG 知识库导出测试（issue #83）。
|
|||
|
|
|
|||
|
|
覆盖：
|
|||
|
|
1. kb.resin.template.yaml 清单声明的 9 篇文档全部导出且非空（validate_export）；
|
|||
|
|
2. resin_kb_loader 按标题返回文本（含 GB/T 标题规范化）；
|
|||
|
|
3. loader 可直接用于 RagKnowledgeBase（from_template_config + loader）；
|
|||
|
|
4. 未知标题 → 空串。
|
|||
|
|
"""
|
|||
|
|
import os
|
|||
|
|
import sys
|
|||
|
|
import unittest
|
|||
|
|
|
|||
|
|
_RAGKB_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|||
|
|
sys.path.insert(0, _RAGKB_DIR)
|
|||
|
|
from loader import ( # noqa: E402
|
|||
|
|
DEFAULT_KB_CONFIG,
|
|||
|
|
resin_kb_loader,
|
|||
|
|
validate_export,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class TestExportCompleteness(unittest.TestCase):
|
|||
|
|
"""导出完整性：清单 ↔ 文档文件一致。"""
|
|||
|
|
|
|||
|
|
def test_all_listed_documents_exported(self):
|
|||
|
|
problems = validate_export()
|
|||
|
|
self.assertEqual(problems, [], f"导出不完整: {problems}")
|
|||
|
|
|
|||
|
|
def test_document_count(self):
|
|||
|
|
docs = os.listdir(os.path.join(_RAGKB_DIR, "documents"))
|
|||
|
|
self.assertEqual(len(docs), 8)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class TestLoader(unittest.TestCase):
|
|||
|
|
"""加载器行为。"""
|
|||
|
|
|
|||
|
|
def test_process_doc(self):
|
|||
|
|
text = resin_kb_loader("吸附树脂合成工艺规范")
|
|||
|
|
self.assertIn("悬浮聚合法", text)
|
|||
|
|
|
|||
|
|
def test_gbt_title_normalization(self):
|
|||
|
|
# 标题含 / → 文件名用 -（GB/T 5475 → GB-T-5475）
|
|||
|
|
text = resin_kb_loader("GB/T 5475 离子交换树脂取样方法")
|
|||
|
|
self.assertIn("取样", text)
|
|||
|
|
|
|||
|
|
def test_unknown_title_empty(self):
|
|||
|
|
self.assertEqual(resin_kb_loader("不存在的文档"), "")
|
|||
|
|
|
|||
|
|
def test_kb_config_loads(self):
|
|||
|
|
import yaml
|
|||
|
|
with open(DEFAULT_KB_CONFIG, "r", encoding="utf-8") as fh:
|
|||
|
|
raw = yaml.safe_load(fh)
|
|||
|
|
self.assertEqual(raw["template"], "resin")
|
|||
|
|
|
|||
|
|
|
|||
|
|
class TestRagKbIntegration(unittest.TestCase):
|
|||
|
|
"""与内核 rag-kb 集成（loader 直接可用）。"""
|
|||
|
|
|
|||
|
|
def test_build_kb_from_export(self):
|
|||
|
|
import importlib.util
|
|||
|
|
# 挂载 core/rag-kb（目录含连字符）
|
|||
|
|
repo_root = os.path.dirname(os.path.dirname(os.path.dirname(
|
|||
|
|
os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
|
|||
|
|
rag_kb_dir = os.path.join(repo_root, "core", "rag-kb")
|
|||
|
|
if "rag_kb" not in sys.modules:
|
|||
|
|
spec = importlib.util.spec_from_file_location(
|
|||
|
|
"rag_kb", os.path.join(rag_kb_dir, "__init__.py"),
|
|||
|
|
submodule_search_locations=[rag_kb_dir])
|
|||
|
|
_kb = importlib.util.module_from_spec(spec)
|
|||
|
|
sys.modules["rag_kb"] = _kb
|
|||
|
|
spec.loader.exec_module(_kb)
|
|||
|
|
|
|||
|
|
from rag_kb import RagKnowledgeBase, load_kb_config # noqa: E402
|
|||
|
|
|
|||
|
|
kb = RagKnowledgeBase.from_template_config(
|
|||
|
|
load_kb_config(DEFAULT_KB_CONFIG), loader=resin_kb_loader)
|
|||
|
|
self.assertGreater(kb.doc_count, 0)
|
|||
|
|
hits = kb.search("交换容量测定", top_k=3, categories=None)
|
|||
|
|
self.assertGreater(len(hits), 0)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
unittest.main()
|