Agent-System/scripts/seed_ten_projects_chat.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
seed_ten_projects_chat.py
将 10 个评审项目的 30 个会话与 240 条种子对话消息完整写入 PostgreSQL 数据库
"""

import json
import subprocess
import sys
from datetime import datetime, timedelta

def main():
    seed_path = "web/src/data/seeds/tenProjectsSeed.json"
    with open(seed_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    sessions = data.get("sessions", [])
    messages = data.get("messages", [])

    print(f"读取到 {len(sessions)} 个会话，{len(messages)} 条种子消息。")

    sql_lines = [
        "-- ==============================================================================",
        "-- 迁移脚本: 10 个评审项目会话与种子消息持久化导入",
        "-- ==============================================================================",
        "BEGIN;",
    ]

    # 1. 导入 chat_sessions
    sess_user_map = {}
    for s in sessions:
        sid = s["sessionId"]
        pid = s.get("projectId", 6001)
        uid = s.get("userId", 1)
        sess_user_map[sid] = uid
        title = s.get("title", "项目协同会话").replace("'", "''")
        stage = s.get("stage", "V0.5")

        sql_lines.append(f"""
INSERT INTO chat_sessions (session_id, project_id, user_id, title, current_stage, pending_exp, status, created_at, updated_at)
VALUES ('{sid}', {pid}, {uid}, '{title}', '{stage}', '{{}}', 'active', '2026-09-11 08:00:00+08', '2026-09-11 08:00:00+08')
ON CONFLICT (session_id) DO UPDATE SET
    project_id = EXCLUDED.project_id,
    user_id = EXCLUDED.user_id,
    title = EXCLUDED.title,
    current_stage = EXCLUDED.current_stage,
    updated_at = EXCLUDED.updated_at;
""")

    # 2. 清理这 30 个 session 的旧测试数据以保持干净
    sids_list = "('" + "', '".join(sess_user_map.keys()) + "')"
    sql_lines.append(f"DELETE FROM chat_messages WHERE session_id IN {sids_list};")
    sql_lines.append(f"DELETE FROM chat_read_positions WHERE session_id IN {sids_list};")

    # 3. 导入 240 条消息
    base_time = datetime(2026, 9, 11, 9, 0, 0)
    # 按会话分组
    sess_msgs = {}
    for m in messages:
        sid = m["sessionId"]
        sess_msgs.setdefault(sid, []).append(m)

    last_msg_id_map = {}

    for sid, mlist in sess_msgs.items():
        uid = sess_user_map.get(sid, 1)
        curr_time = base_time

        for idx, m in enumerate(mlist):
            role = m.get("role", "user")
            content = m.get("content", "").replace("'", "''")
            intent = m.get("intentTag", "").replace("'", "''")
            category = m.get("category", "general").replace("'", "''")
            
            # 解析原始时间戳如 "09:14"
            ts = m.get("timestamp", "")
            if ts and ":" in ts:
                try:
                    parts = ts.split(":")
                    curr_time = datetime(2026, 9, 11, int(parts[0]), int(parts[1]), idx % 60)
                except Exception:
                    curr_time += timedelta(seconds=15)
            else:
                curr_time += timedelta(seconds=15)

            time_str = curr_time.strftime("%Y-%m-%d %H:%M:%S+08")
            day_key = "2026-09-11"

            meta = {
                "author_name": m.get("authorName", "项目成员"),
                "client_id": m.get("id"),
            }
            if m.get("attachments"):
                meta["attachments"] = m["attachments"]
            if m.get("interactiveCard"):
                meta["interactive_card"] = m["interactiveCard"]
            if m.get("phaseEvents"):
                meta["phase_events"] = m["phaseEvents"]

            meta_json = json.dumps(meta, ensure_ascii=False).replace("'", "''")
            msg_uid = f"{uid}" if role == "user" else "NULL"

            sql_lines.append(f"""
INSERT INTO chat_messages (session_id, user_id, role, content, intent, branch, metadata, day_key, category, created_at, content_kind, export_policy, content_revision)
VALUES ('{sid}', {msg_uid}, '{role}', '{content}', '{intent}', '', '{meta_json}'::jsonb, '{day_key}', '{category}', '{time_str}', 'text', 'full', 1);
""")

    # 4. 初始化这 30 个 session 的已读位点至最新消息
    for sid, uid in sess_user_map.items():
        sql_lines.append(f"""
INSERT INTO chat_read_positions (session_id, user_id, last_read_message_id, dismissed_unread_ids, updated_at)
SELECT '{sid}', {uid}, COALESCE(MAX(id), 0), '{{}}'::bigint[], now()
FROM chat_messages
WHERE session_id = '{sid}'
ON CONFLICT (session_id, user_id) DO UPDATE SET
    last_read_message_id = EXCLUDED.last_read_message_id,
    dismissed_unread_ids = '{{}}'::bigint[],
    updated_at = now();
""")

    sql_lines.append("COMMIT;")

    out_file = "deploy/postgres/migrations/2026-09-11-seed-ten-projects-chat.sql"
    with open(out_file, "w", encoding="utf-8") as f:
        f.write("\n".join(sql_lines))

    print(f"成功生成 SQL: {out_file} (行数: {len(sql_lines)})")

    # 执行 SQL
    print("正在向 PostgreSQL (shuimu-context-pg-dev) 执行导入...")
    cmd = [
        "docker", "exec", "-i", "shuimu-context-pg-dev",
        "psql", "-U", "app", "-d", "agent_system_dev", "-f", "-"
    ]
    proc = subprocess.run(cmd, input="\n".join(sql_lines).encode("utf-8"), capture_output=True)
    if proc.returncode != 0:
        print("执行失败:", proc.stderr.decode("utf-8"))
        sys.exit(1)
    else:
        print("成功导入数据库！")

if __name__ == "__main__":
    main()