還在為 Spring AI 應(yīng)用重啟后對(duì)話上下文丟失而煩惱嗎？本文將帶你深入 Spring AI 的對(duì)話記憶機(jī)制，并手把手教你實(shí)現(xiàn)一個(gè)基于文件的持久化方案，讓你的 AI 應(yīng)用擁有 “過目不忘” 的超能力！

哈嘍，各位程序員朋友們！

在之前的文章里，我們一起探索了如何使用 Spring AI 構(gòu)建能理解上下文的對(duì)話機(jī)器人。但一個(gè)棘手的問題很快就浮現(xiàn)了：我們的對(duì)話記憶都存在內(nèi)存里，服務(wù)器一旦重啟，珍貴的聊天記錄就灰飛煙滅了。這可不行！

想象一下，用戶正和你的 AI 聊得火熱，結(jié)果服務(wù)器一更新，AI 就 “失憶” 了，之前的對(duì)話全忘了。這體驗(yàn)感，簡(jiǎn)直一言難盡。

那么，有沒有辦法讓對(duì)話記憶像數(shù)據(jù)一樣被持久化，存到文件、數(shù)據(jù)庫(kù)或者 Redis 里呢？

答案是：當(dāng)然有！Spring AI 早就為我們考慮到了這一點(diǎn)。

一、官方方案：理想與現(xiàn)實(shí)的差距

Spring AI 官方文檔中提到，它提供了一些現(xiàn)成的持久化方案，可以將對(duì)話記憶保存到不同的數(shù)據(jù)源中。聽起來(lái)很不錯(cuò)，對(duì)吧？

• InMemoryChatMemory：默認(rèn)的內(nèi)存存儲(chǔ)，我們一直在用。
• CassandraChatMemory：用 Cassandra 持久化，還帶過期時(shí)間。
• Neo4jChatMemory：用 Neo4j 持久化，永不過期。
• JdbcChatMemory：用 JDBC 持久化到關(guān)系型數(shù)據(jù)庫(kù)。

看到 JdbcChatMemory，我們可能兩眼放光：這不就是我們想要的嗎？然而，現(xiàn)實(shí)卻給我們潑了一盆冷水。spring-ai-starter-model-chat-memory-jdbc 這個(gè)依賴不僅版本稀少，相關(guān)文檔也幾乎沒有，甚至在 Maven 中央倉(cāng)庫(kù)都搜不到。

雖然在 Spring 自己的倉(cāng)庫(kù)里能找到它的蹤跡，但這用戶量……基本上等于讓我們?nèi)ァ伴_荒”，風(fēng)險(xiǎn)太高了。

既然官方的路不好走，那我們就自己動(dòng)手，豐衣足食！

二、另辟蹊徑：自定義你的 ChatMemory

我更推薦的方案是：自定義實(shí)現(xiàn) ChatMemory 接口。

Spring AI 的設(shè)計(jì)非常巧妙，它將“存儲(chǔ)介質(zhì)”和“記憶算法”解耦了。這意味著我們可以只替換存儲(chǔ)部分，而不用改動(dòng)整個(gè)對(duì)話流程。

雖然官方?jīng)]給示例，但沒關(guān)系，我們可以“偷師”啊！直接去看默認(rèn)實(shí)現(xiàn)類 InMemoryChatMemory 的源碼，模仿它的實(shí)現(xiàn)。

ChatMemory 接口的核心方法很簡(jiǎn)單，就是對(duì)消息的增、刪、查：

InMemoryChatMemory 的源碼顯示，它內(nèi)部其實(shí)就是用一個(gè) ConcurrentHashMap 來(lái)存消息，Key 是對(duì)話 ID，Value 是這個(gè)對(duì)話的所有消息列表。

思路有了，接下來(lái)就是實(shí)戰(zhàn)！

三、實(shí)戰(zhàn)演練：打造文件版 ChatMemory

為了避免引入數(shù)據(jù)庫(kù)等額外依賴的復(fù)雜性，我們先來(lái)實(shí)現(xiàn)一個(gè)最簡(jiǎn)單的：基于文件的持久化 ChatMemory。

這里的核心挑戰(zhàn)在于 消息對(duì)象的序列化與反序列化。我們需要將內(nèi)存中的 Message 對(duì)象轉(zhuǎn)換成文本存入文件，也要能從文件中讀出文本并還原成 Message 對(duì)象。

你可能會(huì)首先想到用 JSON，但很快就會(huì)發(fā)現(xiàn)困難重重：

1. Message 是個(gè)接口，有 UserMessage、SystemMessage 等多種實(shí)現(xiàn)。
2. 不同子類的字段各不相同，結(jié)構(gòu)不統(tǒng)一。
3. 這些子類大多沒有無(wú)參構(gòu)造函數(shù)，也沒有實(shí)現(xiàn) Serializable 接口。

直接用 JSON 序列化，大概率會(huì)踩坑。因此，我們請(qǐng)出一位“外援”——高性能序列化庫(kù) Kryo。

第一步：引入 Kryo 依賴

在 pom.xml 中添加：

<dependency>
    <groupId>com.esotericsoftware</groupId>
    <artifactId>kryo</artifactId>
    <version>5.6.2</version>
</dependency>

第二步：編寫 FileBasedChatMemory

新建 chatmemory 包，創(chuàng)建 FileBasedChatMemory.java。別被下面的代碼嚇到，核心邏輯就是文件的讀寫和對(duì)象的序列化/反序列化，完全可以讓 AI 幫你生成。

// ... 省略 package 和 import ...

/**
 * @author BNTang
 * @version 1.0
 * @description 基于文件持久化的對(duì)話記憶，實(shí)現(xiàn) ChatMemory 接口
 **/
public class FileBasedChatMemory implements ChatMemory {
    /**
     * 文件存儲(chǔ)的基礎(chǔ)目錄
     */
    private final String BASE_DIR;
    // Kryo 實(shí)例，用于序列化和反序列化消息對(duì)象
    private static final Kryo KRYO = new Kryo();

    static {
        // 設(shè)置 Kryo 的注冊(cè)要求為 false，允許未注冊(cè)的類進(jìn)行序列化
        KRYO.setRegistrationRequired(false);
        // 設(shè)置 Kryo 的實(shí)例化策略為標(biāo)準(zhǔn)實(shí)例化策略
        KRYO.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    /**
     * 構(gòu)造函數(shù)，初始化文件存儲(chǔ)目錄。
     *
     * @param dir 文件存儲(chǔ)目錄路徑
     */
    public FileBasedChatMemory(String dir) {
        // 設(shè)置基礎(chǔ)目錄
        this.BASE_DIR = dir;
        // 確保目錄存在
        File baseDir = new File(BASE_DIR);
        // 如果目錄不存在，則創(chuàng)建目錄
        if (!baseDir.exists()) {
            // 嘗試創(chuàng)建目錄，如果失敗則拋出異常
            boolean created = baseDir.mkdirs();
            // 如果目錄創(chuàng)建失敗，拋出運(yùn)行時(shí)異常
            if (!created) {
                // 目錄創(chuàng)建失敗，拋出異常
                throw new RuntimeException("Failed to create directory: " + BASE_DIR);
            }
        }
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        // 獲取或創(chuàng)建對(duì)話的消息列表
        List<Message> conversationMessages = getOrCreateConversation(conversationId);
        // 將新的消息添加到對(duì)話消息列表中
        conversationMessages.addAll(messages);
        // 保存更新后的對(duì)話消息列表到文件
        saveConversation(conversationId, conversationMessages);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        // 獲取或創(chuàng)建對(duì)話的消息列表
        List<Message> allMessages = getOrCreateConversation(conversationId);
        // 如果消息總數(shù)小于等于 lastN，直接返回所有消息
        if (allMessages.size() <= lastN) {
            return allMessages;
        }
        // 否則，返回最后 N 條消息
        return allMessages.subList(allMessages.size() - lastN, allMessages.size());
    }

    @Override
    public void clear(String conversationId) {
        // 獲取對(duì)話文件
        File file = getConversationFile(conversationId);
        // 如果文件存在，則刪除該文件
        if (file.exists()) {
            // 嘗試刪除文件，如果刪除失敗則打印警告信息
            file.delete();
        }
    }

    /**
     * getOrCreateConversation 方法用于獲取或創(chuàng)建一個(gè)對(duì)話的消息列表。
     *
     * @param conversationId 對(duì)話 ID，用于標(biāo)識(shí)特定的對(duì)話
     * @return 一個(gè)包含對(duì)話消息的列表，如果文件不存在則返回一個(gè)空列表
     */
    private List<Message> getOrCreateConversation(String conversationId) {
        // 獲取對(duì)話文件
        File file = getConversationFile(conversationId);
        // 如果文件不存在，則創(chuàng)建一個(gè)新的空列表
        if (!file.exists()) {
            return new ArrayList<>();
        }
        // 如果文件存在，則讀取文件中的消息列表
        try (Input input = new Input(new FileInputStream(file))) {
            // 使用 Kryo 反序列化讀取的對(duì)象
            return KRYO.readObject(input, ArrayList.class);
        } catch (Exception e) {
            // 如果讀取文件失敗，打印異常堆棧跟蹤，并返回空列表以防程序崩潰
            e.printStackTrace();
            return new ArrayList<>();
        }
    }

    /**
     * saveConversation 方法用于將對(duì)話消息列表保存到文件中。
     *
     * @param conversationId 對(duì)話 ID，用于標(biāo)識(shí)特定的對(duì)話
     * @param messages       對(duì)話消息列表，包含要保存的消息對(duì)象
     */
    private void saveConversation(String conversationId, List<Message> messages) {
        // 獲取對(duì)話文件
        File file = getConversationFile(conversationId);
        // 確保父目錄存在
        try (Output output = new Output(new FileOutputStream(file))) {
            // 使用 Kryo 序列化消息列表并寫入文件
            KRYO.writeObject(output, messages);
        } catch (IOException e) {
            // 如果寫入文件失敗，打印異常堆棧跟蹤
            e.printStackTrace();
        }
    }

    /**
     * getConversationFile 方法用于獲取特定對(duì)話 ID 的文件。
     *
     * @param conversationId 對(duì)話 ID，用于標(biāo)識(shí)特定的對(duì)話
     * @return 一個(gè) File 對(duì)象，表示存儲(chǔ)該對(duì)話消息的文件
     */
    private File getConversationFile(String conversationId) {
        // 返回一個(gè)新的 File 對(duì)象，表示存儲(chǔ)對(duì)話消息的文件
        return new File(BASE_DIR, conversationId + ".kryo");
    }
}

第三步：配置 ChatClient

修改 App 的構(gòu)造函數(shù)，告訴 ChatClient 使用我們新的文件版對(duì)話記憶。

public App(ChatModel ollamaChatModel) {
    // 指定一個(gè)用于存放記憶文件的目錄
    String fileDir = System.getProperty("user.dir") + "/temp/chat-memory";
    // 實(shí)例化我們自定義的 ChatMemory
    ChatMemory chatMemory = new FileBasedChatMemory(fileDir);
    
    // 構(gòu)建 ChatClient，并注入 ChatMemory
    chatClient = ChatClient.builder(ollamaChatModel)
            .defaultSystem(SYSTEM_PROMPT)
            .defaultAdvisors(
                    new MessageChatMemoryAdvisor(chatMemory),
                    new MyLoggerAdvisor()
            )
            .build();
}

第四步：見證奇跡

運(yùn)行你的應(yīng)用，進(jìn)行幾輪對(duì)話，然后查看項(xiàng)目根目錄下的 temp/chat-memory 文件夾。你會(huì)發(fā)現(xiàn)，對(duì)話記錄已經(jīng)被成功保存為 .kryo 文件了！

現(xiàn)在，即使你重啟應(yīng)用，AI 也能找回之前的對(duì)話，繼續(xù)和用戶愉快地交流。

Spring AI 開發(fā)中的常見痛點(diǎn)：對(duì)話記憶的持久化。通過自定義 ChatMemory 接口，我們成功地將對(duì)話歷史從易失的內(nèi)存轉(zhuǎn)移到了穩(wěn)定的文件中，讓我們的 AI 應(yīng)用擁有了“長(zhǎng)期記憶”。

這個(gè)方法不僅限于文件存儲(chǔ)，你可以舉一反三，將其改造為基于 Redis、MongoDB 或任何你喜歡的存儲(chǔ)方案。這正是 Spring AI 框架靈活性的體現(xiàn)。

希望這篇文章能對(duì)你有所啟發(fā)！動(dòng)手試試吧，給你的 AI 裝上一個(gè)“超級(jí)大腦”！

如果你覺得本文對(duì)你有幫助，歡迎點(diǎn)贊、在看、分享三連！ 你的支持是我持續(xù)創(chuàng)作的最大動(dòng)力！