cotest/reset.py at main · 508liang/cotest · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
"""
重置 IMM / SMM JSON 数据。

用途：测试前清空历史状态，避免不同测试用户互相影响。
默认行为：
1) 删除 jl/ 下所有 imm_*.json（保留 PDF）
2) 基于 .env 的 yh* 映射重建空白 imm_yh*.json
3) 重置 jl/smm_shared_models.json 为 {}

用法：
  conda run -n python_class python .\\reset_mm_json.py
"""

from __future__ import annotations

import json
import os
import re
from pathlib import Path

from config import settings  # noqa: F401  # 触发 .env 自动加载


def _atomic_write_json(path: Path, payload: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    tmp.replace(path)


def main() -> None:
    root = Path(__file__).resolve().parent
    jl_dir = root / "jl"
    jl_dir.mkdir(parents=True, exist_ok=True)

    removed = 0
    for path in jl_dir.glob("imm_*.json"):
        try:
            path.unlink()
            removed += 1
        except Exception as e:
            print(f"[reset_mm_json] 删除失败 {path.name}: {e}")

    mapping: dict[str, str] = {}
    for key, value in os.environ.items():
        k = str(key or "").strip().lower()
        v = str(value or "").strip()
        if re.fullmatch(r"yh\d+", k) and v:
            mapping[k] = v

    created = 0
    for alias, user_id in sorted(mapping.items(), key=lambda x: x[0]):
        path = jl_dir / f"imm_{alias}.json"
        payload = {
            "user_id": user_id,
            "user_name": alias,
            "个人画像": {
                "姓名": alias,
                "专业领域": "",
                "核心专长": [],
                "历史项目经验": [],
            },
            "个人领域知识库": {
                "提取术语": [],
                "术语解释": [],
            },
            "认知盲区 (未涉及知识)": [],
            "个人任务认知 (Task Stance)": {
                "期望研究方向": "",
                "提议研究方法": "",
                "预期实验流程": "",
            },
            "last_confirmed_ts": 0.0,
            "updated_at": 0.0,
        }
        _atomic_write_json(path, payload)
        created += 1

    _atomic_write_json(jl_dir / "smm_shared_models.json", {})

    print(f"[reset_mm_json] 删除IMM文件: {removed}")
    print(f"[reset_mm_json] 新建IMM文件: {created}")
    print("[reset_mm_json] 已重置 SMM: smm_shared_models.json")


if __name__ == "__main__":
    main()