"""修复：按种子定义给 Agent 设置 kb_ids（库名字 → id）"""
import sys
sys.path.insert(0, r'C:\Users\xinmiclaw\Desktop\雅文内容创造工作台\server')
from sqlmodel import Session, select
from app.db import engine
from app.models import Agent, KnowledgeBase

# 与 seed.py AGENTS 定义一致
AGENT_KBS = {
    "主控 Agent": [],
    "开场段 Agent": ["风格档案库"],
    "心法段 Agent": ["风格档案库", "训练营框架库"],
    "技法段 Agent": ["风格档案库", "直播原文库"],
    "产品销售 Agent": ["产品知识库", "风格档案库"],
    "收尾段 Agent": ["风格档案库"],
    "钩子段 Agent": ["训练营框架库"],
    "结构质检": ["质检规则库"],
    "内容质检": ["质检规则库", "产品知识库"],
    "风格尺度质检": ["质检规则库"],
    "跨稿质检": ["备稿库", "训练营框架库"],
    "组合 Agent": [],
    "风格顾问": ["风格档案库"],
    "产品顾问": ["产品知识库"],
}

with Session(engine) as s:
    kbs = {k.name: k.id for k in s.exec(select(KnowledgeBase)).all()}
    print("现有库:", kbs)
    for a in s.exec(select(Agent)).all():
        names = AGENT_KBS.get(a.name, [])
        ids = [kbs[n] for n in names if n in kbs]
        missing = [n for n in names if n not in kbs]
        a.kb_ids = ids
        s.add(a)
        print(f"{a.name}: {ids}{'（缺库:' + str(missing) + '）' if missing else ''}")
    s.commit()
    print("完成")