Agent/examples/create/tool/search_person_tree.py

"""
人设树常量点检索工具 - 根据人设名称提取所有常量点

用于 Agent 执行时自主提取人设树中的常量点。
"""

import json
from pathlib import Path
from typing import Any, Dict, List

from agent.tools import tool, ToolResult

# 人设数据基础路径
PERSONA_DATA_BASE_PATH = Path(__file__).parent.parent / "data"


def _extract_constant_points(tree_data: Dict[str, Any], dimension: str, path_prefix: str = "") -> List[Dict[str, Any]]:
    """
    递归提取树中所有常量点

    Args:
        tree_data: 树节点数据
        dimension: 维度类型（形式/实质/意图）
        path_prefix: 当前路径前缀

    Returns:
        常量点列表
    """
    constant_points = []

    for node_name, node_data in tree_data.items():
        # 跳过元数据字段
        if node_name.startswith("_"):
            continue

        # 构建当前节点的路径
        current_path = f"{path_prefix}>{node_name}" if path_prefix else node_name

        # 检查是否为常量点：_type 为 "ID" 且 _is_constant 为 true
        if isinstance(node_data, dict):
            node_type = node_data.get("_type")
            is_constant = node_data.get("_is_constant", False)

            if node_type == "ID" and is_constant:
                constant_points.append({
                    "点名称": node_name,
                    "维度": dimension,
                    "路径": current_path,
                    "权重": node_data.get("_persona_weight_score", 0),
                    "帖子数量": node_data.get("_post_count", 0)
                })

            # 递归处理子节点
            if "children" in node_data:
                child_points = _extract_constant_points(
                    node_data["children"],
                    dimension,
                    current_path
                )
                constant_points.extend(child_points)

    return constant_points


def _load_persona_tree(persona_name: str) -> Dict[str, List[Dict[str, Any]]]:
    """
    加载人设树数据并提取所有常量点

    Args:
        persona_name: 人设名称，如 "家有大志"

    Returns:
        包含三个维度常量点的字典
    """
    persona_dir = PERSONA_DATA_BASE_PATH / persona_name / "tree"

    if not persona_dir.exists():
        return {
            "found": False,
            "persona_name": persona_name,
            "error": f"人设目录不存在: {persona_dir}"
        }

    tree_files = {
        "形式": persona_dir / "形式_point_tree_how.json",
        "实质": persona_dir / "实质_point_tree_how.json",
        "意图": persona_dir / "意图_point_tree_how.json"
    }

    all_constant_points = {
        "形式": [],
        "实质": [],
        "意图": []
    }

    missing_files = []

    for dimension, tree_file in tree_files.items():
        if not tree_file.exists():
            missing_files.append(f"{dimension}_point_tree_how.json")
            continue

        try:
            with open(tree_file, 'r', encoding='utf-8') as f:
                tree_data = json.load(f)

            # 提取该维度的所有常量点
            constant_points = _extract_constant_points(tree_data, dimension)
            all_constant_points[dimension] = constant_points

        except Exception as e:
            return {
                "found": False,
                "persona_name": persona_name,
                "error": f"读取文件 {tree_file.name} 失败: {str(e)}"
            }

    if missing_files:
        return {
            "found": False,
            "persona_name": persona_name,
            "error": f"缺失文件: {', '.join(missing_files)}"
        }

    return {
        "found": True,
        "persona_name": persona_name,
        "constant_points": all_constant_points
    }


@tool(
    description="根据人设名称检索该人设树中的所有常量点（_is_constant=true的点）。",
    display={
        "zh": {
            "name": "人设常量点检索",
            "params": {
                "persona_name": "人设名称（如：家有大志）",
            },
        },
    },
)
async def search_person_tree_constants(persona_name: str) -> ToolResult:
    """
    根据人设名称检索该人设树中的所有常量点。

    常量点是指人设树中 _type 为 "ID" 且 _is_constant 为 true 的节点。
    这些点代表了该人设的核心特征和固定属性。

    Args:
        persona_name: 人设名称，如 "家有大志"

    Returns:
        ToolResult: 包含三个维度（形式、实质、意图）的常量点列表
    """
    if not persona_name:
        return ToolResult(
            title="人设常量点检索失败",
            output="",
            error="请提供人设名称",
        )

    try:
        result = _load_persona_tree(persona_name)
    except Exception as e:
        return ToolResult(
            title="人设常量点检索失败",
            output="",
            error=f"检索异常: {str(e)}",
        )

    if not result.get("found"):
        return ToolResult(
            title="人设常量点检索失败",
            output="",
            error=result.get("error", "未知错误"),
        )

    constant_points = result["constant_points"]

    # 统计信息
    total_count = sum(len(points) for points in constant_points.values())

    # 格式化输出
    output_data = {
        "人设名称": persona_name,
        "常量点总数": total_count,
        "形式常量点": constant_points["形式"],
        "实质常量点": constant_points["实质"],
        "意图常量点": constant_points["意图"],
        "统计": {
            "形式点数": len(constant_points["形式"]),
            "实质点数": len(constant_points["实质"]),
            "意图点数": len(constant_points["意图"])
        }
    }

    output = json.dumps(output_data, ensure_ascii=False, indent=2)

    return ToolResult(
        title=f"人设常量点检索 - {persona_name}",
        output=output,
        long_term_memory=f"检索到 {total_count} 个常量点（形式:{len(constant_points['形式'])}，实质:{len(constant_points['实质'])}，意图:{len(constant_points['意图'])}）",
    )