Agent/examples_how/overall_derivation/tools/find_tree_node.py

"""
查找树节点 Tool - 人设树节点查询

功能：
1. 获取人设树的常量节点（全局常量、局部常量）
2. 获取符合条件概率阈值的节点（按条件概率排序返回 topN）
"""

import importlib.util
import json
from pathlib import Path
from typing import Any, Optional

try:
    from agent.tools import tool, ToolResult, ToolContext
except ImportError:
    def tool(*args, **kwargs):
        return lambda f: f
    ToolResult = None  # 仅用 main() 测核心逻辑时可无 agent
    ToolContext = None

# 加载同目录层级的 conditional_ratio_calc（不依赖包结构）
_utils_dir = Path(__file__).resolve().parent.parent / "utils"
_cond_spec = importlib.util.spec_from_file_location(
    "conditional_ratio_calc",
    _utils_dir / "conditional_ratio_calc.py",
)
_cond_mod = importlib.util.module_from_spec(_cond_spec)
_cond_spec.loader.exec_module(_cond_mod)
calc_node_conditional_ratio = _cond_mod.calc_node_conditional_ratio

# 相对本文件：tools -> overall_derivation，input 在 overall_derivation 下
_BASE_INPUT = Path(__file__).resolve().parent.parent / "input"


def _tree_dir(account_name: str) -> Path:
    """人设树目录：../input/{account_name}/原始数据/tree/"""
    return _BASE_INPUT / account_name / "原始数据" / "tree"


def _load_trees(account_name: str) -> list[tuple[str, dict]]:
    """加载该账号下所有维度的人设树。返回 [(维度名, 根节点 dict), ...]。"""
    td = _tree_dir(account_name)
    if not td.is_dir():
        return []
    result = []
    for p in td.glob("*.json"):
        try:
            with open(p, "r", encoding="utf-8") as f:
                data = json.load(f)
            for dim_name, root in data.items():
                if isinstance(root, dict):
                    result.append((dim_name, root))
                break
        except Exception:
            continue
    return result


def _iter_all_nodes(account_name: str):
    """遍历该账号下所有人设树节点，产出 (节点名称, 父节点名称, 节点 dict)。"""
    for dim_name, root in _load_trees(account_name):
        def walk(parent_name: str, node_dict: dict):
            for name, child in (node_dict.get("children") or {}).items():
                if not isinstance(child, dict):
                    continue
                yield (name, parent_name, child)
                yield from walk(name, child)

        yield from walk(dim_name, root)


# ---------------------------------------------------------------------------
# 1. 获取人设树常量节点
# ---------------------------------------------------------------------------

def get_constant_nodes(account_name: str) -> list[dict[str, Any]]:
    """
    获取人设树的常量节点。
    - 全局常量：_is_constant=True
    - 局部常量：_is_local_constant=True 且 _is_constant=False
    返回列表项：节点名称、概率(_ratio)、常量类型。
    """
    result = []
    for node_name, _parent, node in _iter_all_nodes(account_name):
        is_const = node.get("_is_constant") is True
        is_local = node.get("_is_local_constant") is True
        if is_const:
            const_type = "全局常量"
        elif is_local and not is_const:
            const_type = "局部常量"
        else:
            continue
        ratio = node.get("_ratio")
        result.append({
            "节点名称": node_name,
            "概率": ratio,
            "常量类型": const_type,
        })
    result.sort(key=lambda x: (x["概率"] is None, -(x["概率"] or 0)))
    return result


# ---------------------------------------------------------------------------
# 2. 获取符合条件概率阈值的节点
# ---------------------------------------------------------------------------

def get_nodes_by_conditional_ratio(
    account_name: str,
    derived_list: list[tuple[str, str]],
    threshold: float,
    top_n: int,
) -> list[dict[str, Any]]:
    """
    获取人设树中条件概率 >= threshold 的节点，按条件概率降序，返回前 top_n 个。
    derived_list: 已推导列表，每项 (已推导的选题点, 推导来源人设树节点)；为空时使用节点自身的 _ratio 作为条件概率。
    返回列表项：节点名称、条件概率、父节点名称。
    """
    base_dir = _BASE_INPUT
    scored: list[tuple[str, float, str]] = []

    if not derived_list:
        # derived_items 为空：条件概率取节点本身的 _ratio
        for node_name, parent_name, node in _iter_all_nodes(account_name):
            ratio = node.get("_ratio")
            if ratio is None:
                ratio = 0.0
            else:
                ratio = float(ratio)
            if ratio >= threshold:
                scored.append((node_name, ratio, parent_name))
    else:
        node_to_parent: dict[str, str] = {}
        for node_name, parent_name, _ in _iter_all_nodes(account_name):
            node_to_parent[node_name] = parent_name
        for node_name, parent_name in node_to_parent.items():
            ratio = calc_node_conditional_ratio(
                account_name, derived_list, node_name, base_dir=base_dir
            )
            if ratio >= threshold:
                scored.append((node_name, ratio, parent_name))

    scored.sort(key=lambda x: x[1], reverse=True)
    top = scored[:top_n]
    return [
        {"节点名称": name, "条件概率": ratio, "父节点名称": parent}
        for name, ratio, parent in top
    ]


def _parse_derived_list(derived_items: list[dict[str, str]]) -> list[tuple[str, str]]:
    """将 agent 传入的 [{"topic": "x", "source_node": "y"}, ...] 转为 DerivedItem 列表。"""
    out = []
    for item in derived_items:
        if isinstance(item, dict):
            topic = item.get("topic") or item.get("已推导的选题点")
            source = item.get("source_node") or item.get("推导来源人设树节点")
            if topic is not None and source is not None:
                out.append((str(topic).strip(), str(source).strip()))
        elif isinstance(item, (list, tuple)) and len(item) >= 2:
            out.append((str(item[0]).strip(), str(item[1]).strip()))
    return out


# ---------------------------------------------------------------------------
# Agent Tools（参考 glob_tool 封装）
# ---------------------------------------------------------------------------

@tool(description="获取人设树的常量节点（全局常量、局部常量）。输入账号名，返回节点名称、概率、常量类型。")
async def find_tree_constant_nodes(
    account_name: str,
    context: Optional[ToolContext] = None,
) -> ToolResult:
    """
    获取人设树的常量节点。
    读取该账号 input/{account_name}/原始数据/tree/ 下的人设树 JSON，
    筛选 _is_constant=true（全局常量）或 _is_local_constant=true 且 _is_constant=false（局部常量）的节点，
    返回：节点名称、概率(_ratio)、常量类型。
    """
    tree_dir = _tree_dir(account_name)
    if not tree_dir.is_dir():
        return ToolResult(
            title="人设树目录不存在",
            output=f"目录不存在: {tree_dir}",
            error="Directory not found",
        )
    try:
        items = get_constant_nodes(account_name)
        if not items:
            output = "未找到常量节点"
        else:
            lines = [f"- {x['节点名称']}\t概率={x['概率']}\t{x['常量类型']}" for x in items]
            output = "\n".join(lines)
        return ToolResult(
            title=f"常量节点 ({account_name})",
            output=output,
            metadata={"account_name": account_name, "count": len(items), "items": items},
        )
    except Exception as e:
        return ToolResult(
            title="获取常量节点失败",
            output=str(e),
            error=str(e),
        )


@tool(
    description="获取人设树中条件概率不低于阈值的节点，按条件概率从高到低返回 topN。"
    "输入：账号名、已推导选题点列表（可为空）、条件概率阈值、topN。"
    "derived_items 为空时，条件概率使用节点自身的 _ratio。"
)
async def find_tree_nodes_by_conditional_ratio(
    account_name: str,
    derived_items: list[dict[str, str]],
    conditional_ratio_threshold: float,
    top_n: int = 20,
    context: Optional[ToolContext] = None,
) -> ToolResult:
    """
    获取人设树中符合条件概率阈值的节点。
    derived_items：可为空；非空时每项为 {\"topic\": \"已推导选题点\", \"source_node\": \"推导来源人设树节点\"}。
    当 derived_items 为空时，各节点的条件概率取其自身的 _ratio；非空时按已推导帖子集合计算条件概率。
    返回：节点名称、条件概率、父节点名称，按条件概率降序最多 top_n 条。
    """
    tree_dir = _tree_dir(account_name)
    if not tree_dir.is_dir():
        return ToolResult(
            title="人设树目录不存在",
            output=f"目录不存在: {tree_dir}",
            error="Directory not found",
        )
    try:
        derived_list = _parse_derived_list(derived_items or [])
        items = get_nodes_by_conditional_ratio(
            account_name, derived_list, conditional_ratio_threshold, top_n
        )
        if not items:
            output = f"未找到条件概率 >= {conditional_ratio_threshold} 的节点"
        else:
            lines = [
                f"- {x['节点名称']}\t条件概率={x['条件概率']}\t父节点={x['父节点名称']}"
                for x in items
            ]
            output = "\n".join(lines)
        return ToolResult(
            title=f"条件概率节点 ({account_name}, 阈值={conditional_ratio_threshold})",
            output=output,
            metadata={
                "account_name": account_name,
                "threshold": conditional_ratio_threshold,
                "top_n": top_n,
                "count": len(items),
                "items": items,
            },
        )
    except Exception as e:
        return ToolResult(
            title="按条件概率查询节点失败",
            output=str(e),
            error=str(e),
        )


def main() -> None:
    """本地测试：用家有大志账号测常量节点与条件概率节点，有 agent 时再跑一遍 tool 接口。"""
    import asyncio

    account_name = "家有大志"
    derived_items = [
        {"topic": "分享", "source_node": "分享"},
    ]
    conditional_ratio_threshold = 0.1
    top_n = 10

    # 1）常量节点
    constant_nodes = get_constant_nodes(account_name)
    print(f"账号: {account_name} — 常量节点共 {len(constant_nodes)} 个（前 50 个）:")
    for x in constant_nodes[:50]:
        print(f"  - {x['节点名称']}\t概率={x['概率']}\t{x['常量类型']}")
    print()

    # 2）条件概率节点（核心函数）
    derived_list = _parse_derived_list(derived_items)
    ratio_nodes = get_nodes_by_conditional_ratio(
        account_name, derived_list, conditional_ratio_threshold, top_n
    )
    print(f"条件概率节点 阈值={conditional_ratio_threshold}, top_n={top_n}, 共 {len(ratio_nodes)} 个:")
    for x in ratio_nodes:
        print(f"  - {x['节点名称']}\t条件概率={x['条件概率']}\t父节点={x['父节点名称']}")
    print()

    # 3）有 agent 时通过 tool 接口再跑一遍
    if ToolResult is not None:
        async def run_tools():
            r1 = await find_tree_constant_nodes(account_name)
            print("--- find_tree_constant_nodes ---")
            print(r1.output[:200] + "..." if len(r1.output) > 200 else r1.output)
            r2 = await find_tree_nodes_by_conditional_ratio(
                account_name,
                derived_items=derived_items,
                conditional_ratio_threshold=conditional_ratio_threshold,
                top_n=top_n,
            )
            print("\n--- find_tree_nodes_by_conditional_ratio ---")
            print(r2.output)
        asyncio.run(run_tools())


if __name__ == "__main__":
    main()