"""把 _002_ 抽好的中药图谱 / 方剂图谱导入 Neo4j。
读 _002__extract_information/{中药图谱,方剂图谱}/*.json 里的 entities 和 relations，
在 Python 里拼 Cypher，走 common/neo4j_manager.py 的 Neo4jClient 写库。
几个关键取舍：
- **先在内存里去重，再写库。** 49226 条 entity 里只有 18870 个不同节点，50627 条
relation 里只有 48346 条不同边 —— 同一个「咳嗽」在几百个文件里都出现。逐条 MERGE
等于把同样的活干几百遍，所以先聚合成唯一集合再写。
- **节点主键是 (label, name)，不是 name。** 有 356 个名字同时挂在两种类型下
（「热咳」既是 Disease 又是 Symptom），只按 name 去重会把两类实体粘成一个节点。
- **label 和关系类型只能拼进语句，不能当参数**，Cypher 不支持参数化它们。所以这两个
位置都先对白名单做校验再拼字符串，值一律走 $ 参数，不做字符串插值。
- 属性以「实体自己那个文件」为准：一贯煎.json 里的「当归」只是配料，没有性味剂量，
当归.json 里才是全的。同名属性冲突时后者覆盖前者，其余只补空缺。
用法（在项目根目录执行）：
python -X utf8 -m _003_create_neo4j_database.01_import_graph_to_neo4j --limit 5
python -X utf8 -m _003_create_neo4j_database.01_import_graph_to_neo4j --dry-run
python -X utf8 -m _003_create_neo4j_database.01_import_graph_to_neo4j
python -X utf8 -m _003_create_neo4j_database.01_import_graph_to_neo4j --only formula
"""
import argparse
import json
import os
from collections import Counter
from tqdm import tqdm
from common.config import Config
from common.neo4j_manager import Neo4jClient
from common.path_utils import get_file_path
# 白名单 = 唯一允许拼进 Cypher 的 label / 关系类型，和 _002_ 的 Pydantic schema 一致。
# 抽取结果里出现过 TREATS、TREATS_DISE 这种模型笔误，不在名单里就丢掉并记账。
ENTITY_TYPES = {"Symptom", "Disease", "Formula", "Herb", "Effect", "Source"}
RELATION_TYPES = {
"TREATS_DISEASE",
"ALLEVIATES_SYMPTOM",
"HAS_EFFECT",
"HAS_INGREDIENT",
"HAS_SYMPTOM",
"FROM_SOURCE",
}
# FormulaAttributes + HerbAttributes 的全部字段。属性名同样会拼进 SET，所以也过白名单，
# 免得模型多吐一个键就在库里长出一个野属性。
ATTR_KEYS = {
"alias",
"dosage",
"effect",
"indication",
"meridian",
"origin",
"place",
"processing",
"property_flavor",
"taboo",
"traits",
"usage",
}
GRAPH_DIRS = {
"herb": "_002__extract_information/中药图谱",
"formula": "_002__extract_information/方剂图谱",
}
# 一个事务里塞多少条语句。太小则往返次数多，太大则事务占内存、失败回滚的代价高。
BATCH_SIZE = 500
# ======================
# 读盘 + 聚合
# ======================
def iter_graph_files(only=None, limit=0):
"""按 (来源, 路径) 产出要导入的图谱 json 路径。"""
paths = []
for kind, rel_dir in GRAPH_DIRS.items():
if only and kind != only:
continue
directory = get_file_path(rel_dir)
if not os.path.isdir(directory):
print(f"图谱目录不存在，跳过: {directory}")
continue
names = sorted(f for f in os.listdir(directory) if f.endswith(".json"))
if limit:
names = names[:limit]
paths.extend((kind, os.path.join(directory, f)) for f in names)
return paths
def clean_attributes(raw):
"""过滤属性字典：只留白名单里的键，丢掉空值，值统一成字符串。"""
if not isinstance(raw, dict):
return {}
out = {}
for key, value in raw.items():
if key not in ATTR_KEYS or value is None:
continue
text = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False)
text = text.strip()
if text:
out[key] = text
return out
def collect(paths, stats):
"""把所有图谱 json 聚合成「唯一节点」和「唯一关系」两个集合。
nodes: {(label, name): {属性}}
rels: {(主label, 主name, 关系, 宾label, 宾name)}
关系两端的节点也顺手补进 nodes —— 有 177 处 relation 的端点没在同文件的 entities
里列出来，不补的话 MATCH 匹配不上，那条边就凭空丢了。
"""
nodes = {}
rels = set()
for kind, path in tqdm(paths, desc="读取图谱", unit="个", ncols=90):
try:
with open(path, encoding="utf-8") as f:
graph = json.load(f)
except (OSError, json.JSONDecodeError):
stats["json 读取失败"] += 1
continue
# 主体实体的属性最全（同名的配料条目往往只有名字），最后合并时优先它
main_name = graph.get("name")
for entity in graph.get("entities") or []:
label = entity.get("type")
name = (entity.get("name") or "").strip()
if label not in ENTITY_TYPES or not name:
stats["实体不合法丢弃"] += 1
continue
key = (label, name)
attrs = clean_attributes(entity.get("attributes"))
existing = nodes.setdefault(key, {})
if name == main_name:
# 自己那个文件里的定义最权威，直接盖掉别处带来的零碎属性
existing.update(attrs)
else:
# 只补空缺，不覆盖已有的
for k, v in attrs.items():
existing.setdefault(k, v)
for relation in graph.get("relations") or []:
rel_type = relation.get("relation")
s_label, o_label = relation.get("subject_type"), relation.get("object_type")
s_name = (relation.get("subject") or "").strip()
o_name = (relation.get("object") or "").strip()
if rel_type not in RELATION_TYPES:
stats["关系类型不合法丢弃"] += 1
continue
if s_label not in ENTITY_TYPES or o_label not in ENTITY_TYPES:
stats["关系端点类型不合法丢弃"] += 1
continue
if not s_name or not o_name:
stats["关系端点为空丢弃"] += 1
continue
nodes.setdefault((s_label, s_name), {})
nodes.setdefault((o_label, o_name), {})
rels.add((s_label, s_name, rel_type, o_label, o_name))
return nodes, rels
# ======================
# 拼 Cypher
# ======================
def build_constraint_cypher():
"""每类实体按 name 建唯一约束。
顺带把 MERGE 需要的索引也建好了（约束自带索引），不然 18870 次 MERGE 都是全表扫。
IF NOT EXISTS 让脚本可以重复跑。
"""
return [
(
f"CREATE CONSTRAINT {label.lower()}_name IF NOT EXISTS "
f"FOR (n:{label}) REQUIRE n.name IS UNIQUE",
None,
)
for label in sorted(ENTITY_TYPES)
]
def build_node_cypher(nodes):
"""一个节点一条 MERGE。
label 只能拼进语句（Cypher 不支持参数化 label），name 和属性全走 $ 参数。
label 已经过了 ENTITY_TYPES 白名单，不存在注入面。
属性用 SET n += $attrs 而不是 MERGE (n {name:..., effect:...})：属性写进 MERGE
的话，属性一变就 MERGE 不上、变成新建一个节点。ON CREATE SET 也不合适 —— 重跑或
补数据时已存在的节点就永远补不上属性了。
"""
queries = []
for (label, name), attrs in sorted(nodes.items()):
if attrs:
query = f"MERGE (n:{label} {{name: $name}}) SET n += $attrs"
params = {"name": name, "attrs": attrs}
else:
query = f"MERGE (n:{label} {{name: $name}})"
params = {"name": name}
queries.append((query, params))
return queries
def build_relation_cypher(rels):
"""一条关系一条语句：两端 MATCH（节点已经在上一步 MERGE 好了）再 MERGE 边。
关系类型和两端 label 都拼进语句，同样先过了白名单。MERGE 而不是 CREATE，
保证重复跑不会长出平行边。
两个 MATCH 中间夹一个 WITH，而不是写成 `MATCH (s), (o)`：后者是一条语句里两个不
相连的模式，Neo4j 每条都会回一个 cartesian product 警告，5 万条刷得日志没法看。
这里两端都是唯一约束上的等值查找，笛卡尔积本来就是 1×1，换个写法只是让驱动闭嘴。
"""
queries = []
for s_label, s_name, rel_type, o_label, o_name in sorted(rels):
query = (
f"MATCH (s:{s_label} {{name: $s_name}}) "
f"WITH s MATCH (o:{o_label} {{name: $o_name}}) "
f"MERGE (s)-[:{rel_type}]->(o)"
)
queries.append((query, {"s_name": s_name, "o_name": o_name}))
return queries
# ======================
# 写库
# ======================
def run_in_batches(client, queries, desc):
"""分批提交。
run_multiple_cypher 整批是一个事务，5 万条一次性提交内存吃不消、失败也全回滚，
所以切成 BATCH_SIZE 一批。批与批之间不是原子的，但所有语句都是 MERGE，
中断后重跑是幂等的，无所谓。
"""
total = len(queries)
done = 0
for start in range(0, total, BATCH_SIZE):
batch = queries[start : start + BATCH_SIZE]
client.run_multiple_cypher(batch)
done += len(batch)
print(f" {desc}: {done}/{total}")
return done
def run(only=None, limit=0, dry_run=False, wipe=False):
paths = iter_graph_files(only=only, limit=limit)
if not paths:
print("没有找到图谱 json，先跑 _002_ 的抽取脚本。")
return
stats = Counter()
nodes, rels = collect(paths, stats)
label_count = Counter(label for label, _ in nodes)
rel_count = Counter(r[2] for r in rels)
print(f"\n图谱文件 {len(paths)} 个")
print(f"唯一节点 {len(nodes)} 个: {dict(label_count.most_common())}")
print(f"唯一关系 {len(rels)} 条: {dict(rel_count.most_common())}")
for key, value in stats.most_common():
print(f" {key}: {value}")
node_queries = build_node_cypher(nodes)
rel_queries = build_relation_cypher(rels)
if dry_run:
print("\n--dry-run，不写库。示例语句：")
for query, params in node_queries[:2] + rel_queries[:2]:
print(f" {query}\n {params}")
return
conf = Config()
client = Neo4jClient(conf.NEO4J_URI, conf.NEO4J_USER, conf.NEO4J_PASSWORD)
if wipe:
# 只删这个图谱的六类节点，不动库里别的东西。DETACH 连边一起删。
print("\n清空旧数据...")
for label in sorted(ENTITY_TYPES):
client.run_cypher(f"MATCH (n:{label}) DETACH DELETE n")
print("\n建唯一约束...")
client.run_multiple_cypher(build_constraint_cypher())
print("写入节点...")
run_in_batches(client, node_queries, "节点")
print("写入关系...")
run_in_batches(client, rel_queries, "关系")
# 从库里回读一遍计数，别只信自己发了多少条语句
print("\n库内实际计数：")
for label in sorted(ENTITY_TYPES):
rows = client.run_cypher(f"MATCH (n:{label}) RETURN count(n) AS c")
print(f" {label}: {rows[0]['c']}")
for rel_type in sorted(RELATION_TYPES):
rows = client.run_cypher(f"MATCH ()-[r:{rel_type}]->() RETURN count(r) AS c")
print(f" {rel_type}: {rows[0]['c']}")
def main():
parser = argparse.ArgumentParser(description="把 _002_ 抽取的图谱导入 Neo4j")
parser.add_argument(
"--only",
choices=sorted(GRAPH_DIRS),
help="只导入一边：herb=中药图谱，formula=方剂图谱，默认两边都导",
)
parser.add_argument(
"--limit",
type=int,
default=0,
help="每个目录只取前 N 个 json 试跑，默认 0 表示全量",
)
parser.add_argument(
"--dry-run",
action="store_true",
help="只统计并打印示例 Cypher，不连库不写入",
)
parser.add_argument(
"--wipe",
action="store_true",
help="导入前先 DETACH DELETE 这六类节点（会删掉已有图谱数据）",
)
args = parser.parse_args()
run(only=args.only, limit=args.limit, dry_run=args.dry_run, wipe=args.wipe)
if __name__ == "__main__":
main()