import json from datetime import datetime, timezone from typing import Dict, List, Optional, Tuple, Union from neo4j import GraphDatabase class MilitaryKnowledgeGraph: def __init__(self, uri: str, auth: Tuple[str, str], Ontology: List[str] = None): """ 初始化知识图谱 Args: uri: Neo4j 数据库连接地址 auth: 认证信息 (用户名, 密码) Ontology: 节点类型列表，用于创建唯一性约束和类型校验 """ self.driver = GraphDatabase.driver(uri, auth=auth) # 默认节点类型 self.Ontology = Ontology or ["舰艇", "系统", "子系统", "部件", "一级零部件", "二级零部件", "故障模式", "维修方案", "操作项目", "备品备件", "维修手册", "操作使用手册"] self._setup_constraints() def close(self): """关闭数据库连接""" self.driver.close() def _setup_constraints(self): """为每个节点类型设置 名称 字段的唯一性约束""" with self.driver.session() as session: for node_type in self.Ontology: try: session.run(f"CREATE CONSTRAINT IF NOT EXISTS FOR (n:{node_type}) REQUIRE n.名称 IS UNIQUE") except Exception as e: print(f"创建约束失败 {node_type}: {e}") def construct_fact(self, from_name: str, to_name: str, rel_type: str, attributes: Dict) -> str: """构造关系的自然语言描述（fact）""" attr_str = ", ".join([f"{k}：{v}" for k, v in attributes.items() if v]) if attr_str: return f"{from_name} {rel_type} {to_name}，{attr_str}" else: return f"{from_name} {rel_type} {to_name}" # def _resolve_attribute_conflicts(self, existing_attrs: Dict, new_attrs: Dict) -> Dict: # """解决属性冲突（保留最新状态、最早开始时间、最晚结束时间等）""" # resolved = existing_attrs.copy() # current_time = datetime.now(timezone.utc).isoformat() # for key, new_value in new_attrs.items(): # if new_value is None or new_value == "": # continue # if key not in existing_attrs: # resolved[key] = new_value # elif existing_attrs[key] != new_value: # if key in ["当前状态", "执行状态", "完成状态"]: # resolved[key] = new_value # elif key in ["开始时间", "下达时间", "承担时间"]: # if new_value < existing_attrs[key]: # resolved[key] = new_value # elif key in ["结束时间", "执行时限"]: # if new_value > existing_attrs[key]: # resolved[key] = new_value # else: # resolved[key] = new_value # if f"{key}_history" not in resolved: # resolved[f"{key}_history"] = [] # resolved[f"{key}_history"].append({ # "previous_value": existing_attrs[key], # "new_value": new_value, # "update_time": current_time # }) # return resolved def _resolve_attribute_conflicts(self, existing_attrs: Dict, new_attrs: Dict) -> Dict: """ 解决属性冲突： 1. 状态类：直接覆盖 2. 开始时间类：取最早时间 3. 结束时间类：取最晚时间 4. 其他类：直接覆盖 (不再记录历史) """ resolved = existing_attrs.copy() # current_time 在此版本中不再需要，但保留导入以防后续扩展使用 # current_time = datetime.now(timezone.utc).isoformat() for key, new_value in new_attrs.items(): if new_value is None or new_value == "": continue if key not in existing_attrs: # 新属性直接添加 resolved[key] = new_value elif existing_attrs[key] != new_value: # 属性存在且值不同，根据类型处理 if key in ["当前状态", "执行状态", "完成状态"]: # 状态类：直接覆盖为新值 resolved[key] = new_value elif key in ["开始时间", "下达时间", "承担时间"]: # 开始时间类：保留更早的时间 if new_value < existing_attrs[key]: resolved[key] = new_value elif key in ["结束时间", "执行时限"]: # 结束时间类：保留更晚的时间 if new_value > existing_attrs[key]: resolved[key] = new_value else: # 其他普通属性：直接覆盖，不再生成 _history 字段 resolved[key] = new_value return resolved def _create_relationship_with_timeline(self, tx, from_name: str, to_name: str, rel_type: str, attributes: Dict) -> bool: """创建或更新带有时间线的关系""" current_time = datetime.now(timezone.utc).isoformat() clean_attributes = {k: v for k, v in attributes.items() if v is not None and v != ""} fact = self.construct_fact(from_name, to_name, rel_type, attributes) clean_attributes["fact"] = fact existing_rel = tx.run(""" MATCH (a {名称: $from_name})-[r]->(b {名称: $to_name}) WHERE type(r) = $rel_type RETURN r, properties(r) as props """, from_name=from_name, to_name=to_name, rel_type=rel_type).single() if existing_rel: existing_props = existing_rel["props"] updated_props = existing_props.copy() for key, new_value in clean_attributes.items(): if key == "fact": continue if key not in existing_props or existing_props[key] != new_value: timeline_key = f"{key}_timeline" if timeline_key not in updated_props: updated_props[timeline_key] = [] timeline_entry_dict = { "value": new_value, "timestamp": current_time, "source": "text_extraction" } if key in existing_props: timeline_entry_dict["previous_value"] = existing_props[key] timeline_entry_str = json.dumps(timeline_entry_dict, ensure_ascii=False) updated_props[timeline_key].append(timeline_entry_str) updated_props[key] = new_value if "fact" in clean_attributes: fact_timeline_key = "fact_timeline" if fact_timeline_key not in updated_props: updated_props[fact_timeline_key] = [] fact_entry_str = json.dumps({ "value": clean_attributes["fact"], "timestamp": current_time, "source": "text_extraction" }, ensure_ascii=False) updated_props[fact_timeline_key].append(fact_entry_str) updated_props["fact"] = clean_attributes["fact"] updated_props["last_updated"] = current_time query = f""" MATCH (a {{名称: $from_name}})-[r:{rel_type}]->(b {{名称: $to_name}}) SET r = $properties RETURN r """ tx.run(query, from_name=from_name, to_name=to_name, properties=updated_props) print(f"更新关系: {from_name} -[{rel_type}]-> {to_name}") else: relationship_props = clean_attributes.copy() for key, value in clean_attributes.items(): timeline_key = f"{key}_timeline" timeline_entry_dict = { "value": value, "timestamp": current_time, "source": "text_extraction" } timeline_entry_str = json.dumps(timeline_entry_dict, ensure_ascii=False) relationship_props[timeline_key] = [timeline_entry_str] relationship_props["created_at"] = current_time relationship_props["last_updated"] = current_time relationship_props["update_count"] = 1 query = f""" MATCH (a {{名称: $from_name}}), (b {{名称: $to_name}}) CREATE (a)-[r:{rel_type} $properties]->(b) RETURN r """ tx.run(query, from_name=from_name, to_name=to_name, properties=relationship_props) print(f"创建关系: {from_name} -[{rel_type}]-> {to_name}") return True def _create_or_update_node(self, tx, node_type: str, 名称: str, attributes: Dict) -> bool: """ 创建或更新节点，支持 knowledge_source 等复杂字段 Args: tx: Neo4j 事务对象 node_type: 节点类型 名称: 节点名称（唯一标识） attributes: 节点属性字典 """ if not 名称 or not node_type: return False # 安全校验：只允许预定义的节点类型 if node_type not in self.Ontology: print(f"跳过未知节点类型: {node_type}") return False # 预处理 attributes：将 knowledge_source 转为 JSON 字符串 clean_attributes = {"名称": 名称} # 确保名称字段存在 for k, v in attributes.items(): if v is None or v == "": continue if k == "knowledge_source" and isinstance(v, list): clean_attributes[k] = json.dumps(v, ensure_ascii=False) else: clean_attributes[k] = v existing_node = tx.run( f"MATCH (n:{node_type} {{名称: $名称}}) RETURN n, properties(n) as props", 名称=名称 ).single() current_time = datetime.now(timezone.utc).isoformat() if existing_node: existing_props = existing_node["props"] existing_attrs = {k: v for k, v in existing_props.items() if not k.endswith('_timeline') and k not in ['created_at', 'last_updated']} resolved_attributes = self._resolve_attribute_conflicts(existing_attrs, clean_attributes) timeline_data = {k: v for k, v in existing_props.items() if k.endswith('_timeline') or k in ['created_at', 'last_updated']} final_attributes = {**timeline_data, **resolved_attributes} final_attributes["last_updated"] = current_time query = f""" MATCH (n:{node_type} {{名称: $名称}}) SET n = $attributes RETURN n """ tx.run(query, 名称=名称, attributes=final_attributes) print(f"更新节点: {node_type} - {名称}") else: clean_attributes["created_at"] = current_time clean_attributes["last_updated"] = current_time query = f""" CREATE (n:{node_type} $attributes) RETURN n """ tx.run(query, attributes=clean_attributes) print(f"创建节点: {node_type} - {名称}") return True def process_data(self, data: Dict) -> Tuple[int, int]: """ 处理字典数据并写入Neo4j（批量版：按类型分组，减少 Neo4j 往返次数） Args: data: 包含 entities 和 relationships 的字典 Returns: Tuple[int, int]: (实体数量, 关系数量) """ if not data or not isinstance(data, dict): print("无效的数据格式，需要字典类型") return 0, 0 BATCH_SIZE = 50 entity_count = 0 relationship_count = 0 current_time = datetime.now(timezone.utc).isoformat() # ── 1. 预处理实体，按类型分组 ────────────────────────────── entities_by_type: Dict[str, List[dict]] = {} for entity in data.get("entities", []): node_type = entity.get("type", "") if node_type not in self.Ontology: print(f"跳过未知节点类型: {node_type}") continue props = entity.get("properties", {}).copy() if "knowledge_source" in entity: props["knowledge_source"] = entity["knowledge_source"] 名称 = props.pop("名称", None) if not 名称: print(f"跳过无名称节点: {entity}") continue clean: Dict = {"名称": 名称} for k, v in props.items(): if v is None or v == "": continue if k == "knowledge_source" and isinstance(v, list): clean[k] = json.dumps(v, ensure_ascii=False) else: clean[k] = v entities_by_type.setdefault(node_type, []).append(clean) # ── 2. 按类型批量读取已存在节点，Python侧合并属性，批量写入 ── with self.driver.session() as session: for node_type, type_entities in entities_by_type.items(): names = [e["名称"] for e in type_entities] # 一次读取该类型下所有涉及名称 existing_list = session.execute_read( self._batch_fetch_nodes, node_type, names ) existing_map = {r["名称"]: r for r in existing_list} to_create: List[dict] = [] to_update: List[dict] = [] for clean_attrs in type_entities: name = clean_attrs["名称"] if name in existing_map: existing_props = existing_map[name] existing_attrs = { k: v for k, v in existing_props.items() if not k.endswith("_timeline") and k not in ("created_at", "last_updated") } resolved = self._resolve_attribute_conflicts(existing_attrs, clean_attrs) timeline_data = { k: v for k, v in existing_props.items() if k.endswith("_timeline") or k in ("created_at", "last_updated") } final = {**timeline_data, **resolved} final["last_updated"] = current_time to_update.append({"名称": name, "props": final}) else: new_attrs = dict(clean_attrs) new_attrs["created_at"] = current_time new_attrs["last_updated"] = current_time to_create.append(new_attrs) for i in range(0, len(to_create), BATCH_SIZE): batch = to_create[i: i + BATCH_SIZE] try: session.execute_write(self._batch_create_nodes, node_type, batch) entity_count += len(batch) except Exception as e: print(f"批量创建节点失败 ({node_type} 批次{i // BATCH_SIZE}): {e}") for i in range(0, len(to_update), BATCH_SIZE): batch = to_update[i: i + BATCH_SIZE] try: session.execute_write(self._batch_update_nodes, node_type, batch) entity_count += len(batch) except Exception as e: print(f"批量更新节点失败 ({node_type} 批次{i // BATCH_SIZE}): {e}") # ── 3. 预处理关系，按类型分组 ────────────────────────────── rels_by_type: Dict[str, List[dict]] = {} for rel in data.get("relationships", []): rel_type = rel.get("type", "") if not rel_type or not rel.get("from_entity") or not rel.get("to_entity"): continue rels_by_type.setdefault(rel_type, []).append(rel) # ── 4. 按类型批量读取已存在关系，Python侧计算 timeline，批量写入 ── with self.driver.session() as session: for rel_type, type_rels in rels_by_type.items(): pairs = [(r["from_entity"], r["to_entity"]) for r in type_rels] existing_list = session.execute_read( self._batch_fetch_relationships, rel_type, pairs ) existing_map = { (r["from_name"], r["to_name"]): r["props"] for r in existing_list } to_create_r: List[dict] = [] to_update_r: List[dict] = [] for rel in type_rels: from_name = rel["from_entity"] to_name = rel["to_entity"] attributes = rel.get("properties", {}) clean_attrs = {k: v for k, v in attributes.items() if v is not None and v != ""} clean_attrs["fact"] = self.construct_fact(from_name, to_name, rel_type, attributes) key = (from_name, to_name) if key in existing_map: existing_props = existing_map[key] updated_props = dict(existing_props) for k, new_value in clean_attrs.items(): if k == "fact": continue if k not in existing_props or existing_props[k] != new_value: tl_key = f"{k}_timeline" if tl_key not in updated_props: updated_props[tl_key] = [] entry: dict = { "value": new_value, "timestamp": current_time, "source": "text_extraction", } if k in existing_props: entry["previous_value"] = existing_props[k] updated_props[tl_key].append( json.dumps(entry, ensure_ascii=False) ) updated_props[k] = new_value fact_tl_key = "fact_timeline" if fact_tl_key not in updated_props: updated_props[fact_tl_key] = [] updated_props[fact_tl_key].append( json.dumps( {"value": clean_attrs["fact"], "timestamp": current_time, "source": "text_extraction"}, ensure_ascii=False, ) ) updated_props["fact"] = clean_attrs["fact"] updated_props["last_updated"] = current_time to_update_r.append( {"from_name": from_name, "to_name": to_name, "props": updated_props} ) else: rel_props = dict(clean_attrs) for k, v in clean_attrs.items(): rel_props[f"{k}_timeline"] = [ json.dumps( {"value": v, "timestamp": current_time, "source": "text_extraction"}, ensure_ascii=False, ) ] rel_props["created_at"] = current_time rel_props["last_updated"] = current_time rel_props["update_count"] = 1 to_create_r.append( {"from_name": from_name, "to_name": to_name, "props": rel_props} ) for i in range(0, len(to_create_r), BATCH_SIZE): batch = to_create_r[i: i + BATCH_SIZE] try: session.execute_write(self._batch_create_relationships, rel_type, batch) relationship_count += len(batch) except Exception as e: print(f"批量创建关系失败 ({rel_type} 批次{i // BATCH_SIZE}): {e}") for i in range(0, len(to_update_r), BATCH_SIZE): batch = to_update_r[i: i + BATCH_SIZE] try: session.execute_write(self._batch_update_relationships, rel_type, batch) relationship_count += len(batch) except Exception as e: print(f"批量更新关系失败 ({rel_type} 批次{i // BATCH_SIZE}): {e}") print(f"成功处理 {entity_count} 个节点和 {relationship_count} 个关系") return entity_count, relationship_count # ── 批量辅助方法 ──────────────────────────────────────────────── @staticmethod def _batch_fetch_nodes(tx, node_type: str, names: List[str]) -> List[dict]: """一次性读取指定类型下所有涉及名称的节点属性""" result = tx.run( f"MATCH (n:{node_type}) WHERE n.名称 IN $names RETURN properties(n) AS props", names=names, ) return [record["props"] for record in result] @staticmethod def _batch_create_nodes(tx, node_type: str, nodes: List[dict]): """批量创建节点（仅用于确认不存在的新节点）""" tx.run( f"UNWIND $nodes AS n CREATE (x:{node_type}) SET x = n", nodes=nodes, ) @staticmethod def _batch_update_nodes(tx, node_type: str, updates: List[dict]): """批量更新节点属性（完整覆盖，已在Python侧完成属性合并）""" tx.run( f"UNWIND $updates AS u " f"MATCH (n:{node_type} {{名称: u.名称}}) " f"SET n = u.props", updates=updates, ) @staticmethod def _batch_fetch_relationships(tx, rel_type: str, pairs: List[Tuple[str, str]]) -> List[dict]: """一次性读取指定类型下所有涉及节点对的关系属性""" result = tx.run( f"UNWIND $pairs AS pair " f"MATCH (a {{名称: pair[0]}})-[r:{rel_type}]->(b {{名称: pair[1]}}) " f"RETURN pair[0] AS from_name, pair[1] AS to_name, properties(r) AS props", pairs=[[p[0], p[1]] for p in pairs], ) return [ {"from_name": r["from_name"], "to_name": r["to_name"], "props": r["props"]} for r in result ] @staticmethod def _batch_create_relationships(tx, rel_type: str, items: List[dict]): """批量创建关系（仅用于确认不存在的新关系）""" tx.run( f"UNWIND $items AS item " f"MATCH (a {{名称: item.from_name}}), (b {{名称: item.to_name}}) " f"CREATE (a)-[r:{rel_type}]->(b) SET r = item.props", items=items, ) @staticmethod def _batch_update_relationships(tx, rel_type: str, items: List[dict]): """批量更新关系属性（完整覆盖，已在Python侧完成 timeline 合并）""" tx.run( f"UNWIND $items AS item " f"MATCH (a {{名称: item.from_name}})-[r:{rel_type}]->(b {{名称: item.to_name}}) " f"SET r = item.props", items=items, ) def get_relationship_history(self, from_name: str, to_name: str, rel_type: str) -> Optional[Dict]: """获取关系的完整历史记录""" with self.driver.session() as session: result = session.run(""" MATCH (a {名称: $from_name})-[r]->(b {名称: $to_name}) WHERE type(r) = $rel_type RETURN properties(r) as props """, from_name=from_name, to_name=to_name, rel_type=rel_type) record = result.single() return record["props"] if record else None def get_property_timeline(self, from_name: str, to_name: str, rel_type: str, property_name: str) -> Optional[List]: """获取特定属性的时间线（返回字符串列表，需 json.loads 解析）""" history = self.get_relationship_history(from_name, to_name, rel_type) if history: return history.get(f"{property_name}_timeline", []) return None def clear_database(self): """清空数据库（谨慎使用）""" with self.driver.session() as session: session.run("MATCH (n) DETACH DELETE n") print("数据库已清空") # if __name__ == "__main__": # # 连接本地 Neo4j（请确保服务已启动） # try: # URI = "bolt://192.168.0.46:57687" # USERNAME = "neo4j" # PASSWORD = "zdht123@" # with open("/storage01/home/hdf/project/wxkgrag/kgextract/output/output3/neo4j_export_merged1.json", "r", encoding="utf-8") as f: # data = json.load(f) # kg = MilitaryKnowledgeGraph( # uri="bolt://192.168.0.46:57687", # auth=("neo4j", "zdht123@"), # ) # entity_count, rel_count = kg.process_data(data) # finally: # # kg.clear_database() # 取消注释以清理测试数据 # kg.close()