iailab-module-ai/iailab-module-ai-biz/src/main/java/com/iailab/module/ai/controller/admin/chat/AiChatMessageController.java

@@ -7,10 +7,7 @@
import com.iailab.framework.common.pojo.PageResult;
import com.iailab.framework.common.util.collection.MapUtils;
import com.iailab.framework.common.util.object.BeanUtils;
import com.iailab.module.ai.controller.admin.chat.vo.message.AiChatMessagePageReqVO;
import com.iailab.module.ai.controller.admin.chat.vo.message.AiChatMessageRespVO;
import com.iailab.module.ai.controller.admin.chat.vo.message.AiChatMessageSendReqVO;
import com.iailab.module.ai.controller.admin.chat.vo.message.AiChatMessageSendRespVO;
import com.iailab.module.ai.controller.admin.chat.vo.message.*;
import com.iailab.module.ai.dal.dataobject.chat.AiChatConversationDO;
import com.iailab.module.ai.dal.dataobject.chat.AiChatMessageDO;
import com.iailab.module.ai.dal.dataobject.knowledge.AiKnowledgeDocumentDO;
@@ -33,6 +30,7 @@
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
@@ -70,6 +68,29 @@
        return chatMessageService.sendChatMessageStream(sendReqVO, getLoginUserId());
    }

    @Operation(summary = "获得指定对话的消息列表(分页)")
    @GetMapping("/page-list-by-conversation-id")
    public CommonResult<PageResult<AiChatMessageRespVO>> getChatMessagePageListByConversationId(AiChatMessageReqVO queryReqVO) {
        PageResult<AiChatMessageRespVO> pageResult = new PageResult<>();
        AiChatConversationDO conversation = chatConversationService.getChatConversation(queryReqVO.getConversationId());
        if (conversation == null || ObjUtil.notEqual(conversation.getUserId(), getLoginUserId())) {
            return success(pageResult);
        }
        // 1. 获取消息列表
        PageResult<AiChatMessageDO> pageList = chatMessageService.getChatMessagePageByConversationId(queryReqVO);
        List<AiChatMessageDO> messageList = pageList.getList();
        if (CollUtil.isEmpty(messageList)) {
            pageResult.setList(new ArrayList<>());
            pageResult.setTotal(0L);
            return success(pageResult);
        }
        List<AiChatMessageRespVO> messageVOList = BeanUtils.toBean(messageList, AiChatMessageRespVO.class);
        dealKnowledge(messageList, messageVOList);
        pageResult.setList(messageVOList);
        pageResult.setTotal(pageList.getTotal());
        return success(pageResult);
    }

    @Operation(summary = "获得指定对话的消息列表")
    @GetMapping("/list-by-conversation-id")
    @Parameter(name = "conversationId", required = true, description = "对话编号", example = "1024")
@@ -84,32 +105,8 @@
        if (CollUtil.isEmpty(messageList)) {
            return success(Collections.emptyList());
        }

        // 2. 拼接数据，主要是知识库段落信息
        Map<Long, AiKnowledgeSegmentDO> segmentMap = knowledgeSegmentService.getKnowledgeSegmentMap(convertListByFlatMap(messageList,
                message -> CollUtil.isEmpty(message.getSegmentIds()) ? null : message.getSegmentIds().stream()));
        Map<Long, AiKnowledgeDocumentDO> documentMap = knowledgeDocumentService.getKnowledgeDocumentMap(
                convertList(segmentMap.values(), AiKnowledgeSegmentDO::getDocumentId));
        List<AiChatMessageRespVO> messageVOList = BeanUtils.toBean(messageList, AiChatMessageRespVO.class);
        for (int i = 0; i < messageList.size(); i++) {
            AiChatMessageDO message = messageList.get(i);
            if (CollUtil.isEmpty(message.getSegmentIds())) {
                continue;
            }
            // 设置知识库段落信息
            messageVOList.get(i).setSegments(convertList(message.getSegmentIds(), segmentId -> {
                AiKnowledgeSegmentDO segment = segmentMap.get(segmentId);
                if (segment == null) {
                    return null;
                }
                AiKnowledgeDocumentDO document = documentMap.get(segment.getDocumentId());
                if (document == null) {
                    return null;
                }
                return new AiChatMessageRespVO.KnowledgeSegment().setId(segment.getId()).setContent(segment.getContent())
                        .setDocumentId(segment.getDocumentId()).setDocumentName(document.getName());
            }));
        }
        dealKnowledge(messageList, messageVOList);
        return success(messageVOList);
    }

@@ -180,4 +177,30 @@
        return success(true);
    }

    private void dealKnowledge(List<AiChatMessageDO> messageList, List<AiChatMessageRespVO> messageVOList) {
        // 拼接数据，主要是知识库段落信息
        Map<Long, AiKnowledgeSegmentDO> segmentMap = knowledgeSegmentService.getKnowledgeSegmentMap(convertListByFlatMap(messageList,
                message -> CollUtil.isEmpty(message.getSegmentIds()) ? null : message.getSegmentIds().stream()));
        Map<Long, AiKnowledgeDocumentDO> documentMap = knowledgeDocumentService.getKnowledgeDocumentMap(
                convertList(segmentMap.values(), AiKnowledgeSegmentDO::getDocumentId));
        for (int i = 0; i < messageList.size(); i++) {
            AiChatMessageDO message = messageList.get(i);
            if (CollUtil.isEmpty(message.getSegmentIds())) {
                continue;
            }
            // 设置知识库段落信息
            messageVOList.get(i).setSegments(convertList(message.getSegmentIds(), segmentId -> {
                AiKnowledgeSegmentDO segment = segmentMap.get(segmentId);
                if (segment == null) {
                    return null;
                }
                AiKnowledgeDocumentDO document = documentMap.get(segment.getDocumentId());
                if (document == null) {
                    return null;
                }
                return new AiChatMessageRespVO.KnowledgeSegment().setId(segment.getId()).setContent(segment.getContent())
                        .setDocumentId(segment.getDocumentId()).setDocumentName(document.getName());
            }));
        }
    }
}