copy_rag_chain.py 63 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608
  1. #!/usr/bin/env python3
  2. """
  3. RAG 全链路优化系统 — 员工手册智能问答
  4. ============================================
  5. 8 项核心优化:
  6. 1. 使用 XX科技有限公司员工手册.md 作为测试文档
  7. 2. 检索器封装为 Agent 工具,自动判断是否需要检索
  8. 3. 使用本地 BGE-Large-Zh-v1.5 embedding 模型 (1024d)
  9. 4. 语义文本切块 (SemanticChunker, percentile=60)
  10. 5. 多路召回: BM25 / 向量 / 混合(Ensemble) / 多查询(MultiQuery)
  11. 6. 查询增强: 查询重写 + HyDE
  12. 7. BGE-Reranker 精排
  13. 8. RAG 效果评估: 测试集生成 → Hit@K/MRR/Recall → 幻觉率 → Bad Case 分析
  14. 运行方式: python copy_rag_chain.py
  15. """
  16. # ============================================================================
  17. # Section 1: 导入 + 路径常量 + 配置常量
  18. # ============================================================================
  19. from __future__ import annotations
  20. import json
  21. import os
  22. import re
  23. import sys
  24. import shutil
  25. import warnings
  26. from pathlib import Path
  27. from typing import Optional, Union
  28. from collections import defaultdict
  29. # --- 抑制无关日志 ---
  30. warnings.filterwarnings("ignore", category=DeprecationWarning)
  31. os.environ.setdefault("GRPC_VERBOSITY", "NONE")
  32. os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
  33. # --- 加载 .env ---
  34. from dotenv import load_dotenv
  35. load_dotenv(Path(__file__).resolve().parent.parent / ".env")
  36. # --- 路径常量 ---
  37. _PROJECT_ROOT = Path(__file__).resolve().parent.parent
  38. _CURRENT_DIR = Path(__file__).resolve().parent
  39. _MODELS_DIR = _PROJECT_ROOT / "models"
  40. _BGE_PATH = _MODELS_DIR / "BAAI" / "bge-large-zh-v1.5"
  41. _RERANKER_PATH = _MODELS_DIR / "BAAI" / "bge-reranker-base"
  42. _DATA_PATH = _PROJECT_ROOT / "XX科技有限公司员工手册.md"
  43. _CHROMA_DIR = _CURRENT_DIR / "chroma_db_bge"
  44. # --- 配置常量 ---
  45. # 各参数取值理由:经过员工手册(~7600 字、62 个条款)实测调优
  46. CHUNK_PERCENTILE = 45 # 语义切块阈值:45 在粒度和语义完整间取平衡(值越小切块越细)
  47. RETRIEVAL_K = 10 # 召回阶段返回数量:取 10 是精度/速度折中(太小漏召回,太大稀释 Reranker)
  48. DEFAULT_K = 4 # 最终给 LLM 的文档数:4 篇足以覆盖一个问题,再多会撑爆上下文窗口
  49. ENSEMBLE_WEIGHTS = [0.4, 0.6] # 混合检索权重 [BM25, 向量]:政策文档关键词重要,BM25 提到 0.4
  50. MULTI_QUERY_VARIANTS = 3 # 多查询变体数:3 个变体已足够覆盖同义词,再多边际效益递减
  51. N_EVAL_QA = 25 # 评估测试集大小:覆盖全部 16 章的最少问题数
  52. RRF_K = 60 # RRF 平滑参数:业界标准值,避免排名靠前的文档过度霸占分数
  53. # --- 确保根目录在 sys.path 中 ---
  54. sys.path.insert(0, str(_PROJECT_ROOT))
  55. # ============================================================================
  56. # Section 2: 模型初始化
  57. # ============================================================================
  58. # --- 2a. Embedding 模型 ---
  59. class SentenceTransformerEmbeddings:
  60. """将 SentenceTransformer 包装为 LangChain Embeddings 接口。
  61. 支持 BGE 系列模型的 query_prefix 机制:
  62. - embed_query: 查询前加前缀「为这个句子生成表示以用于检索相关文章:」
  63. - embed_documents: 文档不加前缀(BGE 官方推荐用法)
  64. """
  65. def __init__(self, model_path: str, device: str = "cpu",
  66. query_prefix: str = ""):
  67. from sentence_transformers import SentenceTransformer
  68. self._model = SentenceTransformer(model_path, device=device)
  69. self._query_prefix = query_prefix
  70. def embed_documents(self, texts: list[str]) -> list[list[float]]:
  71. # 文档侧不加前缀
  72. embeddings = self._model.encode(
  73. texts,
  74. normalize_embeddings=True,
  75. show_progress_bar=False,
  76. batch_size=32,
  77. )
  78. return embeddings.tolist()
  79. def embed_query(self, text: str) -> list[float]:
  80. # 查询侧加前缀(BGE 模型关键!不加前缀效果掉 10%+)
  81. if self._query_prefix:
  82. text = self._query_prefix + text
  83. embedding = self._model.encode(
  84. [text],
  85. normalize_embeddings=True,
  86. show_progress_bar=False,
  87. )
  88. return embedding[0].tolist()
  89. _embedding_model: Optional[SentenceTransformerEmbeddings] = None
  90. # BGE 系列模型的 query_prefix(官方推荐,不加的话检索精度明显下降)
  91. _BGE_QUERY_PREFIX = "为这个句子生成表示以用于检索相关文章:"
  92. def init_embedding_model() -> SentenceTransformerEmbeddings:
  93. """初始化 BGE-Large-Zh-v1.5 embedding 模型(1024 维)。
  94. BGE-Large-Zh-v1.5 是 BAAI 中文优化版,MTEB 中文榜单榜首。
  95. 相比之前的 Corom (768d),维度更高、中文理解更强。
  96. 关键: 查询时必须加 query_prefix,文档不加。
  97. """
  98. global _embedding_model
  99. if _embedding_model is not None:
  100. return _embedding_model
  101. if not _BGE_PATH.exists():
  102. raise FileNotFoundError(f"BGE 模型未找到: {_BGE_PATH}")
  103. print(f"[Embedding] 加载 BGE-Large-Zh-v1.5: {_BGE_PATH}")
  104. _embedding_model = SentenceTransformerEmbeddings(
  105. str(_BGE_PATH), device="cpu",
  106. query_prefix=_BGE_QUERY_PREFIX,
  107. )
  108. print(f"[Embedding] 加载完成,维度: 1024 (含 query_prefix)")
  109. return _embedding_model
  110. # --- 2b. LLM 客户端 ---
  111. _llm = None
  112. def get_llm():
  113. """获取 LLM 客户端(复用根目录 llm_client.py 的 get_llm())。"""
  114. global _llm
  115. if _llm is not None:
  116. return _llm
  117. from llm_client import get_llm as _root_get_llm
  118. _llm = _root_get_llm()
  119. print(f"[LLM] 客户端初始化完成")
  120. return _llm
  121. # --- 2c. Reranker 模型 ---
  122. _reranker = None
  123. def download_reranker() -> Optional[str]:
  124. """从 ModelScope 下载 BGE-Reranker-Base。
  125. 返回模型路径,如果下载失败返回 None。
  126. 已下载则直接返回路径,跳过下载。
  127. """
  128. config_file = _RERANKER_PATH / "config.json"
  129. if config_file.exists():
  130. print(f"[Reranker] 模型已存在: {_RERANKER_PATH}")
  131. return str(_RERANKER_PATH)
  132. print(f"[Reranker] 正在从 ModelScope 下载 BGE-Reranker-Base (~1.1GB)...")
  133. print(f"[Reranker] 目标路径: {_RERANKER_PATH}")
  134. try:
  135. from modelscope import snapshot_download
  136. _RERANKER_PATH.parent.mkdir(parents=True, exist_ok=True)
  137. downloaded = snapshot_download(
  138. "BAAI/bge-reranker-base",
  139. cache_dir=str(_RERANKER_PATH.parent),
  140. )
  141. # snapshot_download 返回缓存路径,可能不是我们指定的路径
  142. # 如果路径不同,做复制
  143. downloaded_path = Path(downloaded)
  144. if downloaded_path != _RERANKER_PATH:
  145. if _RERANKER_PATH.exists():
  146. shutil.rmtree(str(_RERANKER_PATH))
  147. shutil.copytree(str(downloaded_path), str(_RERANKER_PATH))
  148. print(f"[Reranker] 下载完成: {_RERANKER_PATH}")
  149. return str(_RERANKER_PATH)
  150. except Exception as e:
  151. print(f"[Reranker] 下载失败: {e}")
  152. print(f"[Reranker] 将跳过精排步骤(不影响其他功能)")
  153. return None
  154. def init_reranker() -> Optional[object]:
  155. """初始化 CrossEncoder reranker。
  156. 返回 CrossEncoder 实例,如果不可用返回 None。
  157. """
  158. global _reranker
  159. if _reranker is not None:
  160. return _reranker if _reranker != "NONE" else None
  161. model_path = download_reranker()
  162. if model_path is None:
  163. _reranker = "NONE"
  164. return None
  165. try:
  166. from sentence_transformers import CrossEncoder
  167. print(f"[Reranker] 加载模型: {model_path}")
  168. _reranker = CrossEncoder(model_path, device="cpu")
  169. print(f"[Reranker] 加载完成")
  170. return _reranker
  171. except Exception as e:
  172. print(f"[Reranker] 加载失败: {e}")
  173. _reranker = "NONE"
  174. return None
  175. # ============================================================================
  176. # Section 3: 文档加载 + 数据清洗 + 章节切块 (优化1, 优化4)
  177. # ============================================================================
  178. def clean_markdown_escapes(text: str) -> str:
  179. """清理 Markdown 转义符。
  180. Markdown 文件中常见 \-、\.、\* 等转义符,会污染关键词匹配和
  181. Reranker 相关度计算(例如 "试用期\\-" 和 "试用期" 在 BM25 中是不同 token)。
  182. 清理示例:
  183. 9:00\\-18:00 → 9:00-18:00
  184. 1\\.5小时 → 1.5小时
  185. V2\\.0 → V2.0
  186. """
  187. # 去掉反斜杠转义符:匹配 \<符号> 并只保留 <符号>
  188. # 字符类中的特殊符号需要逐个列出:- . * ? + [ ] ( ) { } ! ^ $ | #
  189. text = re.sub(r'\\([\-.*?+\[\](){}!^$|#])', r'\1', text)
  190. # 合并多余空行(>2 个连续换行):避免 chunks 中出现大段空白
  191. text = re.sub(r'\n{3,}', '\n\n', text)
  192. return text
  193. def load_and_chunk() -> list:
  194. """加载员工手册 Markdown 文件,按章节结构切块。
  195. 优化1: 使用 XX科技有限公司员工手册.md
  196. 优化4: 按 Markdown 章节(## 第X条)切块,保证每条语义完整
  197. 切块策略:
  198. 1. 读取 .md 文件全部文本
  199. 2. **数据清洗**:清理 Markdown 转义符(9:00\\-18:00 → 9:00-18:00)
  200. 3. 跳过目录,从 # 第一章 开始
  201. 4. 按章节标题(## 第X条)切块,每个条款作为一个独立 chunk
  202. 5. 过长的条款二次切分(保留章节上下文前缀)
  203. 为什么不用 SemanticChunker?
  204. → SemanticChunker 在文档结构清晰时(如员工手册有章节)反而把相关内容
  205. 合并过度,导致向量被稀释。按 Markdown 标题切能保证每个 chunk 语义完整。
  206. Returns:
  207. list[Document]: 文档块列表,每个包含 page_content 和 metadata
  208. """
  209. from langchain_core.documents import Document
  210. from langchain_text_splitters import RecursiveCharacterTextSplitter
  211. print(f"\n{'='*60}")
  212. print(f"[文档加载] 读取: {_DATA_PATH.name}")
  213. # Step 1: 读取文件
  214. raw_text = _DATA_PATH.read_text(encoding="utf-8")
  215. print(f"[文档加载] 原始长度: {len(raw_text)} 字符")
  216. # Step 2: 数据清洗 - 清理 Markdown 转义符
  217. cleaned_text = clean_markdown_escapes(raw_text)
  218. print(f"[文档加载] 清洗后长度: {len(cleaned_text)} 字符")
  219. # Step 3: 跳过目录,从正文开始
  220. body_start = cleaned_text.find("# 第一章")
  221. if body_start == -1:
  222. body_start = cleaned_text.find("# 第")
  223. if body_start > 0:
  224. body_text = cleaned_text[body_start:]
  225. else:
  226. body_text = cleaned_text
  227. # Step 4: 按章节结构切块
  228. # 每个 "## 第X条" 作为一个独立的 chunk,保证语义完整
  229. chunks = []
  230. current_chapter = ""
  231. current_section = ""
  232. current_section_lines = []
  233. def flush_section():
  234. """保存当前 section 为一个 Document。"""
  235. if not current_section or not current_section_lines:
  236. return
  237. content = "\n".join(current_section_lines).strip()
  238. if len(content) < 15:
  239. return
  240. # 跳过纯符号行
  241. if re.match(r"^[\s\-—#*=_.]+$", content):
  242. return
  243. # 附加上下文前缀:让 embedding 知道这段文字属于哪个章节
  244. full_content = f"【{current_chapter} - {current_section}】\n{content}"
  245. chunks.append(Document(
  246. page_content=full_content,
  247. metadata={
  248. "source": _DATA_PATH.name,
  249. "chapter": current_chapter,
  250. "section": current_section,
  251. }
  252. ))
  253. for line in body_text.split("\n"):
  254. # 检测章标题 (# 第X章 ...) — 一级标题,更新当前所在章
  255. chap_match = re.match(r"^#\s+(第[一二三四五六七八九十\d]+章\s*\S.*)", line)
  256. if chap_match:
  257. flush_section() # 先保存上一节(章切换意味着上一节结束)
  258. current_chapter = chap_match.group(1).strip()
  259. current_section = "" # 进入新章时清空 section,等遇到 ## 才开始记录
  260. current_section_lines = []
  261. continue
  262. # 检测条标题 (## 第X条 ...) — 二级标题,每个条款是一个独立 chunk
  263. sec_match = re.match(r"^##\s+(第[一二三四五六七八九十\d]+条\s*\S.*)", line)
  264. if sec_match:
  265. flush_section() # 先保存上一节(条款切换意味着上一条结束)
  266. current_section = sec_match.group(1).strip()
  267. current_section_lines = []
  268. continue
  269. # 跳过分隔线 (---):Markdown 的水平分隔符,对内容无意义
  270. if re.match(r"^---+\s*$", line):
  271. continue
  272. # 收集正文:只有在已进入某个 section 后才开始收集
  273. # (目录、文档头部说明等非条款内容会被跳过)
  274. if current_section:
  275. stripped = line.strip()
  276. if stripped: # 跳过空行
  277. current_section_lines.append(stripped)
  278. flush_section() # 保存最后一节(循环结束后内存中还残留最后一节未写入)
  279. print(f"[文档加载] 按章节切块完成: {len(chunks)} 个条款")
  280. # Step 5: 对过长的 section 二次切分(保留章节上下文)
  281. # 防止单个条款过长(>1200 字符)稀释向量
  282. final_chunks = []
  283. sub_splitter = RecursiveCharacterTextSplitter(
  284. chunk_size=800,
  285. chunk_overlap=100,
  286. separators=["\n\n", "\n", "。", "!", "?", ";", ",", " ", ""]
  287. )
  288. for c in chunks:
  289. if len(c.page_content) > 1200:
  290. sub_chunks = sub_splitter.split_documents([c])
  291. final_chunks.extend(sub_chunks)
  292. else:
  293. final_chunks.append(c)
  294. print(f"[文档加载] 最终块数量: {len(final_chunks)} 个 "
  295. f"(过长条款已二次切分)")
  296. if len(final_chunks) < 10:
  297. print(f"[文档加载] 警告: 块数量过少 ({len(final_chunks)})")
  298. return final_chunks
  299. # ============================================================================
  300. # Section 4: ChromaDB 向量存储 (优化3)
  301. # ============================================================================
  302. def build_vectorstore(chunks: list) -> object:
  303. """构建 ChromaDB 向量存储。
  304. 使用 BGE-Large-Zh-v1.5 embedding 模型 (1024d),collection 配置余弦相似度。
  305. 每次调用强制重建(开发阶段)。
  306. """
  307. from langchain_community.vectorstores import Chroma
  308. print(f"\n[向量存储] 构建 ChromaDB ({len(chunks)} 个文档)")
  309. embedding_model = init_embedding_model()
  310. # 清理旧数据
  311. if _CHROMA_DIR.exists():
  312. print(f"[向量存储] 清理旧数据库: {_CHROMA_DIR}")
  313. shutil.rmtree(str(_CHROMA_DIR))
  314. vectorstore = Chroma.from_documents(
  315. documents=chunks,
  316. embedding=embedding_model,
  317. collection_metadata={"hnsw:space": "cosine"},
  318. persist_directory=str(_CHROMA_DIR),
  319. )
  320. print(f"[向量存储] 构建完成,持久化到: {_CHROMA_DIR}")
  321. return vectorstore
  322. def get_vectorstore(chunks: list) -> object:
  323. """获取向量存储(优先加载已有,否则重建)。"""
  324. from langchain_community.vectorstores import Chroma
  325. sqlite_file = _CHROMA_DIR / "chroma.sqlite3"
  326. if sqlite_file.exists():
  327. embedding_model = init_embedding_model()
  328. try:
  329. print(f"[向量存储] 从 {_CHROMA_DIR} 加载已有数据库")
  330. return Chroma(
  331. persist_directory=str(_CHROMA_DIR),
  332. embedding_function=embedding_model,
  333. collection_metadata={"hnsw:space": "cosine"},
  334. )
  335. except Exception as e:
  336. print(f"[向量存储] 加载失败 ({e}),将重建")
  337. return build_vectorstore(chunks)
  338. # ============================================================================
  339. # Section 5: 多路召回系统 (优化5)
  340. # ============================================================================
  341. _bm25_retriever = None
  342. _vector_retriever = None
  343. _ensemble_retriever = None
  344. _mq_retriever = None
  345. def _jieba_tokenize(text: str) -> list[str]:
  346. """使用 jieba 对中文文本分词。
  347. BM25Retriever 默认用英文空白分词器,遇到中文时会将整段文本当作一个 token,
  348. 导致完全无法匹配关键词(不同查询返回相同结果)。
  349. 用 jieba 分词后,"试用期是几个月" 会被切成 ["试用期", "是", "几个", "月"],
  350. BM25 才能正常按词频打分。
  351. """
  352. import jieba
  353. # lcut 精确模式(vs cut 全模式):返回 list,每个 token 是一个词
  354. tokens = jieba.lcut(text)
  355. # 过滤规则:去空白 + 丢弃单字
  356. # 单字(如"的""是""一")信息量极低,且大量重复,会拉低 BM25 区分度
  357. return [t for t in tokens if t.strip() and len(t.strip()) > 1]
  358. def create_bm25_retriever(chunks: list):
  359. """创建 BM25 关键词检索器。
  360. BM25 基于词频统计,擅长精确关键词匹配,
  361. 弥补向量检索在专有名词、数字、缩写上的短板。
  362. 注意:必须传入中文分词器 _jieba_tokenize,否则默认英文分词器
  363. 会把整段中文当成一个 token,BM25 退化为按原始顺序返回。
  364. """
  365. global _bm25_retriever
  366. if _bm25_retriever is not None:
  367. return _bm25_retriever
  368. from langchain_community.retrievers import BM25Retriever
  369. _bm25_retriever = BM25Retriever.from_documents(
  370. chunks,
  371. preprocess_func=_jieba_tokenize,
  372. )
  373. _bm25_retriever.k = RETRIEVAL_K
  374. print(f"[检索器] BM25 初始化完成 (k={RETRIEVAL_K}, 中文分词=jieba)")
  375. return _bm25_retriever
  376. def create_vector_retriever(vectorstore):
  377. """创建向量语义检索器。
  378. 基于 BGE-Large-Zh-v1.5 embedding 的余弦相似度检索,
  379. 能在语义层面理解同义词、近义词和概念关联。
  380. """
  381. global _vector_retriever
  382. if _vector_retriever is not None:
  383. return _vector_retriever
  384. _vector_retriever = vectorstore.as_retriever(
  385. search_type="similarity",
  386. search_kwargs={"k": RETRIEVAL_K},
  387. )
  388. print(f"[检索器] 向量检索器初始化完成 (k={RETRIEVAL_K})")
  389. return _vector_retriever
  390. def create_ensemble_retriever(chunks, vectorstore):
  391. """创建混合检索器: BM25 + 向量检索融合。
  392. - BM25 权重 0.3: 关键词精确匹配
  393. - 向量权重 0.7: 语义理解(权重更高,因为语义匹配通常更重要)
  394. 融合方式: 分数归一化后按权重加权求和。
  395. """
  396. global _ensemble_retriever
  397. if _ensemble_retriever is not None:
  398. return _ensemble_retriever
  399. from langchain_classic.retrievers import EnsembleRetriever
  400. bm25 = create_bm25_retriever(chunks)
  401. vec = create_vector_retriever(vectorstore)
  402. _ensemble_retriever = EnsembleRetriever(
  403. retrievers=[bm25, vec],
  404. weights=ENSEMBLE_WEIGHTS,
  405. )
  406. print(f"[检索器] 混合检索器初始化完成 "
  407. f"(BM25:{ENSEMBLE_WEIGHTS[0]} + 向量:{ENSEMBLE_WEIGHTS[1]})")
  408. return _ensemble_retriever
  409. def create_mq_retriever(vectorstore):
  410. """创建多查询检索器 (Multi-Query Retriever)。
  411. LLM 自动生成 3 个查询变体 → 分别检索 → 合并去重。
  412. 适合处理用户问题表述不够精确的场景。
  413. """
  414. global _mq_retriever
  415. if _mq_retriever is not None:
  416. return _mq_retriever
  417. from langchain_classic.retrievers import MultiQueryRetriever
  418. vec = create_vector_retriever(vectorstore)
  419. _mq_retriever = MultiQueryRetriever.from_llm(
  420. retriever=vec,
  421. llm=get_llm(),
  422. )
  423. print(f"[检索器] 多查询检索器初始化完成 (变体数={MULTI_QUERY_VARIANTS})")
  424. return _mq_retriever
  425. def retrieve(query: str, strategy: str = "ensemble",
  426. top_k: int = DEFAULT_K) -> list:
  427. """统一的检索调度接口。
  428. Args:
  429. query: 用户问题
  430. strategy: 检索策略,可选 "bm25" | "vector" | "ensemble" | "multi_query"
  431. top_k: 返回的文档数量
  432. Returns:
  433. list[Document]: 检索到的文档列表
  434. """
  435. vs = _get_vectorstore_ref()
  436. chunks = _get_chunks_ref()
  437. if strategy == "bm25":
  438. ret = create_bm25_retriever(chunks)
  439. elif strategy == "vector":
  440. ret = create_vector_retriever(vs)
  441. elif strategy == "multi_query":
  442. ret = create_mq_retriever(vs)
  443. else:
  444. ret = create_ensemble_retriever(chunks, vs)
  445. # 为不同检索器适配 k 参数
  446. # BM25Retriever / VectorStoreRetriever 用 .k 属性控制返回数量
  447. # MultiQueryRetriever / EnsembleRetriever 没有可变 .k,需要 invoke 后切片
  448. if hasattr(ret, "k"):
  449. # 临时修改 k → 调用 → 恢复原值(防止影响后续其他调用)
  450. old_k = ret.k
  451. ret.k = top_k
  452. docs = ret.invoke(query)
  453. ret.k = old_k
  454. else:
  455. docs = ret.invoke(query)
  456. docs = docs[:top_k]
  457. return docs
  458. # 模块级引用(由主流程设置)
  459. _VS = None
  460. _CHUNKS = None
  461. def _get_vectorstore_ref():
  462. if _VS is None:
  463. raise RuntimeError("向量存储未初始化,请先调用 load_and_chunk() 和 get_vectorstore()")
  464. return _VS
  465. def _get_chunks_ref():
  466. if _CHUNKS is None:
  467. raise RuntimeError("文档块未初始化")
  468. return _CHUNKS
  469. # ============================================================================
  470. # Section 6: 查询增强 (优化6)
  471. # ============================================================================
  472. def rewrite_query(original: str) -> str:
  473. """查询重写:用 LLM 将用户问题改写为更精确的检索查询。
  474. 改写策略:
  475. - 提取核心实体和关键词
  476. - 消除歧义,明确指代
  477. - 将口语化表达转为文档化风格
  478. """
  479. llm = get_llm()
  480. prompt = f"""你是一个查询优化助手。给定用户的原始问题,请将其改写成更适合文档检索的形式。
  481. 要求:
  482. 1. 提取核心实体和关键词
  483. 2. 移除歧义,明确指代
  484. 3. 保持原意不变
  485. 4. 使用文档中可能出现的术语(如"年假""加班费""绩效考核"等正式表述)
  486. 5. 直接输出改写后的查询,不要任何解释
  487. 原始问题:{original}
  488. 改写后的查询:"""
  489. try:
  490. raw = llm.invoke(prompt)
  491. rewritten = raw.content if hasattr(raw, "content") else str(raw)
  492. rewritten = rewritten.strip().strip('"').strip("'")
  493. if len(rewritten) < 2:
  494. return original
  495. return rewritten
  496. except Exception as e:
  497. print(f"[查询重写] 失败 ({e}),使用原始查询")
  498. return original
  499. def hyde_retrieve(query: str, vectorstore, top_k: int = RETRIEVAL_K) -> list:
  500. """HyDE (Hypothetical Document Embeddings) 检索。
  501. 核心思路: LLM 先生成一个「假设性答案」→ 将这个答案向量化 →
  502. 用这个向量去检索 → 假设性答案的语义空间比问题本身更接近真实文档。
  503. 流程:
  504. 1. LLM 生成假设性答案(不需要准确,只需"看起来像"文档内容)
  505. 2. Embedding 模型将假设性答案转为向量
  506. 3. 在 ChromaDB 中检索相似文档
  507. """
  508. llm = get_llm()
  509. prompt = f"""请根据以下问题,生成一个假设性的员工手册回答。这个回答不需要完全准确,
  510. 只需要看起来像是一段真实的公司制度文档内容。用来帮助搜索引擎找到相关文档。
  511. 要求:直接输出内容,不要任何前缀或解释。
  512. 问题:{query}
  513. 假设文档:"""
  514. try:
  515. raw = llm.invoke(prompt)
  516. hypothesis = raw.content if hasattr(raw, "content") else str(raw)
  517. hypothesis = hypothesis.strip()
  518. if len(hypothesis) < 10:
  519. return []
  520. docs = vectorstore.similarity_search(hypothesis, k=top_k)
  521. return docs
  522. except Exception as e:
  523. print(f"[HyDE] 失败 ({e})")
  524. return []
  525. def reciprocal_rank_fusion(doc_lists: list[list], k: int = RRF_K) -> list:
  526. """倒数秩融合 (Reciprocal Rank Fusion)。
  527. 将多个检索结果列表按 RRF 算法合并为一个去重排序列表。
  528. 公式: score(d) = Σ 1/(rank_i(d) + k)
  529. Args:
  530. doc_lists: 多个检索结果列表
  531. k: 平滑参数(默认60,标准设置)
  532. Returns:
  533. 融合后按 RRF 分数降序排列的文档列表
  534. """
  535. scores: dict[int, float] = {}
  536. doc_map: dict[int, object] = {}
  537. # Step 1: 遍历每个检索器返回的列表,按排名累计 RRF 分数
  538. for doc_list in doc_lists:
  539. for rank, doc in enumerate(doc_list): # rank 从 0 开始,表示该文档在该列表中的位置
  540. # 用 page_content 的 hash 作为文档唯一标识(同一段文字会得到相同 hash)
  541. doc_id = hash(doc.page_content)
  542. if doc_id not in scores:
  543. scores[doc_id] = 0.0
  544. doc_map[doc_id] = doc
  545. # RRF 核心公式:score(d) = Σ 1/(rank + k)
  546. # rank 从 0 开始,所以要 +1 才是真正的「第几名」
  547. # k 是平滑常数:避免 rank=0 的文档得到 1.0 满分,让排名靠后的也有贡献
  548. # k=60 是业界标准值,源自论文《Reciprocal Rank Fusion outperforms Condorcet》
  549. scores[doc_id] += 1.0 / (rank + k)
  550. # Step 2: 按 RRF 分数降序排列,输出融合后的文档列表
  551. sorted_ids = sorted(scores.keys(), key=lambda x: scores[x], reverse=True)
  552. return [doc_map[did] for did in sorted_ids]
  553. def enhanced_retrieve(query: str, strategy: str = "ensemble",
  554. top_k: int = DEFAULT_K) -> list:
  555. """增强检索: 查询重写 + HyDE + RRF 融合。
  556. 流程:
  557. 1. LLM 重写查询 → 得到更精确的检索查询
  558. 2. HyDE 生成假设答案 → 从另一个语义角度检索
  559. 3. 标准检索(用重写后的查询)
  560. 4. RRF 融合 HyDE 结果 + 标准结果 → 取 top_k
  561. """
  562. vectorstore = _get_vectorstore_ref()
  563. # Step 1: 查询重写
  564. rewritten = rewrite_query(query)
  565. if rewritten != query:
  566. print(f"[增强检索] 查询重写: 「{query[:40]}...」→「{rewritten[:60]}...」")
  567. # Step 2: HyDE 检索
  568. hyde_docs = hyde_retrieve(query, vectorstore, top_k=RETRIEVAL_K)
  569. print(f"[增强检索] HyDE 召回 {len(hyde_docs)} 个文档")
  570. # Step 3: 标准检索(使用重写后的查询)
  571. standard_docs = retrieve(rewritten, strategy=strategy, top_k=RETRIEVAL_K)
  572. print(f"[增强检索] {strategy} 召回 {len(standard_docs)} 个文档")
  573. # Step 4: RRF 融合
  574. if hyde_docs and standard_docs:
  575. fused = reciprocal_rank_fusion([standard_docs, hyde_docs])
  576. print(f"[增强检索] RRF 融合后 {len(fused)} 个文档")
  577. elif standard_docs:
  578. fused = standard_docs
  579. else:
  580. fused = hyde_docs
  581. return fused[:top_k]
  582. # ============================================================================
  583. # Section 7: Reranker 精排 (优化7)
  584. # ============================================================================
  585. def rerank(query: str, candidates: list, top_k: int = DEFAULT_K) -> list:
  586. """使用 BGE-Reranker 对候选文档精排。
  587. Cross-Encoder 将 query-doc 对联合编码,比 embedding 相似度更精确。
  588. 缺点: 速度较慢,适合对 Top-N 候选做精排(N 通常 10-20)。
  589. Args:
  590. query: 用户问题
  591. candidates: 候选文档列表
  592. top_k: 返回的文档数量
  593. Returns:
  594. 按相关度分数降序排列的文档列表
  595. """
  596. reranker = init_reranker()
  597. if reranker is None or len(candidates) == 0:
  598. return candidates[:top_k]
  599. # 优化:候选数 ≤ top_k 时无需精排,直接返回(省下推理时间)
  600. if len(candidates) <= top_k:
  601. return candidates
  602. try:
  603. # 构建 query-doc 配对:CrossEncoder 是双塔模型的对比
  604. # 与 embedding(Bi-Encoder,query 和 doc 分别编码再算相似度)不同,
  605. # CrossEncoder 把 query 和 doc 拼在一起送入模型,深度交互注意力,
  606. # 因此精度更高但速度更慢(适合 Top-N 精排,N 通常 10-20)
  607. pairs = [(query, doc.page_content) for doc in candidates]
  608. scores = reranker.predict(pairs, show_progress_bar=False)
  609. # 按相关度分数降序排列,取 top_k
  610. # scored 是 [(score, doc), ...],sort 后取前 top_k 个的 doc
  611. scored = list(zip(scores, candidates))
  612. scored.sort(key=lambda x: x[0], reverse=True)
  613. reranked = [doc for _, doc in scored[:top_k]]
  614. print(f"[Reranker] 精排完成: {len(candidates)} → {len(reranked)} "
  615. f"(最高分: {scored[0][0]:.4f})")
  616. return reranked
  617. except Exception as e:
  618. print(f"[Reranker] 精排失败 ({e}),使用原始顺序")
  619. return candidates[:top_k]
  620. # ============================================================================
  621. # Section 8: Agent 工具封装 (优化2)
  622. # ============================================================================
  623. def create_retrieval_tool():
  624. """将检索系统封装为 LangChain Tool。
  625. Agent 可根据问题自动判断是否需要调用此工具检索员工手册。
  626. """
  627. from langchain_core.tools import tool
  628. @tool
  629. def search_employee_handbook(query: str) -> str:
  630. """从《XX科技有限公司员工手册》中检索与 query 相关的内容。
  631. 当用户问到以下内容时使用此工具:
  632. - 公司制度、考勤规定、请假流程
  633. - 薪资福利、社保公积金
  634. - 加班管理、出差报销
  635. - 绩效考核、晋升条件
  636. - 离职手续、保密规定
  637. - 员工关怀、培训发展
  638. - 入职流程、试用期规定
  639. 输入:用户的问题或关键词。
  640. 输出:手册中的相关章节内容。
  641. """
  642. # 使用增强检索 + 精排
  643. docs = enhanced_retrieve(query, strategy="ensemble", top_k=RETRIEVAL_K)
  644. docs = rerank(query, docs, top_k=DEFAULT_K)
  645. parts = []
  646. for i, doc in enumerate(docs):
  647. parts.append(f"【参考条款 {i+1}】\n{doc.page_content}")
  648. return "\n\n---\n\n".join(parts)
  649. return search_employee_handbook
  650. def create_rag_agent():
  651. """创建 RAG Agent。
  652. Agent 能自动判断:
  653. - 需要查手册时 → 调用 search_employee_handbook 工具
  654. - 闲聊/通用问题 → 直接回答
  655. """
  656. from langchain.agents import create_agent
  657. tool = create_retrieval_tool()
  658. llm = get_llm()
  659. agent = create_agent(
  660. model=llm,
  661. tools=[tool],
  662. system_prompt=(
  663. "你是XX科技有限公司的员工手册问答助手。\n\n"
  664. "规则:\n"
  665. "1. 当用户问到公司制度、考勤、请假、薪资、福利、加班、离职、保密等"
  666. "规定时,必须先调用 search_employee_handbook 工具检索手册内容。\n"
  667. "2. 基于检索到的内容回答,不要编造。如果手册中没有相关信息,"
  668. "如实告知用户。\n"
  669. "3. 回答要简洁清晰,引用手册中的具体条款编号。\n"
  670. "4. 对于闲聊或与员工手册无关的问题,直接回答即可。"
  671. ),
  672. )
  673. print(f"[Agent] RAG Agent 创建完成")
  674. return agent
  675. # ============================================================================
  676. # Section 9: 完整 RAG 链路
  677. # ============================================================================
  678. def build_context(docs: list) -> str:
  679. """将检索到的文档拼接为 LLM 上下文。"""
  680. parts = []
  681. for i, doc in enumerate(docs):
  682. parts.append(f"[参考文档 {i+1}]\n{doc.page_content}")
  683. return "\n\n---\n\n".join(parts)
  684. def rag_query(query: str, strategy: str = "ensemble",
  685. use_enhance: bool = True, use_rerank: bool = True) -> dict:
  686. """确定性 RAG 查询流水线(不使用 Agent)。
  687. Args:
  688. query: 用户问题
  689. strategy: 检索策略
  690. use_enhance: 是否使用查询增强(重写 + HyDE)
  691. use_rerank: 是否使用 Reranker 精排
  692. Returns:
  693. dict: {"answer": str, "sources": list[Document], "strategy": str}
  694. """
  695. from langchain_core.prompts import ChatPromptTemplate
  696. from langchain_core.output_parsers import StrOutputParser
  697. # Step 1: 检索(增强检索走 enhanced_retrieve,普通检索走 retrieve)
  698. if use_enhance:
  699. docs = enhanced_retrieve(query, strategy=strategy, top_k=RETRIEVAL_K)
  700. else:
  701. docs = retrieve(query, strategy=strategy, top_k=RETRIEVAL_K)
  702. # Step 2: 精排(可选):对召回的 top-K 候选重新打分排序,取 top-k
  703. if use_rerank:
  704. docs = rerank(query, docs, top_k=DEFAULT_K)
  705. # Step 3: 组装 prompt
  706. # ChatPromptTemplate.from_messages: 支持 system + user 多轮对话格式
  707. # {context} 和 {query} 是占位符,invoke 时填充
  708. context = build_context(docs)
  709. prompt = ChatPromptTemplate.from_messages([
  710. ("system", """你是XX科技有限公司的员工手册问答助手。请根据以下参考资料回答问题。
  711. **规则:**
  712. - 只基于提供的参考资料回答,不要编造
  713. - 如果参考资料中没有相关信息,直接说「根据现有员工手册,我找不到这个问题的答案」
  714. - 回答要简洁清晰,引用具体条款时注明出处
  715. - 如果参考资料不足以回答,可以说明需要查阅手册的哪个章节
  716. **参考资料:**
  717. {context}"""),
  718. ("user", "{query}")
  719. ])
  720. # Step 4: 生成
  721. # LCEL 链式语法:prompt | llm | parser
  722. # - prompt.invoke({...}) → 填充占位符生成最终 prompt
  723. # - llm.invoke(prompt) → 调用 LLM 生成回复
  724. # - StrOutputParser().invoke(resp) → 从 AIMessage 中提取纯文本
  725. # 等价于:parser.invoke(llm.invoke(prompt.invoke({...})))
  726. llm = get_llm()
  727. chain = prompt | llm | StrOutputParser()
  728. answer = chain.invoke({"context": context, "query": query})
  729. return {"answer": answer, "sources": docs, "strategy": strategy}
  730. def agent_query(query: str) -> str:
  731. """Agent 模式查询(自动判断是否需要检索)。
  732. Agent 根据问题内容自动决定:
  733. - 调用 search_employee_handbook 工具检索
  734. - 直接回答(不检索)
  735. """
  736. agent = create_rag_agent()
  737. result = agent.invoke({"messages": [("user", query)]})
  738. # 提取最后一条消息
  739. messages = result.get("messages", [])
  740. if messages:
  741. return messages[-1].content
  742. return "Agent 未返回结果"
  743. # ============================================================================
  744. # Section 10: RAG 效果评估体系 (优化8)
  745. # ============================================================================
  746. #
  747. # 评估策略说明:
  748. # 1. 手动标注测试集: 覆盖员工手册全部 16 章的真实问题 → Section 级检索评估
  749. # 2. 检索交叉验证: LLM 逐条判断检索到的文档能否回答该问题
  750. # 3. LLM 辅助生成评估 + Bad Case 根因分析
  751. # ========================================================================
  752. # 10a. 手动标注真实测试集(覆盖全部 16 章,~30% 文档覆盖率)
  753. # ========================================================================
  754. # 每条测试用例包含:
  755. # question: 员工真实会问的问题
  756. # section_keywords: 期望匹配的章节关键词(用于 Section 级检索评估)
  757. # reference_answer: 从手册原文提取的参考答案(用于生成评估)
  758. _MANUAL_TEST_SET = [
  759. # === 第一章 总则 ===
  760. {
  761. "question": "员工手册制定的法律依据是什么?",
  762. "section_keywords": ["总则", "劳动法", "劳动合同法"],
  763. "reference_answer": "根据《中华人民共和国劳动法》《中华人民共和国劳动合同法》及相关法律法规制定。",
  764. },
  765. # === 第二章 入职与试用期 ===
  766. {
  767. "question": "试用期是多长时间?",
  768. "section_keywords": ["试用期", "入职"],
  769. "reference_answer": "新员工试用期为3个月,试用期包含在劳动合同期限内。",
  770. },
  771. {
  772. "question": "入职需要准备哪些材料?",
  773. "section_keywords": ["入职", "入职流程", "报到"],
  774. "reference_answer": "入职需准备:身份证、学历学位证书、离职证明、体检报告、银行卡。",
  775. },
  776. # === 第三章 工作时间与考勤 ===
  777. {
  778. "question": "公司标准工作时间是怎样的?",
  779. "section_keywords": ["工作时间", "考勤", "上下班"],
  780. "reference_answer": "标准工作时间为周一至周五9:00-18:00,午休1小时,每天工作8小时。",
  781. },
  782. # === 第四章 假期管理 ===
  783. {
  784. "question": "年假没休完可以结转到下一年吗?",
  785. "section_keywords": ["年假", "带薪年假", "假期"],
  786. "reference_answer": "当年未休完的年假最多可结转5天至次年3月31日,逾期作废。",
  787. },
  788. {
  789. "question": "请病假需要提供什么材料?",
  790. "section_keywords": ["病假", "请假", "假期管理"],
  791. "reference_answer": "请病假需提供二级甲等以上医院出具的病假证明和诊断证明。",
  792. },
  793. # === 第五章 薪资福利 ===
  794. {
  795. "question": "社保养老保险个人缴纳比例是多少?",
  796. "section_keywords": ["社保", "社会保险", "五险一金", "薪资福利"],
  797. "reference_answer": "养老保险个人缴纳比例为8%,单位缴纳比例为16%。",
  798. },
  799. # === 第六章 绩效考核 ===
  800. {
  801. "question": "绩效考核结果分几个等级?",
  802. "section_keywords": ["绩效考核", "考核等级", "绩效"],
  803. "reference_answer": "绩效考核分为S/A/B/C/D五个等级,其中S级占比不超过10%。",
  804. },
  805. # === 第七章 培训与发展 ===
  806. {
  807. "question": "公司每年要求的最低培训学时是多少?",
  808. "section_keywords": ["培训", "培训与发展", "学习"],
  809. "reference_answer": "每位员工每年需完成不低于40学时的培训,包括技术培训、管理培训和软技能培训。",
  810. },
  811. # === 第八章 晋升制度 ===
  812. {
  813. "question": "员工晋升需要满足哪些基本条件?",
  814. "section_keywords": ["晋升", "晋升制度", "职业发展"],
  815. "reference_answer": "晋升需满足:连续两次绩效考核B+以上,在现岗位工作满1年以上。",
  816. },
  817. # === 第九章 加班管理 ===
  818. {
  819. "question": "周末加班工资按几倍计算?",
  820. "section_keywords": ["加班", "加班费", "加班工资"],
  821. "reference_answer": "休息日(周末)加班,按200%(2倍)工资支付加班费。",
  822. },
  823. # === 第十章 出差与报销 ===
  824. {
  825. "question": "出差一线城市住宿标准多少钱一天?",
  826. "section_keywords": ["出差", "报销", "住宿标准"],
  827. "reference_answer": "一线城市(北上广深)出差住宿标准为500元/天。",
  828. },
  829. # === 第十一章 劳动合同 ===
  830. {
  831. "question": "第一次签劳动合同的期限是几年?",
  832. "section_keywords": ["劳动合同", "合同期限", "签订"],
  833. "reference_answer": "首次签订劳动合同期限为3年,续签第二次为5年,之后可签无固定期限合同。",
  834. },
  835. # === 第十二章 奖惩制度 ===
  836. {
  837. "question": "员工连续旷工几天会被开除?",
  838. "section_keywords": ["旷工", "处罚", "奖惩", "开除", "严重违纪"],
  839. "reference_answer": "连续旷工3天以上或一年内累计旷工5天以上,属于严重违纪,公司可单方解除劳动合同。",
  840. },
  841. # === 第十三章 保密规定 ===
  842. {
  843. "question": "公司商业机密包括哪些内容?",
  844. "section_keywords": ["保密", "商业机密", "保密规定"],
  845. "reference_answer": "商业机密包括技术资料、客户信息、财务数据、经营决策、未公开的人事信息等。",
  846. },
  847. # === 第十四章 离职管理 ===
  848. {
  849. "question": "离职需要提前多少天申请?",
  850. "section_keywords": ["离职", "离职管理", "解除"],
  851. "reference_answer": "正式员工离职需提前30天书面申请,试用期员工提前3天申请。",
  852. },
  853. # === 第十五章 员工关怀 ===
  854. {
  855. "question": "公司有哪些员工关怀和福利项目?",
  856. "section_keywords": ["员工关怀", "员工福利", "关怀"],
  857. "reference_answer": "公司提供补充商业保险、年度体检、节日礼品、团队建设、生日福利、婚育贺礼、困难补助等。",
  858. },
  859. # === 第十六章 附则 ===
  860. {
  861. "question": "员工手册的解释权归哪个部门?",
  862. "section_keywords": ["附则", "解释权", "人力资源部"],
  863. "reference_answer": "本手册的最终解释权归公司人力资源部所有。",
  864. },
  865. {
  866. "question": "员工手册从什么时候开始生效?",
  867. "section_keywords": ["附则", "生效", "施行"],
  868. "reference_answer": "本手册自2025年1月1日起正式生效施行。",
  869. },
  870. ]
  871. def get_manual_test_set() -> list[dict]:
  872. """获取手动标注的真实测试集。
  873. 这些测试用例是基于员工手册原文内容手动编写的,用于客观评估 RAG 系统。
  874. 与 LLM 自动生成的测试集不同,这里的问题和答案都是人工确定的,
  875. 不存在"同一 LLM 既当选手又当裁判"的问题。
  876. """
  877. return _MANUAL_TEST_SET
  878. # ========================================================================
  879. # 10b. Section 级检索评估(核心指标)
  880. # ========================================================================
  881. def evaluate_retrieval_section_level(test_set: list[dict]) -> dict:
  882. """Section 级检索评估 — 客观评估检索质量。
  883. 评估方法:
  884. 对每个问题,使用各检索策略检索 Top-K 文档,检查检索到的文档
  885. 是否来自员工手册的正确章节(通过章节关键词匹配)。
  886. 为什么不用"同一个 chunk"检测?
  887. → 那是循环验证:LLM 从 chunk A 生成问题 → 检查是否检索到 chunk A
  888. → 毫无意义,因为 embedding 向量天然相近。
  889. 为什么用 Section 级匹配?
  890. → 比如问"试用期",理应检索到"第二章 入职与试用期"的内容
  891. → 如果检索到了"第四章 假期管理",说明检索跑偏了
  892. → 这能真正区分不同检索策略的效果
  893. Returns:
  894. dict: 每种策略的 {"precision@3", "precision@5", "section_recall", "details"}
  895. """
  896. print(f"\n{'='*60}")
  897. print(f"[检索评估] Section 级评估 (测试 {len(test_set)} 个真实查询)")
  898. print(f"[检索评估] 衡量检索到的文档是否来自正确章节")
  899. strategies = ["bm25", "vector", "ensemble", "multi_query", "enhanced"]
  900. all_results = {}
  901. for strategy in strategies:
  902. per_query = []
  903. for qa in test_set:
  904. query = qa["question"]
  905. expected_keywords = qa["section_keywords"]
  906. # "enhanced" 用增强检索链(重写+HyDE+RRF,是实际生产 pipeline)
  907. if strategy == "enhanced":
  908. docs = enhanced_retrieve(query, strategy="vector", top_k=10)
  909. else:
  910. docs = retrieve(query, strategy=strategy, top_k=10)
  911. # 对每个检索到的 doc,检查其内容是否包含期望的章节关键词
  912. top3_hits = 0
  913. top5_hits = 0
  914. top3_texts = []
  915. for rank, doc in enumerate(docs):
  916. doc_text = doc.page_content
  917. # 检查文档中是否包含任何期望的章节关键词
  918. matched = any(kw in doc_text for kw in expected_keywords)
  919. if matched:
  920. if rank < 3:
  921. top3_hits += 1
  922. if rank < 5:
  923. top5_hits += 1
  924. if rank < 3:
  925. top3_texts.append(doc_text[:80])
  926. per_query.append({
  927. "query": query[:60],
  928. "expected_sections": expected_keywords,
  929. "top3_hits": top3_hits,
  930. "top5_hits": top5_hits,
  931. "top3_preview": top3_texts,
  932. })
  933. n = len(test_set)
  934. precision_3 = sum(q["top3_hits"] for q in per_query) / (n * 3) # 理想情况是 top3 中每篇都命中
  935. precision_5 = sum(q["top5_hits"] for q in per_query) / (n * 5)
  936. # Section Recall: 至少有一篇文档命中正确章节的查询比例
  937. section_recall = sum(1 for q in per_query if q["top3_hits"] > 0) / n
  938. all_results[strategy] = {
  939. "precision_at_3": precision_3,
  940. "precision_at_5": precision_5,
  941. "section_recall": section_recall,
  942. "details": per_query,
  943. }
  944. print(f"[检索评估] {strategy:<15s} | P@3: {precision_3:.1%} | "
  945. f"P@5: {precision_5:.1%} | Section Recall: {section_recall:.1%}")
  946. return all_results
  947. # ========================================================================
  948. # 10c. 检索交叉验证
  949. # ========================================================================
  950. def evaluate_retrieval_cross_verify(test_set: list[dict]) -> dict:
  951. """检索交叉验证: LLM 逐条判断检索到的文档是否能回答问题。
  952. 与 Section 级匹配互补:
  953. - Section 匹配是"粗粒度"的(文档来自正确章节?)
  954. - 交叉验证是"细粒度"的(文档内容真的能回答这个问题?)
  955. 方法:
  956. 对每个查询,取检索到的 Top-3 文档,让 LLM 逐条判断:
  957. "这篇文档是否包含回答该问题所需的信息?"
  958. → 是(1) / 部分(0.5) / 否(0)
  959. 优势: 不需要人工标注参考答案,LLM 只做"是否相关"的二元判断,
  960. 比对比参考答案更客观。
  961. """
  962. print(f"\n{'='*60}")
  963. print(f"[交叉验证] LLM 逐条判断检索文档是否相关")
  964. strategies = ["bm25", "vector", "ensemble", "multi_query"]
  965. all_results = {}
  966. llm = get_llm()
  967. for strategy in strategies:
  968. doc_scores = []
  969. for qi, qa in enumerate(test_set):
  970. query = qa["question"]
  971. docs = retrieve(query, strategy=strategy, top_k=3)
  972. for rank, doc in enumerate(docs):
  973. verify_prompt = f"""判断以下文档片段是否包含回答该问题所需的信息。
  974. 问题:{query}
  975. 文档:{doc.page_content[:400]}
  976. 只回答数字: 1=包含关键信息 / 0.5=部分相关 / 0=无关"""
  977. try:
  978. raw = llm.invoke(verify_prompt)
  979. ans = raw.content if hasattr(raw, "content") else str(raw)
  980. score = float(ans.strip()[0]) if ans.strip() else 0
  981. score = max(0, min(1, score))
  982. except Exception:
  983. score = 0
  984. doc_scores.append({
  985. "query": query[:50], "strategy": strategy,
  986. "rank": rank + 1, "relevant": score,
  987. })
  988. n = len(test_set) * 3 # 每个查询 3 篇文档
  989. avg_relevance = sum(d["relevant"] for d in doc_scores) / max(1, n)
  990. perfect_ratio = sum(1 for d in doc_scores if d["relevant"] == 1) / max(1, n)
  991. all_results[strategy] = {
  992. "avg_relevance": avg_relevance,
  993. "perfect_ratio": perfect_ratio,
  994. "details": doc_scores,
  995. }
  996. print(f"[交叉验证] {strategy:<15s} | 平均相关度: {avg_relevance:.2%} | "
  997. f"完美相关比: {perfect_ratio:.1%}")
  998. return all_results
  999. # ========================================================================
  1000. # 10d. 生成质量评估 + Bad Case 分析
  1001. # ========================================================================
  1002. def evaluate_generation(test_set: list[dict], sample_size: int = 10) -> dict:
  1003. """生成效果评估: 准确率、完整度、相关度、幻觉率。
  1004. 对每个测试问题:
  1005. 1. RAG 生成答案
  1006. 2. LLM judge 对比"参考答案"和"生成答案"
  1007. 3. 从 4 个维度打分
  1008. 注意: 这个评估的可靠性取决于 LLM judge 的质量,建议作为辅助指标。
  1009. """
  1010. import random
  1011. print(f"\n{'='*60}")
  1012. print(f"[生成评估] LLM Judge 打分(辅助指标)")
  1013. if len(test_set) > sample_size:
  1014. samples = random.sample(test_set, sample_size)
  1015. else:
  1016. samples = test_set
  1017. llm = get_llm()
  1018. judge_results = []
  1019. bad_cases = []
  1020. for i, qa in enumerate(samples):
  1021. query = qa["question"]
  1022. reference = qa["reference_answer"]
  1023. # RAG 生成答案
  1024. rag_result = rag_query(query, use_enhance=True, use_rerank=True)
  1025. generated = rag_result["answer"]
  1026. # 同时获取检索结果,用于 Bad Case 根因分析
  1027. retrieved_docs = retrieve(query, strategy="ensemble", top_k=5)
  1028. # LLM judge
  1029. judge_prompt = f"""你是一个严格但公正的评估助手。请对比参考答案和RAG系统生成的答案。
  1030. 参考答案(来自员工手册原文):{reference}
  1031. RAG生成答案:{generated}
  1032. 请从以下维度严格打分(1-5分,5分最好):
  1033. - accuracy(事实准确性):生成答案的事实是否与参考答案一致?(有事实错误扣分)
  1034. - completeness(完整性):是否覆盖了参考答案的关键信息?(遗漏重要信息扣分)
  1035. - relevance(相关性):是否直接回答了问题?(答非所问扣分)
  1036. - hallucination(幻觉程度):是否编造了参考答案中没有的内容?(5=无编造,1=严重编造)
  1037. 输出严格的JSON:
  1038. {{"accuracy": 整数1-5, "completeness": 整数1-5, "relevance": 整数1-5, "hallucination": 整数1-5}}"""
  1039. try:
  1040. raw_resp = llm.invoke(judge_prompt)
  1041. raw_text = raw_resp.content if hasattr(raw_resp, "content") else str(raw_resp)
  1042. json_match = re.search(r'\{[\s\S]*\}', raw_text)
  1043. if json_match:
  1044. scores = json.loads(json_match.group())
  1045. scores["query"] = query
  1046. scores["reference"] = reference
  1047. scores["generated"] = generated[:300]
  1048. judge_results.append(scores)
  1049. # 识别 Bad Case
  1050. is_bad = (scores.get("accuracy", 5) <= 2 or
  1051. scores.get("hallucination", 5) <= 2)
  1052. if is_bad:
  1053. # 根因分析
  1054. retrieved_texts = [d.page_content[:200] for d in retrieved_docs[:3]]
  1055. # 检查检索结果中是否有与 reference 相关的内容
  1056. retrieval_ok = any(
  1057. any(kw in " ".join(retrieved_texts) for kw in qa.get("section_keywords", []))
  1058. )
  1059. root_cause = "生成幻觉" if retrieval_ok else "检索失败"
  1060. bad_cases.append({
  1061. "query": query,
  1062. "reference": reference,
  1063. "generated": generated[:200],
  1064. "root_cause": root_cause,
  1065. "retrieved_preview": retrieved_texts[:2],
  1066. "scores": scores,
  1067. })
  1068. print(f"[生成评估] [{i+1}/{len(samples)}] "
  1069. f"acc={scores.get('accuracy','?')} "
  1070. f"hall={scores.get('hallucination','?')} "
  1071. f"{'⚠️ BAD' if is_bad else '✅'}")
  1072. except Exception as e:
  1073. print(f"[生成评估] 第 {i+1} 个评判失败: {e}")
  1074. # 汇总指标
  1075. if judge_results:
  1076. avg_acc = sum(r.get("accuracy", 0) for r in judge_results) / len(judge_results)
  1077. avg_comp = sum(r.get("completeness", 0) for r in judge_results) / len(judge_results)
  1078. avg_rel = sum(r.get("relevance", 0) for r in judge_results) / len(judge_results)
  1079. avg_hall = sum(r.get("hallucination", 0) for r in judge_results) / len(judge_results)
  1080. hall_rate = sum(1 for r in judge_results if r.get("hallucination", 5) <= 2) / len(judge_results)
  1081. bad_rate = len(bad_cases) / len(judge_results)
  1082. else:
  1083. avg_acc = avg_comp = avg_rel = avg_hall = hall_rate = bad_rate = 0
  1084. print(f"\n[生成评估] 汇总:")
  1085. print(f" 准确率: {avg_acc:.2f}/5 | 完整度: {avg_comp:.2f}/5")
  1086. print(f" 相关度: {avg_rel:.2f}/5 | 抗幻觉: {avg_hall:.2f}/5")
  1087. print(f" 幻觉率: {hall_rate:.1%} | Bad率: {bad_rate:.1%}")
  1088. return {
  1089. "avg_accuracy": avg_acc,
  1090. "avg_completeness": avg_comp,
  1091. "avg_relevance": avg_rel,
  1092. "avg_anti_hallucination": avg_hall,
  1093. "hallucination_rate": hall_rate,
  1094. "bad_case_rate": bad_rate,
  1095. "bad_cases": bad_cases,
  1096. "details": judge_results,
  1097. }
  1098. def bad_case_analysis(gen_results: dict) -> list[dict]:
  1099. """Bad Case 深度分析与优化建议。
  1100. 对每个 Bad Case 按根因分类,给出针对性优化方向。
  1101. """
  1102. bad_cases = gen_results.get("bad_cases", [])
  1103. print(f"\n{'='*60}")
  1104. print(f"[Bad Case 分析] 根因定位")
  1105. if not bad_cases:
  1106. print("[Bad Case 分析] ✅ 所有案例表现良好!")
  1107. return []
  1108. # 按根因分类
  1109. categories = defaultdict(list)
  1110. for case in bad_cases:
  1111. categories[case["root_cause"]].append(case)
  1112. print(f"[Bad Case 分析] 共 {len(bad_cases)} 个 Bad Case:\n")
  1113. for root_cause, cases in sorted(categories.items()):
  1114. print(f" 📌 [{root_cause}] ({len(cases)} 例)")
  1115. for case in cases[:1]:
  1116. print(f" ❓ {case['query'][:60]}")
  1117. print(f" 🤖 {case['generated'][:80]}...")
  1118. print()
  1119. # 优化建议
  1120. print(f"[Bad Case 分析] 针对性优化建议:")
  1121. if "检索失败" in categories:
  1122. print(f" 🔧 检索失败 ({len(categories['检索失败'])}例):")
  1123. print(f" → 降低 SemanticChunker 的 percentile 值(更细粒度切块)")
  1124. print(f" → 增加 RETRIEVAL_K 召回量")
  1125. print(f" → 考虑换用更大模型或调整 chunk 粒度")
  1126. print(f" → 检查查询重写是否改变了语义")
  1127. if "生成幻觉" in categories:
  1128. print(f" 🔧 生成幻觉 ({len(categories['生成幻觉'])}例):")
  1129. print(f" → 在 prompt 中加强对「不知道就说不知道」的约束")
  1130. print(f" → 降低 LLM temperature(当前默认值)")
  1131. print(f" → 要求逐句标注引用来源")
  1132. print(f" → 考虑接入 Reranker 精排后再生成")
  1133. return bad_cases
  1134. # ========================================================================
  1135. # 10e. 完整评估入口
  1136. # ========================================================================
  1137. def run_full_evaluation() -> dict:
  1138. """运行完整评估流程。
  1139. 三层评估:
  1140. 第一层: Section 级检索评估 → 客观指标(不依赖 LLM judge)
  1141. 第二层: 检索交叉验证 → LLM 逐篇判断相关性
  1142. 第三层: 生成质量评估 → LLM judge 打分 + Bad Case 根因分析
  1143. Returns:
  1144. dict: 完整评估报告
  1145. """
  1146. print(f"\n{'#'*60}")
  1147. print(f"# RAG 效果评估(三层体系)")
  1148. print(f"{'#'*60}")
  1149. # 使用手动标注的真实测试集(覆盖全部 16 章)
  1150. test_set = get_manual_test_set()
  1151. print(f"[评估] 手动标注测试集: {len(test_set)} 条,覆盖全部 16 章")
  1152. # ---- 第一层: Section 级检索评估 ----
  1153. section_results = evaluate_retrieval_section_level(test_set)
  1154. # ---- 第二层: 检索交叉验证 ----
  1155. cross_verify_results = evaluate_retrieval_cross_verify(test_set)
  1156. # ---- 第三层: 生成质量评估 + Bad Case ----
  1157. gen_results = evaluate_generation(test_set, sample_size=min(12, len(test_set)))
  1158. bad_cases = bad_case_analysis(gen_results)
  1159. # ---- 综合报告 ----
  1160. print(f"\n{'='*60}")
  1161. print(f"📊 综合评估报告")
  1162. print(f"{'='*60}")
  1163. print(f"测试集: {len(test_set)} 条手动标注的真实员工问题\n")
  1164. # 检索排名
  1165. print(f"🏆 检索策略排名 (按 Section Recall):")
  1166. sorted_sr = sorted(section_results.items(),
  1167. key=lambda x: x[1]["section_recall"], reverse=True)
  1168. for rank, (name, m) in enumerate(sorted_sr):
  1169. medal = ["🥇", "🥈", "🥉", "4️⃣", "5️⃣"][rank] if rank < 5 else " "
  1170. print(f" {medal} {name:<15s} P@3: {m['precision_at_3']:.1%} "
  1171. f"P@5: {m['precision_at_5']:.1%} "
  1172. f"SectionRecall: {m['section_recall']:.1%}")
  1173. print(f"\n🔍 检索交叉验证 (LLM 逐篇判断相关性):")
  1174. sorted_cv = sorted(cross_verify_results.items(),
  1175. key=lambda x: x[1]["avg_relevance"], reverse=True)
  1176. for rank, (name, m) in enumerate(sorted_cv):
  1177. medal = ["🥇", "🥈", "🥉", "4️⃣", "5️⃣"][rank] if rank < 5 else " "
  1178. print(f" {medal} {name:<15s} 平均相关度: {m['avg_relevance']:.2%} "
  1179. f"完美相关比: {m['perfect_ratio']:.1%}")
  1180. print(f"\n📈 生成质量 (LLM Judge 辅助指标):")
  1181. print(f" 准确率: {gen_results.get('avg_accuracy', 0):.2f}/5")
  1182. print(f" 完整度: {gen_results.get('avg_completeness', 0):.2f}/5")
  1183. print(f" 幻觉率: {gen_results.get('hallucination_rate', 0):.1%}")
  1184. print(f" Bad率: {gen_results.get('bad_case_rate', 0):.1%}")
  1185. print(f" Bad Case: {len(bad_cases)} 个")
  1186. # 确定最佳策略
  1187. best_strategy = sorted_sr[0][0] if sorted_sr else "ensemble"
  1188. best_cv = sorted_cv[0][0] if sorted_cv else "ensemble"
  1189. print(f"\n💡 推荐策略: 检索={best_strategy}, 交叉验证最佳={best_cv}")
  1190. if best_strategy != best_cv:
  1191. print(f" ⚠️ Section评估与交叉验证结论不一致,建议综合参考")
  1192. return {
  1193. "test_set_size": len(test_set),
  1194. "section_retrieval": section_results,
  1195. "cross_verify": cross_verify_results,
  1196. "generation": gen_results,
  1197. "bad_cases": bad_cases,
  1198. "recommended_strategy": best_strategy,
  1199. }
  1200. # ============================================================================
  1201. # Section 11: 主流程
  1202. # ============================================================================
  1203. def _print_banner():
  1204. print("""
  1205. ╔══════════════════════════════════════════════════════════╗
  1206. ║ RAG 全链路优化系统 — 员工手册智能问答 ║
  1207. ║ ║
  1208. ║ 1.员工手册.md 2.Agent工具 3.BGE-Large-Zh(1024d) ║
  1209. ║ 4.语义切块 5.多路召回 6.查询增强(HyDE) ║
  1210. ║ 7.Reranker精排 8.三层评估体系 ║
  1211. ╚══════════════════════════════════════════════════════════╝
  1212. """)
  1213. def main():
  1214. """主流程: 初始化 → 示例验证 → 交互测试。
  1215. 优化:默认跳过完整评估(200+ 次 LLM 调用,耗时 10+ 分钟),
  1216. 改为运行 1 个示例查询快速验证 RAG 链路可用,再进入交互模式。
  1217. 完整评估通过命令行参数 `--eval` 或交互命令 `eval` 触发。
  1218. """
  1219. global _VS, _CHUNKS
  1220. # 命令行参数: --eval 启动时直接跑完整评估;--demo 只跑示例不进入交互
  1221. run_eval_at_startup = "--eval" in sys.argv
  1222. demo_only = "--demo" in sys.argv
  1223. _print_banner()
  1224. # ---- Step 1: 核心初始化(仅加载必要组件)----
  1225. print("[初始化] 加载模型和文档...\n")
  1226. init_embedding_model() # embedding 模型(必须)
  1227. _CHUNKS = load_and_chunk() # 文档切块(必须)
  1228. _VS = get_vectorstore(_CHUNKS) # 向量存储(必须)
  1229. create_ensemble_retriever(_CHUNKS, _VS) # 默认策略检索器(必须)
  1230. # 以下组件改为按需初始化,避免启动时加载不必要的模型
  1231. # - create_mq_retriever(): 仅 multi_query 策略用到,首次调用时再初始化
  1232. # - init_reranker(): CrossEncoder 1.1GB,加载慢,rerank() 内部已按需加载
  1233. # ---- Step 2: 快速示例验证(1 个查询,验证 RAG 链路可用)----
  1234. print(f"\n{'='*60}")
  1235. print(f"[示例] 快速验证 RAG 链路(不启用查询增强/精排,约 3 秒)")
  1236. print(f"{'='*60}")
  1237. demo_q = "试用期是几个月?"
  1238. print(f"问题: {demo_q}")
  1239. result = rag_query(demo_q, strategy="ensemble",
  1240. use_enhance=False, use_rerank=False)
  1241. print(f"答案: {result['answer'][:200]}")
  1242. print(f"来源: {len(result['sources'])} 个文档 "
  1243. f"(top-1: {result['sources'][0].metadata.get('section', '?')})")
  1244. # ---- Step 3: 可选完整评估(默认跳过)----
  1245. if run_eval_at_startup:
  1246. eval_report = run_full_evaluation()
  1247. else:
  1248. print(f"\n[提示] 完整评估已跳过(耗时较长)。")
  1249. print(f" 如需运行评估:交互模式输入 'eval',"
  1250. f"或重启时加 --eval 参数")
  1251. # ---- Step 4: 交互模式 ----
  1252. if demo_only:
  1253. print(f"\n[--demo] 示例完成,不进入交互模式")
  1254. return
  1255. print(f"\n{'='*60}")
  1256. print(f"💬 交互模式")
  1257. print(f"{'='*60}")
  1258. print(f" 输入问题开始查询")
  1259. print(f" 命令: 'strategy <name>' 切换检索策略")
  1260. print(f" 策略: bm25 | vector | ensemble | multi_query")
  1261. print(f" 命令: 'agent <query>' 使用 Agent 模式")
  1262. print(f" 命令: 'enhance on/off' 切换查询增强(重写+HyDE)")
  1263. print(f" 命令: 'rerank on/off' 切换精排")
  1264. print(f" 命令: 'eval' 运行完整评估(耗时)")
  1265. print(f" 命令: 'exit' 退出")
  1266. print(f"{'='*60}\n")
  1267. current_strategy = "ensemble"
  1268. use_enhance = False # 默认关闭(每次查询省 2 次 LLM 调用)
  1269. use_rerank = False # 默认关闭(首次调用时加载 Reranker 模型)
  1270. while True:
  1271. try:
  1272. user_input = input("🔍 > ").strip()
  1273. except (EOFError, KeyboardInterrupt):
  1274. print("\n再见!")
  1275. break
  1276. if not user_input:
  1277. continue
  1278. if user_input.lower() in ("exit", "quit", "q"):
  1279. print("再见!")
  1280. break
  1281. # 命令处理
  1282. if user_input.lower().startswith("strategy "):
  1283. new_strategy = user_input.split(" ", 1)[1].strip()
  1284. if new_strategy in ("bm25", "vector", "ensemble", "multi_query"):
  1285. current_strategy = new_strategy
  1286. print(f"✅ 已切换到: {current_strategy}")
  1287. else:
  1288. print(f"❌ 未知策略: {new_strategy}")
  1289. continue
  1290. if user_input.lower().startswith("agent "):
  1291. query = user_input.split(" ", 1)[1].strip()
  1292. print(f"🤖 Agent 模式(自动判断是否需要检索)...\n")
  1293. answer = agent_query(query)
  1294. print(f"\n📝 {answer}\n")
  1295. continue
  1296. if user_input.lower().startswith("enhance "):
  1297. arg = user_input.split(" ", 1)[1].strip().lower()
  1298. use_enhance = arg in ("on", "true", "1", "yes")
  1299. print(f"✅ 查询增强: {'开启' if use_enhance else '关闭'}")
  1300. continue
  1301. if user_input.lower().startswith("rerank "):
  1302. arg = user_input.split(" ", 1)[1].strip().lower()
  1303. use_rerank = arg in ("on", "true", "1", "yes")
  1304. print(f"✅ 精排: {'开启' if use_rerank else '关闭'}")
  1305. continue
  1306. if user_input.lower() == "eval":
  1307. eval_report = run_full_evaluation()
  1308. continue
  1309. # 正常查询
  1310. enhance_tag = "查询增强" if use_enhance else "无增强"
  1311. rerank_tag = "精排" if use_rerank else "无精排"
  1312. print(f"🔎 检索策略: {current_strategy} | {enhance_tag} | {rerank_tag}")
  1313. result = rag_query(user_input, strategy=current_strategy,
  1314. use_enhance=use_enhance, use_rerank=use_rerank)
  1315. print(f"\n📝 {result['answer']}\n")
  1316. print(f"📖 参考来源 ({len(result['sources'])} 个文档):")
  1317. for i, doc in enumerate(result["sources"]):
  1318. print(f" [{i+1}] {doc.page_content[:100]}...")
  1319. print()
  1320. if __name__ == "__main__":
  1321. main()