LangGraph4j 学习系列(2)-条件工作流

接上节继续，本节将演示条件工作流如何用langgraph4j实现

经过node-1节点时，根据该节点的输出状态值，来决定流向node-2，还是node-3 ?

 

一、定义节点

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 Node3Action implements NodeAction<AgentState> {
    @Override
    public Map<String, Object> apply(AgentState state) throws Exception {
        System.out.println("current Node: node-3");
        return Map.of("myData", "node3-my-value",
                "node3Key", "node3-value");
    }
}

public class Node1Action implements NodeAction<AgentState> {

    private static final String KEY_NEXT_NODE = "nextNode";
    private static final String DEFAULT_NEXT = "3";

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

    private String resolveNextNode(AgentState state) {
        Optional<Object> routeTo = state.value(KEY_NEXT_NODE);
        if (routeTo.isEmpty()) {
            return DEFAULT_NEXT;
        }
        Object v = routeTo.get();
        String s = v.toString();
        if ("2".equals(s) || "3".equals(s)) {
            return s;
        }
        return DEFAULT_NEXT;
    }
}

 

二、定义条件边

public class RoutingEdgeAction implements EdgeAction<AgentState> {

    private static final String KEY_NEXT_NODE = "nextNode";

    private final Set<String> allowedNodes;
    private final String defaultNode;

    public RoutingEdgeAction(Set<String> allowedNodes, String defaultNode) {
        this.allowedNodes = Set.copyOf(allowedNodes);
        if (!this.allowedNodes.contains(defaultNode)) {
            throw new IllegalArgumentException("defaultNode must be in allowedNodes: " + defaultNode);
        }
        this.defaultNode = defaultNode;
    }

    @Override
    public String apply(AgentState state) throws Exception {
        //根据上1个节点的状态值，执行路由
        System.out.println("current Edge: routing edge");
        String next = resolveNextNode(state);
        if (allowedNodes.contains(next)) {
            return next;
        }
        return defaultNode;
    }

    private String resolveNextNode(AgentState state) {
        Optional<Object> nextOpt = state.value(KEY_NEXT_NODE);
        if (nextOpt.isEmpty()) {
            return defaultNode;
        }
        Object v = nextOpt.get();
        return v.toString();
    }
}

 

三、完整示例

public class ConditionalGraphApplication {

    private static final Set<String> CONDITIONAL_TARGETS = Set.of("2", "3");
    private static final String DEFAULT_TARGET = "2";

    public static void main(String[] args) throws GraphStateException {
        StateGraph<AgentState> conditionalGraph = getConditionalGraph();

        System.out.println(conditionalGraph.getGraph(GraphRepresentation.Type.MERMAID, "conditional Graph", true).content());

        try {
            conditionalGraph.compile()
                    .invoke(Map.of("test", "test-init-value"))
                    .ifPresent(c -> System.out.println(c.data()));
        } catch (GraphStateException e) {
            System.err.println("Graph execution failed: " + e.getMessage());
            throw e;
        } catch (Exception e) {
            System.err.println("Unexpected error: " + e.getMessage());
            throw new RuntimeException(e);
        }
    }

    public static StateGraph<AgentState> getConditionalGraph() throws GraphStateException {
        return new StateGraph<>(AgentState::new)
                .addNode("node-1", node_async(new Node1Action()))
                .addNode("node-2", node_async(new Node2Action()))
                .addNode("node-3", node_async(new Node3Action()))
                .addEdge(START, "node-1")
                .addConditionalEdges("node-1",
                        edge_async(new RoutingEdgeAction(CONDITIONAL_TARGETS, DEFAULT_TARGET)),
                        // 如果返回2，则跳到node-2，如果返回3，则跳到node-3
                        Map.of("2", "node-2",
                                "3", "node-3"))
                .addEdge("node-2", GraphDefinition.END)
                .addEdge("node-3", GraphDefinition.END);
    }
}

 

四、运行结果

current Node: node-1
current Edge: routing edge
current Node: node-3
{node1Key=node1-value, test=test-init-value, nextNode=3, node3Key=node3-value, myData=node3-my-value}

 

文中源码：https://github.com/yjmyzz/langgraph4j-study/tree/main/src/main/java/org/bsc/langgraph4j/agent/_06_conditional