self-constructing_graph/agent/tree_utils.py

import json
from typing import List, Dict
from dataclasses import dataclass

@dataclass
class TreeNodeDTO:
    id: int
    name: str
    pId: int
    icon: str = ""
    iconOpen: str = ""
    iconClose: str = ""

@dataclass
class TreeDTO:
    nodes: List[TreeNodeDTO]

def get_tree_dto(tree_json: str) -> TreeDTO:
    """
    将JSON树结构转换为TreeDTO
    :param tree_json: JSON格式的树结构字符串
    :return: TreeDTO对象
    """
    tree_dto = TreeDTO(nodes=[])
    try:
        tree_obj = json.loads(tree_json)
        if tree_obj:
            # 使用计数器生成唯一ID
            id_counter = 0
            
            # 存储父节点路径到ID的映射
            path_id_map = {}
            
            # 根节点特殊处理
            root_id = id_counter
            id_counter += 1
            path_id_map["根节点"] = root_id
            
            # 从根节点开始构建树
            add_tree(tree_dto.nodes, tree_obj, path_id_map, "根节点", id_counter, "根节点")
    except json.JSONDecodeError:
        pass
    
    return tree_dto

def add_tree(nodes: List[TreeNodeDTO], tree_obj: Dict, path_id_map: Dict[str, int], 
             parent_name: str, id_counter: int, parent_path: str):
    """
    递归添加树节点
    :param nodes: 节点列表
    :param tree_obj: 树对象
    :param path_id_map: 路径到ID的映射
    :param parent_path: 父节点路径
    :param id_counter: ID计数器
    :param current_path: 当前节点路径
    """
    # 构建当前节点的完整路径
    current_path = f"{parent_path}/{tree_obj['name']}"
    
    # 生成当前节点ID
    current_id = id_counter
    id_counter += 1
    path_id_map[current_path] = current_id
    
    # 获取父节点ID
    parent_id = path_id_map[parent_path]
    
    # 判断是否为叶子节点
    is_leaf = not tree_obj.get('sNode', [])

    # 设置图标
    icon = "/images/node.png" if is_leaf else "/images/icon.png"
    icon_open = "" if is_leaf else "/images/iconOpen.png"
    icon_close = "" if is_leaf else "/images/iconClose.png"
    
    # 创建树节点
    node = TreeNodeDTO(
        id=current_id,
        name=tree_obj['name'],
        pId=parent_id,
        icon=icon,
        iconOpen=icon_open,
        iconClose=icon_close
    )
    nodes.append(node)
    
    # 递归处理子节点
    if not is_leaf:
        for child in tree_obj['sNode']:
            id_counter = add_tree(nodes, child, path_id_map, 
                                tree_obj['name'], id_counter, current_path)
    
    return id_counter