retrieval.py 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672
  1. import asyncio
  2. import json
  3. from typing import List, Dict, Any, Optional
  4. from foundation.ai.models.rerank_model import rerank_model
  5. from foundation.observability.monitoring.time_statistics import track_execution_time
  6. from foundation.infrastructure.config.config import config_handler
  7. from foundation.observability.logger.loggering import review_logger
  8. from foundation.database.base.vector.milvus_vector import MilvusVectorManager
  9. class RetrievalManager:
  10. """
  11. 召回管理器,实现多路召回功能
  12. """
  13. def __init__(self):
  14. """
  15. 初始化召回管理器
  16. """
  17. self.vector_manager = MilvusVectorManager()
  18. self.logger = review_logger
  19. self.dense_weight = config_handler.get('hybrid_search', 'DENSE_WEIGHT', 0.7)
  20. self.sparse_weight = config_handler.get('hybrid_search', 'SPARSE_WEIGHT', 0.3)
  21. # 重排序模型配置(从 [model] 部分统一管理)
  22. self.rerank_model_type = config_handler.get('model', 'RERANK_MODEL_TYPE', 'bge_rerank_model')
  23. self.logger.info(f"初始化重排序模型类型: {self.rerank_model_type}")
  24. def set_rerank_model(self, model_type: str):
  25. """
  26. 设置重排序模型类型
  27. Args:
  28. model_type: 配置section名称 ('bge_rerank_model', 'lq_rerank_model', 'silicoflow_rerank_model')
  29. """
  30. valid_models = ['bge_rerank_model', 'lq_rerank_model', 'silicoflow_rerank_model', 'shutian_rerank_model']
  31. if model_type not in valid_models:
  32. raise ValueError(f"model_type 必须是 {valid_models}")
  33. self.rerank_model_type = model_type
  34. self.logger.info(f"重排序模型类型已设置为: {model_type}")
  35. def _clean_document(self, doc: str) -> str:
  36. """
  37. 清理文档文本,移除HTML标签和特殊字符
  38. Args:
  39. doc: 原始文档文本
  40. Returns:
  41. str: 清理后的文档文本
  42. """
  43. if not isinstance(doc, str):
  44. self.logger.debug(f"文档类型转换: {type(doc)} -> str")
  45. return str(doc)
  46. original_length = len(doc)
  47. # 移除HTML标签
  48. import re
  49. doc = re.sub(r'<[^>]+>', '', doc)
  50. # 移除多余的空白字符
  51. doc = re.sub(r'\s+', ' ', doc)
  52. # 更宽松的字符过滤 - 保留更多字符
  53. doc = re.sub(r'[^\u4e00-\u9fff\w\s.,;:!?()()。,;:!?\-\+\=\*/%&@#¥$【】「」""''""\n\r]', '', doc)
  54. # 截断过长的文本
  55. if len(doc) > 8000: # 设置最大长度限制
  56. doc = doc[:8000] + "..."
  57. cleaned_doc = doc.strip()
  58. self.logger.debug(f"文档清理: {original_length} -> {len(cleaned_doc)} 字符")
  59. return cleaned_doc
  60. def _get_rerank_results(self, query_text: str, documents: List[str], top_k: int = None) -> List[Dict[str, Any]]:
  61. """
  62. 根据配置选择重排序模型并执行重排序
  63. Args:
  64. query_text: 查询文本
  65. documents: 文档列表
  66. top_k: 返回结果数量
  67. Returns:
  68. List[Dict]: 重排序后的结果列表
  69. """
  70. try:
  71. # 清理和验证文档列表
  72. cleaned_documents = []
  73. valid_original_docs = []
  74. for doc in documents:
  75. if doc and isinstance(doc, str) and doc.strip():
  76. cleaned_doc = self._clean_document(doc)
  77. if cleaned_doc and len(cleaned_doc) > 3:
  78. cleaned_documents.append(cleaned_doc)
  79. valid_original_docs.append(doc)
  80. if not cleaned_documents:
  81. return []
  82. # 根据配置section名称路由到对应的reranker方法
  83. if self.rerank_model_type == 'lq_rerank_model':
  84. self.logger.info("使用本地 Qwen3-Reranker-8B (lq_rerank_model) 进行重排序")
  85. rerank_results = rerank_model.lq_rerank(query_text, cleaned_documents, top_k)
  86. elif self.rerank_model_type == 'silicoflow_rerank_model':
  87. self.logger.info("使用硅基流动 Qwen3-Reranker-8B (silicoflow_rerank_model) 进行重排序")
  88. rerank_results = rerank_model.qwen3_rerank(query_text, cleaned_documents, top_k)
  89. elif self.rerank_model_type == 'shutian_rerank_model':
  90. self.logger.info("使用蜀天云算力 Qwen3-Reranker-8B (shutian_rerank_model) 进行重排序")
  91. rerank_results = rerank_model.shutian_rerank(query_text, cleaned_documents, top_k)
  92. else: # bge_rerank_model (默认)
  93. self.logger.info("使用 BGE Reranker (bge_rerank_model) 进行重排序")
  94. rerank_results = rerank_model.bge_rerank(query_text, cleaned_documents, top_k)
  95. # 将清理后的文本映射回原始文本(所有reranker都需要)
  96. for result in rerank_results:
  97. cleaned_text = result.get('text', '')
  98. # 查找原始文本
  99. for i, cleaned in enumerate(cleaned_documents):
  100. if cleaned == cleaned_text:
  101. result['text'] = valid_original_docs[i]
  102. break
  103. # 统一字段名:将 relevance_score 转换为 score
  104. if 'relevance_score' in result and 'score' not in result:
  105. result['score'] = float(result['relevance_score'])
  106. return rerank_results
  107. except Exception as e:
  108. self.logger.error(f"重排序失败,模型类型: {self.rerank_model_type}, 错误: {str(e)}")
  109. # 返回原始顺序作为fallback
  110. return [{"text": doc, "score": 0.0} for i, doc in enumerate(documents[:top_k])]
  111. @track_execution_time
  112. async def entity_recall(self, main_entity: str, assisted_search_entity: list,
  113. recall_top_k: int = 5, max_results: int = None) -> List[str]:
  114. """
  115. 执行实体召回
  116. Args:
  117. main_entity: 主查询实体
  118. assisted_search_entity: 辅助搜索实体列表
  119. recall_top_k: 每次单实体召回返回的数量(默认5)
  120. max_results: 最终返回的最大数量,如果为None则返回所有召回结果(默认None)
  121. Returns:
  122. List[str]: 实体文本内容列表
  123. Note:
  124. 实际返回数量 = min(max_results, 主实体召回数 + 所有辅助实体召回数)
  125. 如果不设置max_results,可能返回较多结果(取决于辅助实体数量)
  126. """
  127. self.logger.info(f"[entity_recall] 开始召回, recall_top_k={recall_top_k}, max_results={max_results}, 主实体='{main_entity}', 辅助实体数量={len(assisted_search_entity)}")
  128. collection_name = config_handler.get('rag_collections', 'ENTITY_COLLECTION', 'first_bfp_collection_entity')
  129. # 主实体搜索 - 使用异步方法
  130. entity_result = await self.async_multi_stage_recall(
  131. collection_name=collection_name,
  132. query_text=main_entity,
  133. hybrid_top_k=10, # 降低召回数量,减少搜索耗时和上下文量
  134. top_k=recall_top_k
  135. )
  136. self.logger.info(f"[entity_recall] 主实体召回完成, 返回 {len(entity_result)} 个结果")
  137. assist_tasks = [
  138. self.async_multi_stage_recall(
  139. collection_name=collection_name,
  140. query_text=assisted_search_entity,
  141. hybrid_top_k=10, # 降低召回数量,减少搜索耗时和上下文量
  142. top_k=recall_top_k
  143. ) for assisted_search_entity in assisted_search_entity
  144. ]
  145. # 辅助搜索,异步并发
  146. assist_results_list = await asyncio.gather(*assist_tasks,return_exceptions=True)
  147. assist_results = []
  148. for res in assist_results_list:
  149. if isinstance(res, Exception):
  150. self.logger.error(f"辅助实体召回失败: {str(res)}")
  151. else:
  152. assist_results.extend(res)
  153. all_results = entity_result + assist_results
  154. # if self.rerank_model_type == 'silicoflow_rerank_model':
  155. # with open("temp\entity_bfp_recall\silicoflow_rerank_model.json", "w", encoding="utf-8") as f:
  156. # json.dump(all_results, f, ensure_ascii=False, indent=4)
  157. # elif self.rerank_model_type == 'lq_rerank_model':
  158. # with open("temp\entity_bfp_recall\lq_rerank_model.json", "w", encoding="utf-8") as f:
  159. # json.dump(all_results, f, ensure_ascii=False, indent=4)
  160. # 去重并提取文本内容
  161. entity_list = list(set([item['text_content'] for item in all_results]))
  162. # 如果设置了max_results,进行截断
  163. if max_results is not None and len(entity_list) > max_results:
  164. entity_list = entity_list[:max_results]
  165. self.logger.info(f"[entity_recall] 结果截断到 max_results={max_results}")
  166. self.logger.info(f"entity_list_len:{len(entity_list)}")
  167. return entity_list
  168. @track_execution_time
  169. async def async_bfp_recall(self, entity_list: List[str],background: str ,
  170. top_k: int = 3,) -> List[Dict[str, Any]]:
  171. """
  172. 混合搜索召回 - 向量+BM25召回
  173. Args:
  174. entity_list: 实体列表
  175. background: 背景/上下文信息,用于二次重排
  176. top_k: 返回结果数量
  177. """
  178. import time
  179. start_time = time.time()
  180. self.logger.info(f"[async_bfp_recall] 开始召回, top_k={top_k}, 实体数量={len(entity_list)}, 背景='{background[:50]}...'")
  181. # 异步并发召回编制依据
  182. collection_name = config_handler.get('rag_collections', 'CHILDREN_COLLECTION', 'rag_children_hybrid')
  183. gather_start = time.time()
  184. # 优化:降低hybrid_top_k参数从50到20,减少混合搜索时间
  185. bfp_tasks = [
  186. self.async_multi_stage_recall(
  187. collection_name=collection_name,
  188. query_text=entity,
  189. hybrid_top_k=10, # 从50降到20,减少60%的混合搜索时间
  190. top_k=top_k
  191. ) for entity in entity_list
  192. ]
  193. bfp_tasks_list = await asyncio.gather(*bfp_tasks,return_exceptions=True)
  194. gather_end = time.time()
  195. bfp_results = []
  196. for res in bfp_tasks_list:
  197. if isinstance(res, Exception):
  198. self.logger.error(f"辅助实体召回失败: {str(res)}")
  199. else:
  200. bfp_results.extend(res)
  201. self.logger.info(f"[async_bfp_recall] 第一阶段召回完成, 共召回 {len(bfp_results)} 个文档")
  202. # BFP召回结果已经通过multi_stage_recall进行了重排序,保持原有顺序
  203. # 只对第一次重排序得分大于0.8的文档进行二次重排序
  204. high_score_results = [item for item in bfp_results if (item.get('rerank_score') or 0) > 0.8]
  205. low_score_results = [item for item in bfp_results if (item.get('rerank_score') or 0) <= 0.8]
  206. self.logger.info(f"筛选结果:高分文档(>0.8) {len(high_score_results)} 个,低分文档(≤0.8) {len(low_score_results)} 个")
  207. # 如果没有高分文档,直接返回top_k个结果(按hybrid_similarity排序)
  208. if not high_score_results:
  209. self.logger.info(f"没有得分大于0.8的文档,跳过二次重排序,返回top_k={top_k}个结果(按hybrid_similarity排序)")
  210. # 按 hybrid_similarity 降序排序,返回 top_k 个
  211. sorted_results = sorted(bfp_results, key=lambda x: x.get('hybrid_similarity') or 0, reverse=True)
  212. return sorted_results[:top_k]
  213. # 检查background是否为空,如果为空则跳过二次重排序
  214. if not background or not background.strip():
  215. self.logger.warning("background为空,跳过二次重排序,直接返回高分文档")
  216. return high_score_results
  217. # 提取高分文档的文本内容用于二次重排(保持顺序去重)
  218. seen_texts = set()
  219. high_score_text_content = []
  220. for item in high_score_results:
  221. text = item['text_content']
  222. if text not in seen_texts:
  223. seen_texts.add(text)
  224. high_score_text_content.append(text)
  225. self.logger.info(f"提取高分文档文本内容,共 {len(high_score_text_content)} 个,准备二次重排")
  226. # 二次重排 - 使用配置的重排序模型
  227. rerank_start = time.time()
  228. # 使用传入的 top_k 参数,而不是硬编码为5
  229. bfp_rerank_result = self._get_rerank_results(background, high_score_text_content, top_k)
  230. rerank_end = time.time()
  231. self.logger.info(f"二次重排序耗时: {rerank_end - rerank_start:.3f}秒, top_k={top_k}")
  232. # 根据重排结果重新组织数据
  233. reorganize_start = time.time()
  234. final_results = []
  235. # 构建 text_content -> 原始文档列表 的映射(保留所有匹配项)
  236. text_to_items = {}
  237. for item in high_score_results:
  238. text = item['text_content']
  239. if text not in text_to_items:
  240. text_to_items[text] = []
  241. text_to_items[text].append(item)
  242. # 处理二次重排序的高分文档
  243. added_texts = set() # 用于跟踪已添加的文本,避免重复
  244. for rerank_item in bfp_rerank_result:
  245. text = rerank_item.get('text', '')
  246. parent_id = rerank_item.get('parent_id', '')
  247. score = rerank_item.get('score', 0.0)
  248. if text in text_to_items and text not in added_texts:
  249. # 获取该文本的所有候选文档,选择 rerank_score 最高的
  250. candidates = text_to_items[text]
  251. best_candidate = max(candidates, key=lambda x: x.get('rerank_score', 0.0))
  252. result_item = best_candidate.copy()
  253. result_item['bfp_rerank_score'] = score
  254. result_item['bfp_rerank_parent_id'] = parent_id
  255. final_results.append(result_item)
  256. added_texts.add(text) # 标记该文本已添加
  257. reorganize_end = time.time()
  258. total_time = reorganize_end - start_time
  259. self.logger.info(f"结果重组耗时: {reorganize_end - reorganize_start:.3f}秒")
  260. self.logger.info(f"二次重排完成,返回 {len(final_results)} 个高分文档(top_k={top_k}),丢弃 {len(low_score_results)} 个低分文档")
  261. self.logger.info(f"[async_bfp_recall] 总耗时: {total_time:.3f}秒 (召回: {gather_end-gather_start:.3f}s + 重排: {rerank_end-rerank_start:.3f}s + 其他: {total_time-(gather_end-gather_start)-(rerank_end-rerank_start):.3f}s)")
  262. return final_results
  263. def hybrid_search_recall(self, collection_name: str, query_text: str,
  264. top_k: int = 10 , ranker_type: str = "weighted",
  265. dense_weight: float = 0.7, sparse_weight: float = 0.3) -> List[Dict[str, Any]]:
  266. """
  267. 混合搜索召回 - 向量+BM25召回
  268. Args:
  269. collection_name: 集合名称
  270. query_text: 查询文本
  271. top_k: 返回结果数量
  272. ranker_type: 重排序类型 "weighted" 或 "rrf"
  273. dense_weight: 密集向量权重
  274. sparse_weight: 稀疏向量权重
  275. Returns:
  276. List[Dict]: 搜索结果列表
  277. """
  278. try:
  279. self.logger.info(f"开始混合检索")
  280. param = {'collection_name': collection_name}
  281. # 直接调用同步的混合搜索(在同步方法中)
  282. results = self.vector_manager.hybrid_search(
  283. param=param,
  284. query_text=query_text,
  285. top_k=top_k,
  286. ranker_type=ranker_type,
  287. dense_weight=dense_weight,
  288. sparse_weight=sparse_weight
  289. )
  290. # 详细记录混合搜索结果
  291. self.logger.info(f"混合搜索召回返回 {len(results)} 个结果")
  292. return results
  293. except Exception as e:
  294. self.logger.error(f"混合搜索召回失败: {str(e)}")
  295. return []
  296. def rerank_recall(self, candidates_with_metadata: List[Dict[str, Any]], query_text: str,
  297. top_k: int = None ) -> List[Dict[str, Any]]:
  298. """
  299. 重排序召回 - 使用配置的重排序模型对候选文档重新排序
  300. Args:
  301. candidates_with_metadata: 候选文档列表,包含文本内容和元数据
  302. query_text: 查询文本
  303. top_k: 返回结果数量
  304. Returns:
  305. List[Dict]: 重排序后的结果列表,包含原始索引信息
  306. """
  307. try:
  308. # 第一步:基于文本内容+元数据的组合去重
  309. unique_candidates = []
  310. original_indices_map = [] # 记录每个去重后的候选文档对应的原始索引列表
  311. unique_combinations = set() # 记录已见过的文本+元数据组合
  312. for original_index, candidate in enumerate(candidates_with_metadata):
  313. text_content = candidate.get('text_content', '')
  314. metadata = candidate.get('metadata', {})
  315. # 处理嵌套的metadata字符串
  316. title = ''
  317. file = ''
  318. if 'metadata' in metadata and isinstance(metadata['metadata'], str):
  319. import json
  320. try:
  321. # 解析JSON格式的metadata
  322. inner_metadata = json.loads(metadata['metadata'])
  323. title = inner_metadata.get('title', '')
  324. file = inner_metadata.get('file', '')
  325. except (json.JSONDecodeError, TypeError):
  326. pass
  327. else:
  328. title = metadata.get('title', '')
  329. file = metadata.get('file', '')
  330. # 创建组合键:文本内容 + 关键元数据
  331. combination_key = (text_content, title, file)
  332. if combination_key not in unique_combinations:
  333. # 新的唯一组合
  334. unique_candidates.append(candidate)
  335. original_indices_map.append([original_index])
  336. unique_combinations.add(combination_key)
  337. else:
  338. # 找到对应的唯一候选并添加索引
  339. for unique_idx, unique_candidate in enumerate(unique_candidates):
  340. if unique_candidate.get('text_content', '') == text_content:
  341. # 解析唯一候选的元数据
  342. unique_metadata = unique_candidate.get('metadata', {})
  343. unique_title = ''
  344. unique_file = ''
  345. if 'metadata' in unique_metadata and isinstance(unique_metadata['metadata'], str):
  346. import json
  347. try:
  348. inner_metadata = json.loads(unique_metadata['metadata'])
  349. unique_title = inner_metadata.get('title', '')
  350. unique_file = inner_metadata.get('file', '')
  351. except (json.JSONDecodeError, TypeError):
  352. pass
  353. else:
  354. unique_title = unique_metadata.get('title', '')
  355. unique_file = unique_metadata.get('file', '')
  356. if unique_title == title and unique_file == file:
  357. original_indices_map[unique_idx].append(original_index)
  358. break
  359. # 提取唯一候选文档的文本内容用于重排序
  360. unique_texts = [candidate.get('text_content', '') for candidate in unique_candidates]
  361. # 使用配置的重排序模型进行重排序
  362. rerank_results = self._get_rerank_results(query_text, unique_texts, top_k)
  363. # 转换结果格式,使用索引映射来处理原始索引
  364. scored_docs = []
  365. for i, api_result in enumerate(rerank_results):
  366. rerank_text = api_result.get('text', '')
  367. rerank_score = float(api_result.get('score', '0.0'))
  368. # 根据 rerank_text 在 unique_candidates 中查找匹配项
  369. # (rerank 会改变顺序,不能直接用索引 i)
  370. found_index = None
  371. original_candidate = None
  372. for idx, candidate in enumerate(unique_candidates):
  373. if candidate.get('text_content', '') == rerank_text:
  374. found_index = idx
  375. original_candidate = candidate
  376. break
  377. if original_candidate is None:
  378. self.logger.warning(f"[rerank_recall] 未找到匹配的候选文档,跳过: {rerank_text[:50]}...")
  379. continue
  380. # 使用找到的索引获取原始索引映射
  381. original_index = original_indices_map[found_index][0] if found_index < len(original_indices_map) else i
  382. # 获取原始混合搜索的评分信息
  383. hybrid_distance = original_candidate.get('distance', 0.0)
  384. hybrid_similarity = original_candidate.get('similarity', 0.0)
  385. # 解析元数据获取标题用于日志
  386. metadata = original_candidate.get('metadata', {})
  387. title = 'N/A'
  388. if 'metadata' in metadata and isinstance(metadata['metadata'], str):
  389. try:
  390. import json
  391. inner_metadata = json.loads(metadata['metadata'])
  392. title = inner_metadata.get('title', 'N/A')
  393. except:
  394. pass
  395. scored_docs.append({
  396. 'text_content': rerank_text,
  397. 'metadata': original_candidate.get('metadata', {}), # 保留原始元数据
  398. 'rerank_score': rerank_score,
  399. 'original_index': original_index,
  400. 'rerank_rank': i,
  401. 'duplicate_count': len(original_indices_map[i]), # 记录重复数量
  402. 'hybrid_distance': hybrid_distance, # 保留原始混合搜索评分
  403. 'hybrid_similarity': hybrid_similarity
  404. })
  405. return scored_docs
  406. except Exception as e:
  407. self.logger.error(f"重排序召回失败: {str(e)}")
  408. return []
  409. def multi_stage_recall(self, collection_name: str, query_text: str,
  410. hybrid_top_k: int = 50, top_k: int = 10,
  411. ranker_type: str = "weighted") -> List[Dict[str, Any]]:
  412. """
  413. 多路召回 - 先混合搜索召回,再重排序,只返回重排序结果
  414. Args:
  415. collection_name: 集合名称
  416. query_text: 查询文本
  417. hybrid_top_k: 混合搜索召回的文档数量
  418. top_k: 最终返回的文档数量
  419. ranker_type: 混合搜索的重排序类型
  420. Returns:
  421. List[Dict]: 重排序后的结果列表,只包含重排序分数
  422. """
  423. try:
  424. self.logger.info(f"执行多路召回")
  425. # 第一阶段:混合搜索召回(向量+BM25)
  426. hybrid_results = self.hybrid_search_recall(
  427. collection_name=collection_name,
  428. query_text=query_text,
  429. top_k=hybrid_top_k,
  430. ranker_type=ranker_type
  431. )
  432. if not hybrid_results:
  433. self.logger.warning("混合搜索召回无结果,返回空列表")
  434. return []
  435. # 第二阶段:重排序召回,传递完整的混合搜索结果(包含元数据)
  436. rerank_results = self.rerank_recall(
  437. candidates_with_metadata=hybrid_results,
  438. query_text=query_text,
  439. top_k=top_k
  440. )
  441. # 优化重排序结果的元数据结构
  442. final_results = []
  443. for rerank_result in rerank_results:
  444. metadata = rerank_result.get('metadata', {}).copy()
  445. duplicate_count = rerank_result.get('duplicate_count', 1)
  446. # 如果内层有metadata字段,将其提取到外层
  447. if 'metadata' in metadata and isinstance(metadata['metadata'], str):
  448. import json
  449. try:
  450. # 解析JSON格式的metadata
  451. inner_metadata = json.loads(metadata['metadata'])
  452. metadata.update(inner_metadata)
  453. # 移除内层的metadata字符串,避免重复
  454. del metadata['metadata']
  455. except (json.JSONDecodeError, TypeError):
  456. # 如果解析失败,保持原样
  457. pass
  458. # 移除重复的content字段
  459. if 'content' in metadata:
  460. del metadata['content']
  461. # 添加重复计数信息到元数据中
  462. if duplicate_count > 1:
  463. metadata['duplicate_count'] = duplicate_count
  464. # 输出优化后的结果,包含双重评分
  465. final_result = {
  466. 'text_content': rerank_result['text_content'],
  467. 'metadata': metadata,
  468. 'hybrid_similarity': rerank_result.get('hybrid_similarity', 0.0), # 混合搜索相似度
  469. 'rerank_score': rerank_result.get('rerank_score', 0.0) # BGE重排序评分
  470. }
  471. final_results.append(final_result)
  472. self.logger.debug(f"元数据优化完成: 重排序排名{rerank_result.get('rerank_rank')}, 重复数量={duplicate_count}")
  473. return final_results
  474. except Exception as e:
  475. self.logger.error(f"多路召回失败: {str(e)}")
  476. return []
  477. async def async_multi_stage_recall(self, collection_name: str, query_text: str,
  478. hybrid_top_k: int = 50, top_k: int = 10,
  479. ranker_type: str = "weighted") -> List[Dict[str, Any]]:
  480. """
  481. 多路召回 - 先混合搜索召回,再重排序,只返回重排序结果
  482. Args:
  483. collection_name: 集合名称
  484. query_text: 查询文本
  485. hybrid_top_k: 混合搜索召回的文档数量
  486. top_k: 最终返回的文档数量
  487. ranker_type: 混合搜索的重排序类型
  488. Returns:
  489. List[Dict]: 重排序后的结果列表,只包含重排序分数
  490. """
  491. import time
  492. try:
  493. start_time = time.time()
  494. # 第一阶段:混合搜索召回(向量+BM25)
  495. hybrid_results = await asyncio.to_thread(
  496. self.hybrid_search_recall,
  497. collection_name=collection_name,
  498. query_text=query_text,
  499. top_k=hybrid_top_k,
  500. ranker_type=ranker_type
  501. )
  502. if not hybrid_results:
  503. return []
  504. # 第二阶段:重排序召回
  505. rerank_results = self.rerank_recall(
  506. candidates_with_metadata=hybrid_results,
  507. query_text=query_text,
  508. top_k=top_k
  509. )
  510. # 优化重排序结果的元数据结构
  511. final_results = []
  512. for rerank_result in rerank_results:
  513. metadata = rerank_result.get('metadata', {}).copy()
  514. duplicate_count = rerank_result.get('duplicate_count', 1)
  515. # 如果内层有metadata字段,将其提取到外层
  516. if 'metadata' in metadata and isinstance(metadata['metadata'], str):
  517. import json
  518. try:
  519. # 解析JSON格式的metadata
  520. inner_metadata = json.loads(metadata['metadata'])
  521. metadata.update(inner_metadata)
  522. # 移除内层的metadata字符串,避免重复
  523. del metadata['metadata']
  524. except (json.JSONDecodeError, TypeError):
  525. # 如果解析失败,保持原样
  526. pass
  527. # 移除重复的content字段
  528. if 'content' in metadata:
  529. del metadata['content']
  530. # 添加重复计数信息到元数据中
  531. if duplicate_count > 1:
  532. metadata['duplicate_count'] = duplicate_count
  533. # 输出优化后的结果,包含双重评分
  534. final_result = {
  535. 'text_content': rerank_result['text_content'],
  536. 'metadata': metadata,
  537. 'hybrid_similarity': rerank_result.get('hybrid_similarity', 0.0), # 混合搜索相似度
  538. 'rerank_score': rerank_result.get('rerank_score', 0.0) # BGE重排序评分
  539. }
  540. final_results.append(final_result)
  541. self.logger.debug(f"元数据优化完成: 重排序排名{rerank_result.get('rerank_rank')}, 重复数量={duplicate_count}")
  542. return final_results
  543. except Exception as e:
  544. self.logger.error(f"多路召回失败: {str(e)}")
  545. return []
  546. # 创建全局召回管理器实例
  547. retrieval_manager = RetrievalManager()