BJ_ZTEITS-Developers / parking-cloud-yudao

Commit 29e421432d828bad49a147136bec8a7116dcf04d

Authored by xiaoxin
1 parent e39a6c34

【解决todo】AI 写作、脑图：model、systemMessage获取逻辑调整

Inline Side-by-side
Showing 3 changed files with 72 additions and 42 deletions
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/mindmap/AiMindMapDO.java
  View file @29e4214
... ... @@ -12,8 +12,7 @@ import lombok.Data;
12 12 *
13 13 * @author xiaoxin
14 14 */
15   -// TODO @xin：如果没 typehandler 的需求，autoResultMap 可以去掉哈
16   -@TableName(value = "ai_mind_map", autoResultMap = true)
  15 +@TableName(value = "ai_mind_map")
17 16 @Data
18 17 public class AiMindMapDO extends BaseDO {
19 18  
... ... @@ -25,7 +24,7 @@ public class AiMindMapDO extends BaseDO {
25 24  
26 25 /**
27 26 * 用户编号
28   - *
  27 + * <p>
29 28 * 关联 AdminUserDO 的 userId 字段
30 29 */
31 30 private Long userId;
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/mindmap/AiMindMapServiceImpl.java
  View file @29e4214
1 1 package cn.iocoder.yudao.module.ai.service.mindmap;
2 2  
3 3 import cn.hutool.core.collection.CollUtil;
  4 +import cn.hutool.core.lang.Assert;
4 5 import cn.hutool.core.util.StrUtil;
5 6 import cn.iocoder.yudao.framework.ai.core.enums.AiPlatformEnum;
6 7 import cn.iocoder.yudao.framework.ai.core.util.AiUtils;
... ... @@ -57,33 +58,25 @@ public class AiMindMapServiceImpl implements AiMindMapService {
57 58  
58 59 @Override
59 60 public Flux<CommonResult<String>> generateMindMap(AiMindMapGenerateReqVO generateReqVO, Long userId) {
60   - // 1.1 获取脑图模型 尝试获取思维导图助手角色，如果没有则使用默认模型
  61 + // 1 获取脑图模型 尝试获取思维导图助手角色，如果没有则使用默认模型
61 62 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.getSystemMessage();
70   - }
71   -
  63 + // 1.1 获取脑图执行模型
  64 + AiChatModelDO model = getModel(mindMapRole);
  65 + // 1.2 获取角色设定消息
  66 + String systemMessage = Objects.nonNull(mindMapRole) && StrUtil.isNotBlank(mindMapRole.getSystemMessage())
  67 + ? mindMapRole.getSystemMessage() : AiChatRoleEnum.AI_MIND_MAP_ROLE.getSystemMessage();
  68 + // 1.3 校验平台
72 69 AiPlatformEnum platform = AiPlatformEnum.validatePlatform(model.getPlatform());
73 70 ChatModel chatModel = apiKeyService.getChatModel(model.getKeyId());
74 71  
75 72 // 2 插入思维导图信息
76   - AiMindMapDO mindMapDO = BeanUtils.toBean(generateReqVO, AiMindMapDO.class, e -> e.setUserId(userId).setModel(model.getModel()).setPlatform(platform.getPlatform()));
  73 + AiMindMapDO mindMapDO = BeanUtils.toBean(generateReqVO, AiMindMapDO.class,
  74 + mindMap -> mindMap.setUserId(userId).setModel(model.getModel()).setPlatform(platform.getPlatform()));
77 75 mindMapMapper.insert(mindMapDO);
78 76  
79 77 ChatOptions chatOptions = AiUtils.buildChatOptions(platform, model.getModel(), model.getTemperature(), model.getMaxTokens());
80 78 // 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()));
  79 + List<Message> chatMessages = buildMessages(generateReqVO, systemMessage);
87 80 // 3.3 构建提示词
88 81 Prompt prompt = new Prompt(chatMessages, chatOptions);
89 82  
... ... @@ -109,4 +102,28 @@ public class AiMindMapServiceImpl implements AiMindMapService {
109 102  
110 103 }
111 104  
  105 + private static List<Message> buildMessages(AiMindMapGenerateReqVO generateReqVO, String systemMessage) {
  106 + List<Message> chatMessages = new ArrayList<>();
  107 + if (StrUtil.isNotBlank(systemMessage)) {
  108 + // 1.1 角色设定
  109 + chatMessages.add(new SystemMessage(systemMessage));
  110 + }
  111 + // 1.2 用户输入
  112 + chatMessages.add(new UserMessage(generateReqVO.getPrompt()));
  113 + return chatMessages;
  114 + }
  115 +
  116 + // TODO 芋艿：这里脑图、写作都用到了，是不是可以抽哪里去
  117 + private AiChatModelDO getModel(AiChatRoleDO chatRoleDO) {
  118 + AiChatModelDO model = null;
  119 + if (Objects.nonNull(chatRoleDO) && Objects.nonNull(chatRoleDO.getModelId())) {
  120 + model = chatModalService.getChatModel(chatRoleDO.getModelId());
  121 + }
  122 + if (Objects.isNull(model)) {
  123 + model = chatModalService.getRequiredDefaultChatModel();
  124 + }
  125 + Assert.notNull(model, "[AI] 获取不到模型");
  126 + return model;
  127 + }
  128 +
112 129 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/write/AiWriteServiceImpl.java
  View file @29e4214
1 1 package cn.iocoder.yudao.module.ai.service.write;
2 2  
3 3 import cn.hutool.core.collection.CollUtil;
  4 +import cn.hutool.core.lang.Assert;
4 5 import cn.hutool.core.util.StrUtil;
5 6 import cn.iocoder.yudao.framework.ai.core.enums.AiPlatformEnum;
6 7 import cn.iocoder.yudao.framework.ai.core.util.AiUtils;
... ... @@ -67,19 +68,14 @@ public class AiWriteServiceImpl implements AiWriteService {
67 68  
68 69 @Override
69 70 public Flux<CommonResult<String>> generateWriteContent(AiWriteGenerateReqVO generateReqVO, Long userId) {
70   - // 1.1 获取写作模型 尝试获取写作助手角色，如果没有则使用默认模型
  71 + // 1 获取写作模型 尝试获取写作助手角色，没有则使用默认模型
71 72 AiChatRoleDO writeRole = CollUtil.getFirst(chatRoleService.getChatRoleListByName(AiChatRoleEnum.AI_WRITE_ROLE.getName()));
72   - // TODO @xin：如果有 writeRole，但是没 modeId，是不是也可以用 systemMessage 哈？建议的写法是：先通过 modelId 获取 model。如果 model == null，则 chatModalService.getRequiredDefaultChatModel()；如果还是 null，则抛出异常；。。。。。。。。。。。。。。然后，systemMessage = writeRole != null && writeRole.systemPrompt != "" 这样处理。
73   - AiChatModelDO model;
74   - String systemMessage;
75   - if (Objects.nonNull(writeRole) && Objects.nonNull(writeRole.getModelId())) {
76   - model = chatModalService.getChatModel(writeRole.getModelId());
77   - systemMessage = writeRole.getSystemMessage();
78   - } else {
79   - model = chatModalService.getRequiredDefaultChatModel();
80   - systemMessage = AiChatRoleEnum.AI_WRITE_ROLE.getSystemMessage();
81   - }
82   - // 1.2 校验平台
  73 + // 1.1 获取写作执行模型
  74 + AiChatModelDO model = getModel(writeRole);
  75 + // 1.2 获取角色设定消息
  76 + String systemMessage = Objects.nonNull(writeRole) && StrUtil.isNotBlank(writeRole.getSystemMessage())
  77 + ? writeRole.getSystemMessage() : AiChatRoleEnum.AI_WRITE_ROLE.getSystemMessage();
  78 + // 1.3 校验平台
83 79 AiPlatformEnum platform = AiPlatformEnum.validatePlatform(model.getPlatform());
84 80 StreamingChatModel chatModel = apiKeyService.getChatModel(model.getKeyId());
85 81  
... ... @@ -90,16 +86,11 @@ public class AiWriteServiceImpl implements AiWriteService {
90 86  
91 87 // 3. 调用大模型，写作生成
92 88 ChatOptions chatOptions = AiUtils.buildChatOptions(platform, model.getModel(), model.getTemperature(), model.getMaxTokens());
93   - // 3.1 角色设定
94   - // TODO @xin：要不把 90 到 97 这部分，合并到一个方法里。目的是：让这个方法的主干更明确
95   - List<Message> chatMessages = new ArrayList<>();
96   - if (StrUtil.isNotBlank(systemMessage)) {
97   - chatMessages.add(new SystemMessage(systemMessage));
98   - }
99   - // 3.2 用户输入
100   - chatMessages.add(new UserMessage(buildWritingPrompt(generateReqVO)));
101   - // 3.3 构建提示词
  89 + // 3.1 构建消息列表
  90 + List<Message> chatMessages = buildMessages(generateReqVO, systemMessage);
  91 + // 3.2 构建提示词
102 92 Prompt prompt = new Prompt(chatMessages, chatOptions);
  93 + // 3.3 流式调用
103 94 Flux<ChatResponse> streamResponse = chatModel.stream(prompt);
104 95  
105 96 // 4. 流式返回
... ... @@ -122,6 +113,29 @@ public class AiWriteServiceImpl implements AiWriteService {
122 113 }).onErrorResume(error -> Flux.just(error(ErrorCodeConstants.WRITE_STREAM_ERROR)));
123 114 }
124 115  
  116 + private AiChatModelDO getModel(AiChatRoleDO writeRole) {
  117 + AiChatModelDO model = null;
  118 + if (Objects.nonNull(writeRole) && Objects.nonNull(writeRole.getModelId())) {
  119 + model = chatModalService.getChatModel(writeRole.getModelId());
  120 + }
  121 + if (Objects.isNull(model)) {
  122 + model = chatModalService.getRequiredDefaultChatModel();
  123 + }
  124 + Assert.notNull(model, "[AI] 获取不到模型");
  125 + return model;
  126 + }
  127 +
  128 + private List<Message> buildMessages(AiWriteGenerateReqVO generateReqVO, String systemMessage) {
  129 + List<Message> chatMessages = new ArrayList<>();
  130 + if (StrUtil.isNotBlank(systemMessage)) {
  131 + // 1.1 角色设定
  132 + chatMessages.add(new SystemMessage(systemMessage));
  133 + }
  134 + // 1.2 用户输入
  135 + chatMessages.add(new UserMessage(buildWritingPrompt(generateReqVO)));
  136 + return chatMessages;
  137 + }
  138 +
125 139 private String buildWritingPrompt(AiWriteGenerateReqVO generateReqVO) {
126 140 String format = dictDataApi.getDictDataLabel(DictTypeConstants.AI_WRITE_FORMAT, generateReqVO.getFormat());
127 141 String tone = dictDataApi.getDictDataLabel(DictTypeConstants.AI_WRITE_TONE, generateReqVO.getTone());
... ...