""" 图谱检索工具模块 包含图谱搜索、故障图谱检索、操作图谱检索等功能 """ 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] = 40, 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() 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" } 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: return {"success": False, "error": f"请求超时: {str(e)}"} 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 {} try: graph_data = [] results_text = "" if isinstance(data, dict): raw_nodes = data.get("nodes", []) _filter_searchable_labels(raw_nodes) # 过滤 nodes 的 properties 字段，仅保留"名称"属性 for node in raw_nodes: if isinstance(node, dict): props = node.get("properties", {}) or node.get("props", {}) name_value = props.get("名称", props.get("name", "")) node.pop("properties", None) node.pop("props", None) if name_value: node["properties"] = {"名称": name_value} 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) # 过滤 nodes 的 properties 字段，仅保留"名称"属性 for node in raw_nodes: if isinstance(node, dict): props = node.get("properties", {}) or node.get("props", {}) name_value = props.get("名称", props.get("name", "")) node.pop("properties", None) node.pop("props", None) if name_value: node["properties"] = {"名称": name_value} 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: Optional[str] = None, device_name: str = "", fault_symptom: str = "", ) -> Dict[str, Any]: """ 图谱检索工具 """ 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: device_key = str(device_name).strip() ship_id = None ship_node = None # 有舷号时，先定位舰艇节点 if ship_number and str(ship_number).strip(): ship_key = str(ship_number).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 ship_candidates: ship_node = ship_candidates[0] ship_id = ship_node["id"] if ship_id is not None: # 有舰艇节点：在该舰艇范围内搜索设备 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 ) else: # 无舷号：全局搜索设备节点 device_query = """ MATCH (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 200 """ dev_records = await session.run(device_query) device_candidates = [_node_record_to_dict(r) async for r in dev_records] if not device_candidates: return {"success": False, "error": "未找到匹配的设备节点"} 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", } if ship_id is not None: 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], ) else: # 无舷号：不查舰艇路径，只查设备→故障模式→维修项目 graph_nodes_query = """ MATCH (real_e:设备) WHERE id(real_e) = $real_device_id MATCH path1 = (real_e)-[:发生故障使用维修工作]-(w:维修工作)-[:表征]-(fm:故障模式) WHERE id(fm) = $fault_mode_id OPTIONAL MATCH (fm)<-[:修理]-(rp:维修项目) WHERE id(rp) IN $repair_ids WITH path1, collect(DISTINCT rp) AS rps WITH 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, 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 {} # 仅保留"名称"属性 name_value = props.get("名称", props.get("name", "")) properties = {"名称": name_value} if name_value else {} nodes_g.append( { "id": nid, "name": name, "labels": labels, "properties": properties, } ) if ship_id is not None: 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], ) else: # 无舷号：只查设备→故障模式的边 graph_links_query = """ MATCH (real_e:设备) WHERE id(real_e) = $real_device_id MATCH path1 = (real_e)-[:发生故障使用维修工作]-(w:维修工作)-[:表征]-(fm:故障模式) WHERE id(fm) = $fault_mode_id WITH 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, 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: Optional[str] = None, 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 = 40) -> 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