Spring AI Alibaba 59｜记忆管理 · 使用 ModelHook 管理长期记忆

下面的示例展示了如何使用 ModelHook 在模型调用前后自动加载和保存长期记忆。

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.agent.hook.JumpTo;
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.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;

// 创建记忆拦截器
ModelHook memoryInterceptor = new ModelHook() {
  @Override
  public String getName() {
      return "memory_interceptor";
  }

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

  @Override
  public CompletableFuture<Map<String, Object>> beforeModel(OverAllState state, RunnableConfig config) {
      // 从配置中获取用户ID
      String userId = (String) config.metadata("user_id").orElse(null);
      if (userId == null) {
          return CompletableFuture.completedFuture(Map.of());
      }

      Store store = config.store();
      // 从记忆存储中加载用户画像
      Optional<StoreItem> itemOpt = store.getItem(List.of("user_profiles"), userId);
      if (itemOpt.isPresent()) {
          Map<String, Object> profile = itemOpt.get().getValue();

          // 将用户上下文注入系统消息
          String userContext = String.format(
              "用户信息：姓名=%s, 年龄=%s, 邮箱=%s, 偏好=%s",
              profile.get("name"),
              profile.get("age"),
              profile.get("email"),
              profile.get("preferences")
          );

          // 获取消息列表
          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() + "

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

          // 构建新的消息列表
          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));
      }

      return CompletableFuture.completedFuture(Map.of());
  }

  @Override
  public CompletableFuture<Map<String, Object>> afterModel(OverAllState state, RunnableConfig config) {
      // 可以在这里实现对话后的记忆保存逻辑
      return CompletableFuture.completedFuture(Map.of());
  }
};

// 创建带有记忆拦截器的Agent
ReactAgent agent = ReactAgent.builder()
  .name("memory_agent")
  .model(chatModel)
  .hooks(memoryInterceptor)
  .saver(new MemorySaver())
  .build();

// 创建内存存储
MemoryStore memoryStore = new MemoryStore();

// 模拟数据，预先填充用户画像
Map<String, Object> profileData = new HashMap<>();
profileData.put("name", "王小明");
profileData.put("age", 28);
profileData.put("email", "wang@example.com");
profileData.put("preferences", List.of("喜欢咖啡", "喜欢阅读"));

StoreItem profileItem = StoreItem.of(List.of("user_profiles"), "user_001", profileData);
memoryStore.putItem(profileItem);

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

// Agent会自动加载用户画像信息
agent.invoke("请介绍一下我的信息。", config);

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