BJ_ZTEITS-Developers / parking-cloud-yudao

Commit 8b1e1c047bf6c6334b3bd6ed17c1dc20d7677b83

Authored by cherishsince
1 parent 92ee6659

【解决todo】 处理 dall 异步调用，采用 @Async

Inline Side-by-side
Showing 1 changed file with 82 additions and 99 deletions
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/image/AiImageServiceImpl.java
  View file @8b1e1c0
1 1 package cn.iocoder.yudao.module.ai.service.image;
2 2  
3 3 import cn.hutool.core.util.IdUtil;
4   -import cn.hutool.core.util.StrUtil;
5 4 import cn.hutool.http.HttpUtil;
6 5 import cn.iocoder.yudao.framework.ai.core.enums.OpenAiImageModelEnum;
7 6 import cn.iocoder.yudao.framework.ai.core.enums.OpenAiImageStyleEnum;
8 7 import cn.iocoder.yudao.framework.ai.core.exception.AiException;
9 8 import cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil;
10 9 import cn.iocoder.yudao.framework.common.pojo.PageResult;
11   -import cn.iocoder.yudao.framework.common.util.json.JsonUtils;
  10 +import cn.iocoder.yudao.framework.common.util.object.BeanUtils;
12 11 import cn.iocoder.yudao.framework.mybatis.core.query.LambdaQueryWrapperX;
13   -import cn.iocoder.yudao.framework.security.core.util.SecurityFrameworkUtils;
  12 +import cn.iocoder.yudao.module.ai.AiCommonConstants;
14 13 import cn.iocoder.yudao.module.ai.ErrorCodeConstants;
15 14 import cn.iocoder.yudao.module.ai.controller.admin.image.vo.*;
16   -import cn.iocoder.yudao.module.ai.convert.AiImageConvert;
17 15 import cn.iocoder.yudao.module.ai.dal.dataobject.image.AiImageDO;
18 16 import cn.iocoder.yudao.module.ai.dal.mysql.image.AiImageMapper;
  17 +import cn.iocoder.yudao.module.ai.enums.AiImagePublicStatusEnum;
19 18 import cn.iocoder.yudao.module.ai.enums.AiImageStatusEnum;
20 19 import cn.iocoder.yudao.module.infra.api.file.FileApi;
  20 +import com.google.common.collect.ImmutableMap;
21 21 import jakarta.annotation.PostConstruct;
22   -import lombok.AllArgsConstructor;
  22 +import jakarta.annotation.Resource;
23 23 import lombok.extern.slf4j.Slf4j;
24 24 import org.springframework.ai.image.ImageGeneration;
25 25 import org.springframework.ai.image.ImagePrompt;
26 26 import org.springframework.ai.image.ImageResponse;
27 27 import org.springframework.ai.models.midjourney.api.MidjourneyInteractionsApi;
28   -import org.springframework.ai.models.midjourney.api.req.ReRollReq;
29 28 import org.springframework.ai.models.midjourney.webSocket.MidjourneyWebSocketStarter;
30   -import org.springframework.ai.models.midjourney.webSocket.WssNotify;
31 29 import org.springframework.ai.openai.OpenAiImageClient;
32 30 import org.springframework.ai.openai.OpenAiImageOptions;
  31 +import org.springframework.scheduling.annotation.Async;
33 32 import org.springframework.stereotype.Service;
34 33 import org.springframework.transaction.annotation.Transactional;
35 34  
36   -import java.util.Collections;
37 35 import java.util.List;
38   -import java.util.concurrent.LinkedBlockingQueue;
39   -import java.util.concurrent.ThreadPoolExecutor;
40   -import java.util.concurrent.TimeUnit;
41 36  
42 37 import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
43 38  
44 39  
45 40 // TODO @fan：注释优化下哈
  41 +
46 42 /**
47   - * ai 作图
  43 + * AI 绘画(接入 dall2/dall3、midjourney)
48 44 *
49 45 * @author fansili
50 46 * @time 2024/4/25 15:51
51 47 * @since 1.0
52 48 */
53   -@AllArgsConstructor
54 49 @Service
55 50 @Slf4j
56 51 public class AiImageServiceImpl implements AiImageService {
... ... @@ -58,102 +53,92 @@ public class AiImageServiceImpl implements AiImageService {
58 53 // TODO @fan：使用 @Resource 注入
59 54  
60 55 // TODO @fan：imageMapper
61   - private final AiImageMapper aiImageMapper;
62   -
63   - private final FileApi fileApi;
64   -
65   - private final OpenAiImageClient openAiImageClient;
66   -
67   - private final MidjourneyWebSocketStarter midjourneyWebSocketStarter;
68   -
69   - private final MidjourneyInteractionsApi midjourneyInteractionsApi;
70   -
71   - private static ThreadPoolExecutor EXECUTOR = new ThreadPoolExecutor(
72   - 3, 5, 1, TimeUnit.HOURS, new LinkedBlockingQueue<>(32));
  56 + @Resource
  57 + private AiImageMapper imageMapper;
  58 + @Resource
  59 + private FileApi fileApi;
  60 + @Resource
  61 + private OpenAiImageClient openAiImageClient;
  62 + @Resource
  63 + private MidjourneyWebSocketStarter midjourneyWebSocketStarter;
  64 + @Resource
  65 + private MidjourneyInteractionsApi midjourneyInteractionsApi;
73 66  
74 67 // TODO @fan：接 mj proxy
75 68 @PostConstruct
76 69 public void startMidjourney() {
77   - log.info("midjourney web socket starter...");
78   - midjourneyWebSocketStarter.start(new WssNotify() {
79   - @Override
80   - public void notify(int code, String message) {
81   - log.info("code: {}, message: {}", code, message);
82   - if (message.contains("Authentication failed")) {
83   - // TODO 芋艿，这里看怎么处理，token无效的时候会认证失败！
84   - // 认证失败
85   - log.error("midjourney socket 认证失败，检查token是否失效!");
86   - }
87   - }
88   - });
  70 + // todo @fan 暂时注释掉
  71 +// log.info("midjourney web socket starter...");
  72 +// midjourneyWebSocketStarter.start(new WssNotify() {
  73 +// @Override
  74 +// public void notify(int code, String message) {
  75 +// log.info("code: {}, message: {}", code, message);
  76 +// if (message.contains("Authentication failed")) {
  77 +// // TODO 芋艿，这里看怎么处理，token无效的时候会认证失败！
  78 +// // 认证失败
  79 +// log.error("midjourney socket 认证失败，检查token是否失效!");
  80 +// }
  81 +// }
  82 +// });
89 83 }
90 84  
91   - // TODO @fan：1）分页，然后 loginUser 通过参数传入，这样 Service 无状态；2）另外，返回 DO；VO 的翻译，交给 Controller；3：还有，使用 BeanUtils 替代哈
92 85 @Override
93   - public PageResult<AiImageListRespVO> list(AiImageListReqVO req) {
94   - // 获取登录用户
95   - Long loginUserId = SecurityFrameworkUtils.getLoginUserId();
  86 + public PageResult<AiImageDO> getImagePageMy(Long loginUserId, AiImageListReqVO req) {
96 87 // 查询当前用户下所有的绘画记录
97   - PageResult<AiImageDO> pageResult = aiImageMapper.selectPage(req,
  88 + return imageMapper.selectPage(req,
98 89 new LambdaQueryWrapperX<AiImageDO>()
99 90 .eq(AiImageDO::getUserId, loginUserId)
100   - .orderByDesc(AiImageDO::getId)
101   - );
102   - // 转换 PageResult<AiImageListRespVO> 返回
103   - PageResult<AiImageListRespVO> result = new PageResult<>();
104   - result.setTotal(pageResult.getTotal());
105   - result.setList(AiImageConvert.INSTANCE.convertAiImageListRespVO(pageResult.getList()));
106   - return result;
  91 + .orderByDesc(AiImageDO::getId));
107 92 }
108 93  
109   - // TODO @fan：1）返回 DO；VO 的翻译，交给 Controller；2）还有，使用 BeanUtils 替代哈
110 94 @Override
111   - public AiImageListRespVO getMy(Long id) {
112   - AiImageDO aiImageDO = aiImageMapper.selectById(id);
113   - return AiImageConvert.INSTANCE.convertAiImageListRespVO(aiImageDO);
  95 + public AiImageDO getMy(Long id) {
  96 + return imageMapper.selectById(id);
114 97 }
115 98  
116   - // TODO @fan：1）loginUserId 通过 controller 传入；
117 99 @Override
118   - public AiImageDallRespVO dall(AiImageDallReqVO req) {
119   - Long loginUserId = SecurityFrameworkUtils.getLoginUserId();
  100 + public Long dall(Long loginUserId, AiImageDallReqVO req) {
120 101 // 保存数据库
121   - // TODO @fan：1）使用 BeanUtils；2）使用链式调用哈；
122   - AiImageDO aiImageDO = AiImageConvert.INSTANCE.convertAiImageDO(req);
123   - aiImageDO.setStatus(AiImageStatusEnum.IN_PROGRESS.getStatus());
124   - aiImageDO.setUserId(loginUserId);
125   - aiImageMapper.insert(aiImageDO);
  102 + AiImageDO aiImageDO = BeanUtils.toBean(req, AiImageDO.class)
  103 + .setUserId(loginUserId)
  104 + .setWidth(req.getWidth())
  105 + .setHeight(req.getHeight())
  106 + .setDrawRequest(ImmutableMap.of(AiCommonConstants.DRAW_REQ_KEY_STYLE, req.getStyle()))
  107 + .setPublicStatus(AiImagePublicStatusEnum.PRIVATE.getStatus())
  108 + .setStatus(AiImageStatusEnum.IN_PROGRESS.getStatus());
  109 + imageMapper.insert(aiImageDO);
126 110 // 异步执行
127   - // TODO @fan：使用 @Async 去调用哈；
128   - EXECUTOR.execute(() -> {
129   - try {
130   - // 获取 model
131   - OpenAiImageModelEnum openAiImageModelEnum = OpenAiImageModelEnum.valueOfModel(req.getModel());
132   - OpenAiImageStyleEnum openAiImageStyleEnum = OpenAiImageStyleEnum.valueOfStyle(req.getStyle());
133   -
134   - // 转换openai 参数
135   - // TODO @fan：需要考虑，不同平台，参数不同；
136   - OpenAiImageOptions openAiImageOptions = new OpenAiImageOptions();
137   - openAiImageOptions.setModel(openAiImageModelEnum.getModel());
138   - openAiImageOptions.setStyle(openAiImageStyleEnum.getStyle());
139   - openAiImageOptions.setSize(req.getSize());
140   - ImageResponse imageResponse = openAiImageClient.call(new ImagePrompt(req.getPrompt(), openAiImageOptions));
141   - // 发送
142   - ImageGeneration imageGeneration = imageResponse.getResult();
143   - // 图片保存到服务器
144   - String filePath = fileApi.createFile(HttpUtil.downloadBytes(imageGeneration.getOutput().getUrl()));
145   - // 更新数据库
146   - aiImageMapper.updateById(new AiImageDO().setId(aiImageDO.getId()).setStatus(AiImageStatusEnum.COMPLETE.getStatus())
147   - .setPicUrl(filePath).setOriginalPicUrl(imageGeneration.getOutput().getUrl()));
148   - } catch (AiException aiException) {
149   - // TODO @fan：错误日志，也打印下哈；因为 aiException.getMessage() 比较精简；
150   - aiImageMapper.updateById(new AiImageDO().setId(aiImageDO.getId()).setStatus(AiImageStatusEnum.FAIL.getStatus())
151   - .setErrorMessage(aiException.getMessage()));
152   - }
153   - });
154   - // TODO @fan：返回 id 就可以啦
  111 + doDall(aiImageDO, req);
155 112 // 转换 AiImageDallDrawingRespVO
156   - return AiImageConvert.INSTANCE.convertAiImageDallDrawingRespVO(aiImageDO);
  113 + return aiImageDO.getId();
  114 + }
  115 +
  116 + @Async
  117 + public void doDall(AiImageDO aiImageDO, AiImageDallReqVO req) {
  118 + try {
  119 + // 获取 model
  120 + OpenAiImageModelEnum openAiImageModelEnum = OpenAiImageModelEnum.valueOfModel(req.getModel());
  121 + OpenAiImageStyleEnum openAiImageStyleEnum = OpenAiImageStyleEnum.valueOfStyle(req.getStyle());
  122 +
  123 + // 转换openai 参数
  124 + // TODO @fan：需要考虑，不同平台，参数不同；
  125 + OpenAiImageOptions openAiImageOptions = new OpenAiImageOptions();
  126 + openAiImageOptions.setModel(openAiImageModelEnum.getModel());
  127 + openAiImageOptions.setStyle(openAiImageStyleEnum.getStyle());
  128 + openAiImageOptions.setSize(String.format(AiCommonConstants.DALL_SIZE_TEMPLATE, req.getWidth(), req.getHeight()));
  129 + ImageResponse imageResponse = openAiImageClient.call(new ImagePrompt(req.getPrompt(), openAiImageOptions));
  130 + // 发送
  131 + ImageGeneration imageGeneration = imageResponse.getResult();
  132 + // 图片保存到服务器
  133 + String filePath = fileApi.createFile(HttpUtil.downloadBytes(imageGeneration.getOutput().getUrl()));
  134 + // 更新数据库
  135 + imageMapper.updateById(new AiImageDO().setId(aiImageDO.getId()).setStatus(AiImageStatusEnum.COMPLETE.getStatus())
  136 + .setPicUrl(filePath).setOriginalPicUrl(imageGeneration.getOutput().getUrl()));
  137 + } catch (AiException aiException) {
  138 + // TODO @fan：错误日志，也打印下哈；因为 aiException.getMessage() 比较精简；
  139 + imageMapper.updateById(new AiImageDO().setId(aiImageDO.getId()).setStatus(AiImageStatusEnum.FAIL.getStatus())
  140 + .setErrorMessage(aiException.getMessage()));
  141 + }
157 142 }
158 143  
159 144 @Override
... ... @@ -198,18 +183,16 @@ public class AiImageServiceImpl implements AiImageService {
198 183 // );
199 184 }
200 185  
201   - // TODO @fan：1）需要校验存在；2）需要校验属于我；
202 186 @Override
203 187 public void deleteMy(Long id, Long userId) {
204   - // 校验记录是否存在
205   - // TODO @fan：aiImageDO 这种命名 image 就 ok 拉，更简洁
206   - // TODO @fan：下面这个，可以返回图片不存在
207   - AiImageDO aiImageDO = validateExists(id);
208   - if (!aiImageDO.getUserId().equals(userId)) {
209   - throw exception(ErrorCodeConstants.AI_IMAGE_NOT_CREATE_USER);
  188 + // 校验是否存在，并获取 image
  189 + AiImageDO image = validateExists(id);
  190 + // 是否属于当前用户
  191 + if (!image.getUserId().equals(userId)) {
  192 + throw exception(ErrorCodeConstants.AI_IMAGE_NOT_EXISTS);
210 193 }
211 194 // 删除记录
212   - aiImageMapper.deleteById(id);
  195 + imageMapper.deleteById(id);
213 196 }
214 197  
215 198 private void validateMessageId(String mjMessageId, String messageId) {
... ... @@ -237,7 +220,7 @@ public class AiImageServiceImpl implements AiImageService {
237 220 }
238 221  
239 222 private AiImageDO validateExists(Long id) {
240   - AiImageDO aiImageDO = aiImageMapper.selectById(id);
  223 + AiImageDO aiImageDO = imageMapper.selectById(id);
241 224 if (aiImageDO == null) {
242 225 throw ServiceExceptionUtil.exception(ErrorCodeConstants.AI_MIDJOURNEY_IMAGINE_FAIL);
243 226 }
... ...