From ababc914bd5f6ba65b1348d60b2b78d24ac9f37c Mon Sep 17 00:00:00 2001
From: xiaoxin <718949661@qq.com>
Date: Wed, 10 Jul 2024 13:18:17 +0800
Subject: [PATCH] =?UTF-8?q?=E3=80=90=E6=96=B0=E5=A2=9E=E3=80=91AI=20?=
 =?UTF-8?q?=E8=84=91=E5=9B=BE?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../admin/mindmap/AiMindMapController.java    |  35 +++++
 .../mindmap/vo/AiMindMapGenerateReqVO.java    |  13 ++
 .../dal/dataobject/mindmap/AiMindMapDO.java   |  57 +++++++
 .../ai/dal/mysql/mindmap/AiMindMapMapper.java |  14 ++
 .../ai/service/mindmap/AiMindMapService.java  |  23 +++
 .../service/mindmap/AiMindMapServiceImpl.java | 143 ++++++++++++++++++
 6 files changed, 285 insertions(+)
 create mode 100644 yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/mindmap/AiMindMapController.java
 create mode 100644 yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/mindmap/vo/AiMindMapGenerateReqVO.java
 create mode 100644 yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/mindmap/AiMindMapDO.java
 create mode 100644 yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/mysql/mindmap/AiMindMapMapper.java
 create mode 100644 yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/mindmap/AiMindMapService.java
 create mode 100644 yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/mindmap/AiMindMapServiceImpl.java

diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/mindmap/AiMindMapController.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/mindmap/AiMindMapController.java
new file mode 100644
index 000000000..452628fa9
--- /dev/null
+++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/mindmap/AiMindMapController.java
@@ -0,0 +1,35 @@
+package cn.iocoder.yudao.module.ai.controller.admin.mindmap;
+
+import cn.iocoder.yudao.framework.common.pojo.CommonResult;
+import cn.iocoder.yudao.module.ai.controller.admin.mindmap.vo.AiMindMapGenerateReqVO;
+import cn.iocoder.yudao.module.ai.service.mindmap.AiMindMapService;
+import io.swagger.v3.oas.annotations.Operation;
+import io.swagger.v3.oas.annotations.tags.Tag;
+import jakarta.annotation.Resource;
+import jakarta.annotation.security.PermitAll;
+import jakarta.validation.Valid;
+import org.springframework.http.MediaType;
+import org.springframework.web.bind.annotation.PostMapping;
+import org.springframework.web.bind.annotation.RequestBody;
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RestController;
+import reactor.core.publisher.Flux;
+
+import static cn.iocoder.yudao.framework.security.core.util.SecurityFrameworkUtils.getLoginUserId;
+
+@Tag(name = "管理后台 - AI 思维导图")
+@RestController
+@RequestMapping("/ai/mind-map")
+public class AiMindMapController {
+
+    @Resource
+    private AiMindMapService mindMapService;
+
+    @PostMapping(value = "/generate-stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
+    @Operation(summary = "发送消息(流式)", description = "流式返回,响应较快")
+    @PermitAll  // 解决 SSE 最终响应的时候,会被 Access Denied 拦截的问题
+    public Flux<CommonResult<String>> generateMindMap(@RequestBody @Valid AiMindMapGenerateReqVO generateReqVO) {
+        return mindMapService.generateMindMap(generateReqVO, getLoginUserId());
+    }
+
+}
diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/mindmap/vo/AiMindMapGenerateReqVO.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/mindmap/vo/AiMindMapGenerateReqVO.java
new file mode 100644
index 000000000..adc47b8ea
--- /dev/null
+++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/mindmap/vo/AiMindMapGenerateReqVO.java
@@ -0,0 +1,13 @@
+package cn.iocoder.yudao.module.ai.controller.admin.mindmap.vo;
+
+import io.swagger.v3.oas.annotations.media.Schema;
+import jakarta.validation.constraints.NotBlank;
+import lombok.Data;
+
+@Schema(description = "管理后台 - AI 思维导图生成 Request VO")
+@Data
+public class AiMindMapGenerateReqVO {
+    @Schema(description = "思维导图内容提示", example = "Java 学习路线")
+    @NotBlank(message = "思维导图内容提示不能为空")
+    private String prompt;
+}
\ No newline at end of file
diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/mindmap/AiMindMapDO.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/mindmap/AiMindMapDO.java
new file mode 100644
index 000000000..b6b87e92e
--- /dev/null
+++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/mindmap/AiMindMapDO.java
@@ -0,0 +1,57 @@
+package cn.iocoder.yudao.module.ai.dal.dataobject.mindmap;
+
+import cn.iocoder.yudao.framework.ai.core.enums.AiPlatformEnum;
+import cn.iocoder.yudao.framework.mybatis.core.dataobject.BaseDO;
+import com.baomidou.mybatisplus.annotation.IdType;
+import com.baomidou.mybatisplus.annotation.TableId;
+import com.baomidou.mybatisplus.annotation.TableName;
+import lombok.Data;
+
+/**
+ * AI 思维导图 DO
+ *
+ * @author xiaoxin
+ */
+@TableName(value = "ai_mind_map", autoResultMap = true)
+@Data
+public class AiMindMapDO extends BaseDO {
+
+    /**
+     * 编号
+     */
+    @TableId(type = IdType.AUTO)
+    private Long id;
+
+    /**
+     * 用户编号
+     */
+    private Long userId;
+
+    /**
+     * 模型
+     */
+    private String model;
+
+    /**
+     * 平台
+     * <p>
+     * 枚举 {@link AiPlatformEnum}
+     */
+    private String platform;
+
+    /**
+     * 生成内容提示
+     */
+    private String prompt;
+
+    /**
+     * 生成的内容
+     */
+    private String generatedContent;
+
+    /**
+     * 错误信息
+     */
+    private String errorMessage;
+
+}
\ No newline at end of file
diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/mysql/mindmap/AiMindMapMapper.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/mysql/mindmap/AiMindMapMapper.java
new file mode 100644
index 000000000..54fa7235a
--- /dev/null
+++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/mysql/mindmap/AiMindMapMapper.java
@@ -0,0 +1,14 @@
+package cn.iocoder.yudao.module.ai.dal.mysql.mindmap;
+
+import cn.iocoder.yudao.framework.mybatis.core.mapper.BaseMapperX;
+import cn.iocoder.yudao.module.ai.dal.dataobject.mindmap.AiMindMapDO;
+import org.apache.ibatis.annotations.Mapper;
+
+/**
+ * AI 音乐 Mapper
+ *
+ * @author xiaoxin
+ */
+@Mapper
+public interface AiMindMapMapper extends BaseMapperX<AiMindMapDO> {
+}
diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/mindmap/AiMindMapService.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/mindmap/AiMindMapService.java
new file mode 100644
index 000000000..2eb1f1b1a
--- /dev/null
+++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/mindmap/AiMindMapService.java
@@ -0,0 +1,23 @@
+package cn.iocoder.yudao.module.ai.service.mindmap;
+
+import cn.iocoder.yudao.framework.common.pojo.CommonResult;
+import cn.iocoder.yudao.module.ai.controller.admin.mindmap.vo.AiMindMapGenerateReqVO;
+import reactor.core.publisher.Flux;
+
+/**
+ * AI 思维导图 Service 接口
+ *
+ * @author xiaoxin
+ */
+public interface AiMindMapService {
+
+    /**
+     * 生成思维导图内容
+     *
+     * @param generateReqVO 请求参数
+     * @param userId        用户编号
+     * @return 生成结果
+     */
+    Flux<CommonResult<String>> generateMindMap(AiMindMapGenerateReqVO generateReqVO, Long userId);
+
+}
diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/mindmap/AiMindMapServiceImpl.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/mindmap/AiMindMapServiceImpl.java
new file mode 100644
index 000000000..84090d396
--- /dev/null
+++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/mindmap/AiMindMapServiceImpl.java
@@ -0,0 +1,143 @@
+package cn.iocoder.yudao.module.ai.service.mindmap;
+
+import cn.hutool.core.collection.CollUtil;
+import cn.hutool.core.util.StrUtil;
+import cn.iocoder.yudao.framework.ai.core.enums.AiPlatformEnum;
+import cn.iocoder.yudao.framework.ai.core.util.AiUtils;
+import cn.iocoder.yudao.framework.common.pojo.CommonResult;
+import cn.iocoder.yudao.framework.common.pojo.PageResult;
+import cn.iocoder.yudao.framework.common.util.object.BeanUtils;
+import cn.iocoder.yudao.framework.tenant.core.util.TenantUtils;
+import cn.iocoder.yudao.module.ai.controller.admin.mindmap.vo.AiMindMapGenerateReqVO;
+import cn.iocoder.yudao.module.ai.controller.admin.model.vo.chatRole.AiChatRolePageReqVO;
+import cn.iocoder.yudao.module.ai.dal.dataobject.mindmap.AiMindMapDO;
+import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatModelDO;
+import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatRoleDO;
+import cn.iocoder.yudao.module.ai.dal.mysql.mindmap.AiMindMapMapper;
+import cn.iocoder.yudao.module.ai.enums.ErrorCodeConstants;
+import cn.iocoder.yudao.module.ai.service.model.AiApiKeyService;
+import cn.iocoder.yudao.module.ai.service.model.AiChatModelService;
+import cn.iocoder.yudao.module.ai.service.model.AiChatRoleService;
+import jakarta.annotation.Resource;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.ai.chat.messages.Message;
+import org.springframework.ai.chat.messages.SystemMessage;
+import org.springframework.ai.chat.messages.UserMessage;
+import org.springframework.ai.chat.model.ChatModel;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.prompt.ChatOptions;
+import org.springframework.ai.chat.prompt.Prompt;
+import org.springframework.stereotype.Service;
+import reactor.core.publisher.Flux;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+
+import static cn.iocoder.yudao.framework.common.pojo.CommonResult.error;
+import static cn.iocoder.yudao.framework.common.pojo.CommonResult.success;
+
+/**
+ * AI 写作 Service 实现类
+ *
+ * @author xiaoxin
+ */
+@Service
+@Slf4j
+public class AiMindMapServiceImpl implements AiMindMapService {
+
+    @Resource
+    private AiApiKeyService apiKeyService;
+    @Resource
+    private AiChatModelService chatModalService;
+    @Resource
+    private AiChatRoleService chatRoleService;
+
+    @Resource
+    private AiMindMapMapper mindMapMapper;
+
+    private static final String DEFAULT_SYSTEM_MESSAGE = """
+             你是一位非常优秀的思维导图助手,你会把用户的所有提问都总结成思维导图,然后以 Markdown 格式输出。markdown 只需要输出一级标题,二级标题,三级标题,四级标题,最多输出四级,除此之外不要输出任何其他 markdown 标记。下面是一个合格的例子:
+             # Geek-AI 助手
+             
+             ## 完整的开源系统
+             ### 前端开源
+             ### 后端开源
+                      
+             ## 支持各种大模型
+             ### OpenAI
+             ### Azure
+             ### 文心一言
+             ### 通义千问
+                        
+             ## 集成多种收费方式
+             ### 支付宝
+             ### 微信
+                       
+             另外,除此之外不要任何解释性语句。
+            """;
+
+    @Override
+    public Flux<CommonResult<String>> generateMindMap(AiMindMapGenerateReqVO generateReqVO, Long userId) {
+        // 1.1 获取脑图模型 尝试获取思维导图助手角色,如果没有则使用默认模型
+        AiChatRoleDO mindMapRole = selectOneMindMapRole();
+        AiChatModelDO model;
+        String systemMessage;
+        if (Objects.nonNull(mindMapRole)) {
+            model = chatModalService.getChatModel(mindMapRole.getModelId());
+            systemMessage = mindMapRole.getSystemMessage();
+        } else {
+            model = chatModalService.getRequiredDefaultChatModel();
+            systemMessage = DEFAULT_SYSTEM_MESSAGE;
+        }
+
+        AiPlatformEnum platform = AiPlatformEnum.validatePlatform(model.getPlatform());
+        ChatModel chatModel = apiKeyService.getChatModel(model.getKeyId());
+
+        // 1.2 插入思维导图信息
+        AiMindMapDO mindMapDO = BeanUtils.toBean(generateReqVO, AiMindMapDO.class, e -> e.setUserId(userId).setModel(model.getModel()).setPlatform(platform.getPlatform()));
+        mindMapMapper.insert(mindMapDO);
+
+        ChatOptions chatOptions = AiUtils.buildChatOptions(platform, model.getModel(), model.getTemperature(), model.getMaxTokens());
+        // 2.1 角色设定
+        List<Message> chatMessages = new ArrayList<>();
+        if (StrUtil.isNotBlank(systemMessage)) {
+            chatMessages.add(new SystemMessage(systemMessage));
+        }
+        // 2.2 用户输入
+        chatMessages.add(new UserMessage(generateReqVO.getPrompt()));
+        // 2.3 构建提示词
+        Prompt prompt = new Prompt(chatMessages, chatOptions);
+
+        Flux<ChatResponse> streamResponse = chatModel.stream(prompt);
+        // 2.4 流式返回
+        StringBuffer contentBuffer = new StringBuffer();
+        return streamResponse.map(chunk -> {
+            String newContent = chunk.getResult() != null ? chunk.getResult().getOutput().getContent() : null;
+            newContent = StrUtil.nullToDefault(newContent, ""); // 避免 null 的 情况
+            contentBuffer.append(newContent);
+            // 响应结果
+            return success(newContent);
+        }).doOnComplete(() -> {
+            // 忽略租户,因为 Flux 异步无法透传租户
+            TenantUtils.executeIgnore(() ->
+                    mindMapMapper.updateById(new AiMindMapDO().setId(mindMapDO.getId()).setGeneratedContent(contentBuffer.toString())));
+        }).doOnError(throwable -> {
+            log.error("[generateWriteContent][generateReqVO({}) 发生异常]", generateReqVO, throwable);
+            // 忽略租户,因为 Flux 异步无法透传租户
+            TenantUtils.executeIgnore(() ->
+                    mindMapMapper.updateById(new AiMindMapDO().setId(mindMapDO.getId()).setErrorMessage(throwable.getMessage())));
+        }).onErrorResume(error -> Flux.just(error(ErrorCodeConstants.WRITE_STREAM_ERROR)));
+
+    }
+
+    private AiChatRoleDO selectOneMindMapRole() {
+        AiChatRoleDO chatRoleDO = null;
+        PageResult<AiChatRoleDO> mindMapRolePage = chatRoleService.getChatRolePage(new AiChatRolePageReqVO().setName("思维导图助手"));
+        List<AiChatRoleDO> list = mindMapRolePage.getList();
+        if (CollUtil.isNotEmpty(list)) {
+            chatRoleDO = list.get(0);
+        }
+        return chatRoleDO;
+    }
+}