Commit b8a443fbe09e7b742ae7fe33044fa223875c106f

Authored by 芋道源码
Committed by Gitee
2 parents 4c21ae32 bcdb23b8

!13 【AI 】1.支持脑图内容生成 2.AI 写作:做角色设定,提高准确率

Merge pull request !13 from 小新/master-jdk21-ai
Showing 13 changed files with 355 additions and 33 deletions
yudao-module-ai/yudao-module-ai-api/src/main/java/cn/iocoder/yudao/module/ai/enums/AiChatRoleEnum.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai.enums;
  2 +
  3 +import cn.iocoder.yudao.framework.common.core.IntArrayValuable;
  4 +import lombok.AllArgsConstructor;
  5 +import lombok.Getter;
  6 +
  7 +import java.util.Arrays;
  8 +
  9 +/**
  10 + * AI 写作类型的枚举
  11 + *
  12 + * @author xiaoxin
  13 + */
  14 +@AllArgsConstructor
  15 +@Getter
  16 +public enum AiChatRoleEnum implements IntArrayValuable {
  17 +
  18 + AI_WRITE_ROLE(1, "写作助手", """
  19 + 你是一位出色的写作助手,能够帮助用户生成创意和灵感,并在用户提供场景和提示词时生成对应的回复。你的任务包括:
  20 + 1. 撰写建议:根据用户提供的主题或问题,提供详细的写作建议、情节发展方向、角色设定以及背景描写,确保内容结构清晰、有逻辑。
  21 + 2. 回复生成:根据用户提供的场景和提示词,生成合适的对话或文字回复,确保语气和风格符合场景需求。
  22 + 除此之外不需要除了正文内容外的其他回复,如标题、开头、任何解释性语句或道歉。
  23 + """),
  24 + AI_MIND_MAP_ROLE(2, "脑图助手", """
  25 + 你是一位非常优秀的思维导图助手,你会把用户的所有提问都总结成思维导图,然后以 Markdown 格式输出。markdown 只需要输出一级标题,二级标题,三级标题,四级标题,最多输出四级,除此之外不要输出任何其他 markdown 标记。下面是一个合格的例子:
  26 + # Geek-AI 助手
  27 + ## 完整的开源系统
  28 + ### 前端开源
  29 + ### 后端开源
  30 + ## 支持各种大模型
  31 + ### OpenAI
  32 + ### Azure
  33 + ### 文心一言
  34 + ### 通义千问
  35 + ## 集成多种收费方式
  36 + ### 支付宝
  37 + ### 微信
  38 + 除此之外不要任何解释性语句。
  39 + """);
  40 +
  41 +
  42 + /**
  43 + * 角色
  44 + */
  45 + private final Integer role;
  46 + /**
  47 + * 角色名
  48 + */
  49 + private final String name;
  50 +
  51 + /**
  52 + * 角色设定
  53 + */
  54 + private final String prompt;
  55 +
  56 + public static final int[] ARRAYS = Arrays.stream(values()).mapToInt(AiChatRoleEnum::getRole).toArray();
  57 +
  58 + @Override
  59 + public int[] array() {
  60 + return ARRAYS;
  61 + }
  62 +
  63 +}
yudao-module-ai/yudao-module-ai-api/src/main/java/cn/iocoder/yudao/module/ai/enums/write/AiWriteTypeEnum.java
1 package cn.iocoder.yudao.module.ai.enums.write; 1 package cn.iocoder.yudao.module.ai.enums.write;
2 2
3 -import cn.hutool.core.util.ArrayUtil;  
4 -import cn.hutool.core.util.StrUtil;  
5 import cn.iocoder.yudao.framework.common.core.IntArrayValuable; 3 import cn.iocoder.yudao.framework.common.core.IntArrayValuable;
6 import lombok.AllArgsConstructor; 4 import lombok.AllArgsConstructor;
7 import lombok.Getter; 5 import lombok.Getter;
@@ -41,9 +39,4 @@ public enum AiWriteTypeEnum implements IntArrayValuable { @@ -41,9 +39,4 @@ public enum AiWriteTypeEnum implements IntArrayValuable {
41 return ARRAYS; 39 return ARRAYS;
42 } 40 }
43 41
44 - public static void validateType(Integer type) {  
45 - if (ArrayUtil.contains(ARRAYS, type)) return;  
46 - throw new IllegalArgumentException(StrUtil.format("未知写作类型({})", type));  
47 - }  
48 -  
49 } 42 }
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/mindmap/AiMindMapController.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai.controller.admin.mindmap;
  2 +
  3 +import cn.iocoder.yudao.framework.common.pojo.CommonResult;
  4 +import cn.iocoder.yudao.module.ai.controller.admin.mindmap.vo.AiMindMapGenerateReqVO;
  5 +import cn.iocoder.yudao.module.ai.service.mindmap.AiMindMapService;
  6 +import io.swagger.v3.oas.annotations.Operation;
  7 +import io.swagger.v3.oas.annotations.tags.Tag;
  8 +import jakarta.annotation.Resource;
  9 +import jakarta.annotation.security.PermitAll;
  10 +import jakarta.validation.Valid;
  11 +import org.springframework.http.MediaType;
  12 +import org.springframework.web.bind.annotation.PostMapping;
  13 +import org.springframework.web.bind.annotation.RequestBody;
  14 +import org.springframework.web.bind.annotation.RequestMapping;
  15 +import org.springframework.web.bind.annotation.RestController;
  16 +import reactor.core.publisher.Flux;
  17 +
  18 +import static cn.iocoder.yudao.framework.security.core.util.SecurityFrameworkUtils.getLoginUserId;
  19 +
  20 +@Tag(name = "管理后台 - AI 思维导图")
  21 +@RestController
  22 +@RequestMapping("/ai/mind-map")
  23 +public class AiMindMapController {
  24 +
  25 + @Resource
  26 + private AiMindMapService mindMapService;
  27 +
  28 + @PostMapping(value = "/generate-stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
  29 + @Operation(summary = "脑图生成(流式)", description = "流式返回,响应较快")
  30 + @PermitAll // 解决 SSE 最终响应的时候,会被 Access Denied 拦截的问题
  31 + public Flux<CommonResult<String>> generateMindMap(@RequestBody @Valid AiMindMapGenerateReqVO generateReqVO) {
  32 + return mindMapService.generateMindMap(generateReqVO, getLoginUserId());
  33 + }
  34 +
  35 +}
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/mindmap/vo/AiMindMapGenerateReqVO.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai.controller.admin.mindmap.vo;
  2 +
  3 +import io.swagger.v3.oas.annotations.media.Schema;
  4 +import jakarta.validation.constraints.NotBlank;
  5 +import lombok.Data;
  6 +
  7 +@Schema(description = "管理后台 - AI 思维导图生成 Request VO")
  8 +@Data
  9 +public class AiMindMapGenerateReqVO {
  10 + @Schema(description = "思维导图内容提示", example = "Java 学习路线")
  11 + @NotBlank(message = "思维导图内容提示不能为空")
  12 + private String prompt;
  13 +}
0 \ No newline at end of file 14 \ No newline at end of file
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/write/vo/AiWriteGenerateReqVO.java
@@ -11,7 +11,7 @@ import lombok.Data; @@ -11,7 +11,7 @@ import lombok.Data;
11 public class AiWriteGenerateReqVO { 11 public class AiWriteGenerateReqVO {
12 12
13 @Schema(description = "写作类型", requiredMode = Schema.RequiredMode.REQUIRED, example = "1") 13 @Schema(description = "写作类型", requiredMode = Schema.RequiredMode.REQUIRED, example = "1")
14 - @InEnum(AiWriteTypeEnum.class) 14 + @InEnum(value = AiWriteTypeEnum.class, message = "写作类型必须是 {value}")
15 private Integer type; 15 private Integer type;
16 16
17 @Schema(description = "写作内容提示", example = "1.撰写:田忌赛马;2.回复:不批") 17 @Schema(description = "写作内容提示", example = "1.撰写:田忌赛马;2.回复:不批")
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/mindmap/AiMindMapDO.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai.dal.dataobject.mindmap;
  2 +
  3 +import cn.iocoder.yudao.framework.ai.core.enums.AiPlatformEnum;
  4 +import cn.iocoder.yudao.framework.mybatis.core.dataobject.BaseDO;
  5 +import com.baomidou.mybatisplus.annotation.IdType;
  6 +import com.baomidou.mybatisplus.annotation.TableId;
  7 +import com.baomidou.mybatisplus.annotation.TableName;
  8 +import lombok.Data;
  9 +
  10 +/**
  11 + * AI 思维导图 DO
  12 + *
  13 + * @author xiaoxin
  14 + */
  15 +@TableName(value = "ai_mind_map", autoResultMap = true)
  16 +@Data
  17 +public class AiMindMapDO extends BaseDO {
  18 +
  19 + /**
  20 + * 编号
  21 + */
  22 + @TableId(type = IdType.AUTO)
  23 + private Long id;
  24 +
  25 + /**
  26 + * 用户编号
  27 + */
  28 + private Long userId;
  29 +
  30 + /**
  31 + * 模型
  32 + */
  33 + private String model;
  34 +
  35 + /**
  36 + * 平台
  37 + * <p>
  38 + * 枚举 {@link AiPlatformEnum}
  39 + */
  40 + private String platform;
  41 +
  42 + /**
  43 + * 生成内容提示
  44 + */
  45 + private String prompt;
  46 +
  47 + /**
  48 + * 生成的内容
  49 + */
  50 + private String generatedContent;
  51 +
  52 + /**
  53 + * 错误信息
  54 + */
  55 + private String errorMessage;
  56 +
  57 +}
0 \ No newline at end of file 58 \ No newline at end of file
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/mysql/mindmap/AiMindMapMapper.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai.dal.mysql.mindmap;
  2 +
  3 +import cn.iocoder.yudao.framework.mybatis.core.mapper.BaseMapperX;
  4 +import cn.iocoder.yudao.module.ai.dal.dataobject.mindmap.AiMindMapDO;
  5 +import org.apache.ibatis.annotations.Mapper;
  6 +
  7 +/**
  8 + * AI 音乐 Mapper
  9 + *
  10 + * @author xiaoxin
  11 + */
  12 +@Mapper
  13 +public interface AiMindMapMapper extends BaseMapperX<AiMindMapDO> {
  14 +}
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/mysql/model/AiChatRoleMapper.java
@@ -4,9 +4,7 @@ import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum; @@ -4,9 +4,7 @@ import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
4 import cn.iocoder.yudao.framework.common.pojo.PageResult; 4 import cn.iocoder.yudao.framework.common.pojo.PageResult;
5 import cn.iocoder.yudao.framework.mybatis.core.mapper.BaseMapperX; 5 import cn.iocoder.yudao.framework.mybatis.core.mapper.BaseMapperX;
6 import cn.iocoder.yudao.framework.mybatis.core.query.LambdaQueryWrapperX; 6 import cn.iocoder.yudao.framework.mybatis.core.query.LambdaQueryWrapperX;
7 -import cn.iocoder.yudao.framework.mybatis.core.query.QueryWrapperX;  
8 import cn.iocoder.yudao.module.ai.controller.admin.model.vo.chatRole.AiChatRolePageReqVO; 7 import cn.iocoder.yudao.module.ai.controller.admin.model.vo.chatRole.AiChatRolePageReqVO;
9 -import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatModelDO;  
10 import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatRoleDO; 8 import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatRoleDO;
11 import org.apache.ibatis.annotations.Mapper; 9 import org.apache.ibatis.annotations.Mapper;
12 10
@@ -47,4 +45,10 @@ public interface AiChatRoleMapper extends BaseMapperX&lt;AiChatRoleDO&gt; { @@ -47,4 +45,10 @@ public interface AiChatRoleMapper extends BaseMapperX&lt;AiChatRoleDO&gt; {
47 .groupBy(AiChatRoleDO::getCategory)); 45 .groupBy(AiChatRoleDO::getCategory));
48 } 46 }
49 47
  48 + default List<AiChatRoleDO> selectListByName(String name) {
  49 + return selectList(new LambdaQueryWrapperX<AiChatRoleDO>()
  50 + .likeIfPresent(AiChatRoleDO::getName, name)
  51 + .orderByAsc(AiChatRoleDO::getSort));
  52 + }
  53 +
50 } 54 }
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/mindmap/AiMindMapService.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai.service.mindmap;
  2 +
  3 +import cn.iocoder.yudao.framework.common.pojo.CommonResult;
  4 +import cn.iocoder.yudao.module.ai.controller.admin.mindmap.vo.AiMindMapGenerateReqVO;
  5 +import reactor.core.publisher.Flux;
  6 +
  7 +/**
  8 + * AI 思维导图 Service 接口
  9 + *
  10 + * @author xiaoxin
  11 + */
  12 +public interface AiMindMapService {
  13 +
  14 + /**
  15 + * 生成思维导图内容
  16 + *
  17 + * @param generateReqVO 请求参数
  18 + * @param userId 用户编号
  19 + * @return 生成结果
  20 + */
  21 + Flux<CommonResult<String>> generateMindMap(AiMindMapGenerateReqVO generateReqVO, Long userId);
  22 +
  23 +}
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/mindmap/AiMindMapServiceImpl.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai.service.mindmap;
  2 +
  3 +import cn.hutool.core.collection.CollUtil;
  4 +import cn.hutool.core.util.StrUtil;
  5 +import cn.iocoder.yudao.framework.ai.core.enums.AiPlatformEnum;
  6 +import cn.iocoder.yudao.framework.ai.core.util.AiUtils;
  7 +import cn.iocoder.yudao.framework.common.pojo.CommonResult;
  8 +import cn.iocoder.yudao.framework.common.util.object.BeanUtils;
  9 +import cn.iocoder.yudao.framework.tenant.core.util.TenantUtils;
  10 +import cn.iocoder.yudao.module.ai.controller.admin.mindmap.vo.AiMindMapGenerateReqVO;
  11 +import cn.iocoder.yudao.module.ai.dal.dataobject.mindmap.AiMindMapDO;
  12 +import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatModelDO;
  13 +import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatRoleDO;
  14 +import cn.iocoder.yudao.module.ai.dal.mysql.mindmap.AiMindMapMapper;
  15 +import cn.iocoder.yudao.module.ai.enums.AiChatRoleEnum;
  16 +import cn.iocoder.yudao.module.ai.enums.ErrorCodeConstants;
  17 +import cn.iocoder.yudao.module.ai.service.model.AiApiKeyService;
  18 +import cn.iocoder.yudao.module.ai.service.model.AiChatModelService;
  19 +import cn.iocoder.yudao.module.ai.service.model.AiChatRoleService;
  20 +import jakarta.annotation.Resource;
  21 +import lombok.extern.slf4j.Slf4j;
  22 +import org.springframework.ai.chat.messages.Message;
  23 +import org.springframework.ai.chat.messages.SystemMessage;
  24 +import org.springframework.ai.chat.messages.UserMessage;
  25 +import org.springframework.ai.chat.model.ChatModel;
  26 +import org.springframework.ai.chat.model.ChatResponse;
  27 +import org.springframework.ai.chat.prompt.ChatOptions;
  28 +import org.springframework.ai.chat.prompt.Prompt;
  29 +import org.springframework.stereotype.Service;
  30 +import reactor.core.publisher.Flux;
  31 +
  32 +import java.util.ArrayList;
  33 +import java.util.List;
  34 +import java.util.Objects;
  35 +
  36 +import static cn.iocoder.yudao.framework.common.pojo.CommonResult.error;
  37 +import static cn.iocoder.yudao.framework.common.pojo.CommonResult.success;
  38 +
  39 +/**
  40 + * AI 写作 Service 实现类
  41 + *
  42 + * @author xiaoxin
  43 + */
  44 +@Service
  45 +@Slf4j
  46 +public class AiMindMapServiceImpl implements AiMindMapService {
  47 +
  48 + @Resource
  49 + private AiApiKeyService apiKeyService;
  50 + @Resource
  51 + private AiChatModelService chatModalService;
  52 + @Resource
  53 + private AiChatRoleService chatRoleService;
  54 +
  55 + @Resource
  56 + private AiMindMapMapper mindMapMapper;
  57 +
  58 + @Override
  59 + public Flux<CommonResult<String>> generateMindMap(AiMindMapGenerateReqVO generateReqVO, Long userId) {
  60 + // 1.1 获取脑图模型 尝试获取思维导图助手角色,如果没有则使用默认模型
  61 + AiChatRoleDO mindMapRole = CollUtil.getFirst(chatRoleService.getChatRoleListByName(AiChatRoleEnum.AI_MIND_MAP_ROLE.getName()));
  62 + AiChatModelDO model;
  63 + String systemMessage;
  64 + if (Objects.nonNull(mindMapRole) && Objects.nonNull(mindMapRole.getModelId())) {
  65 + model = chatModalService.getChatModel(mindMapRole.getModelId());
  66 + systemMessage = mindMapRole.getSystemMessage();
  67 + } else {
  68 + model = chatModalService.getRequiredDefaultChatModel();
  69 + systemMessage = AiChatRoleEnum.AI_MIND_MAP_ROLE.getPrompt();
  70 + }
  71 +
  72 + AiPlatformEnum platform = AiPlatformEnum.validatePlatform(model.getPlatform());
  73 + ChatModel chatModel = apiKeyService.getChatModel(model.getKeyId());
  74 +
  75 + // 2 插入思维导图信息
  76 + AiMindMapDO mindMapDO = BeanUtils.toBean(generateReqVO, AiMindMapDO.class, e -> e.setUserId(userId).setModel(model.getModel()).setPlatform(platform.getPlatform()));
  77 + mindMapMapper.insert(mindMapDO);
  78 +
  79 + ChatOptions chatOptions = AiUtils.buildChatOptions(platform, model.getModel(), model.getTemperature(), model.getMaxTokens());
  80 + // 3.1 角色设定
  81 + List<Message> chatMessages = new ArrayList<>();
  82 + if (StrUtil.isNotBlank(systemMessage)) {
  83 + chatMessages.add(new SystemMessage(systemMessage));
  84 + }
  85 + // 3.2 用户输入
  86 + chatMessages.add(new UserMessage(generateReqVO.getPrompt()));
  87 + // 3.3 构建提示词
  88 + Prompt prompt = new Prompt(chatMessages, chatOptions);
  89 +
  90 + Flux<ChatResponse> streamResponse = chatModel.stream(prompt);
  91 + // 3.4 流式返回
  92 + StringBuffer contentBuffer = new StringBuffer();
  93 + return streamResponse.map(chunk -> {
  94 + String newContent = chunk.getResult() != null ? chunk.getResult().getOutput().getContent() : null;
  95 + newContent = StrUtil.nullToDefault(newContent, ""); // 避免 null 的 情况
  96 + contentBuffer.append(newContent);
  97 + // 响应结果
  98 + return success(newContent);
  99 + }).doOnComplete(() -> {
  100 + // 忽略租户,因为 Flux 异步无法透传租户
  101 + TenantUtils.executeIgnore(() ->
  102 + mindMapMapper.updateById(new AiMindMapDO().setId(mindMapDO.getId()).setGeneratedContent(contentBuffer.toString())));
  103 + }).doOnError(throwable -> {
  104 + log.error("[generateWriteContent][generateReqVO({}) 发生异常]", generateReqVO, throwable);
  105 + // 忽略租户,因为 Flux 异步无法透传租户
  106 + TenantUtils.executeIgnore(() ->
  107 + mindMapMapper.updateById(new AiMindMapDO().setId(mindMapDO.getId()).setErrorMessage(throwable.getMessage())));
  108 + }).onErrorResume(error -> Flux.just(error(ErrorCodeConstants.WRITE_STREAM_ERROR)));
  109 +
  110 + }
  111 +
  112 +}
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/model/AiChatRoleService.java
@@ -118,4 +118,11 @@ public interface AiChatRoleService { @@ -118,4 +118,11 @@ public interface AiChatRoleService {
118 */ 118 */
119 List<String> getChatRoleCategoryList(); 119 List<String> getChatRoleCategoryList();
120 120
  121 + /**
  122 + * 根据名字获得聊天角色
  123 + * @param name 名字
  124 + * @return 聊天角色列表
  125 + */
  126 + List<AiChatRoleDO> getChatRoleListByName(String name);
  127 +
121 } 128 }
122 \ No newline at end of file 129 \ No newline at end of file
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/model/AiChatRoleServiceImpl.java
@@ -137,5 +137,10 @@ public class AiChatRoleServiceImpl implements AiChatRoleService { @@ -137,5 +137,10 @@ public class AiChatRoleServiceImpl implements AiChatRoleService {
137 return convertList(list, AiChatRoleDO::getCategory, role -> role != null && StrUtil.isNotBlank(role.getCategory())); 137 return convertList(list, AiChatRoleDO::getCategory, role -> role != null && StrUtil.isNotBlank(role.getCategory()));
138 } 138 }
139 139
  140 + @Override
  141 + public List<AiChatRoleDO> getChatRoleListByName(String name) {
  142 + return chatRoleMapper.selectListByName(name);
  143 + }
  144 +
140 } 145 }
141 146
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/write/AiWriteServiceImpl.java
@@ -5,15 +5,14 @@ import cn.hutool.core.util.StrUtil; @@ -5,15 +5,14 @@ import cn.hutool.core.util.StrUtil;
5 import cn.iocoder.yudao.framework.ai.core.enums.AiPlatformEnum; 5 import cn.iocoder.yudao.framework.ai.core.enums.AiPlatformEnum;
6 import cn.iocoder.yudao.framework.ai.core.util.AiUtils; 6 import cn.iocoder.yudao.framework.ai.core.util.AiUtils;
7 import cn.iocoder.yudao.framework.common.pojo.CommonResult; 7 import cn.iocoder.yudao.framework.common.pojo.CommonResult;
8 -import cn.iocoder.yudao.framework.common.pojo.PageResult;  
9 import cn.iocoder.yudao.framework.common.util.object.BeanUtils; 8 import cn.iocoder.yudao.framework.common.util.object.BeanUtils;
10 import cn.iocoder.yudao.framework.tenant.core.util.TenantUtils; 9 import cn.iocoder.yudao.framework.tenant.core.util.TenantUtils;
11 -import cn.iocoder.yudao.module.ai.controller.admin.model.vo.chatRole.AiChatRolePageReqVO;  
12 import cn.iocoder.yudao.module.ai.controller.admin.write.vo.AiWriteGenerateReqVO; 10 import cn.iocoder.yudao.module.ai.controller.admin.write.vo.AiWriteGenerateReqVO;
13 import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatModelDO; 11 import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatModelDO;
14 import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatRoleDO; 12 import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatRoleDO;
15 import cn.iocoder.yudao.module.ai.dal.dataobject.write.AiWriteDO; 13 import cn.iocoder.yudao.module.ai.dal.dataobject.write.AiWriteDO;
16 import cn.iocoder.yudao.module.ai.dal.mysql.write.AiWriteMapper; 14 import cn.iocoder.yudao.module.ai.dal.mysql.write.AiWriteMapper;
  15 +import cn.iocoder.yudao.module.ai.enums.AiChatRoleEnum;
17 import cn.iocoder.yudao.module.ai.enums.DictTypeConstants; 16 import cn.iocoder.yudao.module.ai.enums.DictTypeConstants;
18 import cn.iocoder.yudao.module.ai.enums.ErrorCodeConstants; 17 import cn.iocoder.yudao.module.ai.enums.ErrorCodeConstants;
19 import cn.iocoder.yudao.module.ai.enums.write.AiWriteTypeEnum; 18 import cn.iocoder.yudao.module.ai.enums.write.AiWriteTypeEnum;
@@ -23,6 +22,9 @@ import cn.iocoder.yudao.module.ai.service.model.AiChatRoleService; @@ -23,6 +22,9 @@ import cn.iocoder.yudao.module.ai.service.model.AiChatRoleService;
23 import cn.iocoder.yudao.module.system.api.dict.DictDataApi; 22 import cn.iocoder.yudao.module.system.api.dict.DictDataApi;
24 import jakarta.annotation.Resource; 23 import jakarta.annotation.Resource;
25 import lombok.extern.slf4j.Slf4j; 24 import lombok.extern.slf4j.Slf4j;
  25 +import org.springframework.ai.chat.messages.Message;
  26 +import org.springframework.ai.chat.messages.SystemMessage;
  27 +import org.springframework.ai.chat.messages.UserMessage;
26 import org.springframework.ai.chat.model.ChatResponse; 28 import org.springframework.ai.chat.model.ChatResponse;
27 import org.springframework.ai.chat.model.StreamingChatModel; 29 import org.springframework.ai.chat.model.StreamingChatModel;
28 import org.springframework.ai.chat.prompt.ChatOptions; 30 import org.springframework.ai.chat.prompt.ChatOptions;
@@ -30,6 +32,7 @@ import org.springframework.ai.chat.prompt.Prompt; @@ -30,6 +32,7 @@ import org.springframework.ai.chat.prompt.Prompt;
30 import org.springframework.stereotype.Service; 32 import org.springframework.stereotype.Service;
31 import reactor.core.publisher.Flux; 33 import reactor.core.publisher.Flux;
32 34
  35 +import java.util.ArrayList;
33 import java.util.List; 36 import java.util.List;
34 import java.util.Objects; 37 import java.util.Objects;
35 38
@@ -61,13 +64,15 @@ public class AiWriteServiceImpl implements AiWriteService { @@ -61,13 +64,15 @@ public class AiWriteServiceImpl implements AiWriteService {
61 @Override 64 @Override
62 public Flux<CommonResult<String>> generateWriteContent(AiWriteGenerateReqVO generateReqVO, Long userId) { 65 public Flux<CommonResult<String>> generateWriteContent(AiWriteGenerateReqVO generateReqVO, Long userId) {
63 // 1.1 获取写作模型 尝试获取写作助手角色,如果没有则使用默认模型 66 // 1.1 获取写作模型 尝试获取写作助手角色,如果没有则使用默认模型
64 - AiChatRoleDO writeRole = selectOneWriteRole(); 67 + AiChatRoleDO writeRole = CollUtil.getFirst(chatRoleService.getChatRoleListByName(AiChatRoleEnum.AI_WRITE_ROLE.getName()));
65 AiChatModelDO model; 68 AiChatModelDO model;
66 - // TODO @xin:writeRole.getModelId 可能为空。所以,最好是先通过 chatRole 拿。如果它没拿到,通过 getRequiredDefaultChatModel 再拿。  
67 - if (Objects.nonNull(writeRole)) { 69 + String systemMessage;
  70 + if (Objects.nonNull(writeRole) && Objects.nonNull(writeRole.getModelId())) {
68 model = chatModalService.getChatModel(writeRole.getModelId()); 71 model = chatModalService.getChatModel(writeRole.getModelId());
  72 + systemMessage = writeRole.getSystemMessage();
69 } else { 73 } else {
70 model = chatModalService.getRequiredDefaultChatModel(); 74 model = chatModalService.getRequiredDefaultChatModel();
  75 + systemMessage = AiChatRoleEnum.AI_WRITE_ROLE.getPrompt();
71 } 76 }
72 // 1.2 校验平台 77 // 1.2 校验平台
73 AiPlatformEnum platform = AiPlatformEnum.validatePlatform(model.getPlatform()); 78 AiPlatformEnum platform = AiPlatformEnum.validatePlatform(model.getPlatform());
@@ -77,9 +82,16 @@ public class AiWriteServiceImpl implements AiWriteService { @@ -77,9 +82,16 @@ public class AiWriteServiceImpl implements AiWriteService {
77 AiWriteDO writeDO = BeanUtils.toBean(generateReqVO, AiWriteDO.class, e -> e.setUserId(userId).setModel(model.getModel()).setPlatform(platform.getPlatform())); 82 AiWriteDO writeDO = BeanUtils.toBean(generateReqVO, AiWriteDO.class, e -> e.setUserId(userId).setModel(model.getModel()).setPlatform(platform.getPlatform()));
78 writeMapper.insert(writeDO); 83 writeMapper.insert(writeDO);
79 84
80 - // 3.1 构建提示词  
81 ChatOptions chatOptions = AiUtils.buildChatOptions(platform, model.getModel(), model.getTemperature(), model.getMaxTokens()); 85 ChatOptions chatOptions = AiUtils.buildChatOptions(platform, model.getModel(), model.getTemperature(), model.getMaxTokens());
82 - Prompt prompt = new Prompt(buildWritingPrompt(generateReqVO), chatOptions); 86 + // 3.1 角色设定
  87 + List<Message> chatMessages = new ArrayList<>();
  88 + if (StrUtil.isNotBlank(systemMessage)) {
  89 + chatMessages.add(new SystemMessage(systemMessage));
  90 + }
  91 + // 3.2 用户输入
  92 + chatMessages.add(new UserMessage(buildWritingPrompt(generateReqVO)));
  93 + // 3.3 构建提示词
  94 + Prompt prompt = new Prompt(chatMessages, chatOptions);
83 Flux<ChatResponse> streamResponse = chatModel.stream(prompt); 95 Flux<ChatResponse> streamResponse = chatModel.stream(prompt);
84 96
85 // 3.2 流式返回 97 // 3.2 流式返回
@@ -102,24 +114,8 @@ public class AiWriteServiceImpl implements AiWriteService { @@ -102,24 +114,8 @@ public class AiWriteServiceImpl implements AiWriteService {
102 }).onErrorResume(error -> Flux.just(error(ErrorCodeConstants.WRITE_STREAM_ERROR))); 114 }).onErrorResume(error -> Flux.just(error(ErrorCodeConstants.WRITE_STREAM_ERROR)));
103 } 115 }
104 116
105 - // TODO @xin:chatRoleService 增加一个 getChatRoleListByName;  
106 - private AiChatRoleDO selectOneWriteRole() {  
107 - AiChatRoleDO chatRoleDO = null;  
108 - // TODO @xin:"写作助手" 枚举下。  
109 - PageResult<AiChatRoleDO> writeRolePage = chatRoleService.getChatRolePage(new AiChatRolePageReqVO().setName("写作助手"));  
110 - List<AiChatRoleDO> list = writeRolePage.getList();  
111 - // TODO @xin:CollUtil.getFirst 简化下  
112 - if (CollUtil.isNotEmpty(list)) {  
113 - chatRoleDO = list.get(0);  
114 - }  
115 - return chatRoleDO;  
116 - }  
117 -  
118 private String buildWritingPrompt(AiWriteGenerateReqVO generateReqVO) { 117 private String buildWritingPrompt(AiWriteGenerateReqVO generateReqVO) {
119 - // 校验写作类型是否合法  
120 Integer type = generateReqVO.getType(); 118 Integer type = generateReqVO.getType();
121 - // TODO @xin:这里可以搞到 validator 的校验。InEnum  
122 - AiWriteTypeEnum.validateType(type);  
123 String format = dictDataApi.getDictDataLabel(DictTypeConstants.AI_WRITE_FORMAT, generateReqVO.getFormat()); 119 String format = dictDataApi.getDictDataLabel(DictTypeConstants.AI_WRITE_FORMAT, generateReqVO.getFormat());
124 String tone = dictDataApi.getDictDataLabel(DictTypeConstants.AI_WRITE_TONE, generateReqVO.getTone()); 120 String tone = dictDataApi.getDictDataLabel(DictTypeConstants.AI_WRITE_TONE, generateReqVO.getTone());
125 String language = dictDataApi.getDictDataLabel(DictTypeConstants.AI_WRITE_LANGUAGE, generateReqVO.getLanguage()); 121 String language = dictDataApi.getDictDataLabel(DictTypeConstants.AI_WRITE_LANGUAGE, generateReqVO.getLanguage());