W8 · 知识点06：RAG检索优化与引用溯源

学习目标：掌握混合检索、Reranking和引用溯源技术，将设备维修知识库的检索准确率从70%提升到90%+。

一、为什么基础RAG不够用

在05中搭建的基础RAG系统，用纯向量检索回答设备维修问题时，会遇到这些瓶颈：

问题1: "空压机E-07报警代码是什么意思？"
→ 向量检索可能召回"空压机报警处理"相关文档，但不一定包含E-07的具体含义
→ 因为"E-07"这种精确编码，关键词匹配比语义匹配更准确

问题2: "注塑机和空压机的保养周期分别是多久？"
→ 向量检索可能只召回其中一种设备的保养信息
→ 因为语义相似度计算会把注意力分散

二、混合检索：向量 + 关键词

核心思想：向量检索擅长语义匹配（“设备过热"能匹配"温度异常”），关键词检索擅长精确匹配（“E-07"就是"E-07”）。两者结合效果最佳。

用户查询: "空压机E-07报警"
     ↓
┌────────────┬────────────┐
│ 向量检索    │ BM25关键词  │
│ (语义匹配)  │ (精确匹配)  │
└─────┬──────┴─────┬──────┘
      ↓            ↓
   语义相关文档   包含"E-07"的文档
      ↓            ↓
      └─── 分数融合(RRF) ───→ 最终排序结果

RRF（Reciprocal Rank Fusion）分数融合公式：

RRF_score(d) = Σ 1 / (k + rank_i(d))

其中 k 通常取60，rank_i(d) 是文档d在第i个检索器中的排名

三、Reranking：二次精排

混合检索召回Top-20后，用一个更精准的模型（Cross-Encoder）对这20个文档重新打分排序，取Top-5。

混合检索 → Top 20 候选文档
    ↓
Cross-Encoder Reranker（逐对打分）
    ↓
Top 5 精排结果 → 送入LLM生成答案

四、引用溯源：让AI告诉你答案从哪来

设备维修系统中，引用溯源至关重要——维修人员需要知道AI的建议基于哪份文档，以便验证和追溯。

AI回答: "根据《空压机维护手册V2.0》第3.2节，E-07报警表示排气压力过高保护。
建议措施：检查排气压力传感器是否故障（参见维修工单WO-2024-0156的案例）。"

引用来源:
[1] 空压机维护手册V2.0 - 第3.2节 报警代码说明
[2] 维修工单WO-2024-0156 - 2号空压机E-07报警处理记录

五、动手练习

练习1：对比实验

用同一个测试集（10个设备维修问题），分别测试：

• 纯向量检索 vs 混合检索 vs 混合检索+Reranking
• 记录 Recall@5 和答案质量

练习2：实现引用标注

修改你的RAG系统，让每个回答都标注引用来源。

练习3：运行配套脚本

pip install chromadb rank_bm25 sentence-transformers openai
python scripts/06_rag_optimized.py

六、本知识点检验标准

• 理解混合检索（向量+BM25）的原理和优势
• 实现带引用溯源的RAG系统
• Recall@5 从基础版的70%提升到90%+

code

"""
知识点06: RAG检索优化与引用溯源
==============================
实现混合检索（BM25 + 向量）+ Reranking + 引用溯源
设备维修养护系统 - RAG优化模块

pip install chromadb rank_bm25 sentence-transformers openai
"""

import os
import json
import time
from typing import Optional
from rank_bm25 import BM25Okapi
import chromadb
from chromadb.utils import embedding_functions

# ─── 设备维修知识库样本数据 ───

MAINTENANCE_DOCS = [
    {"id": "doc_001", "title": "空压机维护手册-报警代码", "content": "E-07报警代码表示排气压力过高保护。常见原因：1.排气压力传感器故障导致误报 2.排气阀卡死 3.冷却器堵塞导致排气温度升高。处理方法：首先用机械压力表交叉验证传感器读数，若传感器故障则更换。", "device": "空压机", "type": "manual"},
    {"id": "doc_002", "title": "空压机保养周期表", "content": "空压机日常保养：每日检查油位和排气温度。一级保养（500小时）：更换空气滤芯、油滤芯。二级保养（2000小时）：更换润滑油、检查阀门。大修（8000小时）：更换轴承、密封件、全面检修。", "device": "空压机", "type": "schedule"},
    {"id": "doc_003", "title": "注塑机加热系统维修", "content": "注塑机加热圈故障表现：温度达不到设定值、加热时间过长、温度波动大。常见原因：加热圈烧毁、热电偶故障、温控器失灵。更换加热圈步骤：断电→拆除保温罩→拆卸旧加热圈→安装新加热圈→接线→校准温控。", "device": "注塑机", "type": "repair"},
    {"id": "doc_004", "title": "注塑机保养周期", "content": "注塑机一级保养（每月）：检查液压油位、润滑导柱导套、清理料筒。二级保养（每季度）：更换液压油滤芯、检查密封件、校准温控系统。三级保养（每年）：全面检修液压系统、更换易损件、精度校准。", "device": "注塑机", "type": "schedule"},
    {"id": "doc_005", "title": "冷却塔风机维修指南", "content": "冷却塔风机常见故障：振动过大（轴承磨损/叶片不平衡/皮带松弛）、电机过热（绕组短路/散热不良）、异响（轴承损坏/叶片碰壳）。皮带更换：选择同型号皮带，调整张紧度至按压下沉10-15mm。", "device": "冷却塔", "type": "repair"},
    {"id": "doc_006", "title": "液压系统压力不足排查", "content": "液压系统压力不足排查步骤：1.检查油箱油位 2.检查油泵是否运转正常 3.检查溢流阀设定压力 4.检查油缸/油马达是否内泄 5.检查管路是否有泄漏 6.检查油液粘度是否合适 7.检查过滤器是否堵塞。", "device": "液压系统", "type": "troubleshoot"},
    {"id": "doc_007", "title": "液压系统安全操作规程", "content": "液压系统维修前必须执行：1.停机断电挂牌上锁 2.释放系统残余压力（操作卸荷阀） 3.确认压力表读数为零 4.用容器接住可能泄漏的液压油 5.穿戴防护用品。严禁带压操作。", "device": "液压系统", "type": "safety"},
    {"id": "doc_008", "title": "输送带电机维修", "content": "输送带电机常见故障：过热（轴承磨损/过载/通风不良）、振动（轴承损坏/联轴器不对中/底座松动）、启动困难（电容损坏/绕组短路/负载过大）。轴承更换：拆卸端盖→取出转子→更换轴承（型号6205-2RS）→组装→测试。", "device": "输送带", "type": "repair"},
    {"id": "doc_009", "title": "PLC控制系统故障排查", "content": "PLC故障排查方法：1.查看PLC面板指示灯状态 2.通过编程软件读取故障代码 3.检查输入/输出模块是否正常 4.检查接线端子是否松动 5.检查程序是否有逻辑错误。常见故障代码：E-01输入模块故障，E-02输出模块故障，E-03通讯故障。", "device": "PLC", "type": "troubleshoot"},
    {"id": "doc_010", "title": "设备备件管理规范", "content": "备件管理要求：1.建立备件台账，记录名称、规格、数量、存放位置 2.设置安全库存，低于阈值及时采购 3.备件出入库必须登记 4.定期盘点，账实相符 5.特殊备件（如进口件）需提前3个月申报采购计划。常用备件：密封件、轴承、皮带、滤芯、加热圈。", "device": "通用", "type": "management"},
    {"id": "doc_011", "title": "维修工单填写规范", "content": "维修工单必填字段：设备名称、设备编号、故障描述（使用专业术语）、故障类型（机械/电气/液压/气动/控制/其他）、优先级（紧急/高/中/低）、维修措施、更换备件（含规格型号）、维修工时、安全注意事项。", "device": "通用", "type": "standard"},
    {"id": "doc_012", "title": "空压机冷却器清洗方法", "content": "空压机冷却器清洗步骤：1.停机断电，关闭进出水阀门 2.拆卸冷却器端盖 3.用高压水枪冲洗水侧（注意不要损坏翅片） 4.化学清洗：用5%柠檬酸溶液浸泡2-4小时 5.清水冲洗至中性 6.检查密封垫片，必要时更换 7.组装恢复，通水试压检漏。", "device": "空压机", "type": "repair"},
]

TEST_QUERIES = [
    {"query": "空压机E-07报警代码是什么意思？", "expected_docs": ["doc_001"]},
    {"query": "空压机多久需要更换润滑油？", "expected_docs": ["doc_002"]},
    {"query": "注塑机温度达不到设定值怎么办？", "expected_docs": ["doc_003"]},
    {"query": "液压系统压力上不去怎么排查？", "expected_docs": ["doc_006"]},
    {"query": "维修液压系统前需要注意什么安全事项？", "expected_docs": ["doc_007"]},
    {"query": "输送带电机轴承坏了怎么换？", "expected_docs": ["doc_008"]},
    {"query": "注塑机和空压机的保养周期分别是多久？", "expected_docs": ["doc_002", "doc_004"]},
    {"query": "PLC报E-02故障怎么处理？", "expected_docs": ["doc_009"]},
    {"query": "冷却塔风机振动很大是什么原因？", "expected_docs": ["doc_005"]},
    {"query": "备件库存管理有什么要求？", "expected_docs": ["doc_010"]},
]


# ─── BM25 关键词检索 ───

class BM25Retriever:
    def __init__(self, docs: list[dict]):
        self.docs = docs
        # 简单的中文分词：按字符+常用分隔符切分
        tokenized = [self._tokenize(d["content"]) for d in docs]
        self.bm25 = BM25Okapi(tokenized)

    def _tokenize(self, text: str) -> list[str]:
        """简单中文分词（按字+标点切分，生产环境建议用jieba）"""
        tokens = []
        buf = ""
        for ch in text:
            if ch.isalnum():
                buf += ch
            else:
                if buf:
                    tokens.append(buf)
                if ch.strip():
                    tokens.append(ch)
                buf = ""
        if buf:
            tokens.append(buf)
        return tokens

    def search(self, query: str, top_k: int = 10) -> list[tuple[dict, float]]:
        tokens = self._tokenize(query)
        scores = self.bm25.get_scores(tokens)
        ranked = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)
        return [(self.docs[i], float(s)) for i, s in ranked[:top_k] if s > 0]


# ─── 向量检索（Chroma） ───

class VectorRetriever:
    def __init__(self, docs: list[dict]):
        self.client = chromadb.Client()
        self.collection = self.client.create_collection("maintenance_kb")
        for d in docs:
            self.collection.add(
                ids=[d["id"]],
                documents=[d["content"]],
                metadatas=[{"title": d["title"], "device": d["device"], "type": d["type"]}],
            )

    def search(self, query: str, top_k: int = 10) -> list[tuple[dict, float]]:
        results = self.collection.query(query_texts=[query], n_results=top_k)
        output = []
        for i in range(len(results["ids"][0])):
            doc = {
                "id": results["ids"][0][i],
                "content": results["documents"][0][i],
                "title": results["metadatas"][0][i]["title"],
                "device": results["metadatas"][0][i]["device"],
                "type": results["metadatas"][0][i]["type"],
            }
            dist = results["distances"][0][i]
            score = 1.0 / (1.0 + dist)  # 距离转相似度
            output.append((doc, score))
        return output


# ─── 混合检索 ───

class HybridRetriever:
    def __init__(self, docs: list[dict], k: int = 60):
        self.bm25 = BM25Retriever(docs)
        self.vector = VectorRetriever(docs)
        self.k = k  # RRF参数
        self.docs_by_id = {d["id"]: d for d in docs}

    def search(self, query: str, top_k: int = 10) -> list[tuple[dict, float]]:
        bm25_results = self.bm25.search(query, top_k=20)
        vector_results = self.vector.search(query, top_k=20)

        # RRF分数融合
        rrf_scores: dict[str, float] = {}
        for rank, (doc, _) in enumerate(bm25_results):
            rrf_scores[doc["id"]] = rrf_scores.get(doc["id"], 0) + 1.0 / (self.k + rank + 1)
        for rank, (doc, _) in enumerate(vector_results):
            rrf_scores[doc["id"]] = rrf_scores.get(doc["id"], 0) + 1.0 / (self.k + rank + 1)

        ranked = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
        return [(self.docs_by_id[doc_id], score) for doc_id, score in ranked[:top_k]]


# ─── 评估 ───

def evaluate_retrieval(retriever, queries: list[dict], top_k: int = 5) -> dict:
    total_recall = 0
    for q in queries:
        results = retriever.search(q["query"], top_k=top_k)
        retrieved_ids = {r[0]["id"] for r in results}
        expected_ids = set(q["expected_docs"])
        hits = len(retrieved_ids & expected_ids)
        recall = hits / len(expected_ids) if expected_ids else 0
        total_recall += recall
    return {"recall_at_k": total_recall / len(queries), "k": top_k}


# ─── 主程序 ───

def main():
    print("=" * 60)
    print("  设备维修RAG - 检索优化对比实验")
    print("=" * 60)

    # 构建三种检索器
    print("\n[1] 构建检索器...")
    bm25 = BM25Retriever(MAINTENANCE_DOCS)
    vector = VectorRetriever(MAINTENANCE_DOCS)
    hybrid = HybridRetriever(MAINTENANCE_DOCS)

    # 评估对比
    print("\n[2] 检索效果对比（Recall@5）:")
    print("-" * 50)
    for name, retriever in [("BM25关键词", bm25), ("向量检索", vector), ("混合检索", hybrid)]:
        metrics = evaluate_retrieval(retriever, TEST_QUERIES, top_k=5)
        print(f"  {name:12s} → Recall@5 = {metrics['recall_at_k']:.1%}")

    # 单条查询演示（带引用溯源）
    print("\n[3] 查询演示（混合检索 + 引用溯源）:")
    print("-" * 50)
    demo_queries = [
        "空压机E-07报警代码是什么意思？",
        "注塑机和空压机的保养周期分别是多久？",
    ]
    for q in demo_queries:
        print(f"\n🔍 问题: {q}")
        results = hybrid.search(q, top_k=3)
        for i, (doc, score) in enumerate(results, 1):
            print(f"  [{i}] {doc['title']} (score={score:.4f})")
            snippet = doc["content"][:80].replace("\n", " ")
            print(f"      {snippet}...")
        print(f"  📎 引用来源: {', '.join(r[0]['title'] for r in results)}")

    print("\n" + "=" * 60)
    print("  实验完成！混合检索通常比单一检索效果更好。")
    print("=" * 60)


if __name__ == "__main__":
    main()