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

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

1 public class Node3Action implements NodeAction<AgentState> {
2     @Override
3     public Map<String, Object> apply(AgentState state) throws Exception {
4         System.out.println("current Node: node-3");
5         return Map.of("myData", "node3-my-value",
6                 "node3Key", "node3-value");
7     }
8 }

 1 public class Node1Action implements NodeAction<AgentState> {
 2 
 3     private static final String KEY_NEXT_NODE = "nextNode";
 4     private static final String DEFAULT_NEXT = "3";
 5 
 6     @Override
 7     public Map<String, Object> apply(AgentState state) throws Exception {
 8         System.out.println("current Node: node-1");
 9         String next = resolveNextNode(state);
10         return Map.of(
11                 "myData", "node1-my-value",
12                 "node1Key", "node1-value",
13                 KEY_NEXT_NODE, next);
14     }
15 
16     private String resolveNextNode(AgentState state) {
17         Optional<Object> routeTo = state.value(KEY_NEXT_NODE);
18         if (routeTo.isEmpty()) {
19             return DEFAULT_NEXT;
20         }
21         Object v = routeTo.get();
22         String s = v.toString();
23         if ("2".equals(s) || "3".equals(s)) {
24             return s;
25         }
26         return DEFAULT_NEXT;
27     }
28 }

 1 public class RoutingEdgeAction implements EdgeAction<AgentState> {
 2 
 3     private static final String KEY_NEXT_NODE = "nextNode";
 4 
 5     private final Set<String> allowedNodes;
 6     private final String defaultNode;
 7 
 8     public RoutingEdgeAction(Set<String> allowedNodes, String defaultNode) {
 9         this.allowedNodes = Set.copyOf(allowedNodes);
10         if (!this.allowedNodes.contains(defaultNode)) {
11             throw new IllegalArgumentException("defaultNode must be in allowedNodes: " + defaultNode);
12         }
13         this.defaultNode = defaultNode;
14     }
15 
16     @Override
17     public String apply(AgentState state) throws Exception {
18         //根据上1个节点的状态值，执行路由
19         System.out.println("current Edge: routing edge");
20         String next = resolveNextNode(state);
21         if (allowedNodes.contains(next)) {
22             return next;
23         }
24         return defaultNode;
25     }
26 
27     private String resolveNextNode(AgentState state) {
28         Optional<Object> nextOpt = state.value(KEY_NEXT_NODE);
29         if (nextOpt.isEmpty()) {
30             return defaultNode;
31         }
32         Object v = nextOpt.get();
33         return v.toString();
34     }
35 }

 1 public class ConditionalGraphApplication {
 2 
 3     private static final Set<String> CONDITIONAL_TARGETS = Set.of("2", "3");
 4     private static final String DEFAULT_TARGET = "2";
 5 
 6     public static void main(String[] args) throws GraphStateException {
 7         StateGraph<AgentState> conditionalGraph = getConditionalGraph();
 8 
 9         System.out.println(conditionalGraph.getGraph(GraphRepresentation.Type.MERMAID, "conditional Graph", true).content());
10 
11         try {
12             conditionalGraph.compile()
13                     .invoke(Map.of("test", "test-init-value"))
14                     .ifPresent(c -> System.out.println(c.data()));
15         } catch (GraphStateException e) {
16             System.err.println("Graph execution failed: " + e.getMessage());
17             throw e;
18         } catch (Exception e) {
19             System.err.println("Unexpected error: " + e.getMessage());
20             throw new RuntimeException(e);
21         }
22     }
23 
24     public static StateGraph<AgentState> getConditionalGraph() throws GraphStateException {
25         return new StateGraph<>(AgentState::new)
26                 .addNode("node-1", node_async(new Node1Action()))
27                 .addNode("node-2", node_async(new Node2Action()))
28                 .addNode("node-3", node_async(new Node3Action()))
29                 .addEdge(START, "node-1")
30                 .addConditionalEdges("node-1",
31                         edge_async(new RoutingEdgeAction(CONDITIONAL_TARGETS, DEFAULT_TARGET)),
32                         // 如果返回2，则跳到node-2，如果返回3，则跳到node-3
33                         Map.of("2", "node-2",
34                                 "3", "node-3"))
35                 .addEdge("node-2", GraphDefinition.END)
36                 .addEdge("node-3", GraphDefinition.END);
37     }
38 }

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