BJ_ZTEITS-Developers / parking-cloud-yudao

Commit 03b4460eae5144eb1e6cdbc29a3b405517cce013

Authored by cherishsince
1 parent ae934e84

【增加】对接 Midjourney，增加nonce传递，更新Midjourney image 状态

Inline Side-by-side
Showing 8 changed files with 140 additions and 15 deletions
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/AiImageController.java
  View file @03b4460
... ... @@ -4,7 +4,6 @@ import cn.iocoder.yudao.framework.common.pojo.CommonResult;
4 4 import cn.iocoder.yudao.module.ai.service.AiImageService;
5 5 import cn.iocoder.yudao.module.ai.vo.AiImageDallDrawingReq;
6 6 import cn.iocoder.yudao.module.ai.vo.AiImageMidjourneyReq;
7   -import cn.iocoder.yudao.module.ai.vo.AiImageMidjourneyRes;
8 7 import io.swagger.v3.oas.annotations.Operation;
9 8 import io.swagger.v3.oas.annotations.tags.Tag;
10 9 import lombok.AllArgsConstructor;
... ... @@ -42,7 +41,8 @@ public class AiImageController {
42 41  
43 42 @Operation(summary = "midjourney", description = "midjourney图片绘画流程：1、提交任务 2、获取完成的任务 3、选择对应功能 4、获取最终结果")
44 43 @PostMapping("/midjourney")
45   - public CommonResult<AiImageMidjourneyRes> midjourney(@Validated @RequestBody AiImageMidjourneyReq req) {
46   - return CommonResult.success(aiImageService.midjourney(req));
  44 + public CommonResult<Void> midjourney(@Validated @RequestBody AiImageMidjourneyReq req) {
  45 + aiImageService.midjourney(req);
  46 + return CommonResult.success(null);
47 47 }
48 48 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/AiImageService.java
  View file @03b4460
... ... @@ -28,5 +28,5 @@ public interface AiImageService {
28 28 * @param req
29 29 * @return
30 30 */
31   - AiImageMidjourneyRes midjourney(AiImageMidjourneyReq req);
  31 + void midjourney(AiImageMidjourneyReq req);
32 32 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/impl/AiImageServiceImpl.java
  View file @03b4460
... ... @@ -95,18 +95,15 @@ public class AiImageServiceImpl implements AiImageService {
95 95  
96 96 @Override
97 97 @Transactional(rollbackFor = Exception.class)
98   - public AiImageMidjourneyRes midjourney(AiImageMidjourneyReq req) {
  98 + public void midjourney(AiImageMidjourneyReq req) {
99 99 // 保存数据库
100   - doSave(req.getPrompt(), null, "midjoureny",
  100 + AiImageDO aiImageDO = doSave(req.getPrompt(), null, "midjoureny",
101 101 null, AiChatDrawingStatusEnum.SUBMIT, null);
102 102 // 提交 midjourney 任务
103   - Boolean imagine = midjourneyInteractionsApi.imagine(req.getPrompt());
  103 + Boolean imagine = midjourneyInteractionsApi.imagine(aiImageDO.getId(), req.getPrompt());
104 104 if (!imagine) {
105 105 throw ServiceExceptionUtil.exception(ErrorCodeConstants.AI_MIDJOURNEY_IMAGINE_FAIL);
106 106 }
107   - //
108   -
109   - return null;
110 107 }
111 108  
112 109 private static void sendSseEmitter(Utf8SseEmitter sseEmitter, Object object) {
... ... @@ -120,7 +117,7 @@ public class AiImageServiceImpl implements AiImageService {
120 117 }
121 118 }
122 119  
123   - private void doSave(String prompt,
  120 + private AiImageDO doSave(String prompt,
124 121 String size,
125 122 String model,
126 123 String imageUrl,
... ... @@ -138,5 +135,6 @@ public class AiImageServiceImpl implements AiImageService {
138 135 aiImageDO.setDrawingStatus(drawingStatusEnum.getStatus());
139 136 aiImageDO.setDrawingError(drawingError);
140 137 aiImageMapper.insert(aiImageDO);
  138 + return aiImageDO;
141 139 }
142 140 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/midjourneyHandler/YuDaoMidjourneyMessageHandler.java
  View file @03b4460
1 1 package cn.iocoder.yudao.module.ai.service.midjourneyHandler;
2 2  
  3 +import cn.hutool.core.collection.CollUtil;
  4 +import cn.hutool.core.util.StrUtil;
3 5 import cn.iocoder.yudao.framework.ai.midjourney.MidjourneyMessage;
  6 +import cn.iocoder.yudao.framework.ai.midjourney.constants.MidjourneyGennerateStatusEnum;
4 7 import cn.iocoder.yudao.framework.ai.midjourney.webSocket.MidjourneyMessageHandler;
  8 +import cn.iocoder.yudao.module.ai.dal.dataobject.AiImageDO;
  9 +import cn.iocoder.yudao.module.ai.enums.AiChatDrawingStatusEnum;
  10 +import cn.iocoder.yudao.module.ai.mapper.AiImageMapper;
  11 +import com.alibaba.fastjson2.JSON;
  12 +import lombok.AllArgsConstructor;
5 13 import lombok.extern.slf4j.Slf4j;
6 14 import org.springframework.stereotype.Component;
7 15  
... ... @@ -14,10 +22,51 @@ import org.springframework.stereotype.Component;
14 22 */
15 23 @Component
16 24 @Slf4j
  25 +@AllArgsConstructor
17 26 public class YuDaoMidjourneyMessageHandler implements MidjourneyMessageHandler {
18 27  
  28 + private final AiImageMapper aiImageMapper;
  29 +
19 30 @Override
20 31 public void messageHandler(MidjourneyMessage midjourneyMessage) {
21   - log.info("yudao-midjourney-midjourney-message-handler", midjourneyMessage);
  32 + log.info("yudao-midjourney-midjourney-message-handler {}", JSON.toJSONString(midjourneyMessage));
  33 + if (midjourneyMessage.getContent() != null) {
  34 + log.info("进度id {} 状态 {} 进度 {}",
  35 + midjourneyMessage.getNonce(),
  36 + midjourneyMessage.getGenerateStatus(),
  37 + midjourneyMessage.getContent().getProgress());
  38 + }
  39 + //
  40 + updateImage(midjourneyMessage);
  41 + }
  42 +
  43 + private void updateImage(MidjourneyMessage midjourneyMessage) {
  44 + // Nonce 不存在不更新
  45 + if (StrUtil.isBlank(midjourneyMessage.getNonce())) {
  46 + return;
  47 + }
  48 + // 获取id
  49 + Long aiImageId = Long.valueOf(midjourneyMessage.getNonce());
  50 + // 获取生成 url
  51 + String imageUrl = null;
  52 + if (CollUtil.isNotEmpty(midjourneyMessage.getAttachments())) {
  53 + imageUrl = midjourneyMessage.getAttachments().get(0).getUrl();
  54 + }
  55 + // 转换状态
  56 + AiChatDrawingStatusEnum drawingStatusEnum = null;
  57 + String generateStatus = midjourneyMessage.getGenerateStatus();
  58 + if (MidjourneyGennerateStatusEnum.COMPLETED.getStatus().equals(generateStatus)) {
  59 + drawingStatusEnum = AiChatDrawingStatusEnum.COMPLETE;
  60 + } else if (MidjourneyGennerateStatusEnum.IN_PROGRESS.getStatus().equals(generateStatus)) {
  61 + drawingStatusEnum = AiChatDrawingStatusEnum.IN_PROGRESS;
  62 + } else if (MidjourneyGennerateStatusEnum.WAITING.getStatus().equals(generateStatus)) {
  63 + drawingStatusEnum = AiChatDrawingStatusEnum.WAITING;
  64 + }
  65 + aiImageMapper.updateById(
  66 + new AiImageDO()
  67 + .setId(aiImageId)
  68 + .setDrawingImageUrl(imageUrl)
  69 + .setDrawingStatus(drawingStatusEnum == null ? null : drawingStatusEnum.getStatus())
  70 + );
22 71 }
23 72 }
... ...
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/MidjourneyMessage.java
  View file @03b4460
... ... @@ -15,6 +15,10 @@ public class MidjourneyMessage {
15 15 */
16 16 private String id;
17 17 /**
  18 + * 提交id(nonce 可能会不存在，系统提示的时候，这个为空)
  19 + */
  20 + private String nonce;
  21 + /**
18 22 * 现在已知：
19 23 * 0：我们发送的消息，和指令
20 24 * 20: mj生成图片发送过程中
... ... @@ -45,6 +49,14 @@ public class MidjourneyMessage {
45 49 * {@link MidjourneyGennerateStatusEnum}
46 50 */
47 51 private String generateStatus;
  52 + /**
  53 + * 一般用于提示信息
  54 + * - 错误
  55 + * - 并发队列满了
  56 + * - 账号违规了、敏感词
  57 + * - 账号被封
  58 + */
  59 + private List<Embed> embeds;
48 60  
49 61 @Data
50 62 @Accessors(chain = true)
... ... @@ -123,4 +135,39 @@ public class MidjourneyMessage {
123 135 private String progress;
124 136 private String status;
125 137 }
  138 +
  139 + /**
  140 + * embed 用于警告、提示、错误
  141 + */
  142 + @Data
  143 + @Accessors(chain = true)
  144 + public static class Embed {
  145 +
  146 + // 内容扫描版本号
  147 + private int contentScanVersion;
  148 +
  149 + // 颜色值，这里用Java的Color类来表示，注意实际使用中可能需要自定义方法来从int转换为Color对象
  150 + private String color;
  151 +
  152 + // 页脚信息，包含文本
  153 + private Footer footer;
  154 +
  155 + // 描述信息
  156 + private String description;
  157 +
  158 + // 消息类型，这里是富文本类型(这个区分不同提示类型)
  159 + private String type;
  160 +
  161 + // 标题
  162 + private String title;
  163 +
  164 + // Footer类，作为嵌套类存在，用来表示footer部分的JSON对象
  165 + @Data
  166 + @Accessors(chain = true)
  167 + public static class Footer {
  168 + // 页脚文本
  169 + private String text;
  170 + }
  171 +
  172 + }
126 173 }
... ...
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/api/MidjourneyInteractionsApi.java
  View file @03b4460
... ... @@ -38,11 +38,13 @@ public class MidjourneyInteractionsApi extends MidjourneyInteractions {
38 38 this.url = midjourneyConfig.getServerUrl().concat(midjourneyConfig.getApiInteractions());
39 39 }
40 40  
41   - public Boolean imagine(String prompt) {
  41 + public Boolean imagine(Long id, String prompt) {
  42 + String nonce = String.valueOf(id);
42 43 // 获取请求模板
43 44 String requestTemplate = midjourneyConfig.getRequestTemplates().get("imagine");
44 45 // 设置参数
45 46 HashMap<String, String> requestParams = getDefaultParams();
  47 + requestParams.put("nonce", nonce);
46 48 requestParams.put("prompt", prompt);
47 49 // 解析 template 参数占位符
48 50 String requestBody = MidjourneyUtil.parseTemplate(requestTemplate, requestParams);
... ...
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/constants/MidjourneyConstants.java
  View file @03b4460
... ... @@ -7,6 +7,10 @@ public final class MidjourneyConstants {
7 7 */
8 8 public static final String MSG_ID = "id";
9 9 /**
  10 + * 用于区分操作唯一性
  11 + */
  12 + public static final String MSG_NONCE = "nonce";
  13 + /**
10 14 * 消息 - 类型
11 15 * 现在已知：
12 16 * 0：我们发送的消息，和指令
... ... @@ -32,6 +36,10 @@ public final class MidjourneyConstants {
32 36 * 附件(生成中比较模糊的图片)
33 37 */
34 38 public static final String MSG_ATTACHMENTS = "attachments";
  39 + /**
  40 + * 一般用于提示
  41 + */
  42 + public static final String MSG_EMBEDS = "embeds";
35 43  
36 44  
37 45 //
... ...
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/listener/MidjourneyMessageListener.java
  View file @03b4460
... ... @@ -42,12 +42,14 @@ public class MidjourneyMessageListener {
42 42 if (ignoreAndLogMessage(data, messageType)) {
43 43 return;
44 44 }
  45 + log.info("socket message: {}", raw);
45 46 // 转换几个重要的信息
46 47 MidjourneyMessage mjMessage = new MidjourneyMessage();
47   - mjMessage.setId(data.getString(MidjourneyConstants.MSG_ID));
  48 + mjMessage.setId(getString(data, MidjourneyConstants.MSG_ID, ""));
  49 + mjMessage.setNonce(getString(data, MidjourneyConstants.MSG_NONCE, ""));
48 50 mjMessage.setType(data.getInt(MidjourneyConstants.MSG_TYPE));
49 51 mjMessage.setRawData(StrUtil.str(raw.toJson(), "UTF-8"));
50   - mjMessage.setContent(MidjourneyUtil.parseContent(data.getString(MidjourneyConstants.MSG_CONTENT)));
  52 + mjMessage.setContent(MidjourneyUtil.parseContent(data.getString(MidjourneyConstants.MSG_CONTENT)));
51 53 // 转换 components
52 54 if (!data.getArray(MidjourneyConstants.MSG_COMPONENTS).isEmpty()) {
53 55 String componentsJson = StrUtil.str(data.getArray(MidjourneyConstants.MSG_COMPONENTS).toJson(), "UTF-8");
... ... @@ -60,6 +62,12 @@ public class MidjourneyMessageListener {
60 62 List<MidjourneyMessage.Attachment> attachments = JsonUtils.parseArray(attachmentsJson, MidjourneyMessage.Attachment.class);
61 63 mjMessage.setAttachments(attachments);
62 64 }
  65 + // 转换 embeds 提示信息
  66 + if (!data.getArray(MidjourneyConstants.MSG_EMBEDS).isEmpty()) {
  67 + String embedJson = StrUtil.str(data.getArray(MidjourneyConstants.MSG_EMBEDS).toJson(), "UTF-8");
  68 + List<MidjourneyMessage.Embed> embeds = JsonUtils.parseArray(embedJson, MidjourneyMessage.Embed.class);
  69 + mjMessage.setEmbeds(embeds);
  70 + }
63 71 // 转换状态
64 72 convertGenerateStatus(mjMessage);
65 73 // message handler 调用
... ... @@ -68,7 +76,20 @@ public class MidjourneyMessageListener {
68 76 }
69 77 }
70 78  
  79 + private String getString(DataObject data, String key, String defaultValue) {
  80 + if (!data.hasKey(key)) {
  81 + return defaultValue;
  82 + }
  83 + return data.getString(key);
  84 + }
  85 +
71 86 private void convertGenerateStatus(MidjourneyMessage mjMessage) {
  87 + //
  88 + // tip：提示、警告、异常 content是没有内容的
  89 + // tip: 一般错误信息在 Embeds 只要 Embeds有值，content就没信息。
  90 + if (CollUtil.isNotEmpty(mjMessage.getEmbeds())) {
  91 + return;
  92 + }
72 93 if (mjMessage.getType() == 20 && mjMessage.getContent().getStatus().contains("Waiting")) {
73 94 mjMessage.setGenerateStatus(MidjourneyGennerateStatusEnum.WAITING.getStatus());
74 95 } else if (mjMessage.getType() == 20 && !StrUtil.isBlank(mjMessage.getContent().getProgress())) {
... ...