""" =========================================== 图册检索模块 - atlas_retrieval =========================================== 功能：根据系统、子系统、设备，使用混合检索匹配节点 然后找到关联的图册节点，提取图册信息和PDF URL =========================================== """ import json import os import logging import time from typing import Any, Dict, List, Optional from dotenv import load_dotenv from neo4j import AsyncGraphDatabase, GraphDatabase load_dotenv() logger = logging.getLogger(__name__) NEO4J_CONFIG = { "uri": os.getenv("NEO4J_URI", "bolt://192.168.0.46:57687"), "user": os.getenv("NEO4J_USER", "neo4j"), "password": os.getenv("NEO4J_PASSWORD", "zdht123@"), } class EmbeddingWrapper: """嵌入模型包装器""" def embed_documents(self, texts: List[str]) -> List[List[float]]: """批量获取嵌入向量（一次性处理所有文本）""" from modelsAPI.model_api import OpenaiAPI if not texts: return [] try: result = OpenaiAPI.batch_embeddings(texts) if result and isinstance(result, list): embedding_list = [item.get('embedding') if isinstance(item, dict) else item for item in result] if embedding_list and len(embedding_list) == len(texts): return embedding_list else: logger.error(f"EmbeddingWrapper: 批量向量化返回数量不匹配，期望: {len(texts)}, 实际: {len(embedding_list) if embedding_list else 0}") return [None] * len(texts) else: logger.error(f"EmbeddingWrapper: 批量向量化返回格式异常") return [None] * len(texts) except Exception as e: logger.error(f"EmbeddingWrapper: 批量向量化失败: {e}", exc_info=True) return [None] * len(texts) async def _search_all_nodes_with_hybrid( session, query: str, sync_driver, top_k: int = 20 ) -> List[Dict[str, Any]]: """ 直接对所有节点进行混合检索（embedding+全文），找最匹配的节点 (与 quick_graph_search.py 保持一致的逻辑) 使用 NodeRetrieval 替代 HybridRetriever，避免兼容性问题 """ from graph_search.node_retrieval import NodeRetrieval try: # 对指定标签进行检索 all_labels = ["设备", "系统", "子系统"] logger.info(f"[混合检索] 检索标签: {all_labels}") # 使用 NodeRetrieval 进行检索 embeddings = EmbeddingWrapper() node_retrieval = NodeRetrieval(sync_driver, embeddings) # 为每个标签构造路由结果 route_res = [{"label": label, "entity": query} for label in all_labels] # 执行检索 retrieved_nodes = await node_retrieval.retrieve_nodes(route_res, top_k=top_k) # 汇总结果 all_results = [] for label, nodes in retrieved_nodes.items(): for node in nodes: all_results.append({ "name": node.get("name"), "score": node.get("score", 0.0), "node": node.get("node", {}), "label": label }) # 按score排序，取top_k all_results.sort(key=lambda x: x.get("score", 0), reverse=True) all_results = all_results[:top_k] logger.info(f"[混合检索] 找到 {len(all_results)} 个候选节点") # 获取完整节点信息 final_nodes = [] for item in all_results: node_data = item["node"] label = item["label"] name = item["name"] get_node_query = f""" MATCH (n:{label}) WHERE n.名称 = $name OR n.name = $name RETURN elementId(n) AS id, n.名称 AS name, labels(n) AS labels, properties(n) AS props LIMIT 1 """ result = await session.run(get_node_query, name=name) record = await result.single() if record: final_nodes.append({ "id": record.get("id"), "name": record.get("name"), "labels": record.get("labels"), "props": record.get("props", {}), "score": item.get("score", 0.0) }) return final_nodes except Exception as e: logger.error(f"混合检索失败: {e}", exc_info=True) return [] async def _rerank_nodes( query, nodes ): """ 使用 rerank 对节点进行重排序 (与 quick_graph_search.py 保持一致) """ if not nodes: return [] from modelsAPI.model_api import OpenaiAPI # 格式化节点文本 documents = [] for node in nodes: props = node.get("props", {}) name = node.get("name", "") labels = ", ".join(node.get("labels", [])) doc_text = f"名称: {name}; 类型: {labels}" documents.append(doc_text) # Rerank try: rerank_results = OpenaiAPI.rerank_query(query, documents, top_n=len(documents)) # 按score排序 sorted_rerank = sorted(rerank_results, key=lambda x: x.get('scores', 0), reverse=True) # 重新组织节点 reranked_nodes = [] for rr in sorted_rerank: idx = rr.get('index', 0) if 0 <= idx < len(nodes): node = nodes[idx].copy() node['rerank_score'] = rr.get('scores', 0) reranked_nodes.append(node) return reranked_nodes except Exception as e: logger.error(f"Rerank失败: {e}", exc_info=True) # 失败时返回原顺序 for node in nodes: node['rerank_score'] = 0 return nodes async def _match_nodes_with_retrieval( session, node_names: List[str], sync_driver, top_k: int = 10 ) -> List[Dict[str, Any]]: """ 使用混合检索+Rerank匹配节点（与 quick_graph_search.py 保持一致的逻辑） Args: session: Neo4j 异步会话 node_names: 节点名称列表 sync_driver: Neo4j 同步 driver top_k: 检索返回的节点数量上限 Returns: 匹配到的节点信息列表 """ try: all_matched_nodes = [] # 对每个节点名称分别进行检索+rerank for node_name in node_names: if not node_name: continue logger.info(f"[图册检索] 开始检索节点: '{node_name}'") # 1. 混合检索 candidate_nodes = await _search_all_nodes_with_hybrid( session, node_name, sync_driver, top_k=20 ) if not candidate_nodes: logger.warning(f"[图册检索] 未找到节点: '{node_name}'") continue logger.info(f"[图册检索] 混合检索找到 {len(candidate_nodes)} 个候选节点") # 2. Rerank 重排序 reranked_nodes = await _rerank_nodes(node_name, candidate_nodes) if not reranked_nodes: continue # 3. 取最匹配的节点 best_node = reranked_nodes[0] logger.info(f"[图册检索] 选择节点: {best_node.get('name', '未知')}, rerank_score={best_node.get('rerank_score', 0)}") all_matched_nodes.append(best_node) # 去重（按节点ID） seen_ids = set() unique_nodes = [] for node in all_matched_nodes: nid = node.get("id") if nid and nid not in seen_ids: seen_ids.add(nid) unique_nodes.append(node) logger.info(f"[图册检索] 匹配到 {len(unique_nodes)} 个唯一节点") return unique_nodes except Exception as e: logger.error(f"_match_nodes_with_retrieval 失败: {e}", exc_info=True) return [] async def find_atlas_by_nodes( node_names: Optional[List[str]] = None, top_k: int = 10, driver=None, ) -> Dict[str, Any]: """ 根据节点名称列表检索关联的图册节点 Args: node_names: 节点名称列表（可以是系统、子系统或设备） top_k: 混合检索返回的候选节点数量 driver: Neo4j driver 实例（可选，如果不提供则创建新连接） Returns: Dict: 检索结果，层级结构为：节点名称 -> 图册名称 -> PDF URL """ start_time = time.time() node_names = node_names or [] if not node_names: return {"success": False, "error": "至少需要提供一个节点名称"} uri = NEO4J_CONFIG.get("uri") user = NEO4J_CONFIG.get("user") password = NEO4J_CONFIG.get("password") if not uri or not user or not password: return {"success": False, "error": "Neo4j 配置不完整，请检查 NEO4J_CONFIG"} try: driver = AsyncGraphDatabase.driver(uri, auth=(user, password)) async with driver.session() as session: logger.info(f"[图册检索] 开始检索 - 节点名称: {node_names}") sync_driver = GraphDatabase.driver(uri, auth=(user, password)) try: matched_nodes = await _match_nodes_with_retrieval( session, node_names, sync_driver, top_k=top_k ) finally: sync_driver.close() if not matched_nodes: return {"success": False, "error": "未找到匹配的节点"} logger.info(f"[图册检索] 匹配到 {len(matched_nodes)} 个源节点") node_ids = [node["id"] for node in matched_nodes] atlas_query = """ MATCH (source) WHERE elementId(source) IN $node_ids MATCH (source)-[*1..1]-(atlas:图册) RETURN elementId(source) AS source_id, source.名称 AS source_name, elementId(atlas) AS atlas_id, atlas.名称 AS atlas_name, properties(atlas) AS atlas_props """ result = await session.run(atlas_query, node_ids=node_ids) records = await result.data() print(11111111111111111111111) print(records) if not records: return { "success": False, "error": "未找到关联的图册节点", "matched_nodes": matched_nodes } result_data = {} for record in records: source_name = record.get("source_name") atlas_name = record.get("atlas_name") atlas_props = record.get("atlas_props", {}) pdf_url = None for key, value in atlas_props.items(): if "图" in key or "pdf" in key.lower() or "url" in key.lower() or "文件" in key: if isinstance(value, str) and value.startswith(("http", "/")): pdf_url = value break if source_name not in result_data: result_data[source_name] = {} if atlas_name not in result_data[source_name]: result_data[source_name][atlas_name] = [] if pdf_url and pdf_url not in result_data[source_name][atlas_name]: result_data[source_name][atlas_name].append(pdf_url) total_atlas_count = sum(len(atlas_dict) for atlas_dict in result_data.values()) total_pdf_count = sum(len(pdf_list) for atlas_dict in result_data.values() for pdf_list in atlas_dict.values()) logger.info(f"[图册检索] 找到 {total_atlas_count} 个图册节点，共 {total_pdf_count} 个PDF URL") elapsed_ms = (time.time() - start_time) * 1000 logger.info(f"[时间统计] 图册检索完成，耗时: {elapsed_ms:.2f}ms") return { "success": True, "data": result_data } except Exception as e: logger.error(f"[图册检索] 检索失败: {e}", exc_info=True) return {"success": False, "error": f"检索失败: {str(e)}"} finally: if driver is not None: await driver.close()