Spring AI Alibaba 74｜Graph Core · 在节点操作中整合流式输出

Spring AI Alibaba Graph 内置了对流式处理的原生支持，框架统一是使用 Flux 来在框架中定义和传递流，与 Spring 生态的流式处理保持一致。以下是从 Graph 运行中流式返回输出的不同方式。

.stream() 是一个用于从图运行中流式返回输出的方法。它返回一个 Flux，请记住由于 Flux 流式的特性，流返回后并不会立即出触发图引擎的执行，你需要执行类似 Flux.subscribe() 的操作才能真正启动流引擎。

目前 Flux 返回的是 NodeOutput 类的实例，该类基本上报告执行的节点名称和结果状态。

流的组合（嵌入和组合）

Flux 支持多个流的合并、转换、组合等操作，具备非常强大的能力，这在处理图中多个流式节点时会非常有用。具体使用方式可搜索 Spring Reactor 学习。

在 Spring AI Alibaba Graph 中，您可以在节点操作中直接返回 Flux 对象，框架会自动处理流式输出。

流式节点实现

import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatResponse;

import java.util.Map;

import reactor.core.publisher.Flux;

/**
* 流式节点实现 - 使用 GraphFluxGenerator 处理流式响应
*/
public static class StreamingNode implements NodeAction {

  private final ChatClient chatClient;
  private final String nodeId;

  public StreamingNode(ChatClient.Builder chatClientBuilder, String nodeId) {
      this.chatClient = chatClientBuilder.build();
      this.nodeId = nodeId;
  }

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

      // 获取流式响应
      Flux<ChatResponse> chatResponseFlux = chatClient.prompt()
              .user(query)
              .stream()
              .chatResponse();

      // 将流式响应存储在状态中
      return Map.of("messages", chatResponseFlux);
  }
}

处理流式输出的节点

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

import java.util.Map;

/**
* 处理流式输出的节点 - 接收并处理流式响应
*/
public static class ProcessStreamingNode implements NodeAction {

  @Override
  public Map<String, Object> apply(OverAllState state) {
       // 请注意，虽然上一个节点返回的是Flux对象，但是在引擎运行到当前节点时，
	// 框架已经完成了对上一个节点Flux对象的自动订阅与消费，并将最终的结果汇总后添加到了 messages key 中（基于 AppendStrategy）
	Object messages = state.value("messages").orElse("");
	String result = "流式响应已处理完成: " + messages;
	return Map.of("result", result);
  }
}

完整示例：使用流式输出的图

import com.alibaba.cloud.ai.graph.CompileConfig;
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.RunnableConfig;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.action.AsyncNodeAction;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.state.strategy.AppendStrategy;
import com.alibaba.cloud.ai.graph.streaming.StreamingOutput;

import org.springframework.ai.chat.client.ChatClient;

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;

/**
* 使用 StateGraph 实现流式输出的完整示例
*/
public static void streamLLMTokens(ChatClient.Builder chatClientBuilder) throws GraphStateException {
  // 定义状态策略
  KeyStrategyFactory keyStrategyFactory = () -> {
      Map<String, KeyStrategy> keyStrategyMap = new HashMap<>();
      keyStrategyMap.put("query", new AppendStrategy());
      keyStrategyMap.put("messages", new AppendStrategy());
      keyStrategyMap.put("result", new AppendStrategy());
      return keyStrategyMap;
  };

  // 创建流式节点
  StreamingNode streamingNode = new StreamingNode(chatClientBuilder, "streaming_node");

  // 创建处理节点
  ProcessStreamingNode processNode = new ProcessStreamingNode();

  // 构建图
  StateGraph stateGraph = new StateGraph(keyStrategyFactory)
          .addNode("streaming_node", AsyncNodeAction.node_async(streamingNode))
          .addNode("process_node", AsyncNodeAction.node_async(processNode))
          .addEdge(START, "streaming_node")
          .addEdge("streaming_node", "process_node")
          .addEdge("process_node", END);

  // 编译图
  CompiledGraph graph = stateGraph.compile(
          CompileConfig.builder()
                  .build()
  );

  // 创建配置
  RunnableConfig config = RunnableConfig.builder()
          .threadId("streaming_thread")
          .build();

  // 使用流式方式执行图
  System.out.println("开始流式输出...
");

  graph.stream(Map.of("query", "请用一句话介绍 Spring AI"), config)
          .doOnNext(output -> {
              // 处理流式输出
              if (output instanceof StreamingOutput<?> streamingOutput) {
                  // 流式输出块
                  String chunk = streamingOutput.chunk();
                  if (chunk != null && !chunk.isEmpty()) {
                      System.out.print(chunk); // 实时打印流式内容
                  }
              }
              else {
                  // 普通节点输出
                  String nodeId = output.node();
                  Map<String, Object> state = output.state().data();
                  System.out.println("
节点 '" + nodeId + "' 执行完成");
                  if (state.containsKey("result")) {
                      System.out.println("最终结果: " + state.get("result"));
                  }
              }
          })
          .doOnComplete(() -> {
              System.out.println("

流式输出完成");
          })
          .doOnError(error -> {
              System.err.println("流式输出错误: " + error.getMessage());
          })
          .blockLast(); // 阻塞等待流完成
}

假设我们有一个调用 LLM 的流式节点：

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatResponse;

import reactor.core.publisher.Flux;

ChatClient chatClient = chatClientBuilder.build();

Flux<ChatResponse> chatResponseFlux = chatClient.prompt()
      .user("给我讲一个笑话")
      .stream()
      .chatResponse();

chatResponseFlux.subscribe(
      response -> {
          // 处理每个 token
          String content = response.getResult().getOutput().getText();
          System.out.print(content);
      },
      error -> {
          // 处理错误
          System.err.println("错误: " + error.getMessage());
      },
      () -> {
          // 完成处理
          System.out.println("
流式处理完成");
      }
);

本文整理自 Spring AI Alibaba 官方文档《流式输出》：https://java2ai.com/docs/frameworks/graph-core/core/streaming