Agent/examples_how/overall_derivation/data_process/extract_simple_tree_node.py

"""
从「处理后数据/tree」下的人设树 JSON 提取节点名称与 _ratio，输出为带缩进的纯文本，
便于大模型阅读。输出：input/{account_name}/处理后数据/simple_tree/simple_tree.txt
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any


def _format_ratio_paren(ratio: Any) -> str:
    """概率两位小数，拼在节点名后的括号内；无 _ratio 时为 —。"""
    if ratio is None:
        return "—"
    try:
        return f"{float(ratio):.3f}"
    except (TypeError, ValueError):
        return str(ratio)


def _walk_node(name: str, node: dict[str, Any], depth: int, lines: list[str]) -> None:
    if not isinstance(node, dict):
        return
    indent = "  " * depth
    r = node.get("_ratio")
    lines.append(f"{indent}{name}({_format_ratio_paren(r)})")
    children = node.get("children")
    if not isinstance(children, dict):
        return
    for child_name, child in children.items():
        if isinstance(child, dict):
            _walk_node(str(child_name), child, depth + 1, lines)


def _tree_json_to_lines(data: Any) -> list[str]:
    lines: list[str] = []
    if not isinstance(data, dict):
        return lines
    for root_name, root_node in data.items():
        if isinstance(root_node, dict):
            _walk_node(str(root_name), root_node, 0, lines)
    return lines


def extract_simple_tree_for_account(account_name: str) -> Path:
    base = Path(__file__).resolve().parents[1]
    tree_dir = base / "input" / account_name / "处理后数据" / "tree"
    out_dir = base / "input" / account_name / "处理后数据" / "simple_tree"
    out_file = out_dir / "simple_tree.txt"

    if not tree_dir.is_dir():
        raise FileNotFoundError(f"目录不存在: {tree_dir}")

    json_paths = sorted(tree_dir.glob("*.json"))
    if not json_paths:
        raise FileNotFoundError(f"未找到 JSON 文件: {tree_dir}")

    out_dir.mkdir(parents=True, exist_ok=True)

    chunks: list[str] = []
    chunks.append("说明: 每行「节点名(概率)」；概率表示该节点在账号下出现的频率；缩进表示层级；root节点无概率用 — 表示。\n分三颗树：\n实质-内容的核心主题/对象\n形式-内容的呈现形式\n意图-内容的目标/用户意图\n")

    for jp in json_paths:
        chunks.append("\n" + "=" * 72 + "\n")
        with jp.open(encoding="utf-8") as f:
            data = json.load(f)
        lines = _tree_json_to_lines(data)
        chunks.append("\n".join(lines))
        chunks.append("\n")

    text = "".join(chunks)
    out_file.write_text(text, encoding="utf-8")
    return out_file


def main(account_name) -> None:
    # p = argparse.ArgumentParser(description="人设树节点简化：节点名与概率 节点(0.00)")
    # p.add_argument("account_name", help="账号目录名，对应 input/{account_name}/")
    # args = p.parse_args()
    try:
        out = extract_simple_tree_for_account(account_name)
        print(f"已写入: {out}")
    except FileNotFoundError as e:
        print(str(e), file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main(account_name="家有大志")