500 lines
18 KiB
Python
500 lines
18 KiB
Python
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 _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"]
|
||
|
||
# 保留时间线和时间戳字段
|
||
timeline_data = {k: v for k, v in existing_props.items()
|
||
if k.endswith('_timeline') or k in ['created_at', 'last_updated']}
|
||
|
||
# 直接使用新属性（不调用 _resolve_attribute_conflicts）
|
||
final_attributes = {**timeline_data, **clean_attributes}
|
||
# 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
|
||
|
||
Args:
|
||
data: 包含 entities 和 relationships 的字典
|
||
|
||
Returns:
|
||
Tuple[int, int]: (实体数量, 关系数量)
|
||
"""
|
||
if not data or not isinstance(data, dict):
|
||
print("无效的数据格式，需要字典类型")
|
||
return 0, 0
|
||
|
||
entity_count = 0
|
||
relationship_count = 0
|
||
|
||
with self.driver.session() as session:
|
||
# 创建/更新实体节点
|
||
for entity in data.get("entities", []):
|
||
try:
|
||
# 提取 properties，并注入 knowledge_source（如果存在）
|
||
props = entity.get("properties", {}).copy()
|
||
if "knowledge_source" in entity:
|
||
props["knowledge_source"] = entity["knowledge_source"]
|
||
|
||
# 从 properties 中提取名称
|
||
名称 = props.pop("名称", None)
|
||
if not 名称:
|
||
print(f"跳过无名称节点: {entity}")
|
||
continue
|
||
|
||
success = session.execute_write(
|
||
self._create_or_update_node,
|
||
entity["type"],
|
||
名称,
|
||
props
|
||
)
|
||
if success:
|
||
entity_count += 1
|
||
except Exception as e:
|
||
print(f"处理节点失败: {entity}, 错误: {e}")
|
||
|
||
# 创建/更新关系
|
||
for relationship in data.get("relationships", []):
|
||
try:
|
||
success = session.execute_write(
|
||
self._create_relationship_with_timeline,
|
||
relationship["from_entity"],
|
||
relationship["to_entity"],
|
||
relationship["type"],
|
||
relationship.get("properties", {})
|
||
)
|
||
if success:
|
||
relationship_count += 1
|
||
except Exception as e:
|
||
print(f"处理关系失败: {relationship}, 错误: {e}")
|
||
|
||
print(f"成功处理 {entity_count} 个节点和 {relationship_count} 个关系")
|
||
return entity_count, relationship_count
|
||
|
||
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("数据库已清空")
|
||
|
||
|
||
# 测试数据
|
||
test_data = {
|
||
"entities": [
|
||
{
|
||
"type": "系统",
|
||
|
||
"properties": {
|
||
"名称": "ZZ系统",
|
||
"系统编码": "04",
|
||
"knowledge_source": [
|
||
{
|
||
"filename": "",
|
||
"url": "",
|
||
"info": ""
|
||
}
|
||
],
|
||
}
|
||
},
|
||
{
|
||
"type": "子系统",
|
||
"properties": {
|
||
"名称": "作战系统总论",
|
||
"系统编码": "00",
|
||
"knowledge_source": [
|
||
{
|
||
"filename": "",
|
||
"url": "",
|
||
"info": ""
|
||
}
|
||
],
|
||
}
|
||
},
|
||
{
|
||
"type": "系统",
|
||
"properties": {
|
||
"名称": "DL系统",
|
||
"系统编码": "05",
|
||
"knowledge_source": [
|
||
{
|
||
"filename": "",
|
||
"url": "",
|
||
"info": ""
|
||
}
|
||
],
|
||
}
|
||
},
|
||
{
|
||
"type": "子系统",
|
||
"properties": {
|
||
"名称": "DL子系统",
|
||
"系统编码": "B0",
|
||
"knowledge_source": [
|
||
{
|
||
"filename": "",
|
||
"url": "",
|
||
"info": ""
|
||
}
|
||
],
|
||
}
|
||
},
|
||
{
|
||
"type": "系统",
|
||
"properties": {
|
||
"名称": "HJ系统",
|
||
"系统编码": "B0",
|
||
"knowledge_source": [
|
||
{
|
||
"filename": "",
|
||
"url": "",
|
||
"info": ""
|
||
}
|
||
],
|
||
}
|
||
},
|
||
{
|
||
"type": "子系统",
|
||
"properties": {
|
||
"名称": "HJ子系统",
|
||
"系统编码": "08",
|
||
"knowledge_source": [
|
||
{
|
||
"filename": "",
|
||
"url": "",
|
||
"info": ""
|
||
}
|
||
],
|
||
}
|
||
},
|
||
],
|
||
"relationships": [
|
||
{
|
||
"type": "包含",
|
||
"from_entity": "北京舰",
|
||
"to_entity": "ZZ系统"
|
||
},
|
||
{
|
||
"type": "包含",
|
||
"from_entity": "北京舰",
|
||
"to_entity": "DL系统"
|
||
},
|
||
{
|
||
"type": "包含",
|
||
"from_entity": "ZZ系统",
|
||
"to_entity": "作战系统总论"
|
||
},
|
||
{
|
||
"type": "包含",
|
||
"from_entity": "DL系统",
|
||
"to_entity": "DL子系统"
|
||
},
|
||
{
|
||
"type": "包含",
|
||
"from_entity": "南京舰",
|
||
"to_entity": "HJ系统"
|
||
},
|
||
{
|
||
"type": "包含",
|
||
"from_entity": "HJ系统",
|
||
"to_entity": "HJ子系统"
|
||
}
|
||
]
|
||
}
|
||
|
||
if __name__ == "__main__":
|
||
# 连接本地 Neo4j（请确保服务已启动）
|
||
kg = MilitaryKnowledgeGraph(
|
||
uri="bolt://192.168.0.46:57687",
|
||
auth=("neo4j", "zdht123@"),
|
||
Ontology=[
|
||
"舰艇", "系统", "子系统", "设备", "维修项目", "维修工作", "操作项目",
|
||
"故障模式", "安全警告", "工作原理",
|
||
"组成", "技术指标参数", "布置情况", "保障设施保障设备工具信息汇总表",
|
||
"消耗材料信息汇总表", "维修安全防护及注意事项汇总表", "图册", "安全使用操作",
|
||
"调试", "维护保养", "常见问题分析及排除", "功能", "接口",
|
||
"性能指标", "操作使用","环境条件","设计单位","操作程序","故障现象","维修","器材"
|
||
]
|
||
)
|
||
|
||
try:
|
||
print("开始写入知识图谱...")
|
||
entity_count, rel_count = kg.process_data(test_data)
|
||
print(f"✅ 写入完成：{entity_count} 个实体，{rel_count} 个关系")
|
||
|
||
# 验证读取
|
||
history = kg.get_relationship_history("发动机", "喷油器测试台", "包含")
|
||
if history:
|
||
print("\n🔍 关系 fact:", history.get("fact"))
|
||
|
||
# 检查节点属性
|
||
with kg.driver.session() as session:
|
||
result = session.run("""
|
||
MATCH (e:设备 {名称: '发动机'})
|
||
RETURN e.名称 AS name, e.型号 AS model, e.knowledge_source AS ks
|
||
""")
|
||
record = result.single()
|
||
if record:
|
||
print(f"\n📌 节点验证：{record['name']}, 型号={record['model']}")
|
||
if record['ks']:
|
||
ks_data = json.loads(record['ks'])
|
||
print(f"📄 来源文件: {ks_data[0]['filename']}")
|
||
else:
|
||
print("❌ 未找到节点！")
|
||
|
||
finally:
|
||
# kg.clear_database() # 取消注释以清理测试数据
|
||
kg.close()
|
||
|