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,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)
|
||||
Reference in New Issue
Block a user