Spring AI Alibaba 98｜Graph Core · 子图作为 CompiledGraph · 基本用法

在 Spring AI Alibaba 中，可以先编译 StateGraph 得到 CompiledGraph，然后在其他 Graph 中复用，这种方式性能更好且更灵活。

创建并编译子图

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.AsyncNodeAction.node_async;

/**
* 创建并编译子图
*/
public static CompiledGraph createAndCompileSubGraph() throws GraphStateException {
KeyStrategyFactory subKeyFactory = () -> {
  HashMap<String, KeyStrategy> strategies = new HashMap<>();
  strategies.put("input", new ReplaceStrategy());
  strategies.put("output", new ReplaceStrategy());
  return strategies;
};

// 定义并编译子图
StateGraph subGraphDef = new StateGraph(subKeyFactory)
          .addNode("process", node_async(state -> {
      String input = (String) state.value("input").orElse("");
      String output = "Processed: " + input.toUpperCase();
      return Map.of("output", output);
  }))
          .addEdge(START, "process")
          .addEdge("process", END);

// 编译子图
  return subGraphDef.compile();
}

在节点中使用 CompiledGraph

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

import java.util.Map;

/**
* 在节点中使用 CompiledGraph
*/
public static class CompiledSubGraphNode implements NodeAction {

  private final CompiledGraph compiledGraph;

  public CompiledSubGraphNode(CompiledGraph compiledGraph) {
      this.compiledGraph = compiledGraph;
  }

  @Override
  public Map<String, Object> apply(OverAllState state) {
      // 从父状态提取输入
      String input = (String) state.value("data").orElse("");

      // 执行编译好的子图
      Map<String, Object> subInput = Map.of("input", input);
      OverAllState subResult = compiledGraph.invoke(subInput, RunnableConfig.builder().build()).orElseThrow();

      // 提取子图输出
      String output = (String) subResult.value("output").orElse("");
      return Map.of("result", output);
  }
}

在父图中使用

import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.StateGraph;

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.AsyncNodeAction.node_async;

/**
* 在父图中使用
*/
public static CompiledGraph useInParentGraph(CompiledGraph compiledSubGraph) throws GraphStateException {
KeyStrategyFactory parentKeyFactory = () -> {
  HashMap<String, KeyStrategy> strategies = new HashMap<>();
  strategies.put("data", new ReplaceStrategy());
  strategies.put("result", new ReplaceStrategy());
  return strategies;
};

StateGraph parentGraph = new StateGraph(parentKeyFactory)
          .addNode("prepare", node_async(state ->
      Map.of("data", "hello world")))
          .addNode("subgraph", compiledSubGraph)
          .addNode("finalize", node_async(state -> {
      String result = (String) state.value("result").orElse("");
      return Map.of("final", "Done: " + result);
  }))
          .addEdge(START, "prepare")
  .addEdge("prepare", "subgraph")
  .addEdge("subgraph", "finalize")
          .addEdge("finalize", END);

  return parentGraph.compile();
}

同一个 CompiledGraph 可以在多个地方复用：

/**
* 多个子图复用
*/
public static CompiledGraph reuseMultipleSubGraphs(CompiledGraph dataProcessor) throws GraphStateException {
  KeyStrategyFactory keyFactory = () -> {
      HashMap<String, KeyStrategy> strategies = new HashMap<>();
      strategies.put("data", new ReplaceStrategy());
      strategies.put("result", new ReplaceStrategy());
      return strategies;
  };

  // 在多个节点中复用
  StateGraph mainGraph = new StateGraph(keyFactory)
          .addNode("process1", dataProcessor)
          .addNode("process2", dataProcessor)
          .addNode("process3", dataProcessor)
          .addEdge(START, "process1")
          .addEdge("process1", "process2")
          .addEdge("process2", "process3")
          .addEdge("process3", END);

  return mainGraph.compile();
}

// 在多个节点中复用
StateGraph mainGraph = new StateGraph(keyFactory)
  .addNode("process1", nodeasync(new CompiledSubGraphNode(dataProcessor)))
  .addNode("process2", nodeasync(new CompiledSubGraphNode(dataProcessor)))
  .addNode("process3", nodeasync(new CompiledSubGraphNode(dataProcessor)))
  .addEdge(StateGraph.START, "process1")
  .addEdge("process1", "process2")
  .addEdge("process2", "process3")
  .addEdge("process3", StateGraph.END);

可配置的 CompiledGraph 包装器

public class ConfigurableCompiledSubGraph implements NodeAction {

  private final CompiledGraph compiledGraph;
  private final Map<String, String> inputMapping;
  private final Map<String, String> outputMapping;

  public ConfigurableCompiledSubGraph(
      CompiledGraph compiledGraph,
      Map<String, String> inputMapping,   // parentKey -> subKey
      Map<String, String> outputMapping   // subKey -> parentKey
  ) {
      this.compiledGraph = compiledGraph;
      this.inputMapping = inputMapping;
      this.outputMapping = outputMapping;
  }

  @Override
  public Map<String, Object> apply(OverAllState parentState) {
      // 映射输入
      Map<String, Object> subInput = new HashMap<>();
      inputMapping.forEach((parentKey, subKey) -> {
          parentState.value(parentKey).ifPresent(value ->
              subInput.put(subKey, value));
      });

      // 执行子图
      OverAllState subResult = compiledGraph.invoke(subInput);

      // 映射输出
      Map<String, Object> parentOutput = new HashMap<>();
      outputMapping.forEach((subKey, parentKey) -> {
          subResult.value(subKey).ifPresent(value ->
              parentOutput.put(parentKey, value));
      });

      return parentOutput;
  }
}

// 使用
CompiledGraph processor = createProcessorGraph().compile();

NodeAction node1 = new ConfigurableCompiledSubGraph(
  processor,
  Map.of("userInput", "input"),      // userInput -> input
  Map.of("output", "processedData")  // output -> processedData
);

NodeAction node2 = new ConfigurableCompiledSubGraph(
  processor,
  Map.of("rawData", "input"),        // rawData -> input
  Map.of("output", "cleanedData")    // output -> cleanedData
);

import java.util.concurrent.CompletableFuture;
import java.util.List;

public class ParallelCompiledGraphNode implements NodeAction {

  private final List<CompiledGraph> graphs;

  public ParallelCompiledGraphNode(List<CompiledGraph> graphs) {
      this.graphs = graphs;
  }

  @Override
  public Map<String, Object> apply(OverAllState state) {
      String input = (String) state.value("input").orElse("");

      // 并行执行所有子图
      List<CompletableFuture<OverAllState>> futures = graphs.stream()
          .map(graph -> CompletableFuture.supplyAsync(() ->
              graph.invoke(Map.of("input", input))))
          .toList();

      // 等待所有完成
      List<OverAllState> results = futures.stream()
          .map(CompletableFuture::join)
          .toList();

      // 合并结果
      List<String> outputs = results.stream()
          .map(r -> (String) r.value("output").orElse(""))
          .toList();

      return Map.of("results", outputs);
  }
}

// 使用
CompiledGraph graph1 = createGraph1().compile();
CompiledGraph graph2 = createGraph2().compile();
CompiledGraph graph3 = createGraph3().compile();

NodeAction parallelNode = new ParallelCompiledGraphNode(
  List.of(graph1, graph2, graph3)
);

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