354 lines
15 KiB
Python
354 lines
15 KiB
Python
"""
|
||
===========================================
|
||
操作图谱检索模块 - operation_graph_search
|
||
===========================================
|
||
功能：针对"操作"的固定图谱检索接口
|
||
检索路径：舰艇 -> 系统 -> 子系统 -> 设备 -> 操作使用 -> 操作程序 -> 操作项目
|
||
===========================================
|
||
"""
|
||
|
||
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 _match_node_with_retrieval(
|
||
session,
|
||
label: str,
|
||
node_name: str,
|
||
top_k: int = 10
|
||
) -> Optional[Dict[str, Any]]:
|
||
"""
|
||
使用 NodeRetrieval 匹配节点
|
||
"""
|
||
from graph_search.node_retrieval import NodeRetrieval
|
||
|
||
uri = NEO4J_CONFIG.get("uri")
|
||
user = NEO4J_CONFIG.get("user")
|
||
password = NEO4J_CONFIG.get("password")
|
||
|
||
sync_driver = GraphDatabase.driver(uri, auth=(user, password))
|
||
|
||
try:
|
||
embeddings = EmbeddingWrapper()
|
||
node_retrieval = NodeRetrieval(sync_driver, embeddings)
|
||
|
||
route_res = [{"label": label, "entity": node_name}]
|
||
retrieved_nodes = await node_retrieval.retrieve_nodes(route_res, top_k=top_k)
|
||
|
||
if not retrieved_nodes or label not in retrieved_nodes:
|
||
logger.warning(f"NodeRetrieval 未检索到 {label} 节点: {node_name}")
|
||
return None
|
||
|
||
candidates = retrieved_nodes[label]
|
||
if not candidates:
|
||
return None
|
||
|
||
logger.info(f"NodeRetrieval 检索到 {len(candidates)} 个 {label} 节点候选")
|
||
|
||
for i, node in enumerate(candidates):
|
||
node_data = node.get("node", {})
|
||
candidate_name = node_data.get("名称") or node_data.get("name") or node.get("name", "")
|
||
logger.info(f" 候选 {i+1}: 名称={candidate_name}, score={node.get('score', 0.0)}")
|
||
|
||
best = candidates[0]
|
||
node_data = best.get("node", {})
|
||
best_name = node_data.get("名称") or node_data.get("name") or best.get("name", "")
|
||
|
||
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(query, name=best_name)
|
||
record = await result.single()
|
||
|
||
if record:
|
||
return {
|
||
"id": record.get("id"),
|
||
"name": record.get("name"),
|
||
"labels": record.get("labels"),
|
||
"props": record.get("props", {}),
|
||
"score": best.get("score", 0.0)
|
||
}
|
||
|
||
return None
|
||
finally:
|
||
sync_driver.close()
|
||
|
||
|
||
async def operation_graph_search(
|
||
ship_number: str,
|
||
device_name: str,
|
||
operation_item: str,
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
操作图谱检索接口
|
||
|
||
Args:
|
||
ship_number: 舰艇舷号
|
||
device_name: 设备名称
|
||
operation_item: 操作项目名称
|
||
|
||
Returns:
|
||
Dict: 检索结果，包含 success, xxxx.graph 等字段
|
||
"""
|
||
start_time = time.time()
|
||
|
||
if not ship_number or not str(ship_number).strip():
|
||
return {"success": False, "error": "舷号不能为空"}
|
||
if not device_name or not str(device_name).strip():
|
||
return {"success": False, "error": "设备名称不能为空"}
|
||
if not operation_item or not str(operation_item).strip():
|
||
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:
|
||
async with AsyncGraphDatabase.driver(uri, auth=(user, password)) as driver:
|
||
async with driver.session() as session:
|
||
ship_key = str(ship_number).strip()
|
||
device_key = str(device_name).strip()
|
||
operation_key = str(operation_item).strip()
|
||
|
||
# ========== 1) 匹配舰艇节点 ==========
|
||
ship_query = """
|
||
MATCH (s:舰艇)
|
||
WHERE s.舷号 = $ship_key OR s.舷号 CONTAINS $ship_key OR $ship_key CONTAINS s.舷号
|
||
RETURN elementId(s) AS id, s.名称 AS name, s.舷号 AS ship_number, labels(s) AS labels, properties(s) AS props
|
||
LIMIT 1
|
||
"""
|
||
result = await session.run(ship_query, ship_key=ship_key)
|
||
ship_record = await result.single()
|
||
|
||
if not ship_record:
|
||
return {"success": False, "error": "在图谱中未找到对应舷号的舰艇节点"}
|
||
|
||
ship_id = ship_record.get("id")
|
||
ship_name = ship_record.get("name") or ship_record.get("ship_number")
|
||
logger.info(f"匹配到舰艇节点: {ship_name}, id={ship_id}")
|
||
|
||
# ========== 2) 匹配设备节点 ==========
|
||
device_query = """
|
||
MATCH (ship:舰艇)
|
||
WHERE elementId(ship) = $ship_id
|
||
MATCH (ship)-[:包含*0..8]->(d:设备)
|
||
WHERE d.名称 CONTAINS $device_key OR $device_key CONTAINS d.名称
|
||
OR d.name CONTAINS $device_key OR $device_key CONTAINS d.name
|
||
RETURN elementId(d) AS id, d.名称 AS name, labels(d) AS labels, properties(d) AS props
|
||
LIMIT 1
|
||
"""
|
||
result = await session.run(device_query, ship_id=ship_id, device_key=device_key)
|
||
device_record = await result.single()
|
||
|
||
if not device_record:
|
||
device_node = await _match_node_with_retrieval(session, "设备", device_key)
|
||
if device_node:
|
||
check_query = """
|
||
MATCH (ship:舰艇)
|
||
WHERE elementId(ship) = $ship_id
|
||
MATCH (ship)-[:包含*0..8]->(d:设备)
|
||
WHERE elementId(d) = $device_id
|
||
RETURN true AS in_range
|
||
LIMIT 1
|
||
"""
|
||
check_result = await session.run(check_query, ship_id=ship_id, device_id=device_node["id"])
|
||
check_record = await check_result.single()
|
||
if check_record:
|
||
device_record = device_node
|
||
|
||
if not device_record:
|
||
return {"success": False, "error": "在该舰艇范围内未找到匹配的设备节点"}
|
||
|
||
device_id = device_record.get("id")
|
||
device_name_found = device_record.get("name")
|
||
logger.info(f"匹配到设备节点: {device_name_found}, id={device_id}")
|
||
|
||
# ========== 3) 匹配操作项目节点 ==========
|
||
op_item_query = """
|
||
MATCH (d:设备)
|
||
WHERE elementId(d) = $device_id
|
||
MATCH (d)-[:包含*0..8]-(real_d:设备)
|
||
MATCH (real_d)-[:设备支持操作操作使用]-(op_use:操作使用)
|
||
-[:使用]-(op_proc:操作程序)
|
||
-[:操作时使用]-(op_item:操作项目)
|
||
WHERE op_item.名称 CONTAINS $op_key OR $op_key CONTAINS op_item.名称
|
||
RETURN elementId(op_item) AS id, op_item.名称 AS name, labels(op_item) AS labels, properties(op_item) AS props
|
||
LIMIT 1
|
||
"""
|
||
result = await session.run(op_item_query, device_id=device_id, op_key=operation_key)
|
||
op_item_record = await result.single()
|
||
|
||
if not op_item_record:
|
||
logger.info(f"Cypher查询未找到操作项目，尝试 NodeRetrieval...")
|
||
op_item_node = await _match_node_with_retrieval(session, "操作项目", operation_key)
|
||
if op_item_node:
|
||
logger.info(f"NodeRetrieval 找到操作项目节点: id={op_item_node.get('id')}, name={op_item_node.get('name')}")
|
||
check_query = """
|
||
MATCH (d:设备)
|
||
WHERE elementId(d) = $device_id
|
||
MATCH (d)-[:包含*0..8]-(real_d:设备)
|
||
MATCH (real_d)-[:支持操作]-(op_use:操作使用)
|
||
-[:使用]-(op_proc:操作程序)
|
||
-[:操作时使用]-(op_item:操作项目)
|
||
WHERE elementId(op_item) = $op_item_id
|
||
RETURN true AS in_range
|
||
LIMIT 1
|
||
"""
|
||
check_result = await session.run(check_query, device_id=device_id, op_item_id=op_item_node["id"])
|
||
check_record = await check_result.single()
|
||
if check_record:
|
||
logger.info(f"操作项目节点在设备路径下验证通过")
|
||
op_item_record = op_item_node
|
||
else:
|
||
logger.warning(f"操作项目节点不在设备路径下，尝试直接使用...")
|
||
op_item_record = op_item_node
|
||
|
||
if not op_item_record:
|
||
return {"success": False, "error": "未找到匹配的操作项目节点"}
|
||
|
||
op_item_id = op_item_record.get("id")
|
||
op_item_name = op_item_record.get("name")
|
||
logger.info(f"匹配到操作项目节点: {op_item_name}, id={op_item_id}")
|
||
|
||
# ========== 4) 查找经过这三个节点的链路 ==========
|
||
exclude_keys = {
|
||
"last_updated",
|
||
"created_at",
|
||
"knowledge_source",
|
||
"fulltext",
|
||
"embedding",
|
||
}
|
||
|
||
path_query = """
|
||
MATCH (ship:舰艇)
|
||
WHERE elementId(ship) = $ship_id
|
||
MATCH (device:设备)
|
||
WHERE elementId(device) = $device_id
|
||
MATCH (op_item:操作项目)
|
||
WHERE elementId(op_item) = $op_item_id
|
||
|
||
MATCH path1 = (ship)-[:包含*0..8]->(device)
|
||
MATCH path2 = (device)-[:包含*0..8]-(real_d:设备)
|
||
-[:支持操作]->(op_use:操作使用)
|
||
-[:使用]->(op_proc:操作程序)
|
||
-[:操作时使用]->(op_item)
|
||
|
||
WITH path1, path2
|
||
WITH nodes(path1) + nodes(path2) AS all_nodes,
|
||
relationships(path1) + relationships(path2) AS all_rels
|
||
RETURN all_nodes, all_rels
|
||
LIMIT 1
|
||
"""
|
||
result = await session.run(path_query, ship_id=ship_id, device_id=device_id, op_item_id=op_item_id)
|
||
path_record = await result.single()
|
||
|
||
if not path_record:
|
||
return {"success": False, "error": "未找到连接这三个节点的路径"}
|
||
|
||
# ========== 5) 格式化输出 ==========
|
||
all_nodes = path_record.get("all_nodes", [])
|
||
all_rels = path_record.get("all_rels", [])
|
||
|
||
nodes_g: List[Dict[str, Any]] = []
|
||
node_id_set = set()
|
||
|
||
for node in all_nodes:
|
||
if node is None:
|
||
continue
|
||
nid = node.element_id if hasattr(node, 'element_id') else node.id
|
||
if nid in node_id_set:
|
||
continue
|
||
node_id_set.add(nid)
|
||
|
||
props = dict(node) if hasattr(node, '__iter__') else {}
|
||
props_filtered = {k: v for k, v in props.items() if k not in exclude_keys}
|
||
|
||
name = props.get("名称") or props.get("name") or str(nid)
|
||
labels = list(node.labels) if hasattr(node, 'labels') else []
|
||
|
||
nodes_g.append({
|
||
"id": nid,
|
||
"name": name,
|
||
"labels": labels,
|
||
"properties": props_filtered,
|
||
})
|
||
|
||
links_g: List[Dict[str, Any]] = []
|
||
for rel in all_rels:
|
||
if rel is None:
|
||
continue
|
||
src = rel.start_node.element_id if hasattr(rel.start_node, 'element_id') else rel.start_node.id
|
||
tgt = rel.end_node.element_id if hasattr(rel.end_node, 'element_id') else rel.end_node.id
|
||
|
||
if src not in node_id_set or tgt not in node_id_set:
|
||
continue
|
||
|
||
links_g.append({
|
||
"id": rel.element_id if hasattr(rel, 'element_id') else rel.id,
|
||
"label": rel.type if hasattr(rel, 'type') else str(rel.type),
|
||
"source": src,
|
||
"target": tgt,
|
||
})
|
||
|
||
# ========== 6) 生成 results 文本 ==========
|
||
op_item_props = op_item_record.get("props", {})
|
||
op_info = {k: v for k, v in op_item_props.items() if k not in exclude_keys and v}
|
||
results_text = json.dumps(op_info, ensure_ascii=False) if op_info else ""
|
||
|
||
elapsed_ms = (time.time() - start_time) * 1000
|
||
logger.info(f"[时间统计] 操作图谱检索完成，耗时: {elapsed_ms:.2f}ms")
|
||
|
||
graph_item = {"nodes": nodes_g, "links": links_g, "results": results_text}
|
||
return {"success": True, "xxxx.graph": [graph_item]}
|
||
|
||
except Exception as e:
|
||
logger.error(f"操作图谱检索失败: {e}", exc_info=True)
|
||
return {"success": False, "error": f"Neo4j 查询失败: {str(e)}"}
|