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

Spring AI Alibaba 71|Graph Core · 内存管理 · 添加长期内存

使用长期内存跨对话存储用户特定或应用程序特定的数据。

Spring AI Alibaba 借助 Store 组件来实现记忆的写入或读取管理。Store 是一个抽象接口,可以有不同的实现(如 MemoryStore、RedisStore 等),用于持久化存储跨会话的数据。

使用 Store 存储用户信息

在节点中使用 Store 存储和检索用户特定的长期数据:

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.OverAllState;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import com.alibaba.cloud.ai.graph.store.Store;
import com.alibaba.cloud.ai.graph.store.StoreItem;
import com.alibaba.cloud.ai.graph.store.stores.MemoryStore;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

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

// 在节点中使用 Store 存储用户信息
var userProfileNode = node_async((state, config) -> {
  String userId = (String) state.value("userId").orElse("");

  if (userId.isEmpty()) {
      return Map.of("userProfile", Map.of("name", "Unknown", "preferences", "default"));
  }

  // 从 Store 获取用户配置
  Store store = config.store();
  if (store != null) {
      Optional<StoreItem> itemOpt = store.getItem(List.of("user_profiles"), userId);
      if (itemOpt.isPresent()) {
          Map<String, Object> userProfile = itemOpt.get().getValue();
          return Map.of("userProfile", userProfile);
      }
  }

  // 如果未找到,返回默认值
  Map<String, Object> userProfile = Map.of("name", "User", "preferences", "default");
  return Map.of("userProfile", userProfile);
});

// 创建图
KeyStrategyFactory keyStrategyFactory = () -> {
  Map<String, KeyStrategy> keyStrategyMap = new HashMap<>();
  keyStrategyMap.put("userId", new ReplaceStrategy());
  keyStrategyMap.put("userProfile", new ReplaceStrategy());
  return keyStrategyMap;
};

StateGraph stateGraph = new StateGraph(keyStrategyFactory)
      .addNode("load_profile", userProfileNode)
      .addEdge(START, "load_profile")
      .addEdge("load_profile", END);

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

// 创建长期记忆存储并预填充数据
MemoryStore memoryStore = new MemoryStore();
Map<String, Object> profileData = new HashMap<>();
profileData.put("name", "张三");
profileData.put("preferences", "喜欢编程");
StoreItem profileItem = StoreItem.of(List.of("user_profiles"), "user_001", profileData);
memoryStore.putItem(profileItem);

// 运行图
RunnableConfig config = RunnableConfig.builder()
      .threadId("profile_thread")
      .store(memoryStore)
      .build();

Optional<OverAllState> stateOptional = graph.invoke(Map.of("userId", "user_001"), config);
Map<String, Object> result = stateOptional.get().data();
System.out.println("加载的用户配置: " + result.get("userProfile"));

说明:

  • 使用 AsyncNodeActionWithConfig.node_async 来访问 RunnableConfig,从而获取 Store 实例
  • 通过 config.store() 获取 Store,可能为 null,需要做空值检查
  • 使用 store.getItem(namespace, key) 从 Store 中获取数据,返回 Optional<StoreItem>
  • 使用 StoreItem.of(namespace, key, value) 创建 StoreItem,然后通过 store.putItem(item) 存储数据
  • 在 RunnableConfig 中通过 .store(memoryStore) 指定要使用的 Store 实例

使用 Store 实现缓存

使用 Store 实现缓存机制,避免重复执行耗时操作:

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.OverAllState;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import com.alibaba.cloud.ai.graph.store.Store;
import com.alibaba.cloud.ai.graph.store.StoreItem;
import com.alibaba.cloud.ai.graph.store.stores.MemoryStore;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

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

var cacheNode = node_async((state, config) -> {
  String key = (String) state.value("cacheKey").orElse("");

  if (key.isEmpty()) {
      return Map.of("result", "no_key");
  }

  // 从 Store 获取缓存数据
  Store store = config.store();
  if (store != null) {
      Optional<StoreItem> itemOpt = store.getItem(List.of("cache"), key);
      if (itemOpt.isPresent()) {
          // 缓存命中
          Map<String, Object> cachedData = itemOpt.get().getValue();
          return Map.of("result", cachedData.get("value"));
      }
  }

  // 缓存未命中,执行计算或查询
  Object computedData = performExpensiveOperation(key);

  // 存储到 Store
  if (store != null) {
      Map<String, Object> cacheValue = new HashMap<>();
      cacheValue.put("value", computedData);
      StoreItem cacheItem = StoreItem.of(List.of("cache"), key, cacheValue);
      store.putItem(cacheItem);
  }

  return Map.of("result", computedData);
});

// 创建图
KeyStrategyFactory keyStrategyFactory = () -> {
  Map<String, KeyStrategy> keyStrategyMap = new HashMap<>();
  keyStrategyMap.put("cacheKey", new ReplaceStrategy());
  keyStrategyMap.put("result", new ReplaceStrategy());
  return keyStrategyMap;
};

StateGraph stateGraph = new StateGraph(keyStrategyFactory)
      .addNode("cache", cacheNode)
      .addEdge(START, "cache")
      .addEdge("cache", END);

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

// 创建长期记忆存储
MemoryStore memoryStore = new MemoryStore();

// 第一次调用(缓存未命中)
RunnableConfig config = RunnableConfig.builder()
      .threadId("cache_thread")
      .store(memoryStore)
      .build();

Optional<OverAllState> stateOptional = graph.invoke(Map.of("cacheKey", "expensive_key"), config);
Map<String, Object> result1 = stateOptional.get().data();
System.out.println("第一次调用结果: " + result1.get("result"));

// 第二次调用(缓存命中)
Optional<OverAllState> stateOptional2 = graph.invoke(Map.of("cacheKey", "expensive_key"), config);
Map<String, Object> result2 = stateOptional2.get().data();
System.out.println("第二次调用结果(从缓存): " + result2.get("result"));

辅助方法(在实际代码中定义):

// 模拟耗时操作
private static Object performExpensiveOperation(String key) {
  // 模拟耗时计算
  return "computed_result_for_" + key;
}

说明:

  • 缓存逻辑:首先检查 Store 中是否存在缓存数据,如果存在则直接返回(缓存命中),否则执行耗时操作并将结果存储到 Store(缓存未命中)
  • 使用 List.of(“cache”) 作为命名空间来组织缓存数据
  • 同一个 RunnableConfig 和 Store 实例可以在多次调用中复用,实现跨调用的缓存持久化

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

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

评论 抢沙发

AI项目博客

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

关于本站文章归档

登录

找回密码

注册