300 lines
12 KiB
Python
300 lines
12 KiB
Python
"""
|
||
===========================================
|
||
节点检索模块 - NodeRetrieval
|
||
===========================================
|
||
功能：根据路由结果，使用混合检索获取候选入口节点
|
||
使用HybridRetriever混合检索（向量+全文）
|
||
===========================================
|
||
"""
|
||
|
||
import re
|
||
import asyncio
|
||
import logging
|
||
import time
|
||
from typing import List, Dict, Any, Optional
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# ========== ========== 导入依赖 ========== ==========
|
||
try:
|
||
from neo4j_graphrag.retrievers import HybridRetriever
|
||
import jieba
|
||
# 过滤 jieba 缓存写入权限错误的日志（不影响功能，只是无法写入缓存）
|
||
jieba_logger = logging.getLogger('jieba')
|
||
jieba_logger.setLevel(logging.CRITICAL) # 只显示 CRITICAL 级别，过滤所有 ERROR/WARNING/INFO/DEBUG
|
||
NODE_RETRIEVAL_AVAILABLE = True
|
||
except ImportError:
|
||
NODE_RETRIEVAL_AVAILABLE = False
|
||
|
||
|
||
class NodeRetrieval:
|
||
"""
|
||
节点检索类
|
||
|
||
功能：根据路由结果，检索候选入口节点
|
||
"""
|
||
|
||
def __init__(self, driver, embeddings):
|
||
"""
|
||
初始化节点检索器
|
||
|
||
Args:
|
||
driver: Neo4j driver实例
|
||
embeddings: 嵌入模型实例（需要有embed_documents方法）
|
||
"""
|
||
self.driver = driver
|
||
self.embeddings = embeddings
|
||
|
||
async def retrieve_nodes(
|
||
self,
|
||
route_res: List[Dict[str, str]], # [{"label": "Drug", "entity": "xxx"}]
|
||
top_k: int = 10
|
||
) -> Dict[str, List[Dict[str, Any]]]:
|
||
"""
|
||
节点检索：根据标签和实体，检索入口节点
|
||
|
||
功能：
|
||
使用混合检索（向量+全文）进行检索
|
||
|
||
Args:
|
||
route_res: 路由结果，包含标签和实体信息
|
||
top_k: 检索返回的节点数量上限
|
||
|
||
Returns:
|
||
Dict[str, List[Dict]]: 检索到的节点结果，以标签为键
|
||
"""
|
||
if not NODE_RETRIEVAL_AVAILABLE:
|
||
logger.error("节点检索依赖未安装")
|
||
return {}
|
||
|
||
pairs = [] # 用于存储需要检索的标签-实体对
|
||
retrieved_nodes = {} # 用于存储检索到的节点结果，以标签为键
|
||
|
||
for item in route_res:
|
||
label = item.get("label", "")
|
||
entity = item.get("entity", "")
|
||
|
||
if not entity: # 如果实体为空则跳过
|
||
continue
|
||
|
||
# 将标签和实体作为一个元组添加到pairs列表中
|
||
pairs.append((label, entity))
|
||
|
||
if not pairs: # 如果没有需要检索的标签-实体对，直接返回
|
||
return retrieved_nodes
|
||
|
||
# 将标签-实体对分离成两个独立的列表
|
||
labels, entities = zip(*pairs)
|
||
labels, entities = list(labels), list(entities)
|
||
|
||
# ========== 对每个实体进行中文分词处理，生成全文检索查询文本 ==========
|
||
query_texts = []
|
||
for entity in entities:
|
||
try:
|
||
words = jieba.lcut(entity)
|
||
filtered_words = [
|
||
word.strip()
|
||
for word in words
|
||
if word.strip() and re.fullmatch(r"[a-zA-Z0-9\u4e00-\u9fa5]+", word.strip())
|
||
]
|
||
# 如果过滤后没有词，使用原始实体
|
||
if not filtered_words:
|
||
filtered_words = [entity] if entity.strip() else []
|
||
query_text = " OR ".join(filtered_words) if filtered_words else entity
|
||
query_texts.append(query_text)
|
||
except Exception as e:
|
||
logger.warning(f"实体 '{entity}' 分词失败: {e}，使用原始实体作为查询文本")
|
||
query_texts.append(entity)
|
||
|
||
# ========== 对实体进行向量化处理 ==========
|
||
try:
|
||
query_vectors = self.embeddings.embed_documents(entities)
|
||
except Exception as e:
|
||
logger.error(f"向量化失败: {e}", exc_info=True)
|
||
query_vectors = [None] * len(entities)
|
||
|
||
# ========== 为每个标签创建混合检索任务 ==========
|
||
tasks = []
|
||
for label, entity, query_text, query_vector in zip(labels, entities, query_texts, query_vectors):
|
||
try:
|
||
# vector_index_name = label.lower() + "_vector"
|
||
# fulltext_index_name = label.lower() + "_fulltext"
|
||
vector_index_name = "global_searchable_embedding"
|
||
fulltext_index_name = "global_searchable_fulltext_search"
|
||
retriever = HybridRetriever(
|
||
self.driver,
|
||
vector_index_name=vector_index_name,
|
||
fulltext_index_name=fulltext_index_name
|
||
)
|
||
tasks.append(
|
||
asyncio.to_thread(
|
||
retriever.get_search_results,
|
||
query_text,
|
||
query_vector,
|
||
top_k,
|
||
effective_search_ratio=2,
|
||
)
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"创建HybridRetriever失败 ({label}): {e}", exc_info=True)
|
||
tasks.append(None)
|
||
|
||
# ========== 并发执行所有检索任务 ==========
|
||
valid_tasks = [task for task in tasks if task is not None]
|
||
if valid_tasks:
|
||
results = await asyncio.gather(*valid_tasks, return_exceptions=True)
|
||
else:
|
||
logger.warning("没有有效的检索任务")
|
||
results = []
|
||
|
||
# ========== 处理检索结果 ==========
|
||
valid_results = []
|
||
exception_results = []
|
||
|
||
for i, result in enumerate(results):
|
||
if isinstance(result, Exception):
|
||
label = labels[i] if i < len(labels) else "unknown"
|
||
logger.error(f"检索任务 {i} (标签: {label}) 执行异常: {result}", exc_info=result)
|
||
exception_results.append((i, result))
|
||
else:
|
||
valid_results.append((i, result))
|
||
|
||
# ========== 按标签分组，准备进行RRF融合 ==========
|
||
# 结构: {label: [(entity, result_index, nodes_list), ...]}
|
||
label_results_map = {}
|
||
|
||
for i, result in valid_results:
|
||
label = labels[i] if i < len(labels) else "unknown"
|
||
entity = entities[i] if i < len(entities) else "unknown"
|
||
|
||
if result is None:
|
||
logger.warning(f"标签 '{label}' (实体: '{entity}') 的检索结果为 None")
|
||
continue
|
||
|
||
if not hasattr(result, 'records'):
|
||
logger.warning(f"标签 '{label}' (实体: '{entity}') 的结果没有 'records' 属性，类型: {type(result)}")
|
||
continue
|
||
|
||
|
||
# 提取节点信息（优先使用"名称"属性，兼容"name"）
|
||
nodes_list = [
|
||
{
|
||
"name": record["node"].get("名称") or record["node"].get("name") or str(record["node"]),
|
||
"score": record.get("score", 0.0),
|
||
"node": dict(record["node"])
|
||
}
|
||
for record in result.records
|
||
]
|
||
|
||
if nodes_list:
|
||
if label not in label_results_map:
|
||
label_results_map[label] = []
|
||
label_results_map[label].append((entity, i, nodes_list))
|
||
|
||
# ========== 对每个标签的多个实体结果进行RRF融合 ==========
|
||
# 策略：只对同一标签下相同实体名的结果进行RRF融合
|
||
rrf_fusion_start = time.time()
|
||
for label, entity_results in label_results_map.items():
|
||
# ========== 按实体名分组 ==========
|
||
entity_name_groups = {} # {entity_name: [(entity, result_index, nodes_list), ...]}
|
||
|
||
for entity, result_index, nodes_list in entity_results:
|
||
if entity not in entity_name_groups:
|
||
entity_name_groups[entity] = []
|
||
entity_name_groups[entity].append((entity, result_index, nodes_list))
|
||
|
||
# ========== 对每个实体名的结果进行RRF融合，然后合并 ==========
|
||
all_fused_nodes = []
|
||
|
||
for entity_name, group_results in entity_name_groups.items():
|
||
if len(group_results) == 1:
|
||
# 只有一个结果，不需要RRF融合，直接使用
|
||
_, _, nodes_list = group_results[0]
|
||
all_fused_nodes.extend(nodes_list)
|
||
logger.debug(f"标签 '{label}' 实体 '{entity_name}' 只有一个结果，跳过RRF融合")
|
||
else:
|
||
# 多个结果（相同实体名），使用RRF融合
|
||
fused_nodes = self._rrf_fusion(group_results, top_k=top_k)
|
||
all_fused_nodes.extend(fused_nodes)
|
||
|
||
# ========== 对同一标签内的所有节点按score排序 ==========
|
||
all_fused_nodes.sort(key=lambda x: x.get('score', 0), reverse=True)
|
||
|
||
# ========== 限制每个标签的数量 ==========
|
||
retrieved_nodes[label] = all_fused_nodes[:top_k]
|
||
|
||
# 统计去重后的结果
|
||
dedup_stats = {k: len(v) for k, v in retrieved_nodes.items()}
|
||
logger.info(f"检索结果统计（RRF融合后）: {list(dedup_stats.items())}")
|
||
return retrieved_nodes
|
||
|
||
def _rrf_fusion(
|
||
self,
|
||
entity_results: List[tuple],
|
||
top_k: int = 10,
|
||
k: int = 60
|
||
) -> List[Dict[str, Any]]:
|
||
"""
|
||
使用RRF (Reciprocal Rank Fusion) 融合多个实体的检索结果
|
||
|
||
RRF公式: RRF(d) = Σ 1/(k + rank_i(d))
|
||
其中:
|
||
- d 是文档/节点
|
||
- rank_i(d) 是文档在第i个排序结果中的排名（从1开始）
|
||
- k 是常数，通常为60
|
||
|
||
Args:
|
||
entity_results: [(entity, result_index, nodes_list), ...]
|
||
top_k: 返回的节点数量上限
|
||
k: RRF常数，默认60
|
||
|
||
Returns:
|
||
List[Dict]: 融合后的节点列表，按RRF分数降序排序
|
||
"""
|
||
if not entity_results:
|
||
return []
|
||
|
||
# ========== 1. 为每个节点计算RRF分数 ==========
|
||
node_rrf_scores = {} # {node_name: {"rrf_score": float, "node_data": dict, "ranks": [rank1, rank2, ...]}}
|
||
|
||
for entity, result_index, nodes_list in entity_results:
|
||
for rank, node in enumerate(nodes_list, start=1): # rank从1开始
|
||
node_name = node.get("name")
|
||
if not node_name:
|
||
continue
|
||
|
||
if node_name not in node_rrf_scores:
|
||
node_rrf_scores[node_name] = {
|
||
"rrf_score": 0.0,
|
||
"node_data": node, # 保留第一个出现的节点数据
|
||
"ranks": []
|
||
}
|
||
|
||
# 累加RRF分数: 1/(k + rank)
|
||
rrf_contribution = 1.0 / (k + rank)
|
||
node_rrf_scores[node_name]["rrf_score"] += rrf_contribution
|
||
node_rrf_scores[node_name]["ranks"].append(rank)
|
||
|
||
# ========== 2. 转换为列表并按RRF分数排序 ==========
|
||
fused_nodes = []
|
||
for node_name, data in node_rrf_scores.items():
|
||
node_data = data["node_data"].copy()
|
||
node_data["score"] = data["rrf_score"] # 使用RRF分数替换原始score
|
||
node_data["_rrf_ranks"] = data["ranks"] # 保留排名信息用于调试
|
||
fused_nodes.append(node_data)
|
||
|
||
# 按RRF分数降序排序
|
||
fused_nodes.sort(key=lambda x: x["score"], reverse=True)
|
||
|
||
# ========== 3. 限制返回数量 ==========
|
||
fused_nodes = fused_nodes[:top_k]
|
||
|
||
logger.debug(
|
||
f"RRF融合: {len(entity_results)} 个实体, "
|
||
f"融合前总节点数: {sum(len(nodes) for _, _, nodes in entity_results)}, "
|
||
f"融合后: {len(fused_nodes)} 个节点"
|
||
)
|
||
|
||
return fused_nodes
|
||
|