public class Node1Action implements NodeAction<AgentState> {
    @Override
    public Map<String, Object> apply(AgentState state) throws Exception {
        System.out.println("current Node: node-1");
        return Map.of("myData", "node1-my-value",
                "node1Key", "node1-value");
    }
}

public class Node2Action implements NodeAction<AgentState> {
    @Override
    public Map<String, Object> apply(AgentState state) throws Exception {
        System.out.println("current Node: node-2");
        return Map.of("myData", "node2-my-value",
                "node2Key", "node2-value");
    }
}

public class LoopGraphApplication {

    public static void main(String[] args) throws GraphStateException {
        StateGraph<AgentState> sequenceGraph = getLoopGraph();
        System.out.println(sequenceGraph.getGraph(GraphRepresentation.Type.MERMAID, "loop Graph", true).content());
        sequenceGraph.compile().invoke(Map.of("loopCount", 0L)).ifPresent(c -> {
            System.out.println(c.data());
        });
    }

    private static final int MAX_LOOP_ITERATIONS = 3;

    public static StateGraph<AgentState> getLoopGraph() throws GraphStateException {
        return new StateGraph<>(AgentState::new)
                .addNode("node-1", node_async(new Node1Action()))
                .addNode("node-2", node_async(new Node2Action()))
                .addEdge(GraphDefinition.START, "node-1")
                .addEdge("node-2", "node-1")
                //循环的跳出条件比较简单，3次后退出，这里就不单独定义EdgeAction类了，用lambda表达式
                .addConditionalEdges("node-1", state -> {
                    long count = getLoopCount(state);
                    System.out.println("loop Count: " + count);
                    if (count >= MAX_LOOP_ITERATIONS) {
                        return CompletableFuture.completedFuture("exit");
                    }
                    return CompletableFuture.completedFuture("continue");
                }, Map.of(
                        "exit", GraphDefinition.END,
                        "continue", "node-2"));
    }

    private static long getLoopCount(AgentState state) {
        Optional<Object> loopCount = state.value("loopCount");
        if (loopCount.isEmpty()) {
            return 0L;
        }
        Object v = loopCount.get();
        if (v instanceof Number n) {
            return n.longValue();
        }
        return Long.parseLong(v.toString());
    }


}

current Node: node-1
loop Count: 0
current Node: node-2
current Node: node-1
loop Count: 1
current Node: node-2
current Node: node-1
loop Count: 2
current Node: node-2
current Node: node-1
loop Count: 3
{node1Key=node1-value, loopCount=3, node2Key=node2-value, myData=node1-my-value}