Agent/knowhub/scripts/merge_capabilities.py

#!/usr/bin/env python3
"""
合并 tao_dev_1 版本内的重复 capability，并删除 VCAP 占位条目。

策略：
  1. 对每个 (canonical, members) 簇：
     对 5 个 junction (requirement_capability / capability_tool /
     capability_knowledge / capability_resource / strategy_capability)
     - 先删除会和 canonical 冲突的 member 引用（避免 PK 冲突）
     - 再 UPDATE member → canonical
  2. 删除被合并的 member 行 + VCAP 占位行
  3. 全程打印进度（flush=True），不管道给 tail

用法：
  python knowhub/scripts/merge_capabilities.py --dry-run   # 只打印计划
  python knowhub/scripts/merge_capabilities.py --execute   # 真正执行
"""
import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent.parent))
from knowhub.knowhub_db.pg_capability_store import PostgreSQLCapabilityStore


# canonical_id -> [member_ids to merge into canonical]
MERGE_CLUSTERS = {
    # ── Round C (跨版本：v0 foundation 吸收 tao_dev_1 同义条) ────
    # C1 动画帧/文生视频
    "CAP-009": ["CAP-a95da8d7"],
    # C2 图像上色
    "CAP-007": ["CAP-75ba0eac"],
    # C3 图像细节增强与高清放大
    "CAP-006": ["CAP-767e7848"],
    # C4 草图/3D → 效果图
    "CAP-019": ["CAP-c7c8479a"],
    # C5 图像内文字翻译/重写替换
    "CAP-021": ["CAP-44ada771"],
    # C6 主体/角色一致性（borderline 合并）
    "CAP-003": ["CAP-1c71b52e"],

    # ── Round 2 ──────────────────────────────────────────
    # B1 室内多光源
    "CAP-5fb6dd66": ["CAP-aa75f198", "CAP-4da21eb2"],
    # B2 室内色调
    "CAP-3bf86ae3": ["CAP-20e2e0e6"],
    # B3 室内自然材质
    "CAP-56368e3a": ["CAP-e8b46ca2"],
    # B4 电影级光照
    "CAP-aaaef688": ["CAP-9658a2d9"],
    # B5 虚拟相机运动
    "CAP-49175b92": ["CAP-943f7709"],
    # B6 结构化 Prompt 框架
    "CAP-5b000814": ["CAP-f0d52dd1", "CAP-28983111"],
    # B7 多景别 prompt
    "CAP-4f15a85f": ["CAP-1da1bf7a", "CAP-0e64f839"],
    # B8 微距/极端特写
    "CAP-26100ea8": ["CAP-b93f5cf9", "CAP-2382899f"],
    # B9 图文卡片自动排版
    "CAP-87ba3b7d": ["CAP-25c0eef3", "CAP-a648fb4a"],
    # B10 海报多元素排版
    "CAP-562d91c1": ["CAP-2df55478", "CAP-cee93f12", "CAP-dba33b2d",
                     "CAP-02b573f4"],
    # B11 空间透视景深
    "CAP-3c49ff0a": ["CAP-edb949a3", "CAP-ee3bac1d", "CAP-19e77a40",
                     "CAP-eb5e0ad2"],
    # B12 产品摄影光照
    "CAP-c2c42fc7": ["CAP-be42acba"],
    # B13 材质质感
    "CAP-0dc2a15b": ["CAP-7f63a4f5"],
    # B14 霓虹/光晕/发光
    "CAP-a35e7966": ["CAP-ec490469", "CAP-ef8c3e8b", "CAP-9dea4396"],
    # B15 数据图表生成
    "CAP-3ee6c232": ["CAP-84a60ca0", "CAP-d02af2bd"],
    # B16 复古印刷质感
    "CAP-9359b49f": ["CAP-59414657", "CAP-718daa01"],
    # B17 超现实空间
    "CAP-1f898cd9": ["CAP-d1a4709c"],
    # B18 文字透视融合
    "CAP-bd4828fc": ["CAP-f3d42bb0"],
    # B19 拟人化角色
    "CAP-e962c3ef": ["CAP-93b35d43", "CAP-5c7abd0c"],
    # B20 镜面水面反射
    "CAP-3b51102e": ["CAP-bb547523"],
    # B21 粒子光效
    "CAP-8467736a": ["CAP-a6c6d4fc"],
    # B22 自然光时刻
    "CAP-e8a77f70": ["CAP-7f2d1a72"],
    # B23 经典布光模式
    "CAP-ed4b506e": ["CAP-08c54a3c"],
    # B24 几何蒙版裁剪
    "CAP-47d6893f": ["CAP-7889940a"],
    # B25 批量系列海报
    "CAP-832e80ac": ["CAP-bb063798"],
    # B26 文生视频通用
    "CAP-a95da8d7": ["CAP-523c8623"],

    # ── Round 1（已执行完毕，保留以保障脚本幂等）───────
    # A1 抠图
    "CAP-12d2aa10": ["CAP-b9c2cafc", "CAP-2a40d757", "CAP-88816f1e",
                     "CAP-2ba2bc19", "CAP-f0137cfa", "CAP-ccb3a2fd",
                     "CAP-ad3fd294"],
    # A2 多图拼贴（后处理）
    "CAP-41ac8100": ["CAP-17732b2b", "CAP-409eefd9", "CAP-6f791b59",
                     "CAP-bdc6c7c8", "CAP-c009dcce", "CAP-e7467ebf",
                     "CAP-462897e2", "CAP-a815960c", "CAP-f3d22954"],
    # A2b 大图切割
    "CAP-fddd3349": ["CAP-3cf6ab47", "CAP-336fe318"],
    # A3 鱼眼/广角畸变
    "CAP-0e3d61ca": ["CAP-71de6ed3", "CAP-4e9c99fa", "CAP-d49daa5b"],
    # A4 选择性着色
    "CAP-3178172e": ["CAP-1dcb6702", "CAP-d18b8a24"],
    # A5 丁达尔/体积光
    "CAP-3086677b": ["CAP-c6ee82db", "CAP-f07aa0df", "CAP-81f097e1"],
    # A6 胶片光学
    "CAP-6c14041c": ["CAP-2c85b37d", "CAP-8ce1a9c6", "CAP-c9b3c7a5"],
    # A7 摄影参数
    "CAP-ef0a4c0c": ["CAP-17ffe5ea", "CAP-84d68e1b", "CAP-074db966"],
    # A8 戏剧明暗/Chiaroscuro
    "CAP-1649b549": ["CAP-e06eeb84", "CAP-e707003d", "CAP-aa25bdf9"],
    # A9 唇形同步
    "CAP-98490894": ["CAP-6fe5aecb", "CAP-89b9875b", "CAP-ad6d8e0d"],
    # A10 虚拟试衣/服装迁移
    "CAP-d92ffc99": ["CAP-1fdb00c2", "CAP-35974014", "CAP-b56d72c7",
                     "CAP-f697cb22"],
    # A11 AI 配乐
    "CAP-5f9644fb": ["CAP-9fe2869e"],
    # A12 角色多视图（prompt 驱动）
    "CAP-5342ad19": ["CAP-2aee7861", "CAP-6566ec42", "CAP-faa53945",
                     "CAP-ea7f3b27", "CAP-9edfec88"],
    # A12b 参数化多视角（ControlNet/坐标）
    "CAP-ee7df476": ["CAP-939c3610", "CAP-1374aa64", "CAP-3758e4a5"],
    # A13 逆光/轮廓光
    # CAP-47151d87 单独保留（VCAP 会在最后删除）
    # A14 冷暖对比
    "CAP-a185d6d2": ["CAP-6994f914", "CAP-a6d5ef60"],
    # A15 双重曝光
    "CAP-19e5402a": ["CAP-7841b44d", "CAP-6a7ebaa3"],
    # A16 对话气泡
    "CAP-fc2bd5cf": ["CAP-fd64812b"],
    # A17 文字描边/阴影
    "CAP-16c5174b": ["CAP-c13f0764", "CAP-e67c259b"],
    # A18 LLM 提示词扩写
    "CAP-4d8ba002": ["CAP-dd8c832f", "CAP-cc8d042f", "CAP-eeb71b76"],
    # A18b LLM 故事脚本
    "CAP-da51c2ec": ["CAP-da6ef3eb", "CAP-88752108", "CAP-2880810c",
                     "CAP-3313a654", "CAP-34567777"],
    # A19 360 全景
    "CAP-1b3e966f": ["CAP-fd4786c5"],
    # A20 信息层级排版
    "CAP-067edd94": ["CAP-c321e41d", "CAP-bb78b124", "CAP-5753e9dd"],
    # A21 卡片版式
    "CAP-6e77db54": ["CAP-33a038e1"],
    # A22 单色调锁定（prompt）
    "CAP-2bd87e28": ["CAP-8dfe77e3", "CAP-7f123558", "CAP-a3985fda"],
    # A22b 后处理色调映射
    "CAP-79590b09": ["CAP-39e1bfda", "CAP-eccd1ce8", "CAP-f08eb0eb",
                     "CAP-a0e2c93a", "CAP-dd9e6d34"],
    # A23 低饱和/去色
    "CAP-298dcb55": ["CAP-a5fb0745"],
    # A24 高饱和多色并置
    "CAP-2de278d6": ["CAP-3ba4cb6e", "CAP-76b117fc", "CAP-252b422a"],
    # A26 极端视角
    "CAP-0c30af82": ["CAP-fecf1f7d", "CAP-a0800d7d"],
    # A28 手部修复
    "CAP-0ba3159e": ["CAP-db60d72e"],
    # A29 皮肤真实感
    "CAP-3b0de1ce": ["CAP-ad971785", "CAP-b36560ff"],
    # A30 AI 一次生成多格图
    "CAP-306c15fe": ["CAP-033b21b5", "CAP-ae5276f1", "CAP-611ac7bd",
                     "CAP-e5fbbe2a", "CAP-8e4dbefa", "CAP-a1e8fd4e",
                     "CAP-56175058", "CAP-5bb77d65", "CAP-e861c293",
                     "CAP-07a8b9a2"],
    # A31 多格一致性
    "CAP-e9b763d2": ["CAP-ddb5e870", "CAP-e1e9a807", "CAP-08b34a6b",
                     "CAP-020d2c05", "CAP-6fa6bd25", "CAP-38d84d76"],
    # A32 多情绪表情矩阵
    "CAP-5a1ac59d": ["CAP-92f70ebf", "CAP-164ecc20"],
    # A33 配色方案推荐
    "CAP-689bac61": ["CAP-4bd95998", "CAP-86887ed5", "CAP-d6954059"],
    # A38 轮播分页
    "CAP-20409fa6": ["CAP-716dd15a", "CAP-e4a5c708", "CAP-a9d3293f"],
}

# VCAP 占位 + 空描述（单独删除）
VCAP_AND_EMPTY = [
    "CAP-47310932",
    "CAP-tao_dev_1-00-01", "CAP-tao_dev_1-00-02", "CAP-tao_dev_1-00-03",
    "CAP-tao_dev_1-00-04", "CAP-tao_dev_1-01-01", "CAP-tao_dev_1-01-02",
    "CAP-tao_dev_1-01-03", "CAP-tao_dev_1-02-01", "CAP-tao_dev_1-02-02",
    "CAP-tao_dev_1-02-03", "CAP-tao_dev_1-03-01", "CAP-tao_dev_1-03-02",
    "CAP-tao_dev_1-03-03", "CAP-tao_dev_1-03-04", "CAP-tao_dev_1-03-05",
    "CAP-tao_dev_1-03-06", "CAP-tao_dev_1-04-01", "CAP-tao_dev_1-04-02",
    "CAP-tao_dev_1-04-03", "CAP-tao_dev_1-04-04", "CAP-tao_dev_1-04-05",
]

JUNCTIONS = [
    ("requirement_capability", "capability_id"),
    ("capability_tool", "capability_id"),
    ("capability_knowledge", "capability_id"),
    ("capability_resource", "capability_id"),
    ("strategy_capability", "capability_id"),
]


def sanity_check(cur):
    """校验所有 canonical/member 都真实存在于 DB。"""
    all_ids = set(MERGE_CLUSTERS.keys())
    for members in MERGE_CLUSTERS.values():
        all_ids.update(members)
    all_ids.update(VCAP_AND_EMPTY)

    cur.execute("SELECT id FROM capability WHERE id = ANY(%s)", (list(all_ids),))
    existing = {r['id'] for r in cur.fetchall()}
    missing = all_ids - existing
    if missing:
        print(f"⚠️  以下 ID 在 DB 中不存在（会跳过）：{sorted(missing)}", flush=True)
    return existing


def merge_one_cluster(cur, canonical, members, existing, dry_run):
    """把 members 合并到 canonical。"""
    members = [m for m in members if m in existing]
    if not members:
        print(f"  [skip] {canonical}: 所有 members 都已不存在", flush=True)
        return 0, 0, 0

    junction_updates = 0
    junction_dups_removed = 0
    # 逐个 member 处理 —— 避免 members 之间互相冲突（都指向同一 other_id）
    for member in members:
        for table, col in JUNCTIONS:
            other_col = {
                "requirement_capability": "requirement_id",
                "capability_tool": "tool_id",
                "capability_knowledge": "knowledge_id",
                "capability_resource": "resource_id",
                "strategy_capability": "strategy_id",
            }[table]

            # 删除会和 canonical 当前占用冲突的 member 引用
            del_sql = (
                f"DELETE FROM {table} WHERE {col} = %s "
                f"AND {other_col} IN ("
                f"  SELECT {other_col} FROM {table} WHERE {col} = %s"
                f")"
            )
            if dry_run:
                cur.execute(
                    f"SELECT COUNT(*) AS c FROM {table} WHERE {col} = %s "
                    f"AND {other_col} IN ("
                    f"  SELECT {other_col} FROM {table} WHERE {col} = %s)",
                    (member, canonical),
                )
                dup_n = cur.fetchone()['c']
            else:
                cur.execute(del_sql, (member, canonical))
                dup_n = cur.rowcount or 0
            junction_dups_removed += dup_n

            # UPDATE 剩余 member 引用 → canonical
            upd_sql = f"UPDATE {table} SET {col} = %s WHERE {col} = %s"
            if dry_run:
                cur.execute(
                    f"SELECT COUNT(*) AS c FROM {table} WHERE {col} = %s",
                    (member,),
                )
                upd_n = cur.fetchone()['c']
            else:
                cur.execute(upd_sql, (canonical, member))
                upd_n = cur.rowcount or 0
            junction_updates += upd_n

    # 删除 member 行
    del_caps_sql = "DELETE FROM capability WHERE id = ANY(%s)"
    if dry_run:
        n_del = len(members)
    else:
        cur.execute(del_caps_sql, (members,))
        n_del = cur.rowcount or 0

    print(f"  [{'DRY' if dry_run else 'OK '}] {canonical} ← {len(members)} members | "
          f"junction upd={junction_updates}, dup_removed={junction_dups_removed}, "
          f"cap_deleted={n_del}",
          flush=True)
    return junction_updates, junction_dups_removed, n_del


def delete_vcap(cur, ids_to_delete, existing, dry_run):
    ids = [i for i in ids_to_delete if i in existing]
    if not ids:
        print("  [skip] VCAP: 都已不存在", flush=True)
        return 0

    total_j = 0
    for table, col in JUNCTIONS:
        if dry_run:
            cur.execute(
                f"SELECT COUNT(*) AS c FROM {table} WHERE {col} = ANY(%s)", (ids,)
            )
            n = cur.fetchone()['c']
        else:
            cur.execute(f"DELETE FROM {table} WHERE {col} = ANY(%s)", (ids,))
            n = cur.rowcount or 0
        total_j += n

    if dry_run:
        n_cap = len(ids)
    else:
        cur.execute("DELETE FROM capability WHERE id = ANY(%s)", (ids,))
        n_cap = cur.rowcount or 0
    print(f"  [{'DRY' if dry_run else 'OK '}] VCAP delete: "
          f"{len(ids)} caps, junction_deleted={total_j}, cap_deleted={n_cap}",
          flush=True)
    return n_cap


def main():
    ap = argparse.ArgumentParser()
    g = ap.add_mutually_exclusive_group(required=True)
    g.add_argument("--dry-run", action="store_true")
    g.add_argument("--execute", action="store_true")
    args = ap.parse_args()

    dry = args.dry_run
    print(f"\n{'='*60}\n{'DRY RUN' if dry else 'EXECUTE'} — capability 合并\n"
          f"{'='*60}", flush=True)

    store = PostgreSQLCapabilityStore()
    cur = store._get_cursor()
    try:
        # 前置统计
        cur.execute("SELECT COUNT(*) AS c FROM capability WHERE version='tao_dev_1'")
        before = cur.fetchone()['c']
        print(f"📊 合并前 tao_dev_1 capability 数: {before}\n", flush=True)

        existing = sanity_check(cur)
        print(f"📋 待合并簇: {len(MERGE_CLUSTERS)}，VCAP 待删: "
              f"{len(VCAP_AND_EMPTY)}\n", flush=True)

        print("── 阶段 1: 簇内合并 ─────────────", flush=True)
        tot_upd = tot_dup = tot_del = 0
        for idx, (canonical, members) in enumerate(MERGE_CLUSTERS.items(), 1):
            print(f"\n[{idx}/{len(MERGE_CLUSTERS)}] canonical={canonical}", flush=True)
            u, d, dl = merge_one_cluster(cur, canonical, members, existing, dry)
            tot_upd += u
            tot_dup += d
            tot_del += dl

        print("\n── 阶段 2: 删除 VCAP + 空数据 ─────", flush=True)
        n_vcap = delete_vcap(cur, VCAP_AND_EMPTY, existing, dry)

        # 事后统计
        cur.execute("SELECT COUNT(*) AS c FROM capability WHERE version='tao_dev_1'")
        after = cur.fetchone()['c']
        print(f"\n{'='*60}\n汇总:", flush=True)
        print(f"  junction UPDATE: {tot_upd}", flush=True)
        print(f"  junction 去重 DELETE: {tot_dup}", flush=True)
        print(f"  capability 簇合并删除: {tot_del}", flush=True)
        print(f"  VCAP 删除: {n_vcap}", flush=True)
        print(f"  合并前: {before}  合并{'后(预测)' if dry else '后'}: {after}",
              flush=True)
        if dry:
            print(f"  预期合并后: {before - tot_del - n_vcap}", flush=True)
        print(f"{'='*60}\n", flush=True)
    finally:
        cur.close()
        store.close()


if __name__ == "__main__":
    main()