迁移pro

liyong

13 小时以前 1ca5584d7e3200a9af65a099bd26d3593e2ba702

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
package com.ruoyi.ai.store;
 
import com.ruoyi.ai.mongodbBean.ChatMessages;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import lombok.RequiredArgsConstructor;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Component;
 
import java.util.Date;
import java.util.LinkedList;
import java.util.List;
 
@Component
@RequiredArgsConstructor
public class MongoChatMemoryStore implements ChatMemoryStore {
 
    private final MongoTemplate mongoTemplate;
 
    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        Query query = Query.query(Criteria.where("memoryId").is(memoryIdString(memoryId)));
        ChatMessages chatMessages = mongoTemplate.findOne(query, ChatMessages.class);
        if (chatMessages == null || chatMessages.getContent() == null) {
            return new LinkedList<>();
        }
        return ChatMessageDeserializer.messagesFromJson(chatMessages.getContent());
    }
 
    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        String memoryIdValue = memoryIdString(memoryId);
        Query query = Query.query(Criteria.where("memoryId").is(memoryIdValue));
        Update update = new Update();
        update.set("memoryId", memoryIdValue);
        update.set("content", ChatMessageSerializer.messagesToJson(messages));
        update.set("updateTime", new Date());
        update.setOnInsert("createTime", new Date());
        mongoTemplate.upsert(query, update, ChatMessages.class);
    }
 
    @Override
    public void deleteMessages(Object memoryId) {
        Query query = Query.query(Criteria.where("memoryId").is(memoryIdString(memoryId)));
        mongoTemplate.remove(query, ChatMessages.class);
    }
 
    public void appendMessages(Object memoryId, List<ChatMessage> appendList) {
        List<ChatMessage> messages = new LinkedList<>(getMessages(memoryId));
        messages.addAll(appendList);
        updateMessages(memoryId, messages);
    }
 
    private String memoryIdString(Object memoryId) {
        return memoryId == null ? "" : memoryId.toString();
    }
}