src/model/anchor_tree.py

from __future__ import annotations

from copy import deepcopy

from src.model.anchor_tree_types import AnchorTree, AnchorTreeEdge, AnchorTreeNode, AnchorTreeRelation


def make_anchor_tree(
    *,
    tree_id: str,
    root: AnchorTreeNode,
    nodes: list[AnchorTreeNode],
    edges: list[AnchorTreeEdge],
    domain: str,
    source_kind: str,
) -> AnchorTree:
    all_nodes = {node.node_id: node for node in [root, *nodes]}
    tree = AnchorTree(
        tree_id=tree_id,
        root_id=root.node_id,
        nodes=all_nodes,
        edges=edges,
        domain=domain,
        source_kind=source_kind,
    )
    tree.validate()
    return tree


def clone_tree(tree: AnchorTree) -> AnchorTree:
    return deepcopy(tree)


def attach_child_node(
    tree: AnchorTree,
    *,
    parent_id: str,
    node: AnchorTreeNode,
    relation: AnchorTreeRelation = AnchorTreeRelation.CHILD,
    score: float = 1.0,
) -> AnchorTree:
    updated = clone_tree(tree)
    updated.nodes[node.node_id] = node
    updated.edges.append(
        AnchorTreeEdge(parent_id=parent_id, child_id=node.node_id, relation=relation, score=score)
    )
    updated.validate()
    return updated