graph/graph.py

from langgraph.graph import StateGraph, START, END

from edges.decide_1 import decide_1
from nodes.node1 import node_1
from nodes.node2 import node_2
from nodes.node3 import node_3
from states.state import State


def create_graph_builder() -> StateGraph:
    builder = StateGraph(State)
    builder.add_node("node_1", node_1)
    builder.add_node("decide_1", decide_1)
    builder.add_node("node_2", node_2)
    builder.add_node("node_3", node_3)

    builder.add_edge(START, "node_1")
    builder.add_edge("node_1", "decide_1")
    builder.add_edge("node_2", END)
    builder.add_edge("node_3", END)
    return builder


def create_graph():
    return create_graph_builder().compile()


graph = create_graph()