409 lines
16 KiB
Python
409 lines
16 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"]
|
||
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": "设备",
|
||
# "knowledge_source": [
|
||
# {
|
||
# "filename": "122-06A0014-B01001_发动机-维修手册.pdf",
|
||
# "url": "http://192.168.0.111:8080/api/v1/files/9ac5c3a0-860c-4d30-a199-8daf51bf189b/content#page=5",
|
||
# "info": "主设备，维修对象"
|
||
# }
|
||
# ],
|
||
# "properties": {
|
||
# "名称": "发动机",
|
||
# "型号": "12V280ZJ",
|
||
# "技术责任单位": "中船动力研究院",
|
||
# "生产单位": "沪东重机",
|
||
# "运行条件": "适用于船舶推进系统，额定转速1000 rpm，最大持续功率5200 kW"
|
||
# }
|
||
# },
|
||
# {
|
||
# "type": "设备",
|
||
# "knowledge_source": [
|
||
# {
|
||
# "filename": "122-06A0014-B01001_发动机-维修手册.pdf",
|
||
# "url": "http://192.168.0.111:8080/api/v1/files/9ac5c3a0-860c-4d30-a199-8daf51bf189b/content#page=6",
|
||
# "info": "用于喷油器性能测试的专用设备"
|
||
# }
|
||
# ],
|
||
# "properties": {
|
||
# "名称": "喷油器测试台",
|
||
# "型号": "PTB-300",
|
||
# "生产单位": "沪东重机"
|
||
# }
|
||
# }
|
||
# ],
|
||
# "relationships": [
|
||
# {
|
||
# "type": "包含",
|
||
# "from_entity": "发动机",
|
||
# "to_entity": "喷油器测试台"
|
||
# }
|
||
# ]
|
||
# }
|
||
|
||
# if __name__ == "__main__":
|
||
# # 连接本地 Neo4j（请确保服务已启动）
|
||
# kg = MilitaryKnowledgeGraph(
|
||
# uri="bolt://localhost:7687",
|
||
# auth=("neo4j", "123456hdf"),
|
||
# 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()
|