Spring AI Alibaba 60｜记忆管理 · 结合短期和长期记忆

短期记忆用于存储对话上下文，长期记忆用于存储持久化数据。下面的示例展示了如何同时使用两种记忆。

import com.alibaba.cloud.ai.graph.agent.ReactAgent;
import com.alibaba.cloud.ai.graph.agent.hook.ModelHook;
import com.alibaba.cloud.ai.graph.agent.hook.HookPosition;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.alibaba.cloud.ai.graph.checkpoint.savers.MemorySaver;
import com.alibaba.cloud.ai.graph.store.Store;
import com.alibaba.cloud.ai.graph.store.stores.MemoryStore;
import com.alibaba.cloud.ai.graph.store.StoreItem;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;

import java.util.*;
import java.util.concurrent.CompletableFuture;

// 创建组合记忆Hook
ModelHook combinedMemoryHook = new ModelHook() {
  @Override
  public String getName() {
      return "combined_memory";
  }

  @Override
  public HookPosition[] getHookPositions() {
      return new HookPosition[]{HookPosition.BEFORE_MODEL};
  }

  @Override
  public CompletableFuture<Map<String, Object>> beforeModel(OverAllState state, RunnableConfig config) {
      Optional<Object> userIdOpt = config.metadata("user_id");
      if (userIdOpt.isEmpty()) {
          return CompletableFuture.completedFuture(Map.of());
      }
      String userId = (String) userIdOpt.get();

      Store memoryStore = config.store();
      // 从长期记忆加载
      Optional<StoreItem> profileOpt = memoryStore.getItem(List.of("profiles"), userId);
      if (profileOpt.isEmpty()) {
          return CompletableFuture.completedFuture(Map.of());
      }

      Map<String, Object> profile = profileOpt.get().getValue();
      String contextInfo = String.format("长期记忆：用户 %s, 职业: %s",
          profile.get("name"), profile.get("occupation"));

      // 获取消息列表
      List<Message> messages = (List<Message>) state.value("messages").orElse(new ArrayList<>());
      List<Message> newMessages = new ArrayList<>();

      // 查找是否已存在 SystemMessage
      SystemMessage existingSystemMessage = null;
      int systemMessageIndex = -1;
      for (int i = 0; i < messages.size(); i++) {
          Message msg = messages.get(i);
          if (msg instanceof SystemMessage) {
              existingSystemMessage = (SystemMessage) msg;
              systemMessageIndex = i;
              break;
          }
      }

      // 如果找到 SystemMessage，更新它；否则创建新的
      SystemMessage enhancedSystemMessage;
      if (existingSystemMessage != null) {
          // 更新现有的 SystemMessage
          enhancedSystemMessage = new SystemMessage(
              existingSystemMessage.getText() + "

" + contextInfo
          );
      } else {
          // 创建新的 SystemMessage
          enhancedSystemMessage = new SystemMessage(contextInfo);
      }

      // 构建新的消息列表
      if (systemMessageIndex >= 0) {
          // 如果找到了 SystemMessage，替换它
          for (int i = 0; i < messages.size(); i++) {
              if (i == systemMessageIndex) {
                  newMessages.add(enhancedSystemMessage);
              } else {
                  newMessages.add(messages.get(i));
              }
          }
      } else {
          // 如果没有找到 SystemMessage，在开头添加新的
          newMessages.add(enhancedSystemMessage);
          newMessages.addAll(messages);
      }

      return CompletableFuture.completedFuture(Map.of("messages", newMessages));
  }
};

// 创建Agent
ReactAgent agent = ReactAgent.builder()
  .name("combined_memory_agent")
  .model(chatModel)
  .hooks(combinedMemoryHook)
  .saver(new MemorySaver()) // 短期记忆
  .build();

// 创建记忆存储
MemoryStore memoryStore = new MemoryStore();
// 设置长期记忆
Map<String, Object> userProfile = new HashMap<>();
userProfile.put("name", "李工程师");
userProfile.put("occupation", "软件工程师");

StoreItem profileItem = StoreItem.of(List.of("profiles"), "user_002", userProfile);
memoryStore.putItem(profileItem);

RunnableConfig config = RunnableConfig.builder()
  .threadId("combined_thread")
  .addMetadata("user_id", "user_002")
  .store(memoryStore)
  .build();

// 短期记忆：在对话中记住
agent.invoke("我今天在做一个 Spring 项目。", config);

// 提出需要同时使用两种记忆的问题
agent.invoke("根据我的职业和今天的工作，给我一些建议。", config);
// 响应会同时使用长期记忆（职业）和短期记忆（Spring项目）

本文整理自 Spring AI Alibaba 官方文档《记忆管理》：https://java2ai.com/docs/frameworks/agent-framework/advanced/memory