public class Node1Action implements NodeAction<AgentState> {
    @Override
    public Map<String, Object> apply(AgentState state) throws Exception {
        System.out.println("current Node: node-1");
        //模拟节点耗时
        Thread.sleep(1000);
        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");
        //模拟节点耗时
        Thread.sleep(2000);
        return Map.of("myData", "node2-my-value",
                "node2Key", "node2-value");
    }
}

public class HookSampleApplication {

    public static void main(String[] args) throws GraphStateException {
        runSequenceGraphWithOnlyStaticEdges();
        out.println("\n========== 下面使用带条件边的图，Edge Hook 会执行 ==========");
        runGraphWithConditionalEdge();
    }

    /**
     * 纯静态边：只有 Node Hook 会执行，Edge Hook 不会执行
     */
    private static void runSequenceGraphWithOnlyStaticEdges() throws GraphStateException {
        StateGraph<AgentState> sequenceGraph = getSequenceGraph();

        sequenceGraph.addBeforeCallNodeHook((String node, AgentState data, RunnableConfig config) -> {
            out.println("Before calling node: " + node + ", data: " + data.data());
            return CompletableFuture.completedFuture(data.data());
        });

        sequenceGraph.addAfterCallNodeHook((String node, AgentState data, RunnableConfig config, Map<String, Object> lastResult) -> {
            out.println("After calling node: " + node + ", data: " + data.data() + ", lastResult: " + lastResult);
            return CompletableFuture.completedFuture(lastResult);
        });

        sequenceGraph.addWrapCallNodeHook((String node, AgentState data, RunnableConfig config, AsyncNodeActionWithConfig<AgentState> action) -> {
            out.println("Wrap calling node: " + node + ", data: " + data.data());
            long start = System.currentTimeMillis();
            return action.apply(data, config).whenComplete((result, error) -> {
                var ms = System.currentTimeMillis() - start;
                out.println(String.format("node '%s' took %d ms", node, ms));
            });
        });

        sequenceGraph.addBeforeCallEdgeHook((String sourceId, AgentState state, RunnableConfig config) -> {
            out.println("Before calling edge: " + sourceId);
            return CompletableFuture.completedFuture(new Command(state.data()));
        });

        sequenceGraph.addAfterCallEdgeHook((String sourceId, AgentState state, RunnableConfig config, Command lastResult) -> {
            out.println("After calling edge: " + sourceId);
            return CompletableFuture.completedFuture(lastResult);
        });

        sequenceGraph.addWrapCallEdgeHook((String sourceId, AgentState state, RunnableConfig config, AsyncCommandAction<AgentState> action) -> {
            out.println("Wrap calling edge: " + sourceId);
            long start = System.currentTimeMillis();
            return action.apply(state, config).whenComplete((result, error) -> {
                var ms = System.currentTimeMillis() - start;
                out.println(String.format("source-node '%s' took %d ms", sourceId, ms));
            });

        });

        out.println(sequenceGraph.getGraph(GraphRepresentation.Type.MERMAID, "NodeHook Graph", true).content());

        sequenceGraph.compile().invoke(Map.of("test", "test-init-value")).ifPresent(c -> {
            System.out.println(c.data());
        });
    }

    /**
     * 带条件边：从 node-1 经条件边到 node-2，会触发 Edge Hook
     */
    private static void runGraphWithConditionalEdge() throws GraphStateException {
        StateGraph<AgentState> graph = getGraphWithConditionalEdge();

        graph.addBeforeCallNodeHook((String node, AgentState data, RunnableConfig config) -> {
            out.println("Before calling node: " + node + ", data: " + data.data());
            return CompletableFuture.completedFuture(data.data());
        });
        graph.addAfterCallNodeHook((String node, AgentState data, RunnableConfig config, Map<String, Object> lastResult) -> {
            out.println("After calling node: " + node + ", data: " + data.data() + ", lastResult: " + lastResult);
            return CompletableFuture.completedFuture(lastResult);
        });
        graph.addWrapCallNodeHook((String node, AgentState data, RunnableConfig config, AsyncNodeActionWithConfig<AgentState> action) -> {
            out.println("Wrap calling node: " + node + ", data: " + data.data());
            long start = System.currentTimeMillis();
            return action.apply(data, config).whenComplete((result, error) -> {
                var ms = System.currentTimeMillis() - start;
                out.println(String.format("node '%s' took %d ms", node, ms));
            });
        });

        graph.addBeforeCallEdgeHook((String sourceId, AgentState state, RunnableConfig config) -> {
            out.println("Before calling edge: " + sourceId);
            return CompletableFuture.completedFuture(new Command(state.data()));
        });
        graph.addAfterCallEdgeHook((String sourceId, AgentState state, RunnableConfig config, Command lastResult) -> {
            out.println("After calling edge: " + sourceId);
            return CompletableFuture.completedFuture(lastResult);
        });
        graph.addWrapCallEdgeHook((String sourceId, AgentState state, RunnableConfig config, AsyncCommandAction<AgentState> action) -> {
            out.println("Wrap calling edge: " + sourceId);
            long start = System.currentTimeMillis();
            return action.apply(state, config).whenComplete((result, error) -> {
                var ms = System.currentTimeMillis() - start;
                out.println(String.format("source-node '%s' took %d ms", sourceId, ms));
            });
        });

        out.println(graph.getGraph(GraphRepresentation.Type.MERMAID, "NodeHook And EdgeHook Graph", true).content());

        graph.compile().invoke(Map.of("test", "test-init-value")).ifPresent(c -> System.out.println(c.data()));
    }

    public static StateGraph<AgentState> getSequenceGraph() 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-1", "node-2")
                .addEdge("node-2", GraphDefinition.END);
    }

    /**
     * 含一条条件边：node-1 通过条件边到 node-2，用于演示 Edge Hook 触发
     */
    public static StateGraph<AgentState> getGraphWithConditionalEdge() 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")
                .addConditionalEdges("node-1", state -> CompletableFuture.completedFuture("toNode2"), Map.of("toNode2", "node-2"))
                .addEdge("node-2", GraphDefinition.END);
    }
}

运行结果

Before calling node: node-1, data: {test=test-init-value}
Wrap calling node: node-1, data: {test=test-init-value}
current Node: node-1
node 'node-1' took 1002 ms
After calling node: node-1, data: {test=test-init-value}, lastResult: {myData=node1-my-value, node1Key=node1-value}
Before calling node: node-2, data: {node1Key=node1-value, test=test-init-value, myData=node1-my-value}
Wrap calling node: node-2, data: {node1Key=node1-value, test=test-init-value, myData=node1-my-value}
current Node: node-2
node 'node-2' took 1999 ms
After calling node: node-2, data: {node1Key=node1-value, test=test-init-value, myData=node1-my-value}, lastResult: {myData=node2-my-value, node2Key=node2-value}
{node1Key=node1-value, test=test-init-value, node2Key=node2-value, myData=node2-my-value}

========== 下面使用带条件边的图，Edge Hook 会执行 ==========

Before calling node: node-1, data: {test=test-init-value}
Wrap calling node: node-1, data: {test=test-init-value}
current Node: node-1
node 'node-1' took 1001 ms
After calling node: node-1, data: {test=test-init-value}, lastResult: {myData=node1-my-value, node1Key=node1-value}
Before calling edge: node-1
Wrap calling edge: node-1
source-node 'node-1' took 1 ms
After calling edge: node-1
Before calling node: node-2, data: {node1Key=node1-value, test=test-init-value, myData=node1-my-value}
Wrap calling node: node-2, data: {node1Key=node1-value, test=test-init-value, myData=node1-my-value}
current Node: node-2
node 'node-2' took 2000 ms
After calling node: node-2, data: {node1Key=node1-value, test=test-init-value, myData=node1-my-value}, lastResult: {myData=node2-my-value, node2Key=node2-value}
{node1Key=node1-value, test=test-init-value, node2Key=node2-value, myData=node2-my-value}