kb_structure.py

class Node(BaseModel):
    """
    Node class for the knowledge graph. Each node represents an entity.
    
    Attributes:
        id (int): Unique identifier for the node.
        label (str): Label or name of the node.
        color (str): Color of the node.
        num_targets (int): Number of target nodes connected to this node.
        num_sources (int): Number of source nodes this node is connected from.
        list_target_ids (List[int]): List of unique identifiers of target nodes for which this node is the source node.
        num_edges (int): Total number of edges that this node is a part of, either soruce or target.
        
    """
    
    id: int
    label: str
    color: str
    num_targets: int
    num_sources: int
    list_target_ids: List[int] = Field(default_factory=list)
    num_edges: int = 0

class Edge(BaseModel):
    source: int
    target: int
    label: str
    color: str = "black"

class KnowledgeGraph(BaseModel):
    nodes: List[Node] = Field(default_factory=list)
    edges: List[Edge] = Field(default_factory=list)