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:
2026-08-04 16:31:31 +08:00
parent 9c3e59695a
commit 1fb1d278d5
11 changed files with 1135 additions and 0 deletions
+53
View File
@@ -0,0 +1,53 @@
# RAG KB —— 领域 RAG 知识库模板化接入
对应 PRD 5.4「④ LLM 网关 + RAG」与 Issue #46(EPIC #6 子任务)。
领域 RAG 知识库**按模板配置**(工艺规范 / SOP / 国标 三类知识源),
换行业只改模板资产(`config/kb.template.yaml` + 文档集),内核零改动。
检索返回**命中文档片段 + 来源**(引用溯源,PRD 5.4 强制要求)。
## 模块
| 文件 | 职责 |
|------|------|
| `templating.py` | 知识源分类(process/sop/standard)+ 模板命名推导(向量库 collection / 索引 / namespace)+ 轻量 YAML 配置加载(零第三方依赖) |
| `documents.py` | 文档模型 + 段落抽取分块:每段保留 `文档/章节/段落` 溯源信息 |
| `store.py` | 模板化知识库(`from_template_config` 按配置自动构建)+ 中英混合词频检索 + 类别过滤,返回带来源的命中 |
## 使用示例
```python
from rag_kb import KbTemplateConfig, RagKnowledgeBase
from rag_kb.templating import KbSourceConfig, KnowledgeSourceKind
# 模板 RAG 库配置(通常由 config/kb.template.yaml 加载)
config = KbTemplateConfig(
template="ti-cl4",
sources=[
KbSourceConfig(KnowledgeSourceKind.PROCESS, ["沸腾氯化工艺规范"]),
KbSourceConfig(KnowledgeSourceKind.SOP, ["沸腾氯化炉异常处置SOP"]),
KbSourceConfig(KnowledgeSourceKind.STANDARD, ["GB/T 有关氯气安全标准"]),
],
)
kb = RagKnowledgeBase.from_template_config(
config, loader=lambda title: read_object_storage(title)) # 文档加载器由部署侧注入
hits = kb.search("炉温骤升怎么处理", top_k=3) # 召回 + 溯源
for h in hits:
print(h.source, "→", h.text) # e.g. 沸腾氯化炉异常处置SOP §2.1 ¶3
```
## 验收口径(Issue #46)
- **模板化接入**:三类知识源(工艺规范/SOP/国标)由 `kb.template.yaml` 声明,
换行业只换模板资产,内核零改动(对齐 data-bus 模板化思想)。
- **引用溯源**:检索命中返回 `source`(标题 §章节 ¶段落)——PRD 5.4「答案强制引用溯源」。
- **类别过滤**:可按知识源类别限定检索范围(如只看 SOP)。
- **配置校验**:未知知识源类别 / 缺 template / 文档加载为空 直接报错,杜绝漂移。
## 测试
```bash
cd core/rag-kb
python -m unittest discover -s tests -v
```
+43
View File
@@ -0,0 +1,43 @@
# -*- coding: utf-8 -*-
"""iAOP-Core · 领域 RAG 知识库(RAG KB)—— 模板化接入。
对应 PRD 5.4「④ LLM 网关 + RAG」与 Issue #46(EPIC #6 子任务):
领域 RAG 知识库改为**按模板配置**(工艺规范 / SOP / 国标 三类知识源),
换行业只改模板资产(`config/kb.template.yaml` + 文档集),内核零改动;
检索返回**命中文档片段 + 来源**(引用溯源,PRD 5.4)。
模块:
- templating 知识源分类 + 模板命名推导(collection/索引/namespace)
+ 轻量 YAML 配置加载(零依赖);
- documents 文档模型 + 段落抽取分块(保留 文档/章节/段落 溯源信息);
- store 模板化知识库(from_template_config 构建)+ 检索 + 类别过滤。
测试:`python -m unittest discover -s tests -v`(在 core/rag-kb 目录下执行)。
"""
__version__ = "0.1.0"
from .documents import Chunk, KbDocument, build_document, chunk_document
from .store import RagKnowledgeBase, RetrievalHit
from .templating import (
KbSourceConfig,
KbTemplateConfig,
KbTemplateNaming,
KnowledgeSourceKind,
load_kb_config,
sanitize,
)
__all__ = [
"KnowledgeSourceKind",
"KbTemplateNaming",
"KbTemplateConfig",
"KbSourceConfig",
"load_kb_config",
"sanitize",
"KbDocument",
"Chunk",
"build_document",
"chunk_document",
"RagKnowledgeBase",
"RetrievalHit",
]
+23
View File
@@ -0,0 +1,23 @@
# -*- coding: utf-8 -*-
# 模板 RAG 库资产示例:ti-cl4(氯化车间/海绵钛,Template-Ti 一期)。
#
# 知识源分类固定三类(PRD 7.3 领域RAG知识库一期范围):
# - process 工艺规范(沸腾氯化工艺规范、操作手册);
# - sop 标准作业程序 / 异常处置 SOP;
# - standard 国标 / 行业标准 / 文献(氯、钛相关)。
# 换行业只改本文件 + 对应文档集,内核零改动。
template: ti-cl4
version: 1.0.0
sources:
- kind: process
documents:
- 沸腾氯化工艺规范
- 沸腾氯化炉操作手册
- kind: sop
documents:
- 沸腾氯化炉异常处置SOP
- 交接班报告生成规范
- kind: standard
documents:
- GB/T 氯气安全使用标准
- GB/T 钛及钛合金加工标准
+154
View File
@@ -0,0 +1,154 @@
# -*- coding: utf-8 -*-
"""RAG 文档模型与段落抽取 —— 抽取管线(Issue #46 / PRD 5.4 引用溯源)。
客户文档(工艺规范 / SOP / 国标,格式:文本或 Markdown)经抽取管线转成
**带来源的段落(chunk)**:每个 chunk 保留 `doc_id + 章节号 + 段落号`,
检索命中的段落可精确回指源文档章节(对齐 PRD 5.4「返回命中文档片段+来源」)。
分块规则(简单可靠,零外部依赖):
- 按空行 / Markdown 标题切分段落;标题行单独记录为小节标题;
- 每段按 `max_chars` 阈值再切分(优先在 `。..!?;;` 等句边界断开),
避免超长段落撑爆向量块;
- chunk 携带 `source` 溯源串(如 `沸腾氯化炉异常处置SOP §2.1 ¶3`)。
"""
from __future__ import annotations
import re
from dataclasses import dataclass, field
from typing import List, Optional
# Markdown 标题行(# / ## / ### ...)—— 作为小节起点
_HEADING_RE = re.compile(r"^\s{0,3}#{1,6}\s+(.*)$")
# 句子边界(中文句号/省略号/英文句点/问号/感叹号/分号)
_SENT_BOUNDARY = re.compile(r"(?<=[。..!?;;])")
@dataclass
class KbDocument:
"""一份知识库源文档(标题即唯一 ID,来源类型见 KnowledgeSourceKind)。"""
doc_id: str # 文档唯一 ID(清洗后用于对象键)
title: str # 文档标题(显示与溯源用)
text: str # 原始文本(Markdown 或纯文本)
category: str = "process" # 知识源类别(process/sop/standard)
version: str = "1.0.0" # 文档版本(来源追溯)
def __post_init__(self) -> None:
self.text = (self.text or "").strip()
if not self.title:
raise ValueError("KbDocument.title 不能为空")
if not self.text:
raise ValueError(f"KbDocument {self.title!r} 文本为空")
@dataclass
class Chunk:
"""抽取后的段落(检索最小单元),携带完整来源信息。"""
doc_id: str
title: str
section: str # 小节标题(无标题段落记 `§0 概述` 或空串)
seq: int # 段落序号(文档内 1 起)
text: str
category: str = "process"
@property
def source(self) -> str:
"""溯源串:`标题 §章节 ¶段落`(PRD 5.4 引用溯源返回给调用方)。"""
base = f"{self.title}"
if self.section:
base += f" §{self.section}"
return f"{base}{self.seq}"
def _split_paragraphs(text: str) -> List[str]:
"""按空行 / 标题切分段落;顺带记录标题。返回 (section, para) 对交给调用方。
实现:先按空行粗分块,再在每个块内识别标题行(标题行不入段落文本,
而是成为随后段落的 section 名)。
"""
blocks: List[List[str]] = []
cur: List[str] = []
for raw in text.split("\n"):
line = raw.rstrip()
if not line.strip():
if cur:
blocks.append(cur)
cur = []
continue
cur.append(line)
if cur:
blocks.append(cur)
return ["\n".join(b) for b in blocks]
def _split_by_chars(para: str, max_chars: int) -> List[str]:
"""超长段落按句边界切分,每片不超过 max_chars。"""
if len(para) <= max_chars:
return [para]
pieces: List[str] = []
for sentence in _SENT_BOUNDARY.split(para):
if not sentence.strip():
continue
if pieces and len(pieces[-1]) + len(sentence) <= max_chars:
pieces[-1] += sentence
else:
# 单句仍超长则硬切,避免无限循环
while len(sentence) > max_chars:
pieces.append(sentence[:max_chars])
sentence = sentence[max_chars:]
if sentence:
pieces.append(sentence)
return [p for p in pieces if p.strip()]
def chunk_document(doc: KbDocument, max_chars: int = 500) -> List[Chunk]:
"""文档 → 带来源的段落列表(抽取管线核心,纯函数便于测试)。"""
chunks: List[Chunk] = []
seq = 0
section = ""
for para in _split_paragraphs(doc.text):
lines = para.split("\n")
heading = None
for ln in lines:
m = _HEADING_RE.match(ln)
if m:
heading = m.group(1).strip()
if heading is not None:
section = heading
body_lines = [ln for ln in lines if not _HEADING_RE.match(ln)]
para = "\n".join(body_lines).strip()
if not para:
continue # 纯标题行:仅更新小节名
for piece in _split_by_chars(para, max_chars):
seq += 1
chunks.append(
Chunk(
doc_id=doc.doc_id,
title=doc.title,
section=section,
seq=seq,
text=piece,
category=doc.category,
)
)
return chunks
def build_document(
title: str,
text: str,
category: str = "process",
version: str = "1.0.0",
doc_id: Optional[str] = None,
) -> KbDocument:
"""便捷构造:doc_id 缺省时由标题清洗生成(对齐 templating.sanitize)。"""
from .templating import sanitize # 局部导入避免循环依赖
return KbDocument(
doc_id=doc_id or sanitize(title),
title=title,
text=text,
category=category,
version=version,
)
+206
View File
@@ -0,0 +1,206 @@
# -*- coding: utf-8 -*-
"""RAG 知识库存储与检索(Issue #46 / PRD 5.4)。
`RagKnowledgeBase` 是模板 RAG 库的运行时形态:由模板配置 + 文档加载器
自动构建(`from_template_config`),换行业只换模板资产(YAML + 文档集),
内核零改动。检索返回**命中的文档片段 + 来源串**(`RetrievalHit.source`),
满足 PRD 5.4「RAG 答案强制引用溯源」。
检索实现为零依赖的倒排词频匹配(中英混合分词 + 子串/单词计数打分):
- 向量化/embedding 由部署侧向量库(如 Milvus)接入,本模块保证
检索语义(召回 + 溯源 + 类别过滤)与向量库一致;
- 支持按知识源类别(工艺规范 / SOP / 国标)过滤检索范围。
"""
from __future__ import annotations
import re
from dataclasses import dataclass
from typing import Callable, Dict, List, Optional, Sequence
from .documents import Chunk, KbDocument, chunk_document
from .templating import (
KbTemplateConfig,
KbTemplateNaming,
KnowledgeSourceKind,
SOURCE_KINDS,
)
# 中英混合分词:英文单词/数字 + 中文连续块
_TOKEN_RE = re.compile(r"[a-zA-Z0-9]+|[\u4e00-\u9fff]+")
# 文档加载器:文档标题 → 原始文本(由部署侧提供:读对象存储/本地目录)
DocumentLoader = Callable[[str], str]
def _tokenize(text: str) -> List[str]:
"""分词:英文按单词;中文连续块切**二元组(bigram)**。
bigram 使整句中文查询(如「炉温骤升怎么处置」)与段落中的连续子串
可匹配(共享 bigram 计数),无需外部分词器,零依赖可复现。
"""
tokens: List[str] = []
for t in _TOKEN_RE.findall(text or ""):
t = t.lower()
if t.isascii() or len(t) < 2:
tokens.append(t)
else:
tokens.extend(t[i : i + 2] for i in range(len(t) - 1))
return tokens
def _count_token(text_lower: str, token: str) -> int:
"""chunk 内 token 出现次数:英文按单词边界、中文按子串(查询词原样匹配)。"""
if token.isascii():
return len(re.findall(rf"\b{re.escape(token)}\b", text_lower))
return text_lower.count(token)
@dataclass
class RetrievalHit:
"""一次检索命中:段落文本 + 来源(引用溯源,PRD 5.4)。"""
chunk: Chunk
score: float
@property
def text(self) -> str:
return self.chunk.text
@property
def source(self) -> str:
"""溯源串(如 `沸腾氯化炉异常处置SOP §2.1 ¶3`),随答案返回给用户。"""
return self.chunk.source
@property
def category(self) -> str:
return self.chunk.category
def to_dict(self) -> Dict[str, object]:
return {
"source": self.source,
"category": self.category,
"text": self.chunk.text,
"score": round(self.score, 4),
}
class RagKnowledgeBase:
"""模板化 RAG 知识库(内存实现,零外部依赖)。
用法:
```python
kb = RagKnowledgeBase.from_template_config(
config, loader=lambda title: read_object(title))
hits = kb.search("炉温骤升怎么处理", top_k=3)
for h in hits:
print(h.source, h.text) # 溯源 + 片段
```
"""
def __init__(self, naming: Optional[KbTemplateNaming] = None):
self.naming = naming or KbTemplateNaming("default")
self._docs: Dict[str, KbDocument] = {}
self._chunks: List[Chunk] = []
# ------------------------------------------------------------------
# 构建
# ------------------------------------------------------------------
def add_document(self, doc: KbDocument, max_chars: int = 500) -> int:
"""入库一份文档,返回新增段落数。"""
if doc.doc_id in self._docs:
raise ValueError(f"文档 {doc.doc_id!r} 已存在(同一知识库内 doc_id 唯一)")
self._docs[doc.doc_id] = doc
chunks = chunk_document(doc, max_chars=max_chars)
self._chunks.extend(chunks)
return len(chunks)
def add_documents(self, docs: Sequence[KbDocument], max_chars: int = 500) -> int:
return sum(self.add_document(d, max_chars=max_chars) for d in docs)
@classmethod
def from_template_config(
cls,
config: KbTemplateConfig,
loader: DocumentLoader,
max_chars: int = 500,
) -> "RagKnowledgeBase":
"""按模板配置构建知识库:遍历三类知识源文档清单,经 loader 取文本入库。
换行业只改模板资产(kb.template.yaml + 文档集),内核零改动。
"""
kb = cls(naming=KbTemplateNaming(config.template))
for source in config.sources:
for title in source.documents:
text = loader(title)
if not text or not text.strip():
raise ValueError(
f"文档 {title!r}({source.kind.value})加载为空,无法入库"
)
doc = KbDocument(
doc_id=source.kind.value + ":" + _slug(title),
title=title,
text=text,
category=source.kind.value,
version=config.version,
)
kb.add_document(doc, max_chars=max_chars)
return kb
# ------------------------------------------------------------------
# 检索
# ------------------------------------------------------------------
def search(
self,
query: str,
top_k: int = 5,
categories: Optional[Sequence[KnowledgeSourceKind]] = None,
) -> List[RetrievalHit]:
"""检索:query → 命中文档片段(带来源),按相关度降序取 top_k。
`categories` 限定检索范围(如只看 SOP);缺省检索全部知识源。
"""
tokens = _tokenize(query)
if not tokens or top_k <= 0:
return []
cat_set = {c.value for c in categories} if categories else None
scored: List[RetrievalHit] = []
for chunk in self._chunks:
if cat_set is not None and chunk.category not in cat_set:
continue
text_lower = chunk.text.lower()
freq = sum(_count_token(text_lower, t) for t in tokens)
if freq > 0:
# TF 密度打分:频次 / 段落长度平方根,抑制长段落重复命中偏好
score = freq / (len(chunk.text) ** 0.5)
scored.append(RetrievalHit(chunk=chunk, score=float(score)))
scored.sort(key=lambda h: (-h.score, h.chunk.seq))
return scored[:top_k]
def category_stats(self) -> Dict[str, int]:
"""按知识源类别统计段落数(模板配置台展示用)。"""
stats: Dict[str, int] = {}
for chunk in self._chunks:
stats[chunk.category] = stats.get(chunk.category, 0) + 1
return stats
# ------------------------------------------------------------------
@property
def doc_count(self) -> int:
return len(self._docs)
@property
def chunk_count(self) -> int:
return len(self._chunks)
@property
def template(self) -> str:
return self.naming.template
def _slug(title: str) -> str:
"""文档标题 → 对象键安全 slug(小写 + 非安全字符替换为 `_`)。"""
from .templating import sanitize
return sanitize(title)
+273
View File
@@ -0,0 +1,273 @@
# -*- coding: utf-8 -*-
"""RAG 知识库模板化 —— 换行业只改模板配置,内核零改动(PRD 5.4 / Issue #46)。
领域 RAG 知识库按**模板配置**(对齐 data-bus 的模板化思想,见 PRD 5.4
「知识库范围(由模板RAG库提供)」与 7.3「领域RAG知识库(一期)」):
- 知识源分类固定为三类,模板配置声明每类收录的文档:
- `process` 工艺规范(沸腾氯化工艺规范等);
- `sop` 标准作业程序 / 操作手册 / 异常处置 SOP;
- `standard` 国标 / 行业标准 / 文献(氯、钛相关国标等)。
- 全部外部资源命名由模板名推导(向量库 collection / 索引 / namespace),
命名清洗规则对齐 data-bus(小写 + 非安全字符替换为 `_`),但**保留中文**
(文档对象键 / doc_id 面向中文标题,MinIO 对象键支持 UTF-8)。
配置即「模板 RAG 库」资产(`config/kb.template.yaml`),运行时绑定模板
版本,保证检索结果可复现(对齐 PRD 5.4 Prompt/知识库版本化)。
"""
from __future__ import annotations
import re
from enum import Enum
from typing import Dict, List, Tuple
# 保留字符集:小写字母、数字、`-` `_` `.` 与中文(对齐 data-bus 清洗规则,
# 但保留中文以便文档对象键 / doc_id 直接使用中文标题)
_KEEP = re.compile(r"[^a-z0-9_.\u4e00-\u9fff-]+")
def sanitize(name: str) -> str:
"""通用命名清洗:小写 + 非安全字符(空格/标点等)替换为 `_`。
RAG 命名保留中文(文档标题即对象键),仅清洗 ASCII 不安全字符;
与 data-bus 的 sanitize 不同点即在此。
"""
s = (name or "").strip().lower()
s = _KEEP.sub("_", s)
return s.strip("._") or "kb"
class KnowledgeSourceKind(str, Enum):
"""知识源分类(模板 RAG 库的三类知识范围,PRD 7.3)。"""
PROCESS = "process" # 工艺规范
SOP = "sop" # 标准作业程序 / 操作手册 / 异常处置 SOP
STANDARD = "standard" # 国标 / 行业标准 / 文献
@property
def label(self) -> str:
return {
KnowledgeSourceKind.PROCESS: "工艺规范",
KnowledgeSourceKind.SOP: "SOP/操作手册",
KnowledgeSourceKind.STANDARD: "国标/标准",
}[self]
# 知识源分类注册表:模板配置中类别名 → 枚举(未知类别名直接拒绝,避免拼写漂移)
SOURCE_KINDS: Dict[str, KnowledgeSourceKind] = {k.value: k for k in KnowledgeSourceKind}
class KbTemplateNaming:
"""按模板推导 RAG 知识库全部外部资源命名(向量库 / 索引 / namespace)。
命名对齐 data-bus(`tpl_{tpl}` schema 风格):collection 名用下划线形态,
namespace 用连字符形态(对象存储桶风格)。
"""
def __init__(self, template: str, index_suffix: str = "kb"):
self.template = sanitize(template)
# SQL/向量库标识符用下划线形态(避免 `-` 需引号转义)
self.tpl_sql = self.template.replace("-", "_")
self.index_suffix = sanitize(index_suffix)
def collection(self) -> str:
"""向量库 collection:`{tpl}_kb`(对齐 `tpl_{tpl}` schema 命名)。"""
return f"{self.tpl_sql}_{self.index_suffix}"
def index_name(self) -> str:
"""检索索引名:`{tpl}_kb_idx`。"""
return f"{self.collection()}_idx"
def namespace(self) -> str:
"""对象存储 namespace:`tpl-{template}-kb`(S3 桶风格,允许 `-`)。"""
return f"tpl-{self.template}-{self.index_suffix}"
def doc_object_key(self, doc_id: str, ext: str = "md") -> str:
"""源文档对象键:`kb/{doc_id}.{ext}`(保留中文标题)。"""
return f"kb/{sanitize(doc_id)}.{sanitize(ext) or 'md'}"
class KbSourceConfig:
"""单个知识源的配置(模板 RAG 库声明某类知识源收录哪些文档)。"""
def __init__(self, kind: KnowledgeSourceKind, documents: List[str]):
self.kind = kind
self.documents = list(documents)
def __repr__(self) -> str: # pragma: no cover - 调试辅助
return f"KbSourceConfig({self.kind.value}, docs={len(self.documents)})"
class KbTemplateConfig:
"""模板 RAG 库配置:模板名 + 三类知识源文档清单。"""
def __init__(self, template: str, sources: List[KbSourceConfig], version: str = "1.0.0"):
self.template = template
self.sources = sources
self.version = version
def documents_for(self, kind: KnowledgeSourceKind) -> List[str]:
"""返回某类知识源收录的文档标题列表。"""
for src in self.sources:
if src.kind is kind:
return list(src.documents)
return []
def all_documents(self) -> List[str]:
"""全部知识源文档标题(去重保序)。"""
seen, out = set(), []
for src in self.sources:
for d in src.documents:
if d not in seen:
seen.add(d)
out.append(d)
return out
# ---------------------------------------------------------------------------
# 轻量 YAML 子集解析(零第三方依赖,递归下降):map / list / scalar / 注释。
# 足以解析 `config/kb.template.yaml` 这类模板资产;运行时零依赖可复现。
# ---------------------------------------------------------------------------
def _parse_scalar(text: str) -> str:
"""去掉标量两侧引号与行内注释(`key: value # comment`)。"""
t = text.split(" #", 1)[0].strip()
if len(t) >= 2 and t[0] == t[-1] and t[0] in ("'", '"'):
return t[1:-1]
return t
def _strip_comments(lines: List[str]) -> List[Tuple[str, int]]:
"""剔除空行与整行注释,保留行号(1 起)用于报错定位。"""
out = []
for i, ln in enumerate(lines):
s = ln.strip()
if not s or s.startswith("#"):
continue
out.append((ln, i + 1))
return out
def _parse_node(lines: List[Tuple[str, int]], i: int, indent: int):
"""递归解析从 lines[i] 开始、缩进为 `indent` 的一个节点。
返回 `(value, next_i)`:value 为 dict / list / str,next_i 为下一个
未消费行的下标。
"""
text, no = lines[i]
lead = len(text) - len(text.lstrip(" "))
# ---- list 节点:`- item` 或 `- key: val`(map 项) ----
if text.lstrip(" ").startswith("- "):
items: List[object] = []
while i < len(lines):
t, no2 = lines[i]
stripped = t.lstrip(" ")
if not stripped.startswith("- "):
break
lead_j = len(t) - len(t.lstrip(" "))
if lead_j != indent:
break # 缩进不同的 `-` 项属于外层块,交还上层处理
item_text = stripped[2:].strip()
if not item_text:
raise ValueError(f"kb.yaml 第 {no2} 行:list 项为空")
if ":" in item_text:
# map 项:把 `- key: val` 规范化为缩进 `map_indent` 的 map 首行
map_indent = len(t) - len(t.lstrip(" ")) + 2
lines[i] = (" " * map_indent + item_text, no2)
v, i = _parse_node(lines, i, map_indent)
items.append(v)
else:
items.append(_parse_scalar(item_text))
i += 1
return items, i
# ---- map 节点:`key: value` / `key:`(嵌套值) ----
result: Dict[str, object] = {}
while i < len(lines):
t, no = lines[i]
lead_j = len(t) - len(t.lstrip(" "))
if lead_j < indent or t.lstrip(" ").startswith("- "):
break
if lead_j > indent:
raise ValueError(f"kb.yaml 第 {no} 行缩进异常(期望 {indent},实际 {lead_j})")
if ":" not in t:
raise ValueError(f"kb.yaml 第 {no} 行不是合法键值对:{t!r}")
key, _, rest = t.partition(":")
key = key.strip()
rest = rest.strip()
if rest:
result[key] = _parse_scalar(rest)
i += 1
continue
# `key:` —— 值在后续行嵌套(list 或 map),缩进必须更深
if i + 1 >= len(lines):
raise ValueError(f"kb.yaml 第 {no} 行 {key!r} 缺少值")
sub_indent = len(lines[i + 1][0]) - len(lines[i + 1][0].lstrip(" "))
if sub_indent <= indent:
raise ValueError(f"kb.yaml 第 {no} 行 {key!r} 缺少值(无嵌套内容)")
v, i = _parse_node(lines, i + 1, sub_indent)
result[key] = v
return result, i
def _load_yaml_text(text: str) -> Dict[str, object]:
"""解析 YAML 子集 → 嵌套 dict/list。顶层必须为 map。"""
lines = _strip_comments(text.splitlines())
if not lines:
return {}
top_indent = len(lines[0][0]) - len(lines[0][0].lstrip(" "))
value, next_i = _parse_node(lines, 0, top_indent)
if not isinstance(value, dict):
raise ValueError("kb.yaml 顶层必须是 map")
if next_i < len(lines):
raise ValueError(
f"kb.yaml 第 {lines[next_i][1]} 行:顶层存在多个节点(缩进不一致)"
)
return value
def load_kb_config(path: str) -> KbTemplateConfig:
"""从模板 RAG 库 YAML 资产加载配置(零第三方依赖)。
期望结构(详见 `config/kb.template.yaml`):
```yaml
template: ti-cl4
version: 1.0.0
sources:
- kind: process
documents: [...]
- kind: sop
documents: [...]
- kind: standard
documents: [...]
```
"""
with open(path, "r", encoding="utf-8") as fh:
data = _load_yaml_text(fh.read())
template = str(data.get("template", "")).strip()
if not template:
raise ValueError("kb.yaml 缺少 template 字段")
version = str(data.get("version", "1.0.0")).strip() or "1.0.0"
sources: List[KbSourceConfig] = []
raw_sources = data.get("sources") or []
if not isinstance(raw_sources, list):
raise ValueError("kb.yaml sources 必须是 list")
for item in raw_sources:
if not isinstance(item, dict):
raise ValueError(f"kb.yaml sources 项必须是 map,实际 {item!r}")
kind_name = str(item.get("kind", "")).strip()
if kind_name not in SOURCE_KINDS:
raise ValueError(
f"kb.yaml 未知知识源类别 {kind_name!r}(应为 {sorted(SOURCE_KINDS)})"
)
docs = item.get("documents") or []
if not isinstance(docs, list):
raise ValueError(f"kb.yaml {kind_name} 的 documents 必须是 list")
sources.append(
KbSourceConfig(SOURCE_KINDS[kind_name], [str(d) for d in docs])
)
return KbTemplateConfig(template=template, sources=sources, version=version)
+2
View File
@@ -0,0 +1,2 @@
# -*- coding: utf-8 -*-
"""rag-kb 单元测试包。"""
+16
View File
@@ -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
+91
View File
@@ -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()
+148
View File
@@ -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()
+126
View File
@@ -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()