1077 lines
38 KiB
Python
1077 lines
38 KiB
Python
"""
|
||
图谱检索工具模块
|
||
包含图谱搜索、故障图谱检索、操作图谱检索等功能
|
||
"""
|
||
import re
|
||
import json
|
||
import asyncio
|
||
import httpx
|
||
from typing import Optional, Dict, Any, List
|
||
from neo4j import AsyncGraphDatabase
|
||
from langchain_core.tools import tool
|
||
|
||
from config import ATLAS_CONFIG, GRAPH_SEARCH_CONFIG, NEO4J_CONFIG
|
||
from modelsAPI.model_api import OpenaiAPI
|
||
from utils.function_tracker import track_function_calls
|
||
|
||
|
||
GRAPH_SEARCH_ENDPOINT = GRAPH_SEARCH_CONFIG.get("search_endpoint")
|
||
|
||
|
||
@track_function_calls
|
||
async def graph_search(
|
||
query: str,
|
||
top_k: Optional[int] = 70,
|
||
use_rerank: Optional[bool] = False,
|
||
fulltext_top_k: Optional[int] = 15,
|
||
rerank_top_k: Optional[int] = 15,
|
||
hops: Optional[int] = 3,
|
||
verbose: Optional[bool] = False
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
图谱检索工具
|
||
"""
|
||
if not query or not query.strip():
|
||
return {"success": False, "error": "查询文本不能为空"}
|
||
|
||
query = query.strip()
|
||
|
||
print("图谱query", query)
|
||
|
||
payload = {
|
||
"query": query,
|
||
"use_rerank": use_rerank,
|
||
"top_k": top_k,
|
||
"fulltext_top_k": fulltext_top_k,
|
||
"rerank_top_k": rerank_top_k,
|
||
"hops": hops,
|
||
"verbose": verbose
|
||
}
|
||
|
||
headers = {
|
||
"accept": "application/json",
|
||
"Content-Type": "application/json"
|
||
}
|
||
|
||
max_retries = 1
|
||
retry_delay = 2
|
||
|
||
response = None
|
||
last_exception = None
|
||
data = None
|
||
|
||
for attempt in range(max_retries + 1):
|
||
try:
|
||
async with httpx.AsyncClient(timeout=180.0, verify=False) as client:
|
||
response = await client.post(GRAPH_SEARCH_ENDPOINT, json=payload, headers=headers)
|
||
except (httpx.ReadTimeout, httpx.ConnectTimeout, httpx.RequestError) as e:
|
||
last_exception = e
|
||
if attempt < max_retries:
|
||
print(f"图谱RAG请求超时，第 {attempt+1} 次重试，2秒后重新请求...")
|
||
await asyncio.sleep(retry_delay)
|
||
else:
|
||
return {"success": False, "error": f"请求超时: {str(e)}"}
|
||
continue
|
||
|
||
if response.status_code != 200:
|
||
return {"success": False, "error": f"HTTP {response.status_code}"}
|
||
|
||
result = response.json()
|
||
if result.get("code") != 200:
|
||
msg = result.get("message", "未知错误")
|
||
return {"success": False, "error": msg}
|
||
|
||
data = result.get("data", {}) or {}
|
||
nodes = []
|
||
if isinstance(data, dict):
|
||
nodes = data.get("nodes", [])
|
||
elif isinstance(data, list):
|
||
for item in data:
|
||
if isinstance(item, dict):
|
||
nodes.extend(item.get("nodes", []))
|
||
if not nodes and attempt < max_retries:
|
||
print(f"图谱RAG返回nodes为空，第 {attempt+1} 次重试，2秒后重新请求...")
|
||
await asyncio.sleep(retry_delay)
|
||
continue
|
||
break
|
||
|
||
try:
|
||
graph_data = []
|
||
results_text = ""
|
||
|
||
if isinstance(data, dict):
|
||
raw_nodes = data.get("nodes", [])
|
||
_filter_searchable_labels(raw_nodes)
|
||
graph_item = {
|
||
"nodes": raw_nodes,
|
||
"links": data.get("links", []),
|
||
}
|
||
graph_item = apply_styles(graph_item)
|
||
graph_data.append(graph_item)
|
||
results_text = data.get("results", "")
|
||
elif isinstance(data, list):
|
||
parts = []
|
||
for item in data:
|
||
if isinstance(item, dict):
|
||
raw_nodes = item.get("nodes", [])
|
||
_filter_searchable_labels(raw_nodes)
|
||
graph_item = {
|
||
"nodes": raw_nodes,
|
||
"links": item.get("links", []),
|
||
}
|
||
graph_item = apply_styles(graph_item)
|
||
graph_data.append(graph_item)
|
||
if item.get("results"):
|
||
parts.append(item["results"])
|
||
results_text = "\n".join(parts)
|
||
|
||
for graph_item in graph_data:
|
||
links = graph_item.get("links", [])
|
||
for link in links:
|
||
link.pop("properties", None)
|
||
|
||
return {
|
||
"success": True,
|
||
"xxxx.graph": graph_data,
|
||
"results": results_text,
|
||
}
|
||
|
||
except Exception as e:
|
||
return {"success": False, "error": str(e)}
|
||
|
||
|
||
def _string_similarity(a: str, b: str) -> float:
|
||
"""
|
||
简单字符串相似度：综合长度差、包含关系和公共字符比例。
|
||
0~1 之间，越大越相似。
|
||
"""
|
||
if not a or not b:
|
||
return 0.0
|
||
a = str(a).strip().lower()
|
||
b = str(b).strip().lower()
|
||
if not a or not b:
|
||
return 0.0
|
||
if a == b:
|
||
return 1.0
|
||
if a in b or b in a:
|
||
return 0.9
|
||
|
||
set_a, set_b = set(a), set(b)
|
||
inter = len(set_a & set_b)
|
||
union = len(set_a | set_b) or 1
|
||
jaccard = inter / union
|
||
|
||
len_diff = abs(len(a) - len(b))
|
||
len_penalty = max(0.0, 1.0 - len_diff / max(len(a), len(b), 1))
|
||
|
||
return 0.5 * jaccard + 0.5 * len_penalty
|
||
|
||
|
||
def _pick_best_by_name(candidates: List[Dict[str, Any]], target_name: str) -> Optional[Dict[str, Any]]:
|
||
"""
|
||
在候选节点中，根据名称相似度选择最优节点。
|
||
约定每个节点为：{id, labels, props, display_name}
|
||
"""
|
||
if not candidates or not target_name:
|
||
return None
|
||
|
||
best = None
|
||
best_score = 0.0
|
||
for item in candidates:
|
||
name = (item.get("display_name") or "").strip()
|
||
if not name:
|
||
props = item.get("props") or {}
|
||
name = props.get("name") or props.get("名称") or props.get("设备名称") or ""
|
||
score = _string_similarity(target_name, name)
|
||
if score > best_score:
|
||
best_score = score
|
||
best = item
|
||
|
||
return best if best_score >= 0.3 else None
|
||
|
||
|
||
def _extract_mt_codes(text: str) -> List[str]:
|
||
"""
|
||
从文本中抽取形如 MT007 的维修项目编号（MT + 3位数字）。
|
||
"""
|
||
if not text:
|
||
return []
|
||
codes = re.findall(r"\bMT\d{3}\b", str(text))
|
||
seen = set()
|
||
ordered: List[str] = []
|
||
for c in codes:
|
||
if c not in seen:
|
||
seen.add(c)
|
||
ordered.append(c)
|
||
return ordered
|
||
|
||
|
||
def _node_record_to_dict(record: Any) -> Dict[str, Any]:
|
||
"""
|
||
将 Cypher 返回的 id/labels/props 结构统一为字典。
|
||
"""
|
||
if hasattr(record, "data"):
|
||
data = record.data()
|
||
else:
|
||
data = record or {}
|
||
result = {
|
||
"id": data.get("id"),
|
||
"labels": data.get("labels") or [],
|
||
"props": data.get("props") or {},
|
||
"display_name": data.get("display_name")
|
||
if "display_name" in data
|
||
else None,
|
||
}
|
||
if "embedding" in data:
|
||
result["embedding"] = data.get("embedding")
|
||
if "real_device_id" in data:
|
||
result["real_device_id"] = data.get("real_device_id")
|
||
if "real_device_name" in data:
|
||
result["real_device_name"] = data.get("real_device_name")
|
||
return result
|
||
|
||
|
||
def _to_graph_display_node(node_dict: Dict[str, Any]) -> Dict[str, Any]:
|
||
"""将内部节点格式转为 graph_search 展示用节点：id、name、label。"""
|
||
nid = node_dict.get("id")
|
||
props = node_dict.get("props") or {}
|
||
name = (
|
||
node_dict.get("display_name")
|
||
or props.get("名称")
|
||
or props.get("name")
|
||
or str(nid)
|
||
)
|
||
labels = [l for l in (node_dict.get("labels") or []) if l != "Searchable"]
|
||
label = labels[0] if labels else ""
|
||
return {"id": nid, "name": name, "label": label}
|
||
|
||
|
||
def _build_fault_graph_display(
|
||
ship_node: Optional[Dict[str, Any]],
|
||
device_node: Optional[Dict[str, Any]],
|
||
fault_mode_node: Optional[Dict[str, Any]],
|
||
repair_projects: List[Dict[str, Any]],
|
||
) -> tuple:
|
||
"""
|
||
构建与 graph_search 一致的图谱展示结构：nodes 与 links。
|
||
"""
|
||
nodes: List[Dict[str, Any]] = []
|
||
links: List[Dict[str, Any]] = []
|
||
seen_ids = set()
|
||
|
||
def add_node(nd: Optional[Dict[str, Any]]):
|
||
if not nd or nd.get("id") in seen_ids:
|
||
return
|
||
seen_ids.add(nd["id"])
|
||
nodes.append(_to_graph_display_node(nd))
|
||
|
||
def add_link(source_id: Any, target_id: Any, rel_type: str = ""):
|
||
links.append({"source": source_id, "target": target_id, "type": rel_type})
|
||
|
||
if ship_node:
|
||
add_node(ship_node)
|
||
if device_node:
|
||
add_node(device_node)
|
||
if ship_node:
|
||
add_link(ship_node["id"], device_node["id"], "包含")
|
||
if fault_mode_node:
|
||
add_node(fault_mode_node)
|
||
if device_node:
|
||
add_link(device_node["id"], fault_mode_node["id"], "存在")
|
||
for p in repair_projects or []:
|
||
add_node(p)
|
||
if device_node:
|
||
add_link(device_node["id"], p["id"], "维修时使用")
|
||
|
||
return nodes, links
|
||
|
||
|
||
def _build_fault_graph_results_text(
|
||
ship_node: Optional[Dict[str, Any]],
|
||
device_node: Optional[Dict[str, Any]],
|
||
fault_mode_node: Optional[Dict[str, Any]],
|
||
fault_reason: str,
|
||
repair_plan: str,
|
||
mt_codes: List[str],
|
||
repair_projects: List[Dict[str, Any]],
|
||
ship_number: str,
|
||
device_name: str,
|
||
fault_symptom: str,
|
||
) -> str:
|
||
"""生成与 graph_rag_search 类似的一段话，供展示与下游使用。"""
|
||
ship_name = ""
|
||
if ship_node:
|
||
ship_name = (ship_node.get("display_name") or (ship_node.get("props") or {}).get("名称") or "").strip()
|
||
device_name_d = ""
|
||
if device_node:
|
||
device_name_d = (device_node.get("display_name") or (device_node.get("props") or {}).get("名称") or "").strip()
|
||
fault_name = ""
|
||
if fault_mode_node:
|
||
fault_name = (fault_mode_node.get("display_name") or (fault_mode_node.get("props") or {}).get("名称") or "").strip()
|
||
|
||
parts = [
|
||
f"根据舷号「{ship_number}」、设备「{device_name}」、故障现象「{fault_symptom}」进行图谱检索：",
|
||
f"定位到舰艇「{ship_name or '未知'}」、设备节点「{device_name_d or '未知'}」、故障模式「{fault_name or '未知'}」。",
|
||
]
|
||
if fault_reason:
|
||
parts.append(f"故障原因：{fault_reason}")
|
||
if repair_plan:
|
||
parts.append(f"维修方案摘要：{repair_plan[:500]}{'…' if len(repair_plan) > 500 else ''}")
|
||
if mt_codes:
|
||
parts.append(f"涉及维修项目编号：{'、'.join(mt_codes)}。")
|
||
if repair_projects:
|
||
proj_names = []
|
||
for p in repair_projects:
|
||
nm = (p.get("display_name") or (p.get("props") or {}).get("名称") or (p.get("props") or {}).get("维修项目编号") or "").strip()
|
||
if nm:
|
||
proj_names.append(nm)
|
||
if proj_names:
|
||
parts.append(f"匹配到的维修项目：{'、'.join(proj_names)}。")
|
||
return "\n".join(parts)
|
||
|
||
|
||
@track_function_calls
|
||
async def fault_graph_reason_and_projects(
|
||
ship_number: str,
|
||
device_name: str,
|
||
fault_symptom: str,
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
图谱检索工具
|
||
"""
|
||
|
||
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 fault_symptom or not str(fault_symptom).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()
|
||
|
||
ship_query = """
|
||
MATCH (s:舰艇)
|
||
WHERE s.舷号 IN [$ship_key]
|
||
RETURN DISTINCT
|
||
id(s) AS id,
|
||
labels(s) AS labels,
|
||
properties(s) AS props,
|
||
coalesce(s.名称, s.name, s.舷号) AS display_name
|
||
LIMIT 10
|
||
"""
|
||
ship_records = await session.run(ship_query, ship_key=ship_key)
|
||
ship_candidates = [_node_record_to_dict(r) async for r in ship_records]
|
||
if not ship_candidates:
|
||
return {"success": False, "error": "在图谱中未找到对应舷号的舰艇节点"}
|
||
ship_node = ship_candidates[0]
|
||
ship_id = ship_node["id"]
|
||
|
||
device_query = """
|
||
MATCH (ship:舰艇)
|
||
WHERE id(ship) = $ship_id
|
||
MATCH (ship)-[:包含*0..8]->(n:设备)
|
||
WITH n, coalesce(n.名称, n.name, n.设备名称) AS nm
|
||
WHERE nm IS NOT NULL
|
||
RETURN DISTINCT
|
||
id(n) AS id,
|
||
labels(n) AS labels,
|
||
properties(n) AS props,
|
||
nm AS display_name,
|
||
n.embedding AS embedding
|
||
LIMIT 100
|
||
"""
|
||
dev_records = await session.run(
|
||
device_query, ship_id=ship_id
|
||
)
|
||
device_candidates = [_node_record_to_dict(r) async for r in dev_records]
|
||
|
||
device_node, device_score = await OpenaiAPI.hybrid_match_with_embeddings(
|
||
query_text=device_key,
|
||
candidates=device_candidates,
|
||
name_key="display_name",
|
||
embedding_key="embedding",
|
||
text_weight=0.3,
|
||
semantic_weight=0.7
|
||
)
|
||
if not device_node:
|
||
return {"success": False, "error": f"在该舰艇范围内未找到匹配的设备节点（最高得分: {device_score:.3f}）"}
|
||
device_id = device_node["id"]
|
||
|
||
fault_mode_query = """
|
||
MATCH (start_e:设备)
|
||
WHERE id(start_e) = $device_id
|
||
MATCH (start_e)-[:包含*0..3]-(real_e:设备)
|
||
MATCH (real_e)-[:发生故障使用维修工作]-(w:维修工作)-[:包含故障]-(fm:故障模式)
|
||
RETURN DISTINCT
|
||
id(fm) AS id,
|
||
labels(fm) AS labels,
|
||
properties(fm) AS props,
|
||
fm.名称 AS display_name,
|
||
fm.embedding AS embedding,
|
||
id(real_e) AS real_device_id,
|
||
real_e.名称 AS real_device_name
|
||
"""
|
||
fm_records = await session.run(fault_mode_query, device_id=device_id)
|
||
fault_modes = [_node_record_to_dict(r) async for r in fm_records]
|
||
if not fault_modes:
|
||
return {"success": False, "error": "在该设备节点下未找到故障模式节点"}
|
||
|
||
fault_mode_node, fault_score = await OpenaiAPI.hybrid_match_with_embeddings(
|
||
query_text=fault_symptom,
|
||
candidates=fault_modes,
|
||
name_key="display_name",
|
||
embedding_key="embedding",
|
||
text_weight=0.3,
|
||
semantic_weight=0.7
|
||
)
|
||
if not fault_mode_node:
|
||
return {"success": False, "error": f"未能在故障模式节点中找到与故障现象足够匹配的节点（最高得分: {fault_score:.3f}）"}
|
||
|
||
fault_mode_id = fault_mode_node.get("id")
|
||
if fault_mode_id is None:
|
||
return {"success": False, "error": "故障模式节点缺少 id，无法继续检索维修项目"}
|
||
|
||
real_device_id = fault_mode_node.get("real_device_id")
|
||
|
||
fm_props = fault_mode_node.get("props") or {}
|
||
fault_reason = fm_props.get("故障原因") or ""
|
||
repair_plan = fm_props.get("维修方案") or ""
|
||
mt_codes = _extract_mt_codes(repair_plan)
|
||
|
||
repair_project_query = """
|
||
MATCH (fm:故障模式)
|
||
WHERE id(fm) = $fault_mode_id
|
||
OPTIONAL MATCH (fm)<-[:修理]-(p:维修项目)
|
||
RETURN DISTINCT
|
||
id(p) AS id,
|
||
labels(p) AS labels,
|
||
properties(p) AS props,
|
||
p.名称 AS display_name
|
||
"""
|
||
rp_records = await session.run(
|
||
repair_project_query, fault_mode_id=fault_mode_id
|
||
)
|
||
all_projects = [_node_record_to_dict(r) async for r in rp_records]
|
||
|
||
repair_projects: List[Dict[str, Any]] = []
|
||
if mt_codes:
|
||
for proj in all_projects:
|
||
props = proj.get("props") or {}
|
||
proj_code = (props.get("维修项目编号") or "").strip()
|
||
if proj_code and (
|
||
proj_code in mt_codes or any(c in proj_code for c in mt_codes)
|
||
):
|
||
repair_projects.append(proj)
|
||
elif not proj_code:
|
||
values_text = " ".join(
|
||
[str(v) for v in props.values() if v is not None]
|
||
)
|
||
if any(code in values_text for code in mt_codes):
|
||
repair_projects.append(proj)
|
||
else:
|
||
repair_projects = all_projects
|
||
|
||
repair_ids = [
|
||
p.get("id") for p in repair_projects if p.get("id") is not None
|
||
]
|
||
|
||
exclude_keys = {
|
||
"last_updated",
|
||
"created_at",
|
||
"knowledge_source",
|
||
"fulltext",
|
||
"embedding",
|
||
}
|
||
|
||
graph_nodes_query = """
|
||
MATCH (ship:舰艇)
|
||
WHERE id(ship) = $ship_id
|
||
MATCH (real_e:设备)
|
||
WHERE id(real_e) = $real_device_id
|
||
MATCH path0 = (ship)-[:包含*0..8]->(real_e)
|
||
MATCH path1 = (real_e)-[:发生故障使用维修工作]-(w:维修工作)-[:包含故障]-(fm:故障模式)
|
||
WHERE id(fm) = $fault_mode_id
|
||
OPTIONAL MATCH (fm)<-[:修理]-(rp:维修项目)
|
||
WHERE id(rp) IN $repair_ids
|
||
WITH path0, path1, collect(DISTINCT rp) AS rps
|
||
WITH nodes(path0) + nodes(path1) + rps AS ns
|
||
UNWIND ns AS n
|
||
WITH DISTINCT n
|
||
WHERE n IS NOT NULL
|
||
RETURN elementId(n) AS id,
|
||
labels(n) AS labels,
|
||
coalesce(n.名称, n.name) AS name,
|
||
properties(n) AS props
|
||
"""
|
||
node_records = await session.run(
|
||
graph_nodes_query,
|
||
ship_id=ship_id,
|
||
real_device_id=real_device_id,
|
||
fault_mode_id=fault_mode_id,
|
||
repair_ids=repair_ids or [-1],
|
||
)
|
||
|
||
nodes_g: List[Dict[str, Any]] = []
|
||
async for r in node_records:
|
||
data_r = r.data()
|
||
nid = data_r.get("id")
|
||
if nid is None:
|
||
continue
|
||
labels = data_r.get("labels") or []
|
||
labels = [l for l in labels if l != "Searchable"]
|
||
name = data_r.get("name") or str(nid)
|
||
props = data_r.get("props") or {}
|
||
props_filtered = {
|
||
k: v for k, v in props.items() if k not in exclude_keys
|
||
}
|
||
nodes_g.append(
|
||
{
|
||
"id": nid,
|
||
"name": name,
|
||
"labels": labels,
|
||
"properties": props_filtered,
|
||
}
|
||
)
|
||
|
||
graph_links_query = """
|
||
MATCH (ship:舰艇)
|
||
WHERE id(ship) = $ship_id
|
||
MATCH (real_e:设备)
|
||
WHERE id(real_e) = $real_device_id
|
||
MATCH path0 = (ship)-[:包含*0..8]->(real_e)
|
||
MATCH path1 = (real_e)-[:发生故障使用维修工作]-(w:维修工作)-[:包含故障]-(fm:故障模式)
|
||
WHERE id(fm) = $fault_mode_id
|
||
WITH relationships(path0) + relationships(path1) AS rs
|
||
UNWIND rs AS r
|
||
RETURN elementId(r) AS id,
|
||
type(r) AS label,
|
||
elementId(startNode(r)) AS source,
|
||
elementId(endNode(r)) AS target,
|
||
properties(r) AS props
|
||
UNION
|
||
MATCH (fm:故障模式)
|
||
WHERE id(fm) = $fault_mode_id
|
||
MATCH (fm)<-[r:修理]-(rp:维修项目)
|
||
WHERE id(rp) IN $repair_ids
|
||
RETURN elementId(r) AS id,
|
||
type(r) AS label,
|
||
elementId(startNode(r)) AS source,
|
||
elementId(endNode(r)) AS target,
|
||
properties(r) AS props
|
||
"""
|
||
link_records = await session.run(
|
||
graph_links_query,
|
||
ship_id=ship_id,
|
||
real_device_id=real_device_id,
|
||
fault_mode_id=fault_mode_id,
|
||
repair_ids=repair_ids or [-1],
|
||
)
|
||
|
||
links_g: List[Dict[str, Any]] = []
|
||
async for r in link_records:
|
||
data_r = r.data()
|
||
src = data_r.get("source")
|
||
tgt = data_r.get("target")
|
||
if src is None or tgt is None:
|
||
continue
|
||
links_g.append(
|
||
{
|
||
"id": data_r.get("id"),
|
||
"label": data_r.get("label", ""),
|
||
"source": src,
|
||
"target": tgt,
|
||
}
|
||
)
|
||
|
||
node_ids = {n["id"] for n in nodes_g}
|
||
links_g = [
|
||
l
|
||
for l in links_g
|
||
if l.get("source") in node_ids and l.get("target") in node_ids
|
||
]
|
||
|
||
repair_plans_data: List[Dict[str, Any]] = []
|
||
for p in repair_projects:
|
||
props = p.get("props") or {}
|
||
filtered = {k: v for k, v in props.items() if k not in exclude_keys}
|
||
if filtered:
|
||
repair_plans_data.append(filtered)
|
||
|
||
parts = []
|
||
|
||
if fault_reason and str(fault_reason).strip():
|
||
parts.append(str(fault_reason).strip())
|
||
|
||
if repair_plans_data:
|
||
for plan in repair_plans_data:
|
||
if plan is not None:
|
||
plan_str = str(plan).strip()
|
||
if plan_str:
|
||
parts.append(plan_str)
|
||
|
||
if parts:
|
||
combined_text = "\n".join(parts)
|
||
else:
|
||
combined_text = "暂无故障信息或修复方案"
|
||
|
||
results_obj = combined_text
|
||
|
||
results_text = json.dumps(results_obj, ensure_ascii=False)
|
||
|
||
graph_item = {"nodes": nodes_g, "links": links_g}
|
||
graph_item = apply_styles(graph_item)
|
||
return {"success": True, "xxxx.graph": [graph_item], "results": results_text}
|
||
|
||
except Exception as e:
|
||
return {"success": False, "error": f"Neo4j 查询失败: {str(e)}"}
|
||
|
||
|
||
def add_node_style(node):
|
||
node_type = node.get("labels", [])
|
||
node_type = [l for l in node_type if l != "Searchable"]
|
||
if node_type and len(node_type) > 0:
|
||
node_type = node_type[0]
|
||
else:
|
||
node_type = ""
|
||
|
||
if node_type == "舰艇":
|
||
node["color"] = "#F89A64"
|
||
node["size"] = 86
|
||
elif node_type == "系统":
|
||
node["color"] = "#C25C21"
|
||
node["size"] = 86
|
||
elif node_type == "子系统":
|
||
node["color"] = "#7D2D00"
|
||
node["size"] = 86
|
||
elif node_type == "设备":
|
||
node["color"] = "#540802"
|
||
node["size"] = 86
|
||
elif node_type in ["功能", "环境条件", "接口", "安全警告", "调试", "器材保障", "图册", "设计单位", "维修项目", "操作项目"]:
|
||
node["color"] = "#2261DF"
|
||
node["size"] = 52
|
||
else:
|
||
node["color"] = "#68C3FF"
|
||
node["size"] = 53
|
||
|
||
return node
|
||
|
||
|
||
def add_link_style(link, nodes_dict):
|
||
source_id = link.get("source")
|
||
target_id = link.get("target")
|
||
|
||
source_type = ""
|
||
target_type = ""
|
||
|
||
if source_id and source_id in nodes_dict:
|
||
source_labels = [l for l in nodes_dict[source_id].get("labels", []) if l != "Searchable"]
|
||
if source_labels and len(source_labels) > 0:
|
||
source_type = source_labels[0]
|
||
|
||
if target_id and target_id in nodes_dict:
|
||
target_labels = [l for l in nodes_dict[target_id].get("labels", []) if l != "Searchable"]
|
||
if target_labels and len(target_labels) > 0:
|
||
target_type = target_labels[0]
|
||
|
||
hierarchy_types = ["舰艇", "系统", "子系统", "设备"]
|
||
if source_type in hierarchy_types and target_type in hierarchy_types:
|
||
link["color"] = "#F59E0B"
|
||
else:
|
||
link["color"] = "#FFFFFF"
|
||
|
||
return link
|
||
|
||
|
||
def _filter_searchable_labels(nodes: List[Dict[str, Any]]):
|
||
for node in nodes:
|
||
if isinstance(node, dict) and "labels" in node:
|
||
node["labels"] = [l for l in node["labels"] if l != "Searchable"]
|
||
|
||
|
||
def apply_styles(graph_data):
|
||
nodes = graph_data.get("nodes", [])
|
||
links = graph_data.get("links", [])
|
||
|
||
nodes_dict = {n.get("id"): n for n in nodes if n.get("id")}
|
||
|
||
for node in nodes:
|
||
add_node_style(node)
|
||
|
||
for link in links:
|
||
add_link_style(link, nodes_dict)
|
||
|
||
return graph_data
|
||
|
||
|
||
@tool
|
||
async def fault_graph_rag_search(
|
||
ship_number: str,
|
||
device_name: str,
|
||
fault_symptom: str,
|
||
) -> str:
|
||
"""
|
||
故障诊断图谱检索工具（供工作流调用）：仅返回一段话摘要，与 graph_rag_search 一致。
|
||
根据舷号、设备名称、故障现象在图谱中定位舰艇→设备→故障模式→维修项目，并生成文字摘要。
|
||
工作流中用于与 RAG 结果一起作为上下文，不返回节点/关系（展示用 fault_graph_reason_and_projects）。
|
||
"""
|
||
result = await fault_graph_reason_and_projects(
|
||
ship_number=ship_number,
|
||
device_name=device_name,
|
||
fault_symptom=fault_symptom,
|
||
)
|
||
if not result.get("success", False):
|
||
return result.get("error", "检索失败")
|
||
return result.get("results", "")
|
||
|
||
|
||
@track_function_calls
|
||
async def operation_graph_search_tool(
|
||
ship_number: str,
|
||
device_name: str,
|
||
operation_item: str,
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
图谱检索工具
|
||
"""
|
||
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": "操作项目不能为空"}
|
||
|
||
endpoint = GRAPH_SEARCH_CONFIG.get("operation_search_endpoint")
|
||
timeout = GRAPH_SEARCH_CONFIG.get("graph_timeout")
|
||
|
||
if not endpoint:
|
||
return {"success": False, "error": "操作图谱检索服务配置不完整"}
|
||
|
||
payload = {
|
||
"ship_number": str(ship_number).strip(),
|
||
"device_name": str(device_name).strip(),
|
||
"operation_item": str(operation_item).strip(),
|
||
}
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||
response = await client.post(
|
||
endpoint,
|
||
json=payload,
|
||
headers={"Content-Type": "application/json"},
|
||
)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
|
||
if not result.get("success", False):
|
||
return {"success": False, "error": result.get("error", "检索失败")}
|
||
|
||
graph_list = result.get("xxxx.graph") or []
|
||
results_parts = []
|
||
if graph_list:
|
||
for i, graph_item in enumerate(graph_list):
|
||
if isinstance(graph_item, dict):
|
||
_filter_searchable_labels(graph_item.get("nodes", []))
|
||
if "results" in graph_item:
|
||
r = graph_item.pop("results")
|
||
if r:
|
||
results_parts.append(r)
|
||
graph_list[i] = apply_styles(graph_item)
|
||
result["xxxx.graph"] = graph_list
|
||
|
||
if results_parts:
|
||
result["results"] = "\n".join(results_parts)
|
||
elif "results" not in result:
|
||
result["results"] = ""
|
||
|
||
return result
|
||
|
||
except httpx.ReadTimeout:
|
||
return {"success": False, "error": "操作图谱检索服务超时"}
|
||
except httpx.ConnectError:
|
||
return {"success": False, "error": "无法连接操作图谱检索服务"}
|
||
except httpx.HTTPStatusError as e:
|
||
return {"success": False, "error": f"操作图谱检索服务返回错误: {e.response.status_code}"}
|
||
except Exception as e:
|
||
return {"success": False, "error": f"操作图谱检索失败: {str(e)}"}
|
||
|
||
|
||
@tool
|
||
async def operation_graph_rag_search(
|
||
ship_number: str,
|
||
device_name: str,
|
||
operation_item: str,
|
||
) -> str:
|
||
"""
|
||
操作图谱检索工具（供工作流调用）：仅返回一段话摘要，与 graph_rag_search 一致。
|
||
根据舷号、设备名称、操作项目在图谱中定位舰艇→设备→操作项目，并生成文字摘要。
|
||
工作流中用于与 RAG 结果一起作为上下文，不返回节点/关系（展示用 operation_graph_search_tool）。
|
||
"""
|
||
result = await operation_graph_search_tool(
|
||
ship_number=ship_number,
|
||
device_name=device_name,
|
||
operation_item=operation_item,
|
||
)
|
||
if not result.get("success", False):
|
||
return result.get("error", "检索失败")
|
||
return result.get("results", "")
|
||
|
||
|
||
@tool
|
||
async def graph_rag_search(query: str, entity_type: Optional[str] = None, top_k: int = 70) -> str:
|
||
"""
|
||
GraphRAG检索工具（异步版本）
|
||
在知识图谱中检索与query相关的实体和关系
|
||
|
||
Args:
|
||
query: 检索查询
|
||
entity_type: 实体类型过滤（如：设备、故障、维修项目）- 注意：此参数在当前实现中暂未使用
|
||
top_k: 返回结果数量（保留参数以兼容接口）
|
||
|
||
Returns:
|
||
图谱检索结果文本字符串（旧格式，保持兼容）
|
||
"""
|
||
print(f"graph_rag_query+++++++++++{query}")
|
||
result = await graph_search(query)
|
||
|
||
if not result.get("success", False):
|
||
return ""
|
||
|
||
results_text = result.get("results", "")
|
||
if not results_text:
|
||
return ""
|
||
if len(results_text) > 1000:
|
||
results_text = results_text[:1000] + "..."
|
||
return results_text
|
||
|
||
|
||
|
||
async def atlas_retrieval(
|
||
node_names: List[str],
|
||
top_k: int = 10,
|
||
timeout: int = 30
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
图册检索工具：从图谱中找到节点并关联图册PDF
|
||
|
||
Args:
|
||
node_names: 节点名称列表，如 ["主发动机", "压缩机"]
|
||
top_k: 返回的结果数量
|
||
timeout: 请求超时时间（秒）
|
||
|
||
Returns:
|
||
包含success和data字段的字典
|
||
data格式示例：
|
||
{
|
||
"船舶主发动机": {
|
||
"船舶主发动机的原理图": [
|
||
"http://192.168.0.46:59085/files/163-06A0015-B01001_船舶主发动机原理图.pdf"
|
||
]
|
||
}
|
||
}
|
||
"""
|
||
api_url = ATLAS_CONFIG["endpoint"]
|
||
|
||
if not node_names:
|
||
return {
|
||
"success": False,
|
||
"error": "node_names不能为空",
|
||
"data": {}
|
||
}
|
||
|
||
payload = {
|
||
"node_names": node_names,
|
||
"top_k": top_k
|
||
}
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||
response = await client.post(
|
||
api_url,
|
||
json=payload,
|
||
headers={
|
||
"accept": "application/json",
|
||
"Content-Type": "application/json"
|
||
}
|
||
)
|
||
|
||
if response.status_code == 200:
|
||
result = response.json()
|
||
if result.get("success", False):
|
||
return {
|
||
"success": True,
|
||
"data": result.get("data", {}),
|
||
"error": None
|
||
}
|
||
else:
|
||
return {
|
||
"success": False,
|
||
"error": result.get("message", "API返回失败"),
|
||
"data": {}
|
||
}
|
||
else:
|
||
return {
|
||
"success": False,
|
||
"error": f"HTTP {response.status_code}: {response.text}",
|
||
"data": {}
|
||
}
|
||
|
||
except Exception as e:
|
||
return {
|
||
"success": False,
|
||
"error": str(e),
|
||
"data": {}
|
||
}
|
||
|
||
|
||
def format_atlas_results(atlas_data: Dict[str, Any]) -> str:
|
||
"""
|
||
格式化图册检索结果为文本，便于在提示词中使用
|
||
|
||
Args:
|
||
atlas_data: atlas_retrieval返回的data字段
|
||
|
||
Returns:
|
||
格式化后的文本
|
||
"""
|
||
if not atlas_data:
|
||
return ""
|
||
|
||
parts = []
|
||
for node_name, node_data in atlas_data.items():
|
||
if isinstance(node_data, dict):
|
||
for title, urls in node_data.items():
|
||
if isinstance(urls, list) and urls:
|
||
parts.append(f"\n【{node_name} - {title}】")
|
||
for url in urls:
|
||
# 确保URL后面有明确的换行，避免和其他文字连在一起
|
||
parts.append(f"- {url}\n")
|
||
|
||
return "\n".join(parts) if parts else ""
|
||
|
||
|
||
async def extract_entity_names_from_history(
|
||
history_messages: List[Dict[str, Any]],
|
||
device_name: str = "",
|
||
operation_item: str = "",
|
||
fault: str = "",
|
||
last_generated_scheme: str = ""
|
||
) -> List[str]:
|
||
"""
|
||
从历史记录中抽取部件/设备/零部件名称
|
||
|
||
Args:
|
||
history_messages: 历史对话记录
|
||
device_name: 设备名称
|
||
operation_item: 操作项目名称
|
||
fault: 故障现象
|
||
last_generated_scheme: 最后一次生成的方案（用于重点抽取）
|
||
|
||
Returns:
|
||
抽取的实体名称列表
|
||
"""
|
||
entity_names = []
|
||
|
||
# 添加明确提供的名称
|
||
if device_name:
|
||
entity_names.append(device_name)
|
||
if operation_item:
|
||
entity_names.append(operation_item)
|
||
if fault:
|
||
entity_names.append(fault)
|
||
|
||
# 去重
|
||
entity_names = list(dict.fromkeys(entity_names))
|
||
|
||
# 调用大模型进一步抽取和补充
|
||
if entity_names or last_generated_scheme:
|
||
try:
|
||
extraction_prompt = f"""从以下信息中抽取尽可能多的设备、部件、零部件、机构、系统名称（用于图册检索）。
|
||
注意：
|
||
- 请尽可能多抽取，宁多勿少，准不准没关系
|
||
- 包括所有提及的机械部件、电子设备、机构系统等
|
||
- 即使是重复提到的也可以保留，我们会去重
|
||
- 重点从【生成的方案】中抽取，这里信息最全面
|
||
|
||
【示例 - 输入】
|
||
船舶主发动机（舷号：101）故障分析与维修方案
|
||
1. 故障概况
|
||
设备名称：船舶主发动机
|
||
舷号：101
|
||
故障现象：主机运行中出现异响
|
||
2. 故障原因分析
|
||
主轴承或连杆大端轴承磨损：轴承间隙过大或磨损严重会导致金属撞击声。
|
||
曲轴主轴颈跳动超差：曲轴旋转不平衡或跳动量超出标准。
|
||
飞轮连接螺栓松动：飞轮与曲轴连接处的螺栓松动。
|
||
活塞头烧蚀或活塞环断裂：活塞环断裂或活塞头部损坏可能导致敲缸声。
|
||
排气阀密封面磨损或阀座烧蚀：气门间隙异常或密封失效。
|
||
涡轮增压器轴承损坏：增压器内部轴承损坏。
|
||
3. 维修方案
|
||
执行项目：MT003 检测曲轴组件（含飞轮紧固状态检查）。
|
||
执行项目：MTO03 检测曲轴主轴颈跳动量。
|
||
|
||
【示例 - 输出】
|
||
{{
|
||
"entity_names": ["船舶主发动机", "主轴承", "连杆大端轴承", "曲轴", "飞轮", "飞轮连接螺栓", "活塞头", "活塞环", "排气阀", "排气阀座", "涡轮增压器", "涡轮增压器轴承", "曲轴组件"]
|
||
}}
|
||
|
||
【设备/操作】
|
||
- 设备名: {device_name or '无'}
|
||
- 操作项: {operation_item or '无'}
|
||
- 故障: {fault or '无'}
|
||
|
||
【生成的方案】
|
||
{last_generated_scheme or '无'}
|
||
|
||
请输出JSON格式，格式如下：
|
||
{{
|
||
"entity_names": ["名称1", "名称2", "名称3", "..."]
|
||
}}
|
||
|
||
仅输出JSON，不要其他内容。"""
|
||
|
||
response = await OpenaiAPI.open_api_chat_without_thinking(
|
||
query=extraction_prompt,
|
||
json_output=True,
|
||
system_prompt="你是一个专业的实体抽取助手，负责从文本中抽取尽可能多的设备、部件、零部件、机构、系统名称。请尽量多抽取，宁多勿少。重点从生成的维修方案中抽取。"
|
||
)
|
||
|
||
parsed = _safe_json_extract(response)
|
||
if isinstance(parsed, dict):
|
||
extracted = parsed.get("entity_names", [])
|
||
if extracted:
|
||
entity_names.extend(extracted)
|
||
entity_names = list(dict.fromkeys(entity_names))
|
||
except Exception as e:
|
||
print(f"[图册检索] 实体抽取失败: {str(e)}")
|
||
|
||
# 过滤空值和太短的名称
|
||
entity_names = [name.strip() for name in entity_names if name and len(name.strip()) >= 2]
|
||
return entity_names[:30] # 最多30个，尽量多找图册
|
||
|
||
|
||
def _safe_json_extract(text: str) -> Any:
|
||
"""安全地从文本中提取JSON"""
|
||
try:
|
||
# 尝试直接解析
|
||
return json.loads(text)
|
||
except json.JSONDecodeError:
|
||
# 尝试用正则提取
|
||
match = re.search(r'\{[\s\S]*\}', text)
|
||
if match:
|
||
try:
|
||
return json.loads(match.group(0))
|
||
except json.JSONDecodeError:
|
||
pass
|
||
return None
|
||
|