Files
Novel-Map/backend/graph_builder.py
2026-03-31 17:18:30 +08:00

198 lines
6.9 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""
JSON → Neo4j 导入脚本。
图谱 Schema:
节点: Character, Location, Faction, Event
关系: VISITED, CONTROLS, HAS_MEMBER, LEADS, OCCURRED_AT
"""
import json
from pathlib import Path
from neo4j import Driver
DATA_DIR = Path(__file__).parent.parent / "data"
# ── 工具函数 ──────────────────────────────────────────────
def _split_characters(name: str) -> list[str]:
"""'寇仲 & 徐子陵' → ['寇仲', '徐子陵']"""
return [c.strip() for c in name.split("&") if c.strip()]
def _split_leaders(leader: str) -> list[str]:
"""'翟让/李密' → ['翟让', '李密'];过滤'未提及'"""
parts = [p.strip() for p in leader.split("/") if p.strip()]
return [p for p in parts if p not in ("未提及", "")]
# ── Schema 初始化 ─────────────────────────────────────────
def setup_schema(driver: Driver):
with driver.session() as s:
s.run("CREATE CONSTRAINT IF NOT EXISTS FOR (n:Character) REQUIRE n.name IS UNIQUE")
s.run("CREATE CONSTRAINT IF NOT EXISTS FOR (n:Location) REQUIRE n.id IS UNIQUE")
s.run("CREATE CONSTRAINT IF NOT EXISTS FOR (n:Faction) REQUIRE n.id IS UNIQUE")
s.run("CREATE CONSTRAINT IF NOT EXISTS FOR (n:Event) REQUIRE n.id IS UNIQUE")
s.run("CREATE INDEX IF NOT EXISTS FOR (e:Event) ON (e.vol)")
s.run("CREATE INDEX IF NOT EXISTS FOR ()-[r:VISITED]-() ON (r.vol)")
s.run("CREATE INDEX IF NOT EXISTS FOR ()-[r:CONTROLS]-() ON (r.vol)")
# ── 各类型导入 ────────────────────────────────────────────
def _import_locations(session, locations: list[dict]):
for loc in locations:
session.run(
"""
MERGE (l:Location {id: $id})
SET l.name = $name,
l.type = $type,
l.lat = $lat,
l.lng = $lng
""",
id=loc["id"],
name=loc["name"],
type=loc.get("type", ""),
lat=loc.get("lat"),
lng=loc.get("lng"),
)
def _import_factions(session, factions: list[dict], vol: int):
for f in factions:
session.run(
"""
MERGE (n:Faction {id: $id})
SET n.name = $name, n.type = $type, n.color = $color
""",
id=f["id"], name=f["name"],
type=f.get("type", ""), color=f.get("color", ""),
)
# Faction → CONTROLS → Location
for loc_id in f.get("territory", []):
session.run(
"""
MATCH (fac:Faction {id: $fid})
MATCH (loc:Location {id: $lid})
MERGE (fac)-[:CONTROLS {vol: $vol}]->(loc)
""",
fid=f["id"], lid=loc_id, vol=vol,
)
# Faction → HAS_MEMBER → Character
for figure in f.get("key_figures", []):
if not figure:
continue
session.run(
"""
MERGE (c:Character {name: $name})
WITH c
MATCH (fac:Faction {id: $fid})
MERGE (fac)-[:HAS_MEMBER {vol: $vol}]->(c)
""",
name=figure, fid=f["id"], vol=vol,
)
# Character → LEADS → Faction
for leader_name in _split_leaders(f.get("leader", "")):
session.run(
"""
MERGE (c:Character {name: $name})
WITH c
MATCH (fac:Faction {id: $fid})
MERGE (c)-[:LEADS {vol: $vol}]->(fac)
""",
name=leader_name, fid=f["id"], vol=vol,
)
def _import_routes(session, routes: list[dict], vol: int):
for route in routes:
char_color = route.get("color", "")
char_names = _split_characters(route["character"])
for char_name in char_names:
session.run(
"MERGE (c:Character {name: $name}) SET c.color = $color",
name=char_name, color=char_color,
)
for wp in route.get("route", []):
loc_id = wp.get("location")
if not loc_id:
continue # lat/lng only → 跳过(无命名地点节点)
chapter = wp.get("chapter", 0)
event = wp.get("event", "")
session.run(
"""
MATCH (c:Character {name: $char})
MATCH (l:Location {id: $lid})
MERGE (c)-[v:VISITED {vol: $vol, chapter: $chapter}]->(l)
SET v.event = $event
""",
char=char_name, lid=loc_id,
vol=vol, chapter=chapter, event=event,
)
def _import_events(session, events: list[dict], vol: int):
for i, evt in enumerate(events):
event_id = f"v{vol:02d}_e{i:03d}"
chapter = evt.get("chapter", 0)
description = evt.get("event", "")
session.run(
"""
MERGE (e:Event {id: $id})
SET e.vol = $vol, e.chapter = $chapter, e.description = $description
""",
id=event_id, vol=vol, chapter=chapter, description=description,
)
# 只在有命名地点 id 时建立关系lat/lng 条目跳过)
loc_ref = evt.get("location")
if isinstance(loc_ref, str) and loc_ref:
session.run(
"""
MATCH (e:Event {id: $eid})
MATCH (l:Location {id: $lid})
MERGE (e)-[:OCCURRED_AT]->(l)
""",
eid=event_id, lid=loc_ref,
)
# ── 主入口 ────────────────────────────────────────────────
def build_graph(driver: Driver, clear: bool = False):
if clear:
print("Clearing existing graph data...")
with driver.session() as s:
s.run("MATCH (n) DETACH DELETE n")
print("Setting up schema constraints and indexes...")
setup_schema(driver)
imported = 0
for vol_num in range(1, 64):
filepath = DATA_DIR / f"vol{vol_num:02d}.json"
if not filepath.exists():
continue
with open(filepath, encoding="utf-8") as f:
data = json.load(f)
with driver.session() as session:
_import_locations(session, data.get("locations", []))
_import_factions(session, data.get("factions", []), vol_num)
_import_routes(session, data.get("character_routes", []), vol_num)
_import_events(session, data.get("key_events", []), vol_num)
imported += 1
print(f" [✓] vol{vol_num:02d} imported")
print(f"\nDone. Imported {imported} volumes.")