添加GraphRAG后端
This commit is contained in:
龙澳
197
backend/graph_builder.py
Normal file
197
backend/graph_builder.py
Normal file
@@ -0,0 +1,197 @@
|
||||
"""
|
||||
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.")
|
||||
Reference in New Issue
Block a user