350 lines
12 KiB
Python
350 lines
12 KiB
Python
"""
|
||
===========================================
|
||
设备到系统反向检索模块 - device_to_system
|
||
===========================================
|
||
功能：根据设备名称，反向检索对应的系统节点
|
||
使用混合检索（向量+全文）匹配设备节点
|
||
然后反向遍历2-8级找到对应的系统节点
|
||
===========================================
|
||
"""
|
||
|
||
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_device_with_retrieval(
|
||
session,
|
||
device_name: str,
|
||
sync_driver,
|
||
top_k: int = 10
|
||
) -> Optional[Dict[str, Any]]:
|
||
"""
|
||
使用 NodeRetrieval 混合检索匹配设备节点
|
||
|
||
Args:
|
||
session: Neo4j 异步会话
|
||
device_name: 设备名称
|
||
sync_driver: Neo4j 同步 driver（用于 NodeRetrieval）
|
||
top_k: 检索返回的节点数量上限
|
||
|
||
Returns:
|
||
匹配到的设备节点信息，包含 id, name, labels, props, score
|
||
"""
|
||
from graph_search.node_retrieval import NodeRetrieval
|
||
|
||
try:
|
||
embeddings = EmbeddingWrapper()
|
||
node_retrieval = NodeRetrieval(sync_driver, embeddings)
|
||
|
||
route_res = [{"label": "设备", "entity": device_name}]
|
||
retrieved_nodes = await node_retrieval.retrieve_nodes(route_res, top_k=top_k)
|
||
|
||
if not retrieved_nodes or "设备" not in retrieved_nodes:
|
||
logger.warning(f"NodeRetrieval 未检索到设备节点: {device_name}")
|
||
return None
|
||
|
||
candidates = retrieved_nodes["设备"]
|
||
if not candidates:
|
||
return None
|
||
|
||
logger.info(f"NodeRetrieval 检索到 {len(candidates)} 个设备节点候选")
|
||
|
||
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 = """
|
||
MATCH (n:设备)
|
||
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
|
||
except Exception as e:
|
||
logger.error(f"_match_device_with_retrieval 失败: {e}", exc_info=True)
|
||
return None
|
||
|
||
|
||
async def _do_find_system_by_device(
|
||
session,
|
||
device_name: str,
|
||
uri: str,
|
||
user: str,
|
||
password: str,
|
||
min_hops: int,
|
||
max_hops: int,
|
||
top_k: int,
|
||
start_time: float
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
执行设备到系统的检索逻辑（内部函数）
|
||
"""
|
||
device_key = str(device_name).strip()
|
||
|
||
logger.info(f"[设备到系统检索] 开始检索设备: {device_key}")
|
||
|
||
sync_driver = GraphDatabase.driver(uri, auth=(user, password))
|
||
try:
|
||
device_node = await _match_device_with_retrieval(session, device_key, sync_driver, top_k=top_k)
|
||
finally:
|
||
sync_driver.close()
|
||
|
||
if not device_node:
|
||
return {"success": False, "error": "未找到匹配的设备节点"}
|
||
|
||
device_id = device_node.get("id")
|
||
device_name_found = device_node.get("name")
|
||
logger.info(f"[设备到系统检索] 匹配到设备节点: {device_name_found}, id={device_id}, score={device_node.get('score', 0.0)}")
|
||
|
||
system_query = """
|
||
MATCH (d:设备)
|
||
WHERE elementId(d) = $device_id
|
||
MATCH (d)-[:包含*%d..%d]-(s:系统)
|
||
RETURN DISTINCT elementId(s) AS id, s.名称 AS name, labels(s) AS labels, properties(s) AS props
|
||
ORDER BY s.名称
|
||
""" % (min_hops, max_hops)
|
||
|
||
result = await session.run(system_query, device_id=device_id)
|
||
system_records = await result.data()
|
||
|
||
if not system_records:
|
||
logger.warning(f"[设备到系统检索] 未找到设备对应的系统节点，尝试扩大搜索范围...")
|
||
extended_query = """
|
||
MATCH (d:设备)
|
||
WHERE elementId(d) = $device_id
|
||
MATCH (d)-[:包含*1..10]-(s:系统)
|
||
RETURN DISTINCT elementId(s) AS id, s.名称 AS name, labels(s) AS labels, properties(s) AS props
|
||
ORDER BY s.名称
|
||
"""
|
||
result = await session.run(extended_query, device_id=device_id)
|
||
system_records = await result.data()
|
||
|
||
if not system_records:
|
||
return {
|
||
"success": False,
|
||
"error": "未找到设备对应的系统节点",
|
||
"device": {
|
||
"id": device_id,
|
||
"name": device_name_found,
|
||
"labels": device_node.get("labels", []),
|
||
"properties": device_node.get("props", {})
|
||
}
|
||
}
|
||
|
||
exclude_keys = {
|
||
"last_updated",
|
||
"created_at",
|
||
"knowledge_source",
|
||
"fulltext",
|
||
"embedding",
|
||
}
|
||
|
||
systems = []
|
||
for record in system_records:
|
||
props = record.get("props", {})
|
||
props_filtered = {k: v for k, v in props.items() if k not in exclude_keys}
|
||
|
||
systems.append({
|
||
"id": record.get("id"),
|
||
"name": record.get("name"),
|
||
"labels": record.get("labels", []),
|
||
"properties": props_filtered,
|
||
})
|
||
|
||
logger.info(f"[设备到系统检索] 找到 {len(systems)} 个系统节点")
|
||
|
||
path_query = """
|
||
MATCH (d:设备)
|
||
WHERE elementId(d) = $device_id
|
||
MATCH path = (d)-[:包含*%d..%d]-(s:系统)
|
||
RETURN path
|
||
LIMIT 5
|
||
""" % (min_hops, max_hops)
|
||
|
||
result = await session.run(path_query, device_id=device_id)
|
||
path_records = await result.data()
|
||
|
||
nodes_g: List[Dict[str, Any]] = []
|
||
links_g: List[Dict[str, Any]] = []
|
||
node_id_set = set()
|
||
|
||
for path_record in path_records:
|
||
path = path_record.get("path")
|
||
if not path:
|
||
continue
|
||
|
||
if hasattr(path, 'nodes'):
|
||
for node in path.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,
|
||
})
|
||
|
||
if hasattr(path, 'relationships'):
|
||
for rel in path.relationships:
|
||
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,
|
||
})
|
||
|
||
elapsed_ms = (time.time() - start_time) * 1000
|
||
logger.info(f"[时间统计] 设备到系统检索完成，耗时: {elapsed_ms:.2f}ms")
|
||
|
||
return {
|
||
"success": True,
|
||
"device": {
|
||
"id": device_id,
|
||
"name": device_name_found,
|
||
"labels": device_node.get("labels", []),
|
||
"properties": {k: v for k, v in device_node.get("props", {}).items() if k not in exclude_keys},
|
||
"match_score": device_node.get("score", 0.0)
|
||
},
|
||
"systems": systems,
|
||
"graph": {
|
||
"nodes": nodes_g,
|
||
"links": links_g
|
||
},
|
||
"meta": {
|
||
"elapsed_ms": round(elapsed_ms, 2),
|
||
"system_count": len(systems),
|
||
"min_hops": min_hops,
|
||
"max_hops": max_hops
|
||
}
|
||
}
|
||
|
||
|
||
async def find_system_by_device(
|
||
device_name: str,
|
||
min_hops: int = 2,
|
||
max_hops: int = 8,
|
||
top_k: int = 10,
|
||
driver=None,
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
根据设备名称反向检索对应的系统节点
|
||
|
||
Args:
|
||
device_name: 设备名称
|
||
min_hops: 最小跳数（默认2级）
|
||
max_hops: 最大跳数（默认8级）
|
||
top_k: 混合检索返回的候选节点数量
|
||
driver: Neo4j driver 实例（可选，如果不提供则创建新连接）
|
||
|
||
Returns:
|
||
Dict: 检索结果，包含 success, systems, path 等字段
|
||
"""
|
||
start_time = time.time()
|
||
|
||
if not device_name or not str(device_name).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"}
|
||
|
||
external_driver = driver is not None
|
||
|
||
try:
|
||
if external_driver:
|
||
async with driver.session() as session:
|
||
return await _do_find_system_by_device(
|
||
session, device_name, uri, user, password, min_hops, max_hops, top_k, start_time
|
||
)
|
||
else:
|
||
async with AsyncGraphDatabase.driver(uri, auth=(user, password)) as driver:
|
||
async with driver.session() as session:
|
||
return await _do_find_system_by_device(
|
||
session, device_name, uri, user, password, min_hops, max_hops, top_k, start_time
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"[设备到系统检索] 检索失败: {e}", exc_info=True)
|
||
return {"success": False, "error": f"检索失败: {str(e)}"}
|
||
|