#!/usr/bin/env python3 """ 按舰船号拆分 kg_snapshot.json、kg_output/、files/ 等目录 输出结果放在 OUTPUT_BASE_DIR 下，每艘舰一个子目录 文件名规则与舰船的对应关系只需修改下方「变量配置区」，逻辑不用动。 """ import json import os import re import shutil import sys from collections import defaultdict from typing import Dict, List, Optional, Set # 从统一配置文件读取路径 sys.path.insert(0, os.path.dirname(__file__)) from migration_config import DST_CODE_DIR, DST_FILES_DIR # ===================== 变量配置区（路径已从 migration_config.py 读取） ======================== # 输入路径 SNAPSHOT_PATH = os.path.join(DST_CODE_DIR, "kg_snapshot.json") # 全量快照 KG_OUTPUT_DIR = os.path.join(DST_CODE_DIR, "kg_output") # 抽取结果 JSON 目录 FILES_DIR = DST_FILES_DIR # 原始文档目录 UPLOADS_V1_DIR = os.path.join(DST_FILES_DIR, "uploads_v1") # 高亮 PDF 目录 # 输出根目录（每艘舰一个子目录） OUTPUT_BASE_DIR = os.path.join(DST_CODE_DIR, "migration_split") # -------------------------------------------------------------------------- # 舰船规则表：定义「哪些文件名属于哪艘舰」 # # 每条规则是一个 dict，支持以下匹配方式（可以组合，任意一条命中即匹配）： # prefix : 文件名以此开头（不含路径） # suffix : 文件名以此结尾（含扩展名） # contains : 文件名包含此字符串 # regex : 正则表达式，匹配整个文件名（不含路径） # # 示例： # {"prefix": "122-"} → 以 "122-" 开头 # {"prefix": "163-"} → 以 "163-" 开头 # {"contains": "发动机"} → 包含"发动机" # {"regex": r"^JZ\d+_.+"} → 正则匹配 # -------------------------------------------------------------------------- SHIP_RULES: Dict[str, List[Dict]] = { "101舰": [ {"prefix": "101-"}, ], "122舰": [ {"prefix": "122-"}, ], "163舰": [ {"prefix": "163-"}, ], # 新增舰船示例（取消注释并修改）： # "200舰": [ # {"prefix": "200-"}, # {"contains": "某特殊前缀"}, # ], } # 共享节点（同时属于多艘舰）的处理策略： # "all" : 每艘涉及的舰都保留这个节点（knowledge_source 保持完整） # "majority" : 归属来源文件更多的那艘舰，其他舰不保留 SHARED_NODE_POLICY = "all" # ========================================================================== def match_ship(filename: str) -> Optional[str]: """ 根据 SHIP_RULES 判断文件名属于哪艘舰。 返回舰船名，或 None（没有匹配）。 第一条命中的规则生效（规则之间互斥时注意顺序）。 """ bare = os.path.basename(filename) for ship, rules in SHIP_RULES.items(): for rule in rules: if "prefix" in rule and bare.startswith(rule["prefix"]): return ship if "suffix" in rule and bare.endswith(rule["suffix"]): return ship if "contains" in rule and rule["contains"] in bare: return ship if "regex" in rule and re.match(rule["regex"], bare): return ship return None # -------------------------------------------------------------------------- # 快照拆分 # -------------------------------------------------------------------------- def _get_node_ships(entities: List[Dict]) -> Dict[str, Set[str]]: """返回 {节点名: {归属舰集合}}""" node_ships: Dict[str, Set[str]] = defaultdict(set) for e in entities: name = e.get("properties", {}).get("名称", "") if not name: continue ks_raw = e.get("properties", {}).get("knowledge_source", "") try: ks_list = json.loads(ks_raw) if isinstance(ks_raw, str) else (ks_raw or []) except Exception: ks_list = [] for ks in ks_list: ship = match_ship(ks.get("filename", "")) if ship: node_ships[name].add(ship) return node_ships def split_snapshot(snapshot_path: str, output_base: str): print(f"\n[快照拆分] 读取 {snapshot_path} ...") with open(snapshot_path, "r", encoding="utf-8") as f: data = json.load(f) entities = data.get("entities", []) relationships = data.get("relationships", []) print(f" 总节点: {len(entities)} 总关系: {len(relationships)}") node_ships = _get_node_ships(entities) # 统计共享情况 all_ships = set(SHIP_RULES.keys()) stats = defaultdict(int) for name, ships in node_ships.items(): if len(ships) == 0: stats["无来源"] += 1 elif len(ships) == 1: stats[next(iter(ships))] += 1 else: stats["共享"] += 1 print(f" 节点归属统计: { dict(stats) }") # 为每艘舰生成节点/关系子集 for ship in SHIP_RULES: ship_entities = [] for e in entities: name = e.get("properties", {}).get("名称", "") ships = node_ships.get(name, set()) if ship in ships: ship_entities.append(e) elif not ships and SHARED_NODE_POLICY == "all": # 无来源文件的节点（如顶层舰艇/系统节点）每艘舰都保留 ship_entities.append(e) # 只保留两端节点都在本舰集合里的关系 ship_node_names = {e.get("properties", {}).get("名称", "") for e in ship_entities} ship_rels = [ r for r in relationships if r.get("from_entity") in ship_node_names and r.get("to_entity") in ship_node_names ] out_dir = os.path.join(output_base, ship) os.makedirs(out_dir, exist_ok=True) out_path = os.path.join(out_dir, "kg_snapshot.json") with open(out_path, "w", encoding="utf-8") as f: json.dump({"entities": ship_entities, "relationships": ship_rels}, f, ensure_ascii=False, indent=2) print(f" [{ship}] 节点:{len(ship_entities)} 关系:{len(ship_rels)} → {out_path}") # -------------------------------------------------------------------------- # kg_output 拆分（每个 .json 文件对应一个原始文档） # -------------------------------------------------------------------------- def split_kg_output(kg_output_dir: str, output_base: str): if not os.path.isdir(kg_output_dir): print(f"\n[kg_output 拆分] 目录不存在，跳过: {kg_output_dir}") return print(f"\n[kg_output 拆分] 扫描 {kg_output_dir} ...") files = [f for f in os.listdir(kg_output_dir) if f.endswith(".json")] print(f" 共 {len(files)} 个文件") counts = defaultdict(int) unmatched = [] for fname in files: # kg_output 文件名格式：原始文件名.json（含扩展名，如 122-xxx.pdf.json） # 去掉最后的 .json 拿到原始文件名 original = fname[:-5] if fname.endswith(".json") else fname ship = match_ship(original) if ship: dst_dir = os.path.join(output_base, ship, "kg_output") os.makedirs(dst_dir, exist_ok=True) shutil.copy2(os.path.join(kg_output_dir, fname), os.path.join(dst_dir, fname)) counts[ship] += 1 else: unmatched.append(fname) for ship, cnt in counts.items(): print(f" [{ship}] {cnt} 个 JSON") if unmatched: print(f" [未匹配] {len(unmatched)} 个文件: {unmatched[:5]}") # -------------------------------------------------------------------------- # files/ 原始文档拆分 # -------------------------------------------------------------------------- def split_files(files_dir: str, output_base: str): if not os.path.isdir(files_dir): print(f"\n[files 拆分] 目录不存在，跳过: {files_dir}") return print(f"\n[files 拆分] 扫描 {files_dir} ...") all_files = [f for f in os.listdir(files_dir) if os.path.isfile(os.path.join(files_dir, f))] print(f" 共 {len(all_files)} 个文件") counts = defaultdict(int) unmatched = [] for fname in all_files: # files/ 里有两种格式： # UUID_原文件名.pdf → 去掉 UUID 前缀再匹配 # 原文件名.pdf → 直接匹配 bare = fname # 尝试去掉 UUID 前缀（格式：8-4-4-4-12 字符的十六进制） uuid_pattern = r'^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}_(.+)$' m = re.match(uuid_pattern, fname, re.IGNORECASE) if m: bare = m.group(1) ship = match_ship(bare) if ship: dst_dir = os.path.join(output_base, ship, "files") os.makedirs(dst_dir, exist_ok=True) shutil.copy2(os.path.join(files_dir, fname), os.path.join(dst_dir, fname)) counts[ship] += 1 else: unmatched.append(fname) for ship, cnt in counts.items(): print(f" [{ship}] {cnt} 个文件") if unmatched: print(f" [未匹配] {len(unmatched)} 个: {unmatched[:5]}") # -------------------------------------------------------------------------- # uploads_v1/ 高亮 PDF 拆分（通过快照里的 knowledge_source 反查归属） # -------------------------------------------------------------------------- def split_uploads(snapshot_path: str, uploads_dir: str, output_base: str): if not os.path.isdir(uploads_dir): print(f"\n[uploads 拆分] 目录不存在，跳过: {uploads_dir}") return print(f"\n[uploads 拆分] 通过快照反查高亮 PDF 归属 ...") with open(snapshot_path, "r", encoding="utf-8") as f: data = json.load(f) # 从快照里收集：高亮PDF文件名 → 归属舰船 highlight_ship: Dict[str, str] = {} for e in data.get("entities", []): ks_raw = e.get("properties", {}).get("knowledge_source", "") try: ks_list = json.loads(ks_raw) if isinstance(ks_raw, str) else (ks_raw or []) except Exception: continue for ks in ks_list: url = ks.get("url", "") src_file = ks.get("filename", "") ship = match_ship(src_file) if not ship or not url: continue # URL 格式：http://host/upload/#page=N if "/upload/" in url: highlight_fname = url.split("/upload/")[-1].split("#")[0] # 共享节点下同一个高亮文件只记录第一个匹配舰船 if highlight_fname not in highlight_ship: highlight_ship[highlight_fname] = ship counts = defaultdict(int) missing = [] for fname, ship in highlight_ship.items(): src = os.path.join(uploads_dir, fname) if os.path.exists(src): dst_dir = os.path.join(output_base, ship, "files", "uploads_v1") os.makedirs(dst_dir, exist_ok=True) shutil.copy2(src, os.path.join(dst_dir, fname)) counts[ship] += 1 else: missing.append(fname) for ship, cnt in counts.items(): print(f" [{ship}] {cnt} 个高亮 PDF") if missing: print(f" [本地缺失] {len(missing)} 个（需从服务器 uploads_v1 补全）: {missing[:3]}") # -------------------------------------------------------------------------- # 入口 # -------------------------------------------------------------------------- def main(): print("=" * 50) print(" kgrag 按舰船号拆分工具") print(f" 舰船规则: { {k: v for k, v in SHIP_RULES.items()} }") print(f" 共享节点策略: {SHARED_NODE_POLICY}") print("=" * 50) os.makedirs(OUTPUT_BASE_DIR, exist_ok=True) split_snapshot(SNAPSHOT_PATH, OUTPUT_BASE_DIR) split_kg_output(KG_OUTPUT_DIR, OUTPUT_BASE_DIR) split_files(FILES_DIR, OUTPUT_BASE_DIR) split_uploads(SNAPSHOT_PATH, UPLOADS_V1_DIR, OUTPUT_BASE_DIR) print("\n" + "=" * 50) print(f" 完成！输出目录: {OUTPUT_BASE_DIR}") print(" 每艘舰的子目录结构：") print(" <舰名>/kg_snapshot.json ← 导入新 Neo4j 用") print(" <舰名>/kg_output/ ← 抽取结果 JSON") print(" <舰名>/files/ ← 原始文档") print(" <舰名>/files/uploads_v1/ ← 高亮 PDF") print("=" * 50) if __name__ == "__main__": main()