Browse Source

树结构文件

yuchengwei 1 month ago
parent
commit
0b6fc13b77
2 changed files with 105 additions and 0 deletions
  1. 23 0
      agent/models/db/tree_structure.py
  2. 82 0
      agent/tree_utils.py

+ 23 - 0
agent/models/db/tree_structure.py

@@ -0,0 +1,23 @@
+from sqlalchemy import Column, Integer, String, Text
+from sqlalchemy.ext.declarative import declarative_base
+
+Base = declarative_base()
+
+class TreeStructure(Base):
+    __tablename__ = 'kg_tree_structures'
+    
+    id = Column(Integer, primary_key=True)
+    user_id = Column(Integer, nullable=False)
+    graph_id = Column(Integer, nullable=False)
+    content = Column(Text, nullable=False)
+    
+    def __repr__(self):
+        return f"<TreeStructure(id={self.id}, user_id={self.user_id}, graph_id={self.graph_id})>"
+
+class KgGraphCategory(Base):
+    __tablename__ = 'kg_graph_category'
+    
+    id = Column(Integer, primary_key=True)
+    user_id = Column(Integer)
+    graph_id = Column(Integer)
+    category = Column(String(64))

+ 82 - 0
agent/tree_utils.py

@@ -0,0 +1,82 @@
+import json
+from typing import List, Dict
+from dataclasses import dataclass
+
+@dataclass
+class TreeNodeDTO:
+    id: int
+    name: str
+    pId: int
+
+@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', [])
+    
+    # 创建树节点
+    node = TreeNodeDTO(
+        id=current_id,
+        name=tree_obj['name'],
+        pId=parent_id
+    )
+    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