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

Spring AI Alibaba 72|Graph Core · 内存管理 · 完整示例:结合短期和长期内存

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 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 org.springframework.ai.chat.client.ChatClient;

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.AsyncNodeAction.node_async;
import static com.alibaba.cloud.ai.graph.action.AsyncNodeActionWithConfig.node_async;

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

// 加载用户偏好(长期内存)
var loadUserPreferences = node_async((state, config) -> {
  String userId = (String) state.value("userId").orElse("");

  if (userId.isEmpty()) {
      return Map.of("userPreferences", Map.of("theme", "default", "language", "zh"));
  }

  // 从 Store 加载用户偏好
  Store store = config.store();
  if (store != null) {
      Optional<StoreItem> itemOpt = store.getItem(List.of("user_preferences"), userId);
      if (itemOpt.isPresent()) {
          Map<String, Object> preferences = itemOpt.get().getValue();
          return Map.of("userPreferences", preferences);
      }
  }

  // 如果未找到,返回默认偏好
  Map<String, Object> preferences = Map.of("theme", "dark", "language", "zh");
  return Map.of("userPreferences", preferences);
});

// 聊天节点(使用短期和长期内存)
var chatNode = node_async(state -> {
  List<Map<String, String>> messages =
      (List<Map<String, String>>) state.value("messages").orElse(List.of());
  Map<String, Object> preferences =
      (Map<String, Object>) state.value("userPreferences").orElse(Map.of());

  // 构建包含用户偏好的提示
  String userPrompt = messages.get(messages.size() - 1).get("content");
  String enhancedPrompt = "用户偏好: " + preferences + "
用户问题: " + userPrompt;

  // 调用 AI
  ChatClient chatClient = chatClientBuilder.build();
  String response = chatClient.prompt()
      .user(enhancedPrompt)
      .call()
      .content();

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

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

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

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

// 创建长期记忆存储并预填充用户偏好
MemoryStore memoryStore = new MemoryStore();
Map<String, Object> preferencesData = new HashMap<>();
preferencesData.put("theme", "dark");
preferencesData.put("language", "zh");
preferencesData.put("timezone", "Asia/Shanghai");
StoreItem preferencesItem = StoreItem.of(List.of("user_preferences"), "user_002", preferencesData);
memoryStore.putItem(preferencesItem);

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

// 第一轮对话(加载偏好并开始对话)
graph.invoke(Map.of(
      "userId", "user_002",
      "messages", List.of(Map.of("role", "user", "content", "你好"))
), config);

// 第二轮对话(使用短期和长期记忆)
graph.invoke(Map.of(
      "userId", "user_002",
      "messages", List.of(Map.of("role", "user", "content", "根据我的偏好给我一些建议"))
), config);

通过这种方式,您的应用程序可以同时利用短期内存(对话历史)和长期内存(用户偏好),提供更个性化和上下文感知的体验。

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

赞(0)
未经允许不得转载:AI项目博客 » Spring AI Alibaba 72|Graph Core · 内存管理 · 完整示例:结合短期和长期内存

评论 抢沙发

AI项目博客

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

关于本站文章归档

登录

找回密码

注册