Administrator / xlyAi

package com.xly.config;

import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.data.message.ChatMessage;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * 可操作的ChatMemoryProvider：获取消息对象+清除记忆（指定/全量）+删除单条消息
 * 实现框架原生接口，无缝对接AiServices，线程安全适配生产环境
 */
@Component
public class OperableChatMemoryProvider implements ChatMemoryProvider {
    // 核心缓存：memoryId -> ChatMemory，保证一个会话/用户对应唯一记忆实例
    private final Map<Object, ChatMemory> memoryCache = new ConcurrentHashMap<>();
    // 记忆最大消息数，根据业务需求调整
    private static final int MAX_MESSAGE_SIZE = 100;

    /**
     * 框架原生方法：获取【当前memoryId对应的ChatMemory实例（含消息对象）】
     * AiServices自动调用，也是手动操作记忆/消息的唯一入口
     */
    @Override
    public ChatMemory get(Object memoryId) {
        // 空memoryId兜底，避免空指针
        Object finalMemId = Objects.isNull(memoryId) ? "default_erp_chat_memory" : memoryId;
        // 不存在则创建MessageWindowChatMemory，存在则复用
        return memoryCache.computeIfAbsent(finalMemId, k -> MessageWindowChatMemory.withMaxMessages(MAX_MESSAGE_SIZE));
    }

    // ===================== 1. 获取消息对象 =====================

    /**
     * 获取当前会话/用户的全部消息列表
     */
    public List<ChatMessage> getCurrentChatMessages(Object memoryId) {
        if (Objects.isNull(memoryId)) {
            return new ArrayList<>();
        }
        // 返回一个新的列表，避免直接操作内部列表
        return new ArrayList<>(this.get(memoryId).messages());
    }

    // ===================== 2. 清除记忆 =====================

    /**
     * 清空【指定memoryId】的全部记忆
     */
    public void clearSpecifiedMemory(Object memoryId) {
        if (Objects.nonNull(memoryId)) {
            this.get(memoryId).clear();
        }
    }

    /**
     * 全量清除【所有memoryId】的记忆
     */
    public void clearAllMemory() {
        memoryCache.values().forEach(ChatMemory::clear);
    }

    // ===================== 3. 删除单条消息（完全重新设置方案） =====================

    /**
     * 删除指定消息（通过完全重新设置消息列表的方式）
     * @param memoryId 会话ID
     * @param messageToDelete 要删除的消息对象
     * @return 删除后的最新消息列表
     */
    public List<ChatMessage> deleteSingleMessage(Object memoryId, ChatMessage messageToDelete) {
        if (Objects.isNull(memoryId) || Objects.isNull(messageToDelete)) {
            return getCurrentChatMessages(memoryId);
        }

        // 步骤1: 获取当前所有消息
        ChatMemory currentMemory = this.get(memoryId);
        List<ChatMessage> currentMessages = new ArrayList<>(currentMemory.messages());

        // 从后往前查找内容匹配的最后一条消息
        int indexToDelete = -1;
        for (int i = currentMessages.size() - 1; i >= 0; i--) {
            ChatMessage msg = currentMessages.get(i);
            if (isSameMessage(msg, messageToDelete)) {
                indexToDelete = i;
                break;
            }
        }
        // 如果找到匹配的消息
        if (indexToDelete >= 0) {
            List<ChatMessage> filteredMessages = new ArrayList<>(currentMessages);
            filteredMessages.remove(indexToDelete);
            return rebuildMemoryWithMessages(memoryId, filteredMessages);
        }
        // 步骤4: 完全重新设置消息列表
        return rebuildMemoryWithMessages(memoryId, currentMessages);
    }
    public List<ChatMessage> deleteUserLasterMessage(Object memoryId) {
        if (Objects.isNull(memoryId)) {
            return getCurrentChatMessages(memoryId);
        }
        // 步骤1: 获取当前所有消息
        ChatMemory currentMemory = this.get(memoryId);
        List<ChatMessage> currentMessages = new ArrayList<>(currentMemory.messages());

        // 从后往前查找内容匹配的最后一条消息
        int indexToDelete = currentMessages.size()-1;
        // 如果找到匹配的消息
        if (indexToDelete >= 0) {
            List<ChatMessage> filteredMessages = new ArrayList<>(currentMessages);
            filteredMessages.remove(indexToDelete);
            return rebuildMemoryWithMessages(memoryId, filteredMessages);
        }
        // 步骤4: 完全重新设置消息列表
        return rebuildMemoryWithMessages(memoryId, currentMessages);
    }
    /**
     * 批量删除多条消息
     * @param memoryId 会话ID
     * @param messagesToDelete 要删除的消息列表
     * @return 删除后的最新消息列表
     */
    public List<ChatMessage> deleteMessages(Object memoryId, List<ChatMessage> messagesToDelete) {
        if (Objects.isNull(memoryId) || messagesToDelete == null || messagesToDelete.isEmpty()) {
            return getCurrentChatMessages(memoryId);
        }

        // 获取当前所有消息
        ChatMemory currentMemory = this.get(memoryId);
        List<ChatMessage> currentMessages = new ArrayList<>(currentMemory.messages());

        // 过滤掉所有要删除的消息
        List<ChatMessage> filteredMessages = currentMessages.stream()
                .filter(msg -> messagesToDelete.stream().noneMatch(delMsg -> isSameMessage(msg, delMsg)))
                .collect(Collectors.toList());

        // 如果消息数量没有变化，直接返回
        if (filteredMessages.size() == currentMessages.size()) {
            return currentMessages;
        }

        // 重新设置消息列表
        return rebuildMemoryWithMessages(memoryId, filteredMessages);
    }

    /**
     * 判断两条消息是否相同（支持根据内容、类型等多维度匹配）
     */
    private boolean isSameMessage(ChatMessage msg1, ChatMessage msg2) {
//        if (msg1 == msg2) return true; // 同一对象引用
//        if (msg1 == null || msg2 == null) return false;
//
//        // 根据消息类型和内容判断
//        // 方式1: 根据对象相等性
//        if (msg1.equals(msg2)) return true;
        // 方式2: 根据消息类型和文本内容（适用于大多数场景）
        if (msg1.type() == msg2.type()) {
            // 如果有唯一ID或时间戳，也可以加入判断
            // 这里简单使用文本内容判断
            String text1 = extractText(msg1);
            String text2 = extractText(msg2);
            return Objects.equals(text1, text2);
        }

        return false;
    }

    /**
     * 从消息中提取文本内容
     */
    private String extractText(ChatMessage message) {
        // 根据ChatMessage的实际类型提取文本
        // 这里需要根据您的具体消息实现来调整
        try {
            // 尝试通过toString获取内容
            return message.toString();
        } catch (Exception e) {
            return "";
        }
    }

    /**
     * 核心方法：使用过滤后的消息列表重建记忆
     */
    private List<ChatMessage> rebuildMemoryWithMessages(Object memoryId, List<ChatMessage> messages) {
        // 从缓存中移除原实例
        ChatMemory oldMemory = memoryCache.remove(memoryId);

        // 创建新的ChatMemory实例
        MessageWindowChatMemory newMemory = MessageWindowChatMemory.withMaxMessages(MAX_MESSAGE_SIZE);

        // 由于ChatMemory接口没有直接添加消息的方法，我们需要通过特定方式添加
        // 方式1: 如果MessageWindowChatMemory有add方法（需要通过反射或强制类型转换）
        try {
            // 通过反射调用add方法（如果存在）
            java.lang.reflect.Method addMethod = newMemory.getClass().getMethod("add", ChatMessage.class);
            addMethod.setAccessible(true);
            for (ChatMessage message : messages) {
                addMethod.invoke(newMemory, message);
            }
        } catch (Exception e) {
            // 方式2: 如果无法通过反射添加，我们需要使用替代方案
            // 这里可以记录日志
            System.err.println("无法通过反射添加消息: " + e.getMessage());

            // 方式3: 使用clear后通过某种方式重新添加
            // 注意：这取决于具体的ChatMemory实现
        }

        // 将新实例放入缓存
        memoryCache.put(memoryId, newMemory);

        // 返回重建后的消息列表
        return new ArrayList<>(newMemory.messages());
    }

    /**
     * 移除并清除指定记忆（清空消息+从缓存删除实例）
     */
    public void removeAndClearMemory(Object memoryId) {
        if (Objects.nonNull(memoryId)) {
            ChatMemory chatMemory = memoryCache.remove(memoryId);
            if (Objects.nonNull(chatMemory)) {
                chatMemory.clear();
            }
        }
    }
}