Spring AI Alibaba 102｜Graph Core · 子图作为 StateGraph

在 Spring AI Alibaba 中，可以将一个 StateGraph 嵌入到另一个 StateGraph 中，实现复杂工作流的模块化设计。

将 StateGraph 作为子图使用有以下优势：

• 复用性: 相同的子图可以在多个父图中复用
• 可维护性: 独立开发和测试每个子图
• 清晰性: 层次化结构使工作流更易理解

import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.KeyStrategyFactory;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;

import java.util.HashMap;
import java.util.Map;

import static com.alibaba.cloud.ai.graph.StateGraph.END;
import static com.alibaba.cloud.ai.graph.StateGraph.START;
import static com.alibaba.cloud.ai.graph.action.AsyncEdgeAction.edge_async;
import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;

/**
* 定义子图
*/
public static StateGraph createProcessingSubGraph() throws GraphStateException {
  KeyStrategyFactory keyFactory = () -> {
      HashMap<String, KeyStrategy> strategies = new HashMap<>();
      strategies.put("input", new ReplaceStrategy());
      strategies.put("output", new ReplaceStrategy());
      strategies.put("valid", new ReplaceStrategy());
      return strategies;
  };

  return new StateGraph(keyFactory)
          .addNode("validate", node_async(state -> {
          String input = (String) state.value("input").orElse("");
          boolean isValid = input != null && !input.isEmpty();
          return Map.of("valid", isValid);
      }))
          .addNode("transform", node_async(state -> {
          String input = (String) state.value("input").orElse("");
          String transformed = input.toUpperCase();
          return Map.of("output", transformed);
      }))
          .addEdge(START, "validate")
      .addConditionalEdges("validate",
                  edge_async(state -> {
              Boolean valid = (Boolean) state.value("valid").orElse(false);
              return valid ? "valid" : "invalid";
          }),
          Map.of(
              "valid", "transform",
                          "invalid", END
          ))
          .addEdge("transform", END);
}

方式 1: 直接嵌入

/**
* 在父图中集成子图 - 方式 1: 直接嵌入
*/
public static StateGraph createParentGraphWithDirectEmbedding() throws GraphStateException {
  KeyStrategyFactory keyFactory = () -> {
      HashMap<String, KeyStrategy> strategies = new HashMap<>();
      strategies.put("data", new ReplaceStrategy());
      strategies.put("output", new ReplaceStrategy());
      strategies.put("result", new ReplaceStrategy());
      return strategies;
  };

  StateGraph subGraph = createProcessingSubGraph();

  return new StateGraph(keyFactory)
          .addNode("prepare", node_async(state -> {
          return Map.of("data", "hello world");
      }))
      // 将子图作为节点添加
      .addNode("process", subGraph)
          .addNode("finalize", node_async(state -> {
          String output = (String) state.value("output").orElse("");
          return Map.of("result", "Final: " + output);
      }))
          .addEdge(START, "prepare")
      .addEdge("prepare", "process")
      .addEdge("process", "finalize")
          .addEdge("finalize", END);
}

方式 2: 使用编译后的子图

import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.RunnableConfig;

import java.util.Map;

/**
* 在父图中集成子图 - 方式 2: 使用编译后的子图
*/
public static StateGraph createParentGraphWithCompiledSubGraph() throws GraphStateException {
  KeyStrategyFactory keyFactory = () -> {
      HashMap<String, KeyStrategy> strategies = new HashMap<>();
      strategies.put("data", new ReplaceStrategy());
      strategies.put("output", new ReplaceStrategy());
      strategies.put("processed", new ReplaceStrategy());
      return strategies;
  };

// 先编译子图
CompiledGraph compiledSubGraph = createProcessingSubGraph().compile();

// 在父图中使用
  return new StateGraph(keyFactory)
          .addNode("prepare", node_async(state -> {
      return Map.of("data", "input");
  }))
          .addNode("process", node_async(state -> {
      // 手动调用子图
      Map<String, Object> subInput = Map.of(
          "input", state.value("data").orElse("")
      );
              OverAllState subResult = compiledSubGraph.invoke(subInput, RunnableConfig.builder().build()).orElseThrow();
      return Map.of("processed", subResult.value("output").orElse(""));
  }))
          .addEdge(START, "prepare")
  .addEdge("prepare", "process")
          .addEdge("process", END);
}

共享状态

父子图共享相同的状态键：

// 父子图使用相同的 KeyStrategyFactory
KeyStrategyFactory sharedKeyFactory = () -> {
  HashMap<String, KeyStrategy> strategies = new HashMap<>();
  strategies.put("data", new ReplaceStrategy());
  strategies.put("result", new ReplaceStrategy());
  return strategies;
};

StateGraph subGraph = new StateGraph(sharedKeyFactory)
  .addNode("process", nodeasync(state -> {
      String data = (String) state.value("data").orElse("");
      return Map.of("result", "Processed: " + data);
  }))
  .addEdge(StateGraph.START, "process")
  .addEdge("process", StateGraph.END);

StateGraph parentGraph = new StateGraph(sharedKeyFactory)
  .addNode("prepare", nodeasync(state ->
      Map.of("data", "Input")))
  .addNode("sub", subGraph)  // 子图直接访问 "data" 并写入 "result"
  .addEdge(StateGraph.START, "prepare")
  .addEdge("prepare", "sub")
  .addEdge("sub", StateGraph.END);

状态隔离

子图使用独立的状态空间：

public class IsolatedSubGraphNode implements NodeAction {
  private final CompiledGraph subGraph;

  public IsolatedSubGraphNode(StateGraph subGraphDef) {
      this.subGraph = subGraphDef.compile();
  }

  @Override
  public Map<String, Object> apply(OverAllState parentState) {
      // 提取父状态数据
      String input = (String) parentState.value("input").orElse("");

      // 创建子图独立状态
      Map<String, Object> subState = Map.of("subInput", input);

      // 执行子图
      OverAllState subResult = subGraph.invoke(subState);

      // 将子图结果映射回父状态
      String output = (String) subResult.value("subOutput").orElse("");
      return Map.of("output", output);
  }
}

子图可以包含其他子图：

// Level 2 子图
StateGraph level2SubGraph = new StateGraph(keyFactory)
  .addNode("level2Node", nodeasync(state ->
      Map.of("level2Result", "Level 2")))
  .addEdge(StateGraph.START, "level2Node")
  .addEdge("level2Node", StateGraph.END);

// Level 1 子图，包含 Level 2
StateGraph level1SubGraph = new StateGraph(keyFactory)
  .addNode("level1Node", nodeasync(state ->
      Map.of("level1Result", "Level 1")))
  .addNode("level2", level2SubGraph)
  .addEdge(StateGraph.START, "level1Node")
  .addEdge("level1Node", "level2")
  .addEdge("level2", StateGraph.END);

// 父图，包含 Level 1
StateGraph parentGraph = new StateGraph(keyFactory)
  .addNode("parentNode", nodeasync(state ->
      Map.of("parentResult", "Parent")))
  .addNode("level1", level1SubGraph)
  .addEdge(StateGraph.START, "parentNode")
  .addEdge("parentNode", "level1")
  .addEdge("level1", StateGraph.END);

本文整理自 Spring AI Alibaba 官方文档《子图作为 StateGraph》：https://java2ai.com/docs/frameworks/graph-core/examples/subgraph-as-stategraph