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

AI项目博客






