self-constructing_graph/agent/models/db/tree_structure.py

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))