workbuddy-portal/tools/demo_data.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Wang Chuanli
"""生成**脱敏示例数据**，用于本地体验、界面截图与文档配图。

为什么需要它
------------
`docs/images/` 里的界面截图必须是可公开的，但真实库里的 Prompt 全文、
请求 ID、用量分布与本机路径都属于私有信息。与其手工打码，不如用一份
**完全合成**的数据集重新截图——顺便也让后来者能一键把界面跑起来看。

生成内容
--------
| 表 | 说明 |
|---|---|
| `users`         | 明示例两个账号：`admin`（管理员）与 `demo`（普通账号），各有自己的数据 |
| `usage_records` | 约 900 条合成记录，跨 30 天，含假模型名 / 假 Prompt / 偏斜的积分分布 |
| `collect_runs`  | 采集历史，含 ok / warn / error 三种状态 |
| `settings`      | 走项目默认值（`config.DEFAULTS`），凭证只写**一眼可辨的假值** |
| `audit_log`     | 操作审计（按账号归属） |

刻意造两个账号，是因为「数据按账号隔离」在多用户版里是最该被截进文档的性质：
只有一个账号的话，用户管理页和「我的数据」列都看不出区别。

用法
----
    # 默认写到 data/demo/（该目录在 .gitignore 内，不会误提交）
    python tools/demo_data.py

    # 指定目录与口令，然后起服务看效果
    python tools/demo_data.py --out data/demo --admin-password demo123
    WB_DATA_DIR=$PWD/data/demo python manage.py serve --port 8849 --no-scheduler

注意
----
本脚本**只写 `--out` 指定的目录**，不会读取也不会修改 `data/usage.sqlite`。
已存在的目标库会被拒绝覆盖，除非显式加 `--force`。
"""
from __future__ import annotations

import argparse
import os
import random
import sys
from datetime import datetime, timedelta

BASE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, BASE)

# ---- 合成素材：刻意保持通用，不含任何真实的产品名、Prompt 或业务信息 ----
MODELS = [
    # (模型名, 权重, 积分中位数)
    ("demo-flash", 26, 0.35),
    ("demo-lite", 20, 0.60),
    ("demo-pro", 18, 2.20),
    ("demo-reason", 14, 4.80),
    ("demo-mini", 12, 0.22),
    ("demo-vision", 7, 6.50),
    ("demo-nano", 3, 0.12),
]

CLIENTS = [("vscode", 74), ("webconsole", 18), ("sdk", 8)]

# 写进 settings 的假 Cookie。刻意用重复串，一眼就能看出不是真凭据；
# 作用只是让概览页的健康指示灯是绿的（首装状态是「缺 Cookie」告警）。
# 两个账号给不同的值，这样「各自持有自己的凭证」在截图里看得出来。
DEMO_COOKIE = "wb_demo_session=" + "deadbeef" * 15
DEMO_COOKIE_2 = "wb_demo_session=" + "cafef00d" * 15

# 第二个账号（普通用户）用的模型：刻意与管理员**不重名**，
# 这样「按账号隔离」在按模型的图上立刻可见。
MODELS_2 = [
    ("demo-lite", 30, 0.55),
    ("demo-nano", 26, 0.14),
    ("demo-flash", 22, 0.32),
    ("demo-vision", 12, 5.90),
    ("demo-pro", 10, 2.05),
]

PROMPTS_SHORT = [
    "帮我解释一下这段代码的作用",
    "把这段 SQL 优化一下，避免全表扫描",
    "写一个 Python 脚本，把目录里的 CSV 批量转成 JSON",
    "这个报错是什么意思：connection refused",
    "帮我 review 一下这个接口设计，有什么问题",
    "解释一下 JWT 和 Session 的区别",
    "生成一份周报模板",
    "把下面的需求整理成技术方案",
    "这个正则怎么写：匹配 11 位手机号",
    "Explain the difference between processes and threads",
    "Refactor this function to be more readable",
    "What is the time complexity of this algorithm?",
    "Write unit tests for this module",
    "How do I make this loop faster?",
    "Summarize the key points of this document",
]

PROMPTS_LONG = [
    "你是资深后端工程师。请审查下面的接口实现，重点看：\n"
    "1) 并发写入是否安全；\n2) 异常分支是否都有兜底；\n3) 有没有可以合并的重复查询。\n"
    "请按「问题 / 影响 / 建议」三栏输出。",

    "下面的表结构要支持按天和按模型两个维度聚合，数据量大约千万级。\n"
    "请给出索引设计，并说明每个索引命中的查询模式。",

    "把这段代码从回调风格改成 async/await，保持对外行为不变，"
    "并补充必要的错误处理。改完给出前后对比。",

    "Review the provided module and suggest improvements.\n"
    "Focus on readability, error handling, and testability.\n"
    "Reply as a short bullet list.",
]


def _weighted(rng: random.Random, pairs):
    """按权重取一项（pairs 为 (值, 权重) 列表）。"""
    total = sum(w for _, w in pairs)
    pick = rng.uniform(0, total)
    acc = 0.0
    for val, w in pairs:
        acc += w
        if pick <= acc:
            return val
    return pairs[-1][0]


def _credits(rng: random.Random, median: float) -> float:
    """对数正态分布的积分值，偶尔出现大额。"""
    val = rng.lognormvariate(0.0, 0.85) * median
    if rng.random() < 0.02:                     # 2% 的大额长任务
        val *= rng.uniform(12, 60)
    return round(min(val, 900.0), 2)


def _prompt(rng: random.Random) -> str | None:
    r = rng.random()
    if r < 0.12:                                # 一部分请求不带 Prompt
        return None
    if r < 0.55:
        return rng.choice(PROMPTS_SHORT)
    if r < 0.75:
        return rng.choice(PROMPTS_LONG)
    # 其余用短句拼接，避免重复得过于整齐
    return rng.choice(PROMPTS_SHORT) + "：" + rng.choice(PROMPTS_SHORT)


def build(out_dir: str, days: int, seed: int, admin_password: str,
          admin_user: str) -> str:
    """建库并写入合成数据，返回数据库文件路径。"""
    out_dir = os.path.abspath(out_dir)
    os.makedirs(out_dir, exist_ok=True)
    db_file = os.path.join(out_dir, "usage.sqlite")

    # 必须在 import 项目模块之前设好环境变量：config 在导入时读取它们。
    # 日志目录用 setdefault —— 容器里 WB_LOG_DIR 已由镜像 ENV 指定，
    # 不该在数据卷下再凭空建一个 logs/。
    os.environ["WB_DATA_DIR"] = out_dir
    os.environ.setdefault("WB_LOG_DIR", os.path.join(out_dir, "logs"))
    os.environ["WB_DB"] = db_file

    from workbuddy_portal import db, security              # noqa: E402

    conn = db.connect()
    try:
        db.init_db(conn, create_admin=True, admin_user=admin_user,
                   admin_password=admin_password)

        # ---------- 账号 ----------
        admin_row = db.user_by_name(conn, admin_user)
        admin_uid = admin_row["id"]
        demo_user = "demo"
        if not db.user_by_name(conn, demo_user):
            conn.execute(
                "INSERT INTO users(username,password_hash,display_name,email,is_admin,"
                "status,created_at) VALUES(?,?,?,?,0,'active',?)",
                (demo_user, security.hash_password(admin_password), "演示账号",
                 "demo@example.invalid", db.now_str()))
        db_row = db.user_by_name(conn, demo_user)
        demo_uid = db_row["id"]

        # 写入**明显是假值**的 Cookie：让概览页的健康状态显示为「正常」
        # 而不是首装的「缺 Cookie」告警——演示与截图应当呈现「配置完成」后的样子。
        # 经 set_secret 落库 = 真的走一遍加密，所以示例库里也是密文。
        db.set_secret(conn, "cookie", DEMO_COOKIE, admin_uid)
        db.set_secret(conn, "cookie", DEMO_COOKIE_2, demo_uid)
        # 两个账号各有一套调度时刻，界面上能看出「每人可改自己的」
        db.set_setting(conn, "schedule_times", "09:00,17:00", admin_uid)
        db.set_setting(conn, "schedule_times", "08:30,20:00", demo_uid)

        rng = random.Random(seed)
        now = datetime.now().replace(second=0, microsecond=0)
        today0 = now.replace(hour=0, minute=0, second=0)
        model_pairs = [(m, w) for m, w, _ in MODELS]
        medians = {m: md for m, _, md in MODELS}
        model_pairs_2 = [(m, w) for m, w, _ in MODELS_2]
        medians_2 = {m: md for m, _, md in MODELS_2}
        client_pairs = list(CLIENTS)

        # ---------- usage_records（两个账号各生成一份）----------
        def _gen_records(uid, pairs, med, lo=14, hi=46):
            rows = []
            for d in range(days - 1, -1, -1):
                day0 = today0 - timedelta(days=d)
                for _ in range(rng.randint(lo, hi)):
                    # 工作时间加权：9-19 点更密
                    hour = _weighted(rng, [(h, 6 if 9 <= h <= 19 else 1) for h in range(24)])
                    ts = day0 + timedelta(hours=hour, minutes=rng.randint(0, 59),
                                          seconds=rng.randint(0, 59))
                    if ts > now:
                        continue
                    model = _weighted(rng, pairs)
                    client = _weighted(rng, client_pairs)
                    rid = "req-%s" % "".join(rng.choice("0123456789abcdef") for _ in range(16))
                    stamp = ts.strftime("%Y-%m-%d %H:%M:%S")
                    rows.append((
                        uid, rid, stamp, ts.strftime("%Y-%m-%d"), hour, model, client,
                        _credits(rng, med[model]), _prompt(rng), stamp, stamp, stamp,
                    ))
            return rows

        rows = _gen_records(admin_uid, model_pairs, medians)
        # 普通账号只给大约三分之二的量：列表里一眼能分出主次
        rows += _gen_records(demo_uid, model_pairs_2, medians_2, lo=9, hi=31)
        conn.execute("BEGIN")
        conn.executemany(
            "INSERT OR REPLACE INTO usage_records"
            "(user_id,request_id,ts,day,hour,model,client,credits,prompt,"
            " first_seen,last_seen,cloud_ts) VALUES(?,?,?,?,?,?,?,?,?,?,?,?)", rows)
        conn.execute("COMMIT")

        # ---------- collect_runs ----------
        runs = []
        total = 0
        for i in range(14, 0, -1):
            started = now - timedelta(minutes=i * 37 + rng.randint(0, 9))
            fetched = rng.randint(3, 22)
            dup = rng.randint(0, max(1, fetched - 2))
            added = max(0, fetched - dup)
            total += added
            status = "ok"
            msg = "新增 %d 条，重复 %d 条，存档共 %d 条" % (added, dup, total)
            exit_code = 0
            if i == 9:
                status, exit_code = "warn", 1
                msg = "参数错误：from 不是合法日期；abc（正确写法 2026-09-01）"
            if i == 5:
                status, exit_code = "error", 2
                msg = "云端返回 401 Unauthorized：Cookie 可能已过期，请重新粘贴"
                fetched = dup = added = 0
            trigger = "schedule" if i % 3 else "manual"
            runs.append((
                admin_uid, trigger, status, started.strftime("%Y-%m-%d %H:%M:%S"),
                (started + timedelta(milliseconds=rng.randint(180, 1400))
                 ).strftime("%Y-%m-%d %H:%M:%S"),
                rng.randint(180, 1400),
                (started - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"),
                started.strftime("%Y-%m-%d %H:%M:%S"),
                fetched, added, dup, total, 0, exit_code, msg,
                "[%s] %s" % (status, msg),
            ))
        # 普通账号也给两条，让「日志只显示自己的」在截图里成立
        for k, (st, msg) in enumerate((("ok", "新增 6 条，重复 4 条，存档共 192 条"),
                                       ("ok", "新增 3 条，重复 5 条，存档共 186 条"))):
            at = now - timedelta(hours=5 + k * 9)
            runs.append((demo_uid, "schedule", st, at.strftime("%Y-%m-%d %H:%M:%S"),
                         at.strftime("%Y-%m-%d %H:%M:%S"), 420, 
                         (at - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"),
                         at.strftime("%Y-%m-%d %H:%M:%S"), 10 - k, 6 - k * 3, 4, 192, 0, 0,
                         msg, "[%s] %s" % (st, msg)))
        conn.execute("BEGIN")
        conn.executemany(
            "INSERT INTO collect_runs(user_id,trigger,status,started_at,finished_at,duration_ms,"
            "win_from,win_to,fetched,added,dup,total,conflicts,exit_code,message,detail)"
            " VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)", runs)
        conn.execute("COMMIT")

        # ---------- audit_log ----------
        audits = []
        # IP 用 RFC 5737 的文档专用网段（TEST-NET-1），
        # 保证示例里出现的地址永远不可能是真实主机
        actions = [
            ("login", "登录成功", "192.0.2.10"),
            ("login", "登录成功", "192.0.2.10"),
            ("settings", "修改：schedule_times", "192.0.2.10"),
            ("settings", "修改：page_size", "192.0.2.10"),
            ("maintenance.count", "存档当前 %d 条记录" % total, "192.0.2.10"),
            ("settings_rejected", "参数错误：page_size 必须是数字（条/页）", "192.0.2.10"),
        ]
        for i in range(30):
            act, detail, ip = actions[i % len(actions)]
            at = now - timedelta(hours=i * 3 + rng.randint(0, 2))
            uid = admin_uid if i % 3 else demo_uid
            actor = admin_user if i % 3 else demo_user
            audits.append((uid, at.strftime("%Y-%m-%d %H:%M:%S"), actor, act, detail, ip))
        conn.execute("BEGIN")
        conn.executemany(
            "INSERT INTO audit_log(user_id,at,actor,action,detail,ip)"
            " VALUES(?,?,?,?,?,?)", audits)
        conn.execute("COMMIT")

        # 统计一下，便于打印
        n = conn.execute("SELECT COUNT(*) FROM usage_records").fetchone()[0]
        c = conn.execute("SELECT ROUND(SUM(credits),2) FROM usage_records").fetchone()[0]
        d = conn.execute("SELECT COUNT(DISTINCT day) FROM usage_records").fetchone()[0]
        print("示例库：%s" % db_file)
        print("  记录 %d 条 / 积分 %s / 覆盖 %d 天" % (n, c, d))
        for r in conn.execute(
                "SELECT u.id,u.username,u.is_admin,"
                " (SELECT COUNT(*) FROM usage_records x WHERE x.user_id=u.id) AS n"
                " FROM users u ORDER BY u.id"):
            print("  账号 #%s %-8s %-6s %d 条"
                  % (r["id"], r["username"], "管理员" if r["is_admin"] else "普通", r["n"]))
        print("  口令都是 %s（仅供本地演示）" % admin_password)
        print("  该目录在 .gitignore 内，不会被提交")
    finally:
        conn.close()
    return db_file


def main() -> int:
    ap = argparse.ArgumentParser(description="生成脱敏示例数据（完全合成，不碰真实库）")
    ap.add_argument("--out", default=os.path.join(BASE, "data", "demo"),
                    help="输出目录，默认 data/demo")
    ap.add_argument("--days", type=int, default=30, help="覆盖天数，默认 30")
    ap.add_argument("--seed", type=int, default=20260914, help="随机种子，保证可复现")
    ap.add_argument("--admin-user", default="admin")
    ap.add_argument("--admin-password", default="admin123",
                    help="示例管理员口令，默认 admin123（仅供本地演示）")
    ap.add_argument("--force", action="store_true", help="目标库已存在时覆盖")
    a = ap.parse_args()

    db_file = os.path.join(os.path.abspath(a.out), "usage.sqlite")
    if os.path.exists(db_file) and not a.force:
        print("目标库已存在：%s" % db_file)
        print("如需重建请加 --force")
        return 1
    build(a.out, a.days, a.seed, a.admin_password, a.admin_user)
    return 0


if __name__ == "__main__":
    sys.exit(main())