yunmei
1fb1d278d5
- 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 用例全绿（模板化/分块溯源/检索排序/类别过滤/配置校验）
207 lines
7.4 KiB
Python
207 lines
7.4 KiB
Python
# -*- 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)
|