AI项目博客
记录 AI 工程实践与大模型应用

Spring AI Alibaba 70|Graph Core · 内存管理 · 添加短期内存

AI 应用程序需要支持在同一轮会话的多条消息间共享上下文,或者在不同的会话场景先共享上下文。在 Spring AI Alibaba Graph 中,您可以添加两种类型的内存:

  • 添加短期内存作为智能体状态的一部分,支持与智能体进行多轮聊天对话。
  • 添加长期内存是指跨会话存储的用户特定或应用程序级别的数据。

短期内存(会话级持久化)使智能体能够跟踪多轮对话。要添加短期内存:

// 创建内存检查点器
MemorySaver checkpointer = new MemorySaver();

SaverConfig saverConfig = SaverConfig.builder()
  .register(checkpointer)
  .build();

// 定义状态策略
KeyStrategyFactory keyStrategyFactory = () -> {
  Map<String, KeyStrategy> keyStrategyMap = new HashMap<>();
  keyStrategyMap.put("messages", new AppendStrategy());
  return keyStrategyMap;
};

// 构建图
StateGraph stateGraph = new StateGraph(keyStrategyFactory)
  .addNode("agent", agentNode)
  .addEdge(START, "agent")
  .addEdge("agent", END);

// 使用检查点器编译图
CompiledGraph graph = stateGraph.compile(
  CompileConfig.builder()
      .saverConfig(saverConfig)
      .build()
);

// 使用会话 ID 调用图
RunnableConfig config = RunnableConfig.builder()
  .threadId("user-session-1")
  .build();

Map<String, Object> input = Map.of(
  "messages", List.of(
      Map.of("role", "user", "content", "你好!我是 Bob")
  )
);

graph.invoke(input, config);

生产环境使用

在生产环境中,使用由数据库支持的检查点器:

Redis 检查点器

import com.alibaba.cloud.ai.graph.checkpoint.savers.RedisSaver;
import com.alibaba.cloud.ai.graph.checkpoint.config.SaverConfig;
import com.alibaba.cloud.ai.graph.checkpoint.constant.SaverConstant;

// Redis 配置
String redisHost = "localhost";
int redisPort = 6379;

RedisSaver redisSaver = new RedisSaver(redisHost, redisPort);

SaverConfig saverConfig = SaverConfig.builder()
  .register(SaverConstant.REDIS, redisSaver)
  .build();

CompiledGraph graph = stateGraph.compile(
  CompileConfig.builder()
      .saverConfig(saverConfig)
      .build()
);

完整示例:使用短期内存的多轮对话

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.checkpoint.config.SaverConfig;
import com.alibaba.cloud.ai.graph.checkpoint.savers.MemorySaver;
import com.alibaba.cloud.ai.graph.state.strategy.AppendStrategy;

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

import java.util.HashMap;
import java.util.List;
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;

// 定义状态策略
KeyStrategyFactory keyStrategyFactory = () -> {
  Map<String, KeyStrategy> keyStrategyMap = new HashMap<>();
  keyStrategyMap.put("messages", new AppendStrategy());
  return keyStrategyMap;
};

// 创建聊天节点
var chatNode = node_async(state -> {
  List<Map<String, String>> messages =
      (List<Map<String, String>>) state.value("messages").orElse(List.of());

  // 使用 ChatClient 调用 AI 模型
  ChatClient chatClient = chatClientBuilder.build();
  String response = chatClient.prompt()
      .user(messages.get(messages.size() - 1).get("content"))
      .call()
      .content();

  return Map.of("messages", List.of(
      Map.of("role", "assistant", "content", response)
  ));
});

// 构建图
StateGraph stateGraph = new StateGraph(keyStrategyFactory)
  .addNode("chat", chatNode)
  .addEdge(START, "chat")
  .addEdge("chat", END);

// 配置检查点
SaverConfig saverConfig = SaverConfig.builder()
      .register(new MemorySaver())
  .build();

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

// 第一轮对话
RunnableConfig config = RunnableConfig.builder()
  .threadId("conversation-1")
  .build();

graph.invoke(Map.of("messages", List.of(
  Map.of("role", "user", "content", "你好!我是 Bob")
)), config);

// 第二轮对话(使用相同的 threadId)
graph.invoke(Map.of("messages", List.of(
  Map.of("role", "user", "content", "我的名字是什么?")
)), config);
// AI 将能够记住之前的对话,回答 "Bob"

在子图中使用

如果您的图包含子图,您只需在编译父图时提供检查点器。Spring AI Alibaba Graph 将自动将检查点器传播到子图。

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

// 定义状态
KeyStrategyFactory keyStrategyFactory = () -> {
  Map<String, KeyStrategy> keyStrategyMap = new HashMap<>();
  keyStrategyMap.put("foo", new ReplaceStrategy());
  return keyStrategyMap;
};

// 子图
var subgraphNode = node_async(state -> {
  String foo = (String) state.value("foo").orElse("");
  return Map.of("foo", foo + "bar");
});

StateGraph subgraphBuilder = new StateGraph(keyStrategyFactory)
  .addNode("subgraph_node_1", subgraphNode)
  .addEdge(START, "subgraph_node_1");

// 子图不需要检查点器
CompiledGraph subgraph = subgraphBuilder.compile();

// 父图
StateGraph parentBuilder = new StateGraph(keyStrategyFactory)
  .addNode("node_1", state -> {
      // 调用子图
      return subgraph.invoke(state.data(),
          RunnableConfig.builder().build());
  })
  .addEdge(START, "node_1");

// 只在父图编译时提供检查点器
SaverConfig saverConfig = SaverConfig.builder()
  .register(SaverConstant.MEMORY, new MemorySaver())
  .build();

CompiledGraph graph = parentBuilder.compile(
  CompileConfig.builder()
      .saverConfig(saverConfig)
      .build()
);

如果您希望子图拥有自己的内存,可以使用适当的检查点器选项编译它。这在多智能体系统中很有用,如果您希望智能体跟踪其内部消息历史。

本文整理自 Spring AI Alibaba 官方文档《内存管理》:https://java2ai.com/docs/frameworks/graph-core/core/memory

赞(0)
未经允许不得转载:AI项目博客 » Spring AI Alibaba 70|Graph Core · 内存管理 · 添加短期内存

评论 抢沙发

AI项目博客

专注 AI 工程化落地 —— 覆盖智能体、大模型、Java AI 与 Python 的学习笔记,以及 AI 平台的项目实战记录。

关于本站文章归档

登录

找回密码

注册