Agent/examples/create/tool/search_person_tree.py

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

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

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

from agent.tools import tool, ToolResult

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

GRAPH_FULL_DATA_PATH = os.getenv(
    "GRAPH_FULL_DATA_PATH",
    # str(Path(__file__).parent.parent / "data/library/item_graph/item_graph_full_all_levels.json")
    str(Path(__file__).parent.parent / "data/家有大志/topic_point_data/point_classification_results_trans.json")
)

# 缓存图数据，避免重复加载
_graph_full_cache: Optional[Dict[str, Any]] = None


def _load_graph_full() -> Dict[str, Any]:
    """加载完整图数据（带缓存，包含 edges）"""
    global _graph_full_cache
    if _graph_full_cache is None:
        with open(GRAPH_FULL_DATA_PATH, 'r', encoding='utf-8') as f:
            _graph_full_cache = json.load(f)
    return _graph_full_cache

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

    Args:
        tree_data: 树节点数据

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

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

        # 检查是否为常量点：_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)

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

    return constant_points


def _load_persona_tree(persona_name: str) -> 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)
            all_constant_points.extend(constant_points)

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

    dict = _load_graph_full()

    data = []
    for point in  all_constant_points:
        if point in dict:
            data.append(dict.get(point))
    
    return data



@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:
        constant_points = _load_persona_tree(persona_name)
    except Exception as e:
        return ToolResult(
            title="人设常量点检索失败",
            output="",
            error=f"检索异常: {str(e)}",
        )

    if not constant_points:
        return ToolResult(
            title="人设常量点检索失败",
            output="",
            error="未找到常量点数据",
        )

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

    return ToolResult(
        title=f"人设常量点检索 - {persona_name}",
        output=output,
        long_term_memory=f"检索到 {len(constant_points)} 个常量点",
    )

if __name__ == "__main__":
    print(_load_persona_tree("家有大志"))