feat: 完成 issue #46 RAG 知识库模板化接入(工艺规范/SOP/国标)
- core/rag-kb:领域 RAG 知识库按模板配置(PRD 5.4/7.3,EPIC #6 子任务) - templating.py:知识源三类分类 + 模板命名推导 + 零依赖轻量 YAML 配置加载 - documents.py:文档段落抽取分块,chunk 携带 文档/章节/段落 溯源信息 - store.py:from_template_config 按模板自动构建 + 中英混合词频检索 + 类别过滤 - config/kb.template.yaml:ti-cl4 模板 RAG 库资产示例(工艺规范/SOP/国标) - tests:29 用例全绿(模板化/分块溯源/检索排序/类别过滤/配置校验)
This commit is contained in:
@@ -0,0 +1,2 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""rag-kb 单元测试包。"""
|
||||
@@ -0,0 +1,16 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""测试引导:把 `core/rag-kb` 以包名 `rag_kb` 挂载到 sys.modules。
|
||||
|
||||
目录名 `rag-kb` 含连字符,无法直接以包名 import;挂载后模块内相对导入
|
||||
(`from .templating import ...`)在 unittest 发现机制下可正常解析。
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
|
||||
RAG_KB_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, RAG_KB_DIR)
|
||||
if "rag_kb" not in sys.modules:
|
||||
pkg = types.ModuleType("rag_kb")
|
||||
pkg.__path__ = [RAG_KB_DIR]
|
||||
sys.modules["rag_kb"] = pkg
|
||||
@@ -0,0 +1,91 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""文档抽取(documents)单元测试:分块 / 章节识别 / 引用溯源。"""
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import _bootstrap # noqa: F401
|
||||
|
||||
from rag_kb.documents import ( # noqa: E402
|
||||
KbDocument,
|
||||
build_document,
|
||||
chunk_document,
|
||||
)
|
||||
|
||||
SOP_TEXT = """# 沸腾氯化炉异常处置SOP
|
||||
|
||||
## 2.1 炉温骤升处置
|
||||
当炉温骤升超过阈值时,立即降低氯气流量,并开启冷却循环水。
|
||||
同时通知中控室确认炉压状态。
|
||||
|
||||
## 2.2 炉压异常
|
||||
炉压异常升高时,检查排渣系统是否堵塞,必要时切换备用炉。
|
||||
"""
|
||||
|
||||
|
||||
class ChunkDocumentTest(unittest.TestCase):
|
||||
def test_heading_and_sections(self):
|
||||
doc = KbDocument(
|
||||
doc_id="sop-001",
|
||||
title="沸腾氯化炉异常处置SOP",
|
||||
text=SOP_TEXT,
|
||||
category="sop",
|
||||
)
|
||||
chunks = chunk_document(doc)
|
||||
# 两个小节各至少一个段落
|
||||
self.assertGreaterEqual(len(chunks), 2)
|
||||
sections = {c.section for c in chunks}
|
||||
self.assertIn("2.1 炉温骤升处置", sections)
|
||||
self.assertIn("2.2 炉压异常", sections)
|
||||
|
||||
def test_source_traceability(self):
|
||||
doc = KbDocument(
|
||||
doc_id="sop-001",
|
||||
title="沸腾氯化炉异常处置SOP",
|
||||
text=SOP_TEXT,
|
||||
category="sop",
|
||||
)
|
||||
chunks = chunk_document(doc)
|
||||
first = chunks[0]
|
||||
# 溯源串 = 标题 §章节 ¶序号(PRD 5.4 引用溯源)
|
||||
self.assertTrue(first.source.startswith("沸腾氯化炉异常处置SOP §2.1 炉温骤升处置 ¶1"))
|
||||
self.assertEqual(first.category, "sop")
|
||||
# 段落序号全局递增
|
||||
seqs = [c.seq for c in chunks]
|
||||
self.assertEqual(seqs, sorted(seqs))
|
||||
|
||||
def test_long_paragraph_split_at_sentence_boundary(self):
|
||||
text = "第一句。第二句!第三句?第四句;第五句。" * 10 # 200 字
|
||||
doc = KbDocument(doc_id="d", title="长文档", text=text)
|
||||
chunks = chunk_document(doc, max_chars=120)
|
||||
# 每段不超过阈值(硬切兜底除外)
|
||||
for c in chunks:
|
||||
self.assertLessEqual(len(c.text), 120 + 20) # 句边界允许略超
|
||||
# 拼接后应还原全文(顺序无损)
|
||||
joined = "".join(c.text for c in chunks)
|
||||
self.assertIn("第一句。第二句!", joined)
|
||||
|
||||
def test_pure_heading_skipped(self):
|
||||
doc = KbDocument(doc_id="d", title="只有标题", text="# 第一章\n\n# 第二章\n")
|
||||
chunks = chunk_document(doc)
|
||||
self.assertEqual(chunks, []) # 纯标题行不产生段落
|
||||
|
||||
def test_empty_text_rejected(self):
|
||||
with self.assertRaises(ValueError):
|
||||
KbDocument(doc_id="d", title="空文档", text=" \n ")
|
||||
|
||||
|
||||
class BuildDocumentTest(unittest.TestCase):
|
||||
def test_doc_id_slug(self):
|
||||
doc = build_document("沸腾氯化工艺规范", "正文内容")
|
||||
self.assertEqual(doc.doc_id, "沸腾氯化工艺规范")
|
||||
self.assertEqual(doc.category, "process")
|
||||
|
||||
def test_custom_doc_id(self):
|
||||
doc = build_document("沸腾氯化工艺规范", "正文内容", doc_id="spec-001")
|
||||
self.assertEqual(doc.doc_id, "spec-001")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,148 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""知识库存储与检索(store)单元测试:构建 / 检索 / 类别过滤 / 引用溯源。"""
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import _bootstrap # noqa: F401
|
||||
|
||||
from rag_kb.documents import KbDocument # noqa: E402
|
||||
from rag_kb.store import RagKnowledgeBase # noqa: E402
|
||||
from rag_kb.templating import ( # noqa: E402
|
||||
KbSourceConfig,
|
||||
KbTemplateConfig,
|
||||
KnowledgeSourceKind,
|
||||
load_kb_config,
|
||||
)
|
||||
|
||||
CONFIG_PATH = os.path.join(
|
||||
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
||||
"config", "kb.template.yaml",
|
||||
)
|
||||
|
||||
SPEC_TEXT = """# 沸腾氯化工艺规范
|
||||
|
||||
## 3.1 温度控制
|
||||
沸腾氯化炉炉温应控制在 850-950°C 范围内,防止炉温骤升损坏炉体。
|
||||
温度波动超过 ±20°C 时应立即报告工艺工程师。
|
||||
|
||||
## 3.2 氯气流量
|
||||
氯气流量与加料量保持比例,超压时自动切断氯气供给。
|
||||
"""
|
||||
|
||||
SOP_TEXT = """# 沸腾氯化炉异常处置SOP
|
||||
|
||||
## 2.1 炉温骤升处置
|
||||
当炉温骤升超过阈值时,立即降低氯气流量,并开启冷却循环水。
|
||||
"""
|
||||
|
||||
GB_TEXT = """# GB/T 氯气安全使用标准
|
||||
|
||||
## 5.1 泄漏处置
|
||||
氯气泄漏时应佩戴正压式空气呼吸器,向上风向撤离,并喷水稀释。
|
||||
"""
|
||||
|
||||
|
||||
def _demo_kb() -> RagKnowledgeBase:
|
||||
kb = RagKnowledgeBase()
|
||||
kb.add_documents([
|
||||
KbDocument(doc_id="spec-001", title="沸腾氯化工艺规范",
|
||||
text=SPEC_TEXT, category="process"),
|
||||
KbDocument(doc_id="sop-001", title="沸腾氯化炉异常处置SOP",
|
||||
text=SOP_TEXT, category="sop"),
|
||||
KbDocument(doc_id="gb-001", title="GB/T 氯气安全使用标准",
|
||||
text=GB_TEXT, category="standard"),
|
||||
])
|
||||
return kb
|
||||
|
||||
|
||||
class KnowledgeBaseTest(unittest.TestCase):
|
||||
def test_build_and_stats(self):
|
||||
kb = _demo_kb()
|
||||
self.assertEqual(kb.doc_count, 3)
|
||||
self.assertGreaterEqual(kb.chunk_count, 4) # SPEC 2 段 + SOP 1 段 + GB 1 段
|
||||
stats = kb.category_stats()
|
||||
self.assertIn("process", stats)
|
||||
self.assertIn("sop", stats)
|
||||
self.assertIn("standard", stats)
|
||||
|
||||
def test_duplicate_doc_id_rejected(self):
|
||||
kb = _demo_kb()
|
||||
with self.assertRaises(ValueError):
|
||||
kb.add_document(KbDocument(
|
||||
doc_id="spec-001", title="重复文档", text="内容"))
|
||||
|
||||
def test_search_hit_and_traceability(self):
|
||||
kb = _demo_kb()
|
||||
hits = kb.search("炉温骤升怎么处置", top_k=3)
|
||||
self.assertTrue(hits)
|
||||
top = hits[0]
|
||||
# 命中 SOP 文档且携带完整来源(引用溯源)
|
||||
self.assertEqual(top.category, "sop")
|
||||
self.assertIn("沸腾氯化炉异常处置SOP", top.source)
|
||||
self.assertIn("§", top.source)
|
||||
self.assertIn("炉温骤升", top.text)
|
||||
# 序列化结果(LLM 网关拼 prompt 用)
|
||||
d = top.to_dict()
|
||||
self.assertEqual(d["source"], top.source)
|
||||
self.assertGreaterEqual(d["score"], 0)
|
||||
|
||||
def test_category_filter(self):
|
||||
kb = _demo_kb()
|
||||
hits = kb.search("氯气", top_k=10,
|
||||
categories=[KnowledgeSourceKind.STANDARD])
|
||||
self.assertTrue(hits)
|
||||
for h in hits:
|
||||
self.assertEqual(h.category, "standard")
|
||||
|
||||
def test_empty_query_returns_nothing(self):
|
||||
kb = _demo_kb()
|
||||
self.assertEqual(kb.search(""), [])
|
||||
self.assertEqual(kb.search(" "), [])
|
||||
|
||||
def test_no_match_returns_nothing(self):
|
||||
kb = _demo_kb()
|
||||
self.assertEqual(kb.search("不存在的话题xyz"), [])
|
||||
|
||||
def test_search_scoring_ranked(self):
|
||||
kb = _demo_kb()
|
||||
hits = kb.search("炉温 氯气 流量", top_k=10)
|
||||
scores = [h.score for h in hits]
|
||||
self.assertEqual(scores, sorted(scores, reverse=True))
|
||||
|
||||
|
||||
class FromTemplateConfigTest(unittest.TestCase):
|
||||
def test_load_example_assets(self):
|
||||
# 用真实模板资产(config/kb.template.yaml)+ 内存 loader 构建知识库
|
||||
config = load_kb_config(CONFIG_PATH)
|
||||
texts = {
|
||||
"沸腾氯化工艺规范": SPEC_TEXT,
|
||||
"沸腾氯化炉操作手册": "# 操作手册\n\n启动前检查冷却水。",
|
||||
"沸腾氯化炉异常处置SOP": SOP_TEXT,
|
||||
"交接班报告生成规范": "# 交接班规范\n\n记录炉温与氯气流量。",
|
||||
"GB/T 氯气安全使用标准": GB_TEXT,
|
||||
"GB/T 钛及钛合金加工标准": "# 钛加工标准\n\n控制还原温度。",
|
||||
}
|
||||
kb = RagKnowledgeBase.from_template_config(
|
||||
config, loader=lambda title: texts[title])
|
||||
self.assertEqual(kb.template, "ti-cl4")
|
||||
self.assertEqual(kb.doc_count, 6)
|
||||
self.assertGreaterEqual(kb.chunk_count, 6)
|
||||
# 检索能命中三类知识源
|
||||
hits = kb.search("炉温骤升", top_k=3)
|
||||
self.assertTrue(hits)
|
||||
self.assertEqual(hits[0].category, "sop")
|
||||
|
||||
def test_missing_document_raises(self):
|
||||
config = KbTemplateConfig(
|
||||
template="ti-cl4",
|
||||
sources=[KbSourceConfig(KnowledgeSourceKind.PROCESS, ["不存在的文档"])],
|
||||
)
|
||||
with self.assertRaises(ValueError):
|
||||
RagKnowledgeBase.from_template_config(
|
||||
config, loader=lambda title: "")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,126 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""模板化(templating)单元测试:知识源分类 / 命名推导 / YAML 配置加载。"""
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import _bootstrap # noqa: F401
|
||||
|
||||
from rag_kb.templating import ( # noqa: E402
|
||||
KbTemplateConfig,
|
||||
KbTemplateNaming,
|
||||
KnowledgeSourceKind,
|
||||
SOURCE_KINDS,
|
||||
load_kb_config,
|
||||
sanitize,
|
||||
)
|
||||
|
||||
CONFIG_PATH = os.path.join(
|
||||
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
||||
"config", "kb.template.yaml",
|
||||
)
|
||||
|
||||
|
||||
class SanitizeTest(unittest.TestCase):
|
||||
def test_lower_and_replace_unsafe(self):
|
||||
# 小写 + 空格等非安全字符替换为 `_`;RAG 命名**保留中文**(文档对象键)
|
||||
self.assertEqual(sanitize("Ti-Cl4 模板"), "ti-cl4_模板")
|
||||
self.assertEqual(sanitize("A.B-c_d"), "a.b-c_d")
|
||||
|
||||
def test_chinese_kept(self):
|
||||
self.assertEqual(sanitize("沸腾氯化工艺规范"), "沸腾氯化工艺规范")
|
||||
# `/` 属非安全字符,替换为 `_`
|
||||
self.assertEqual(sanitize("GB/T 氯气安全使用标准"), "gb_t_氯气安全使用标准")
|
||||
|
||||
def test_empty_fallback(self):
|
||||
self.assertEqual(sanitize(""), "kb")
|
||||
self.assertEqual(sanitize("..."), "kb")
|
||||
|
||||
|
||||
class SourceKindTest(unittest.TestCase):
|
||||
def test_kinds_registered(self):
|
||||
# 三类知识源固定注册(PRD 7.3:工艺规范 / SOP / 国标)
|
||||
self.assertEqual(
|
||||
set(SOURCE_KINDS),
|
||||
{"process", "sop", "standard"},
|
||||
)
|
||||
|
||||
def test_labels(self):
|
||||
self.assertEqual(KnowledgeSourceKind.PROCESS.label, "工艺规范")
|
||||
self.assertEqual(KnowledgeSourceKind.SOP.label, "SOP/操作手册")
|
||||
self.assertEqual(KnowledgeSourceKind.STANDARD.label, "国标/标准")
|
||||
|
||||
|
||||
class NamingTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.naming = KbTemplateNaming(template="ti-cl4")
|
||||
|
||||
def test_collection(self):
|
||||
# 向量库 collection:下划线形态(对齐 data-bus `tpl_{tpl}` schema)
|
||||
self.assertEqual(self.naming.collection(), "ti_cl4_kb")
|
||||
|
||||
def test_index_and_namespace(self):
|
||||
self.assertEqual(self.naming.index_name(), "ti_cl4_kb_idx")
|
||||
self.assertEqual(self.naming.namespace(), "tpl-ti-cl4-kb")
|
||||
|
||||
def test_doc_object_key(self):
|
||||
self.assertEqual(
|
||||
self.naming.doc_object_key("沸腾氯化工艺规范"),
|
||||
"kb/沸腾氯化工艺规范.md",
|
||||
)
|
||||
|
||||
def test_default_template(self):
|
||||
self.assertEqual(KbTemplateNaming("").collection(), "kb_kb")
|
||||
|
||||
|
||||
class YamlLoadTest(unittest.TestCase):
|
||||
def test_load_example_config(self):
|
||||
config = load_kb_config(CONFIG_PATH)
|
||||
self.assertEqual(config.template, "ti-cl4")
|
||||
self.assertEqual(config.version, "1.0.0")
|
||||
# 三类知识源各至少声明一份文档
|
||||
self.assertIn(KnowledgeSourceKind.PROCESS, {s.kind for s in config.sources})
|
||||
self.assertIn(KnowledgeSourceKind.SOP, {s.kind for s in config.sources})
|
||||
self.assertIn(KnowledgeSourceKind.STANDARD, {s.kind for s in config.sources})
|
||||
self.assertGreaterEqual(len(config.all_documents()), 3)
|
||||
|
||||
def test_documents_for_kind(self):
|
||||
config = load_kb_config(CONFIG_PATH)
|
||||
docs = config.documents_for(KnowledgeSourceKind.PROCESS)
|
||||
self.assertIn("沸腾氯化工艺规范", docs)
|
||||
self.assertEqual(
|
||||
config.documents_for(KnowledgeSourceKind.STANDARD),
|
||||
["GB/T 氯气安全使用标准", "GB/T 钛及钛合金加工标准"],
|
||||
)
|
||||
|
||||
def test_unknown_kind_rejected(self):
|
||||
# 未知知识源类别必须报错(杜绝配置拼写漂移)
|
||||
import tempfile
|
||||
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False,
|
||||
encoding="utf-8") as fh:
|
||||
fh.write("template: ti-cl4\nsources:\n - kind: hmm\n documents: [a]\n")
|
||||
path = fh.name
|
||||
try:
|
||||
with self.assertRaises(ValueError):
|
||||
load_kb_config(path)
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_missing_template_rejected(self):
|
||||
import tempfile
|
||||
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False,
|
||||
encoding="utf-8") as fh:
|
||||
fh.write("version: 1.0.0\nsources: []\n")
|
||||
path = fh.name
|
||||
try:
|
||||
with self.assertRaises(ValueError):
|
||||
load_kb_config(path)
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user